diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 41198887cab..fb29a92e306 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -1035,9 +1035,9 @@ jobs: run: | diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils else - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))" @@ -1294,7 +1294,7 @@ jobs: - name: Install dependencies if: steps.changes.outputs.decision != 'skip' run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard - name: Cache Prisma binaries if: steps.changes.outputs.decision != 'skip' @@ -1497,9 +1497,9 @@ jobs: RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }} run: | if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils else - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so fi diff --git a/AGENTS.md b/AGENTS.md index dd0561617dd..7bea7a56cbc 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -52,7 +52,6 @@ If you ever write any human-facing text (pull requests, issues, commit messages, - don't use "—". Don't use ";" or ":" unless it would be nonsensical not to. Instead, reach for ",", ".", conjunction words, etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them also feels AI-y. A word cap does not penalize you for adding more sentences - don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc. - unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose -- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a ".". Exception is that docs (i.e., litellm-docs) should have trailing "." - don't use →. Instead, prefer not to use arrows, and if need be, use -> instead - use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure diff --git a/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md index e1f606a2eb7..538a5f8a404 100644 --- a/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md +++ b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md @@ -46,7 +46,43 @@ For a managed workspace rollout, use OpenAI's [gateway configuration guide](http ## Native compatibility and Pi -For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials but does not need a model deployment or resolve model aliases. Restricted keys must allow the upstream model name. +For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials. A ChatGPT model keeps its upstream name and needs no deployment of its own. + +### Mix ChatGPT and LiteLLM models + +On the native route, Codex's model picker shows OpenAI's catalog first and then the other chat and responses models the virtual key can use. Pick one of those and Codex's requests go through LiteLLM's `/v1/responses` with the key's usual routing, budgets, guardrails, and logging. The proxy removes the ChatGPT bearer and account header from that request, so the other provider and your callbacks never receive them. A name is routed this way when another provider serves it as an exact `model_name`, through a router `model_group_alias`, or as a team's public model name. A ChatGPT model, a name that only a wildcard deployment matches, or a name OpenAI's catalog lists for the caller's ChatGPT account keeps the passthrough, so a model OpenAI adds later still uses the subscription even when a deployment of another provider has the same name. The proxy learns the account's list from the catalog Codex fetches, or asks OpenAI with the caller's login. When it cannot get the list, for example because OpenAI is down, the request times out, or the client is not Codex, every request keeps the passthrough for the next minute, as it did before LiteLLM models were added to the picker. When the key can use many models, the picker lists the ones that fit in Codex's catalog size limit, and the proxy logs a warning with how many it left out. + +Codex's picker sends a `none` reasoning effort for a LiteLLM model whose catalog entry lists no effort levels. With `litellm_settings: drop_params: true`, the proxy drops it for a model that does not support `none`, and the provider applies its own default. A model whose cost map entry or `model_info` sets `supports_none_reasoning_effort: true` keeps it. Without `drop_params`, the effort reaches the provider unchanged. + +Codex also sends its ChatGPT app tools with every turn, which on a LiteLLM model adds a few hundred thousand prompt tokens to each turn. Set `apps = false` under `[features]` in Codex's `config.toml` to turn them off. + +Only a key that may call `/v1/responses` gets these models. A key whose `allowed_routes` covers only the native routes sees OpenAI's catalog unchanged, and every request it sends uses the subscription. + +Codex sends earlier reasoning items back with each turn. The proxy tags the reasoning that a LiteLLM model returns on this route, OpenAI and Azure deployments included. When a session switches to a ChatGPT model, the proxy drops those items before forwarding the request to OpenAI, which would reject them. Messages, tool calls, and ChatGPT's own reasoning are kept, and a LiteLLM model still gets its own reasoning back. + +When the catalog adds LiteLLM models for a key, its `ETag` is a digest of OpenAI's ETag and those models, and that key's responses carry the same value in `X-Models-Etag`. Codex keeps its cached catalog while the two match and fetches it again when OpenAI's catalog or the key's models change. A key whose catalog adds no LiteLLM models keeps OpenAI's ETags. + +### Restrict native models with an access group + +A key with a model allowlist needs a rule for ChatGPT names that OpenAI adds over time. Put the subscription wildcards in an access group under `chatgpt/` names: + +```yaml +model_list: + - model_name: chatgpt/gpt-* + litellm_params: + model: chatgpt/gpt-* + model_info: + access_groups: ["chatgpt-native"] + - model_name: chatgpt/codex-* + litellm_params: + model: chatgpt/codex-* + model_info: + access_groups: ["chatgpt-native"] +``` + +Then add `chatgpt-native` to the key's models, next to any other model names it should reach. On the native route, a key that holds only the `chatgpt/` form of a name is checked as that form, so a request for `gpt-5.6-sol` is allowed when the key may call `chatgpt/gpt-5.6-sol`. The key can call every ChatGPT model that matches the wildcards, including ones the proxy configuration does not name. + +On `/v1` these wildcards match only a request that asks for a `chatgpt/` name, so a request for a plain `gpt-` or `codex-` name still goes where it went before, such as a `*` catch-all deployment. A `/v1` request for a `chatgpt/` name without a ChatGPT login falls back to the proxy's own saved login if it has one, so keep that login unset as described above. ### Configure Pi @@ -71,7 +107,7 @@ Keep Pi's subscription login as the provider credential. The header value names Send a prompt and check the completed request in LiteLLM's Logs page. Its gateway user/key and token counts should match the caller. Give each person their own LiteLLM virtual key for separate attribution. The gateway credential and ChatGPT login are independently authenticated; LiteLLM does not establish that both identities belong to the same person. -The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key must allow that name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions. +The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key needs the `chatgpt-native` access group above or that exact name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions. The proxy forwards the ChatGPT bearer and account header to the fixed OpenAI Codex endpoint. It removes gateway credentials and cookies from the upstream request and redacts credential headers in logging callbacks. The client remains responsible for refresh and reauthentication. Native passthrough returns authentication and quota errors without switching subscriptions. Standard model routes retain the caller login across ChatGPT retries; explicitly configured cross-provider fallbacks use their own credentials. An explicit workspace header is preserved for OpenAI to validate. diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 7db2de1f874..64fe03769b8 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.81" +version = "0.1.83" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.81" +version = "0.1.83" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index a12e1e27cfc..1c0c708ae7a 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.112" +version = "0.4.114" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -30,7 +30,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.112" +version = "0.4.114" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm-rust/crates/python-bridge/src/router/host.rs b/litellm-rust/crates/python-bridge/src/router/host.rs index 6ccac4df659..0ded98f27c5 100644 --- a/litellm-rust/crates/python-bridge/src/router/host.rs +++ b/litellm-rust/crates/python-bridge/src/router/host.rs @@ -20,10 +20,10 @@ use std::{ use async_trait::async_trait; use litellm_router_core::{ - Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, Deployment, Error, Host, Output, - Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache, RouterEvent, RouterRequest, - RouterResponse, RoutingContext, SecretRef, SelectStage, ThreadRandomness, - build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo, + Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, CallbackFiltered, Deployment, + Error, Host, Output, Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache, + RouterEvent, RouterRequest, RouterResponse, RoutingContext, SecretRef, SelectStage, + ThreadRandomness, build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo, precall::PreCallFacts, }; use pyo3::{PyResult, Python, types::PyAnyMethods}; @@ -298,8 +298,9 @@ pub(crate) enum HostCallResponse { PreCall(PreCallFacts), /// The candidate models a `routing_plugins` select step reported. Models(Option>), - /// The ids a `callback_filters` select step left, or the error a filter raised. - Filtered(Result, AttemptError>), + /// The ids a `callback_filters` select step left with the one a filter prefers, or the error + /// a filter raised. + Filtered(Result<(Vec, Option), AttemptError>), /// What a pick's select or pre-route step, or a fallback check, raised on the host instead of /// answering: the walk takes it as the try's failure. Raised(AttemptError), @@ -352,8 +353,10 @@ impl HostCallResponse { got: self.kind(), }); } - if let (Self::Select(ids) | Self::Filtered(Ok(ids)), HostCall::Select { candidates, .. }) = - (self, call) + if let ( + Self::Select(ids) | Self::Filtered(Ok((ids, _))), + HostCall::Select { candidates, .. }, + ) = (self, call) && let Some(id) = ids.iter().find(|id| !candidates.contains(id)) { return Err(DriveError::UnknownCandidate { id: id.clone() }); @@ -890,7 +893,7 @@ impl Host for StepHost { &self, ctx: &RoutingContext<'_>, candidates: Vec>, - ) -> Result>, Error> { + ) -> Result { let (scope, group) = ctx.scope.wire(); let joined = |own: &[String], walked: &[String]| [own, walked].concat(); let call = HostCall::Select { @@ -909,8 +912,11 @@ impl Host for StepHost { let HostCallResponse::Filtered(result) = call_host(call).await else { unreachable!("{CHECKED}"); }; - let ids = result.map_err(Error::Attempt)?; - Ok(answered_by_id(&candidates, &ids)) + let (ids, preferred) = result.map_err(Error::Attempt)?; + Ok(CallbackFiltered { + kept: answered_by_id(&candidates, &ids), + preferred, + }) } async fn allow_fallback( diff --git a/litellm-rust/crates/python-bridge/src/router/mod.rs b/litellm-rust/crates/python-bridge/src/router/mod.rs index 8157d0f5a32..6b58d63a714 100644 --- a/litellm-rust/crates/python-bridge/src/router/mod.rs +++ b/litellm-rust/crates/python-bridge/src/router/mod.rs @@ -1181,7 +1181,8 @@ fn host_call_response( stage: Stage::CallbackFilters { .. }, .. } => { - // The remaining ids, or a failure mapping as an attempt reports one. + // The remaining ids, alone or paired with the id a filter prefers, or a failure + // mapping as an attempt reports one. let outcome = outcome.ok_or_else(missing)?; if let Ok(failure) = outcome.cast::() { return match attempt_outcome(failure)? { @@ -1191,7 +1192,10 @@ fn host_call_response( )), }; } - Ok(HostCallResponse::Filtered(Ok(outcome.extract()?))) + if let Ok(answer) = outcome.cast::() { + return Ok(HostCallResponse::Filtered(Ok(answer.extract()?))); + } + Ok(HostCallResponse::Filtered(Ok((outcome.extract()?, None)))) } HostCall::Select { stage: Stage::Strategy { records, .. }, diff --git a/litellm-rust/crates/router-core/src/error.rs b/litellm-rust/crates/router-core/src/error.rs index 926772f435c..1b718ed2c12 100644 --- a/litellm-rust/crates/router-core/src/error.rs +++ b/litellm-rust/crates/router-core/src/error.rs @@ -284,6 +284,7 @@ impl Error { strum::EnumString, strum::IntoStaticStr, strum::VariantArray, + serde_with::SerializeDisplay, )] pub enum FailureKind { #[strum(serialize = "RateLimitError")] @@ -471,6 +472,17 @@ mod tests { assert_eq!(<&'static str>::from(kind), wire); } + #[rstest] + #[case::rate_limit(FailureKind::RateLimit, "RateLimitError")] + #[case::timeout(FailureKind::Timeout, "Timeout")] + #[case::api_connection(FailureKind::ApiConnection, "APIConnectionError")] + fn failure_kind_serializes_as_its_exception_name( + #[case] kind: FailureKind, + #[case] wire: &str, + ) { + assert_eq!(serde_json::to_value(kind).unwrap(), serde_json::json!(wire)); + } + #[rstest] #[case::empty("")] #[case::lowercase("ratelimiterror")] diff --git a/litellm-rust/crates/router-core/src/host.rs b/litellm-rust/crates/router-core/src/host.rs index cfc97480bf1..ae93288b0fa 100644 --- a/litellm-rust/crates/router-core/src/host.rs +++ b/litellm-rust/crates/router-core/src/host.rs @@ -342,7 +342,10 @@ pub struct Output { } /// How a provider blocked a response the call otherwise served. -#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr, strum::EnumString)] +#[derive( + Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, strum::IntoStaticStr, strum::EnumString, +)] +#[serde(rename_all = "snake_case")] #[strum(serialize_all = "snake_case")] pub enum Blocked { /// A chat response whose first choice was content-filtered (or that has no choices). @@ -402,8 +405,16 @@ pub enum SelectStage { OverrideCheck(RoutingStrategy), } -/// Routing events, delivered through [`Host::observe`]. -#[derive(Clone, Debug, PartialEq)] +#[derive(Debug)] +pub struct CallbackFiltered { + pub kept: Vec>, + pub preferred: Option, +} + +/// Routing events, delivered through [`Host::observe`]. They serialize as JSON objects tagged by +/// `event` (the variant name), durations in seconds. +#[derive(Clone, Debug, PartialEq, serde::Serialize)] +#[serde(tag = "event")] pub enum RouterEvent { /// A pick for `model` (the level's group) begins: its route phase opens. PickStarted { @@ -413,6 +424,7 @@ pub enum RouterEvent { model_group: String, deployment_id: String, /// How long the pick took (`_acompletion`'s `async_get_available_deployment` service event). + #[serde(serialize_with = "seconds")] pick: Duration, }, Succeeded { @@ -425,11 +437,13 @@ pub enum RouterEvent { }, CooldownSet { deployment_id: String, + #[serde(serialize_with = "seconds")] cooldown: Duration, }, Retry { model_group: String, retry: u32, + #[serde(serialize_with = "seconds")] delay: Duration, }, Fallback { @@ -584,6 +598,13 @@ impl PickError { } } +/// A pick's error serializes as its message. +impl serde::Serialize for PickError { + fn serialize(&self, serializer: S) -> Result { + serializer.collect_str(&self.0) + } +} + impl PartialEq for PickError { fn eq(&self, other: &Self) -> bool { format!("{:?}", self.0) == format!("{:?}", other.0) @@ -591,7 +612,8 @@ impl PartialEq for PickError { } /// Which error a [`RouterEvent::RetryLogged`] records. -#[derive(Clone, Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq, serde::Serialize)] +#[serde(tag = "cause", content = "error", rename_all = "snake_case")] pub enum RetryCause { /// The last attempt's failure. Attempt, @@ -698,14 +720,19 @@ pub trait Host: Send + Sync { } /// `async_callback_filter_deployments`: the host's routing callbacks narrow `candidates` in - /// turn. A filter that raises fails the try with [`Error::Attempt`], which the call retries and - /// falls back from like any attempt failure. A host without filters keeps every candidate. + /// turn, and may name one deployment they prefer, which the pick takes only when the filters + /// after the callbacks keep it. A filter that raises fails the try with [`Error::Attempt`], + /// which the call retries and falls back from like any attempt failure. A host without + /// filters keeps every candidate and prefers none. async fn callback_filters( &self, _ctx: &RoutingContext<'_>, candidates: Vec>, - ) -> Result>, Error> { - Ok(candidates) + ) -> Result { + Ok(CallbackFiltered { + kept: candidates, + preferred: None, + }) } /// The values of `refs`, one per ref, in order: what the deployment build substitutes for /// each `os.environ/` reference it found in entry `index` ([`crate::build`]). Asked once per @@ -745,3 +772,83 @@ pub trait Host: Send + Sync { Ok(Value::Null) } } + +fn seconds(duration: &Duration, serializer: S) -> Result { + serializer.serialize_f64(duration.as_secs_f64()) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::{Value, json}; + + use super::*; + + #[rstest] + #[case::selected( + RouterEvent::Selected { + model_group: "gpt-4o".to_owned(), + deployment_id: "azure-eastus".to_owned(), + pick: Duration::from_millis(1500), + }, + json!({"event": "Selected", "model_group": "gpt-4o", "deployment_id": "azure-eastus", "pick": 1.5}) + )] + #[case::attempt_failed( + RouterEvent::AttemptFailed { + deployment_id: "openai".to_owned(), + kind: FailureKind::RateLimit, + status: Some(429), + }, + json!({"event": "AttemptFailed", "deployment_id": "openai", "kind": "RateLimitError", "status": 429}) + )] + #[case::cooldown_set( + RouterEvent::CooldownSet { deployment_id: "openai".to_owned(), cooldown: Duration::from_secs(30) }, + json!({"event": "CooldownSet", "deployment_id": "openai", "cooldown": 30.0}) + )] + #[case::retry( + RouterEvent::Retry { model_group: "gpt-4o".to_owned(), retry: 1, delay: Duration::ZERO }, + json!({"event": "Retry", "model_group": "gpt-4o", "retry": 1, "delay": 0.0}) + )] + #[case::blocked( + RouterEvent::Blocked { level: 0, model: "claude".to_owned(), blocked: Blocked::ContentFilter }, + json!({"event": "Blocked", "level": 0, "model": "claude", "blocked": "content_filter"}) + )] + #[case::retry_logged_after_an_attempt( + RouterEvent::RetryLogged { + scope: 0, + level: 0, + model: "gpt-4o".to_owned(), + attempted_retries: None, + cause: RetryCause::Attempt, + }, + json!({ + "event": "RetryLogged", "scope": 0, "level": 0, "model": "gpt-4o", + "attempted_retries": null, "cause": {"cause": "attempt"}, + }) + )] + #[case::retry_logged_after_a_pick( + RouterEvent::RetryLogged { + scope: 1, + level: 1, + model: "backup".to_owned(), + attempted_retries: Some(2), + cause: RetryCause::Pick(PickError(Arc::new(Error::UnknownModel { model: "backup".to_owned() }))), + }, + json!({ + "event": "RetryLogged", "scope": 1, "level": 1, "model": "backup", "attempted_retries": 2, + "cause": {"cause": "pick", "error": "no deployment serves model `backup`"}, + }) + )] + #[case::mock_failed( + RouterEvent::MockFailed { + mock: MockFailure::RateLimit, + level: 0, + model: "gpt-4o".to_owned(), + num_retries: None, + }, + json!({"event": "MockFailed", "mock": "rate_limit", "level": 0, "model": "gpt-4o", "num_retries": null}) + )] + fn an_event_serializes_tagged_by_its_name(#[case] event: RouterEvent, #[case] expected: Value) { + assert_eq!(serde_json::to_value(&event).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs index 80c8f597157..f9b71c689c9 100644 --- a/litellm-rust/crates/router-core/src/lib.rs +++ b/litellm-rust/crates/router-core/src/lib.rs @@ -124,8 +124,8 @@ pub use crate::{ error::{AttemptError, ConfigError, Error, FailureKind, Payload, Unsupported}, health::{Clock, Health, Outcome, SystemClock}, host::{ - Attempt, Blocked, BuildAsk, BuildPath, Built, Host, Output, Provided, RetryCause, - RouterEvent, RouterRequest, RoutingContext, SecretAt, SecretRef, SelectStage, + Attempt, Blocked, BuildAsk, BuildPath, Built, CallbackFiltered, Host, Output, Provided, + RetryCause, RouterEvent, RouterRequest, RoutingContext, SecretAt, SecretRef, SelectStage, }, outcome::{OutcomeCallbacks, OutcomeWrite}, prerouting::{PreRoute, PreRouted}, diff --git a/litellm-rust/crates/router-core/src/mock.rs b/litellm-rust/crates/router-core/src/mock.rs index 898782e6203..1db6c0d8019 100644 --- a/litellm-rust/crates/router-core/src/mock.rs +++ b/litellm-rust/crates/router-core/src/mock.rs @@ -7,7 +7,15 @@ use crate::error::{AttemptError, FailureKind}; /// Which mock failure a request asks for. #[derive( - Clone, Copy, Debug, PartialEq, Eq, Deserialize, strum::IntoStaticStr, strum::EnumString, + Clone, + Copy, + Debug, + PartialEq, + Eq, + Deserialize, + serde::Serialize, + strum::IntoStaticStr, + strum::EnumString, )] #[serde(rename_all = "snake_case")] #[strum(serialize_all = "snake_case")] diff --git a/litellm-rust/crates/router-core/src/select.rs b/litellm-rust/crates/router-core/src/select.rs index c7f3b1f9c31..75198c0a707 100644 --- a/litellm-rust/crates/router-core/src/select.rs +++ b/litellm-rust/crates/router-core/src/select.rs @@ -456,7 +456,7 @@ impl Core { let request_order = request.target_order.filter(|_| !inputs.fallback_hop); // `async_callback_filter_deployments` runs after the cooldown and blocked filters on every // async candidate list but a deployment named by id and a listed pass-through pick. - if settings.callback_filters + let preferred = if settings.callback_filters && !inputs.count_only && !resolution.kind.is_single() && request.pass_through != Some(PassThrough::Listed) @@ -472,8 +472,12 @@ impl Core { ctx.excluded_ids = &inputs.excluded_ids; ctx.skipped_ids = &inputs.skipped_ids; ctx.target_order = inputs.target_order.or(request_order); - healthy = self.host.callback_filters(&ctx, healthy).await?; - } + let filtered = self.host.callback_filters(&ctx, healthy).await?; + healthy = filtered.kept; + filtered.preferred + } else { + None + }; if inputs.count_only { return Ok(Candidates { name: name.to_owned(), @@ -648,8 +652,13 @@ impl Core { healthy.retain(|d| !d.cannot_run_batches()); } - if let Some(pinned) = request.pinned_deployment_id.as_deref() - && healthy.iter().any(|d| d.id == pinned) + if let Some(pinned) = [ + request.pinned_deployment_id.as_deref(), + preferred.as_deref(), + ] + .into_iter() + .flatten() + .find(|id| healthy.iter().any(|d| d.id == *id)) { healthy.retain(|d| d.id == pinned); } diff --git a/litellm-rust/crates/router-core/tests/call_loop.rs b/litellm-rust/crates/router-core/tests/call_loop.rs index a2a741d6135..cd51e53ec6b 100644 --- a/litellm-rust/crates/router-core/tests/call_loop.rs +++ b/litellm-rust/crates/router-core/tests/call_loop.rs @@ -17,7 +17,7 @@ use litellm_router_core::{ }; use parking_lot::Mutex; use rstest::{fixture, rstest}; -use serde_json::json; +use serde_json::{Value, json}; #[derive(Default)] struct ScriptedHost { @@ -1809,3 +1809,55 @@ fn request_for(model: &str) -> RouterRequest { ..RouterRequest::default() } } + +#[rstest] +#[tokio::test] +async fn every_event_of_a_retried_and_fallen_back_call_serializes_tagged_by_its_name( + config: RouterConfig, +) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![failure(FailureKind::RateLimit, 429); 3], + )); + let router = router( + config, + json!({"num_retries": 1, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + let tags: Value = host + .events + .lock() + .iter() + .map(|event| serde_json::to_value(event).unwrap()["event"].clone()) + .collect(); + assert_eq!( + tags, + json!([ + "RetriesOpened", + "PickStarted", + "Selected", + "AttemptFailed", + "CooldownSet", + "RetryBudget", + "RetryLogged", + "Retry", + "PickStarted", + "Selected", + "AttemptFailed", + "CooldownSet", + "RetryLogged", + "FallingBack", + "RetryLogged", + "Fallback", + "RetriesOpened", + "PickStarted", + "Selected", + "Succeeded", + "FallbackDone", + ]) + ); +} diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py index 08e95f99c87..441b326f4d2 100644 --- a/litellm/a2a_protocol/card_resolver.py +++ b/litellm/a2a_protocol/card_resolver.py @@ -4,13 +4,14 @@ Custom A2A Card Resolver for LiteLLM. Extends the A2A SDK's card resolver to support multiple well-known paths. """ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, runtime_checkable from litellm._logging import verbose_logger from litellm.a2a_protocol.exceptions import A2AAgentCardDiscoveryError from litellm.constants import LOCALHOST_URL_PATTERNS +from litellm.litellm_core_utils.core_helpers import has_attributes if TYPE_CHECKING: from a2a.types import AgentCard @@ -130,15 +131,34 @@ def normalize_agent_card_interfaces(agent_card: "AgentCard", protocol_version: s return normalized +class _UrlCarrier(Protocol): + url: str + + +class _InterfacesCarrier(Protocol): + @property + def supported_interfaces(self) -> Sequence[_UrlCarrier]: ... + + +def _url_of(value: object) -> str | None: + return value.url if has_attributes(value, _UrlCarrier, "url") else None + + +def _supported_interfaces(agent_card: object) -> Sequence[_UrlCarrier] | None: + if has_attributes(agent_card, _InterfacesCarrier, "supported_interfaces"): + return agent_card.supported_interfaces + return None + + def get_agent_card_url(agent_card: "AgentCard") -> str | None: """Return the agent endpoint URL from the resolved SDK card.""" - url: Final = getattr(agent_card, "url", None) + url: Final = _url_of(agent_card) if url: return url - interfaces: Final = getattr(agent_card, "supported_interfaces", None) + interfaces: Final = _supported_interfaces(agent_card) if interfaces: - return getattr(interfaces[0], "url", None) + return _url_of(interfaces[0]) return None @@ -168,16 +188,16 @@ def fix_agent_card_url(agent_card: "AgentCard", base_url: str) -> "AgentCard": Returns: The agent card with the URL fixed if necessary """ - card_url: Final = getattr(agent_card, "url", None) + card_url: Final = _url_of(agent_card) if card_url and is_localhost_or_internal_url(card_url): # Normalize base_url to ensure it ends with / fixed_url: Final = base_url.rstrip("/") + "/" agent_card.url = fixed_url - interfaces: Final = getattr(agent_card, "supported_interfaces", None) + interfaces: Final = _supported_interfaces(agent_card) if interfaces: - interface_url: Final = getattr(interfaces[0], "url", None) + interface_url: Final = _url_of(interfaces[0]) if interface_url and is_localhost_or_internal_url(interface_url): interfaces[0].url = base_url.rstrip("/") + "/" @@ -194,6 +214,9 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): - /agentCard/v1.0 """ + if TYPE_CHECKING: + base_url: str + async def get_agent_card( self, relative_card_path: str | None = None, diff --git a/litellm/a2a_protocol/exception_mapping_utils.py b/litellm/a2a_protocol/exception_mapping_utils.py index 1abf32915f5..1fe8b2bf0af 100644 --- a/litellm/a2a_protocol/exception_mapping_utils.py +++ b/litellm/a2a_protocol/exception_mapping_utils.py @@ -4,7 +4,7 @@ A2A Protocol Exception Mapping Utils. Maps A2A SDK exceptions to LiteLLM A2A exception types. """ -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger from litellm.a2a_protocol.card_resolver import ( @@ -21,6 +21,7 @@ from litellm.constants import CONNECTION_ERROR_PATTERNS if TYPE_CHECKING: from a2a.client import Client as A2AClientType + from a2a.types import AgentCard _A2A_KIND_DIALECT_HINT: Final = ( @@ -165,7 +166,7 @@ def map_a2a_exception( async def handle_a2a_localhost_retry( error: A2ALocalhostURLError, - agent_card: Any, + agent_card: "AgentCard | None", a2a_client: "A2AClientType", is_streaming: bool = False, ) -> "A2AClientType": diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index 1e003a6e4bc..06a910a08fd 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -10,7 +10,7 @@ A2A Streaming Events (in order): 4. Status update (kind: "status-update") - Final status "completed" with final=true """ -from collections.abc import AsyncIterator, Callable, Coroutine, Mapping +from collections.abc import AsyncIterator, Mapping from types import MappingProxyType from typing import Any, Final @@ -126,10 +126,7 @@ class A2ACompletionBridgeHandler: @staticmethod async def _acompletion(completion_params: Mapping[str, object]) -> ModelResponse | CustomStreamWrapper: - acompletion_fn: Final[Callable[..., Coroutine[object, object, ModelResponse | CustomStreamWrapper]]] = vars( - litellm - )["acompletion"] - return await acompletion_fn(**completion_params) + return await litellm.acompletion(**completion_params) @staticmethod async def handle_non_streaming( diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index dcbe7a6506e..7d617613b0c 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -14,7 +14,7 @@ import datetime import uuid from collections.abc import AsyncIterator, Coroutine, Mapping from types import MappingProxyType, ModuleType -from typing import TYPE_CHECKING, Final, Optional, cast +from typing import TYPE_CHECKING, Final, Optional, Protocol, cast from typing_extensions import TypedDict, Unpack @@ -24,6 +24,7 @@ from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator from litellm.a2a_protocol.utils import A2ARequestUtils from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -36,7 +37,6 @@ if TYPE_CHECKING: from a2a.client import Client as A2AClientType from a2a.client import ClientCallContext as A2ACallContextType from a2a.compat.v0_3.types import ( - AgentCard, Message, SendMessageRequest, SendMessageResponse, @@ -45,6 +45,7 @@ if TYPE_CHECKING: SendStreamingMessageSuccessResponse, Task, ) + from a2a.types import AgentCard from a2a.types.a2a_pb2 import SendMessageRequest as CoreSendMessageRequest from a2a.types.a2a_pb2 import StreamResponse as CoreStreamResponse @@ -95,6 +96,12 @@ class _SendMessageOptions(TypedDict, total=False, extra_items=object): pass +class _StreamingRequestConverter(Protocol): + def to_core_send_message_request( + self, compat_req: "SendStreamingMessageRequest", / + ) -> "CoreSendMessageRequest": ... + + def _set_usage_on_logging_obj( kwargs: Mapping[str, object], prompt_tokens: int, @@ -375,7 +382,9 @@ async def _stream_messages( "The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk" ) - pb_request: Final[CoreSendMessageRequest] = _a2a_conversions.to_core_send_message_request(request) + pb_request: Final[CoreSendMessageRequest] = require_attributes( + _a2a_conversions, _StreamingRequestConverter, "to_core_send_message_request" + ).to_core_send_message_request(request) async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)): compat_chunk = _to_compat_stream_response(event, request_id=request.id) yield SendStreamingMessageResponse(root=compat_chunk) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 3b7ea66ad6b..b9cef2ea301 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -13,10 +13,10 @@ import json import logging import time import traceback -from collections.abc import Generator, Mapping +from collections.abc import Awaitable, Callable, Generator, Mapping from contextlib import contextmanager from enum import Enum -from typing import Any, Final, Literal, cast +from typing import Any, Final, Literal, Protocol, cast from pydantic import BaseModel, ConfigDict, TypeAdapter @@ -25,6 +25,7 @@ from litellm._internal_context import current_service_target, service_target from litellm._logging import verbose_logger from litellm.constants import CACHED_STREAMING_CHUNK_DELAY from litellm.integrations.otel.runtime import phase_span +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.model_param_helper import ( ModelParamHelper, embeds_thought_signature_in_tool_call_id, @@ -52,6 +53,16 @@ _REDIS_STARTUP_NODES_ADAPTER: Final = TypeAdapter( ) +class _SupportsAsyncPing(Protocol): + @property + def ping(self) -> Callable[[], Awaitable[object]] | None: ... + + +class _SupportsAsyncDeleteCacheKeys(Protocol): + @property + def delete_cache_keys(self) -> Callable[[object], Awaitable[object]] | None: ... + + def _keyed_provider_params_in_extra_body(kwargs: Mapping[str, object]) -> str: raw_extra_body: Final = kwargs.get("extra_body") if not isinstance(raw_extra_body, dict): @@ -1032,14 +1043,16 @@ class Cache: cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs) await self.cache.batch_cache_write(cache_key, cached_data, **kwargs) - async def ping(self): - cache_ping: Final = getattr(self.cache, "ping") + async def ping(self) -> object: + cache_ping: Final = require_attributes(self.cache, _SupportsAsyncPing, "ping").ping if cache_ping: return await cache_ping() return None - async def delete_cache_keys(self, keys): - cache_delete_cache_keys: Final = getattr(self.cache, "delete_cache_keys") + async def delete_cache_keys(self, keys: object) -> object: + cache_delete_cache_keys: Final = require_attributes( + self.cache, _SupportsAsyncDeleteCacheKeys, "delete_cache_keys" + ).delete_cache_keys if cache_delete_cache_keys: return await cache_delete_cache_keys(keys) return None diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 1d80d6952a4..82804d2d707 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -22,6 +22,7 @@ from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError +from typing_extensions import TypedDict import litellm from litellm._internal_context import is_internal_call, post_response_phase @@ -90,6 +91,10 @@ class CachingHandlerResponse(LiteLLMBaseModel): in_memory_cache_obj: Final = InMemoryCache() +class _CacheCallKwargs(TypedDict, total=False, extra_items=object): + pass + + def _drop_logging_obj_from_kwargs(request_kwargs: dict[str, object]) -> dict[str, object]: """ The caching handler is stored on the Logging object @@ -455,7 +460,7 @@ class LLMCachingHandler: if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function): args = args or () # Now that we confirmed caching will happen, prepare kwargs - new_kwargs: Final = kwargs.copy() + new_kwargs: Final = dict[str, object](kwargs) new_kwargs.update( convert_args_to_kwargs( self.original_function, @@ -470,7 +475,7 @@ class LLMCachingHandler: print_verbose("Checking Sync Cache") logging_obj.reset_caching_details() with response_cache_phase("get"): - cached_result = litellm.cache.get_cache(**new_kwargs) + cached_result = litellm.cache.get_cache(**_CacheCallKwargs(**new_kwargs)) if is_response_without_output(cached_result): verbose_logger.debug("LiteLLM Cache: cached response has no output, treating it as a miss") return CachingHandlerResponse(cached_result=None) @@ -480,7 +485,7 @@ class LLMCachingHandler: pass else: call_type = original_function.__name__ - cached_result = self._convert_cached_result_to_model_response( + converted_result: Final = self._convert_cached_result_to_model_response( cached_result=cached_result, call_type=call_type, kwargs=kwargs, @@ -508,16 +513,16 @@ class LLMCachingHandler: logging_obj=logging_obj, model=f"{custom_llm_provider}/{model}", kwargs=kwargs, - cached_result=cached_result, + cached_result=converted_result, is_async=False, custom_llm_provider=custom_llm_provider, ) if not is_internal_call.get() and not _should_defer_streaming_cache_hit_callbacks( - cached_result=cached_result + cached_result=converted_result ): logging_obj.handle_sync_cache_hit_success_callbacks( - result=cached_result, + result=converted_result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, @@ -527,8 +532,8 @@ class LLMCachingHandler: or self.request_kwargs.get("cache_key") or litellm.cache.get_cache_key(**self.request_kwargs) ) - _set_cached_hidden_param(cached_result, "cache_key", cache_key) - return CachingHandlerResponse(cached_result=cached_result) + _set_cached_hidden_param(converted_result, "cache_key", cache_key) + return CachingHandlerResponse(cached_result=converted_result) return CachingHandlerResponse(cached_result=cached_result) _sync_get_cache = sync_get_cache @@ -950,22 +955,14 @@ class LLMCachingHandler: def _convert_cached_result_to_model_response( self, - cached_result: Any, + cached_result: object, call_type: str, kwargs: dict[str, object], logging_obj: LiteLLMLoggingObj, model: str, args: tuple[object, ...], custom_llm_provider: str | None = None, - ) -> ( - ModelResponse - | TextCompletionResponse - | EmbeddingResponse - | RerankResponse - | TranscriptionResponse - | CustomStreamWrapper - | None - ): + ) -> object: """ Internal method to process the cached result @@ -1141,7 +1138,7 @@ class LLMCachingHandler: self, result: object, original_function: Callable, - kwargs: dict[str, Any], + kwargs: Mapping[str, object], args: tuple[object, ...] | None = None, ): """ @@ -1169,7 +1166,7 @@ class LLMCachingHandler: return cache: Final = litellm.cache - new_kwargs: Final = kwargs.copy() + new_kwargs: Final = dict(kwargs) new_kwargs.update( convert_args_to_kwargs( original_function, @@ -1206,7 +1203,7 @@ class LLMCachingHandler: ) ) else: - create_cache_write_task(lambda: cache.async_add_cache(result, **new_kwargs)) + create_cache_write_task(lambda: cache.async_add_cache(result, **_CacheCallKwargs(**new_kwargs))) def sync_set_cache( self, diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 62ee5d83634..4d9d26fc70c 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -40,7 +40,7 @@ from litellm.constants import ( REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, REDIS_TIMEOUT_LOG_INTERVAL, ) -from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs +from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, has_attributes from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.redis.store_settings import RedisStoreSettings, breaker_settings from litellm.types.caching import ( @@ -97,6 +97,22 @@ class _AsyncRedisCommands(Protocol): def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ... +class _AsyncLuaScript(Protocol): + def __call__( + self, keys: Sequence[str], args: Sequence[str | bytes | int | float], client: object + ) -> Awaitable[object]: ... + + +class _SupportsRegisterScript(Protocol): + def register_script(self, script: str, /) -> _AsyncLuaScript: ... + + +class _SupportsScriptLoad(Protocol): + def script_load(self, script: str, /) -> object: ... + + def evalsha(self, sha: object, numkeys: int, /, *keys_and_args: str | bytes | float) -> Awaitable[object]: ... + + _GENERIC_CALLER_MODULES: Final = frozenset( { __name__, @@ -844,6 +860,9 @@ class RedisCache(BaseCache): def _async_commands(self) -> _AsyncRedisCommands: return self.init_async_client() + def _async_client_object(self) -> object: + return self.init_async_client() + def check_and_fix_namespace(self, key: str) -> str: """ Make sure each key starts with the given namespace @@ -1117,8 +1136,8 @@ class RedisCache(BaseCache): Kept separate from async_register_script so each loop caches its own executor; see that method for why the binding must be per loop. """ - _redis_client: Final[Any] = self.init_async_client() - if hasattr(_redis_client, "register_script"): + _redis_client: Final = self._async_client_object() + if has_attributes(_redis_client, _SupportsRegisterScript, "register_script"): registered_script: Final = _redis_client.register_script(script) async def standalone_executor( @@ -1131,7 +1150,7 @@ class RedisCache(BaseCache): return standalone_executor - if hasattr(_redis_client, "script_load"): + if has_attributes(_redis_client, _SupportsScriptLoad, "script_load"): script_sha: Final = _redis_client.script_load(script) async def cluster_executor( diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index 7268511fe7e..c884eb1b001 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -43,6 +43,7 @@ if TYPE_CHECKING: _CACHE_RESPONSE_ADAPTER: Final[TypeAdapter[str | bytes | None]] = TypeAdapter( str | bytes | None, config=ConfigDict(strict=True, hide_input_in_errors=True) ) +_DECODED_CACHE_VALUE: Final = TypeAdapter(object) class RedisSemanticCache(BaseCache): @@ -407,10 +408,10 @@ class RedisSemanticCache(BaseCache): # Convert string representation to Python object try: - return json.loads(response_text) + return _DECODED_CACHE_VALUE.validate_python(json.loads(response_text)) except json.JSONDecodeError: try: - return ast.literal_eval(response_text) + return _DECODED_CACHE_VALUE.validate_python(ast.literal_eval(response_text)) except (ValueError, SyntaxError) as e: print_verbose(f"Error parsing cached response: {e}") return None diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 9902fe1d8fc..fc7a6048538 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -39,7 +39,7 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import ( parse_chat_allowed_tools_tool_choice, responses_allowed_tools_tool_choice, ) -from litellm.litellm_core_utils.core_helpers import drops_unsupported_params +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, iter_items from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params from litellm.litellm_core_utils.prompt_templates.prompt_cache_breakpoint_utils import ( prompt_cache_breakpoint_with_readable_mode, @@ -54,6 +54,7 @@ from litellm.llms.base_llm.bridges.completion_transformation import ( CompletionTransformationBridge, ) from litellm.llms.custom_httpx.sse_lines import split_sse_lines +from litellm.responses.additional_tools import additional_tools_in_input from litellm.responses.litellm_completion_transformation.tool_search import is_anthropic_tool_search_chat_tool from litellm.responses.sse_output_recovery import ( parse_sse_json_chunk, @@ -348,6 +349,9 @@ _TOOL_SEARCH_REPLAY_ITEM_TYPES: Final = frozenset({"reasoning", *_TOOL_SEARCH_IT _TOOL_SEARCH_ITEMS_KEY: Final = "tool_search_items" +_DECODED_JSON: Final = TypeAdapter(object) + + def _json_field(value: object, key: str) -> object: return cast(Mapping[str, object], value).get(key) if isinstance(value, dict) else None # cast-ok: untyped JSON @@ -426,6 +430,17 @@ def replayed_tool_search_tool_names(request: object) -> tuple[str, ...]: return tuple(chain.from_iterable(_tool_definition_name(tool) for tool in replayed_tool_search_tools(request))) +def _top_level_namespaces(request: object) -> Iterator[object]: + tools: Final = _json_list(_json_field(request, "tools")) or () + return (tool for tool in tools if _json_field(tool, "type") == "namespace") + + +def nested_tool_member_names(request: object) -> tuple[str, ...]: + tools: Final = (*_top_level_namespaces(request), *additional_tools_in_input(_json_field(request, "input"))) + members: Final = chain.from_iterable(_namespace_members(tool) for tool in tools) + return tuple(chain.from_iterable(_tool_definition_name(member) for member in members)) + + def _loads_any_tool(message: object, names: frozenset[str]) -> bool: loaded_names: Final = chain.from_iterable(_tool_definition_name(tool) for tool in _message_loaded_tools(message)) return any(name in names for name in loaded_names) @@ -845,8 +860,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Handle message items with output_text content if item_type == "message": - content_list: Final = item.get("content", []) - for content_item in content_list: + content_list: Final = _DECODED_JSON.validate_python(item.get("content", [])) + for content_item in iter_items(content_list): if isinstance(content_item, dict): content_type = content_item.get("type") if content_type == "output_text": diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 3910a554905..5f7d646e150 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -8,7 +8,7 @@ from datetime import datetime from functools import lru_cache from itertools import accumulate from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, cast from httpx import Response from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError @@ -24,7 +24,12 @@ from litellm.constants import ( REALTIME_METERED_USAGE_KEY, ) from litellm.litellm_core_utils.audio_utils.utils import normalized_audio_duration_seconds -from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, billed_usage_cost +from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, + attribute_or_none, + billed_usage_cost, + require_attributes, +) from litellm.litellm_core_utils.cost_discount import resolve_cost_discount from litellm.litellm_core_utils.decisions_pricing import ( decisions_pricing_fields, @@ -1188,6 +1193,10 @@ def _model_contains_known_llm_provider(model: str) -> bool: return _provider_prefix in LlmProvidersSet +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + def _get_response_model(completion_response: object) -> str | None: """ Extract the model name from a completion response object. @@ -1219,7 +1228,7 @@ _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER: Final[dict] = { } -def _map_traffic_type_to_service_tier(traffic_type: str | None) -> str | None: +def _map_traffic_type_to_service_tier(traffic_type: object) -> str | None: """ Map a Gemini usageMetadata.trafficType value to a LiteLLM service_tier string. @@ -1273,11 +1282,13 @@ def _served_service_tier(completion_response: object, usage_object: Usage | None usage_tier: Final = _extract_service_tier(usage_object) if isinstance(usage_tier, str): return usage_tier - hidden_params: Final = getattr(completion_response, "_hidden_params", None) + hidden_params: Final[object] = getattr(completion_response, "_hidden_params", None) if hidden_params is None: return None - provider_specific: Final = hidden_params.get("provider_specific_fields") or {} - raw_traffic_type: Final = provider_specific.get("traffic_type") + provider_specific: Final[object] = ( + require_attributes(hidden_params, _SupportsGet, "get").get("provider_specific_fields") or {} + ) + raw_traffic_type: Final = require_attributes(provider_specific, _SupportsGet, "get").get("traffic_type") if not raw_traffic_type: return None return _map_traffic_type_to_service_tier(raw_traffic_type) or "default" @@ -1341,8 +1352,16 @@ def get_usage_object( return None +class _DumpedUsageFields(TypedDict, total=False, extra_items=object): + pass + + +def _dumped_usage_fields(usage_model: BaseModel) -> Mapping[str, object]: + return usage_model.model_dump() + + def _usage_from_model(usage_model: BaseModel) -> Usage: - return Usage(**usage_model.model_dump()) + return Usage(**_DumpedUsageFields(**_dumped_usage_fields(usage_model))) def _decisions_prompt_tokens_details(usage: DecisionsUsage | OpenAIDecisionUsage) -> PromptTokensDetailsWrapper: diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 6662cf4a2c2..56a5f136b62 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -12,10 +12,11 @@ import datetime import enum from collections.abc import Mapping, Sequence -from typing import Any, Final +from typing import Any, Final, Protocol, SupportsIndex, SupportsInt, runtime_checkable import httpx import openai +from typing_extensions import Buffer import litellm from litellm.types.utils import LiteLLMCommonStrings @@ -1250,6 +1251,19 @@ class BlockedPiiEntityError(Exception): super().__init__(self.message) +@runtime_checkable +class _SupportsTrunc(Protocol): + def __trunc__(self) -> int: ... + + +def _status_code_int(value: object) -> int: + if isinstance(value, (str, Buffer, SupportsInt, SupportsIndex, _SupportsTrunc)): + return int(value) + raise TypeError( + f"int() argument must be a string, a bytes-like object or a real number, not '{type(value).__name__}'" + ) + + class MidStreamFallbackError(ServiceUnavailableError): def __init__( self, @@ -1265,8 +1279,8 @@ class MidStreamFallbackError(ServiceUnavailableError): is_pre_first_chunk: bool = False, emitted_disqualifying_content: bool = False, ): - original_status: Final = getattr(original_exception, "status_code", None) - self.status_code = int(original_status) if original_status is not None else 503 + original_status: Final[object] = getattr(original_exception, "status_code", None) + self.status_code = _status_code_int(original_status) if original_status is not None else 503 self.message = f"litellm.MidStreamFallbackError: {message}" self.model = model self.llm_provider = llm_provider @@ -1309,7 +1323,7 @@ class MidStreamFallbackError(ServiceUnavailableError): ) # Restore the propagated status and original response/request objects - self.status_code = int(original_status) if original_status is not None else 503 + self.status_code = _status_code_int(original_status) if original_status is not None else 503 self.response = _saved_response self.request = _saved_request self.message = _saved_message diff --git a/litellm/harness/endpoint.py b/litellm/harness/endpoint.py index 71e2742fa2e..249bb2cebb7 100644 --- a/litellm/harness/endpoint.py +++ b/litellm/harness/endpoint.py @@ -14,10 +14,11 @@ import contextlib import itertools import json import logging +import math import secrets -from collections.abc import AsyncIterable, AsyncIterator, Mapping +from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping from dataclasses import dataclass -from types import MappingProxyType, ModuleType +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol import httpx @@ -36,6 +37,7 @@ from litellm.constants import ( from litellm.harness.context import GatewayTarget from litellm.harness.errors import HarnessError, HarnessInstallFailed from litellm.harness.types import Harness, Usage +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.custom_http import httpxSpecialProvider @@ -45,7 +47,7 @@ if TYPE_CHECKING: from starlette.requests import Request from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.routing import Route - from uvicorn import Server + from uvicorn import Config, Server verbose_logger: Final = logging.getLogger("LiteLLM") _JSON_VALUE_ADAPTER: Final = TypeAdapter(object, config=ConfigDict(strict=True)) @@ -110,9 +112,22 @@ class _ResponsesModule(Protocol): StreamingResponse: type[StreamingResponse] +class _UvicornModule(Protocol): + Config: type[Config] + Server: type[Server] + + +class _CapturesSignals(Protocol): + capture_signals: Callable[[], contextlib.AbstractContextManager[None]] + + +class _InstallsSignalHandlers(Protocol): + install_signal_handlers: Callable[[], None] + + @dataclass(frozen=True) class _ServerDeps: - uvicorn: ModuleType + uvicorn: _UvicornModule applications: _ApplicationsModule routing: _RoutingModule responses: _ResponsesModule @@ -438,6 +453,14 @@ def _noop() -> None: return None +def _signal_capture(server: object) -> _CapturesSignals | None: + return server if has_attributes(server, _CapturesSignals, "capture_signals") else None + + +def _signal_handler_installer(server: object) -> _InstallsSignalHandlers | None: + return server if has_attributes(server, _InstallsSignalHandlers, "install_signal_handlers") else None + + class ModelEndpoint: """Local HTTP endpoint for one harness session. Use as an async context manager.""" @@ -521,7 +544,7 @@ class ModelEndpoint: await asyncio.sleep(DEFAULT_POLLING_INTERVAL) def _build_server(self, deps: _ServerDeps) -> Server: - config = deps.uvicorn.Config( + config: Final = deps.uvicorn.Config( self._build_app(deps), host=HARNESS_ENDPOINT_HOST, port=0, @@ -529,14 +552,16 @@ class ModelEndpoint: log_level="warning", access_log=False, lifespan="off", - timeout_graceful_shutdown=HARNESS_PROCESS_KILL_GRACE_SECONDS, + timeout_graceful_shutdown=math.ceil(HARNESS_PROCESS_KILL_GRACE_SECONDS), ) - server = deps.uvicorn.Server(config) + server: Final = deps.uvicorn.Server(config) # Never touch the host process's signal handlers. - if hasattr(server, "capture_signals"): - server.capture_signals = contextlib.nullcontext - if hasattr(server, "install_signal_handlers"): - server.install_signal_handlers = _noop + signal_capture: Final = _signal_capture(server) + if signal_capture is not None: + signal_capture.capture_signals = contextlib.nullcontext + handler_installer: Final = _signal_handler_installer(server) + if handler_installer is not None: + handler_installer.install_signal_handlers = _noop return server def _build_app(self, deps: _ServerDeps) -> Starlette: diff --git a/litellm/images/utils.py b/litellm/images/utils.py index d1a2dd76af0..efd37065e33 100644 --- a/litellm/images/utils.py +++ b/litellm/images/utils.py @@ -6,12 +6,25 @@ from pathlib import Path from typing import Any, Final, Protocol, cast, get_type_hints, runtime_checkable import litellm +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.litellm_core_utils.token_counter import get_image_type from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.types.files import FILE_MIME_TYPES, FileType from litellm.types.images.main import ImageEditOptionalRequestParams +class _ReadsBytes(Protocol): + def read(self, size: int, /) -> bytes: ... + + +class _Tells(Protocol): + def tell(self) -> int: ... + + +class _Seeks(Protocol): + def seek(self, offset: int, /) -> object: ... + + class ImageEditRequestUtils: @staticmethod def get_optional_params_image_edit( @@ -114,12 +127,12 @@ class ImageEditRequestUtils: bytes_data = image_data[:100] else: # For other types, try to read if possible - if hasattr(image_data, "read"): - current_pos = getattr(image_data, "tell", lambda: 0)() - if hasattr(image_data, "seek"): + if has_attributes(image_data, _ReadsBytes, "read"): + current_pos = image_data.tell() if has_attributes(image_data, _Tells, "tell") else 0 + if has_attributes(image_data, _Seeks, "seek"): image_data.seek(0) bytes_data = image_data.read(100) - if hasattr(image_data, "seek"): + if has_attributes(image_data, _Seeks, "seek"): image_data.seek(current_pos) else: return FILE_MIME_TYPES[FileType.PNG] # Default fallback diff --git a/litellm/integrations/SlackAlerting/user_spend_alerts.py b/litellm/integrations/SlackAlerting/user_spend_alerts.py index 08799a2f59f..67357cff403 100644 --- a/litellm/integrations/SlackAlerting/user_spend_alerts.py +++ b/litellm/integrations/SlackAlerting/user_spend_alerts.py @@ -53,11 +53,13 @@ async def fetch_user_spend_rows( today: datetime.date, baseline_days: int, ) -> tuple[UserSpendRow, ...]: + from litellm.proxy.utils import query_raw_rows + today_str: Final = today.strftime("%Y-%m-%d") month_start_str: Final = today.replace(day=1).strftime("%Y-%m-%d") baseline_start_str: Final = (today - datetime.timedelta(days=max(baseline_days, 1))).strftime("%Y-%m-%d") return USER_SPEND_ROWS_ADAPTER.validate_python( - await prisma_client.db.query_raw(USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str) + await query_raw_rows(prisma_client.db, USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str) ) diff --git a/litellm/integrations/argilla.py b/litellm/integrations/argilla.py index f873f274ab6..c20b2c332d3 100644 --- a/litellm/integrations/argilla.py +++ b/litellm/integrations/argilla.py @@ -7,11 +7,12 @@ import json import os import random import types -from collections.abc import Mapping -from typing import Final +from collections.abc import Mapping, Sequence +from typing import Final, Protocol import httpx from pydantic import BaseModel +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -29,6 +30,22 @@ from litellm.types.integrations.argilla import ( from litellm.types.utils import StandardLoggingPayload +class _ArgillaDataset(TypedDict): + id: ReadOnly[str] + + +class _ArgillaDatasetsBody(TypedDict, total=False): + items: ReadOnly[Sequence[_ArgillaDataset]] + + +class _ArgillaDatasetsResponse(Protocol): + def json(self) -> _ArgillaDatasetsBody: ... + + +def _datasets_body(response: _ArgillaDatasetsResponse) -> _ArgillaDatasetsBody: + return response.json() + + def is_serializable(value): non_serializable_types: Final = ( types.CoroutineType, @@ -102,7 +119,7 @@ class ArgillaLogger(CustomBatchLogger): url=f"{_credentials_base_url}/api/v1/me/datasets?name={_credentials_dataset_name}", headers={"X-Argilla-Api-Key": _credentials_api_key}, ) - json_response: Final = dataset_response.json() + json_response: Final = _datasets_body(dataset_response) if ( "items" in json_response and isinstance(json_response["items"], list) diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index fd912839c95..ea71300b8a6 100644 --- a/litellm/integrations/arize/_utils.py +++ b/litellm/integrations/arize/_utils.py @@ -1,6 +1,6 @@ import json from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, Protocol from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict, override @@ -10,6 +10,7 @@ from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes impor BaseLLMObsOTELAttributes, safe_set_attribute, ) +from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes, subscript from litellm.litellm_core_utils.redact_messages import ( should_redact_message_logging, ) @@ -34,32 +35,40 @@ from litellm.integrations._types.open_inference import ( _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) +class _SupportsGetWithDefault(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + class ArizeOTELAttributes(BaseLLMObsOTELAttributes): @staticmethod @override - def set_messages(span: "Span", kwargs: dict[str, Any]): + def set_messages(span: "Span", kwargs: Mapping[str, object]): messages: Final = kwargs.get("messages") # for /chat/completions # https://docs.arize.com/arize/large-language-models/tracing/semantic-conventions if messages: - last_message: Final = messages[-1] + last_message: Final = subscript(messages, -1) safe_set_attribute( span, SpanAttributes.INPUT_VALUE, - last_message.get("content", ""), + require_attributes(last_message, _SupportsGetWithDefault, "get").get("content", ""), ) # LLM_INPUT_MESSAGES shows up under `input_messages` tab on the span page. - for idx, msg in enumerate(messages): + for idx, msg in enumerate(iter_items(messages)): prefix = f"{SpanAttributes.LLM_INPUT_MESSAGES}.{idx}" # Set the role per message. - safe_set_attribute(span, f"{prefix}.{MessageAttributes.MESSAGE_ROLE}", msg.get("role")) + safe_set_attribute( + span, + f"{prefix}.{MessageAttributes.MESSAGE_ROLE}", + require_attributes(msg, _SupportsGetWithDefault, "get").get("role"), + ) # Set the content per message. safe_set_attribute( span, f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}", - msg.get("content", ""), + require_attributes(msg, _SupportsGetWithDefault, "get").get("content", ""), ) # Additive: emit structured tool_calls / multimodal content diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index 84dfc770e8f..f8047f97e75 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -18,10 +18,11 @@ import time import traceback from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Final, TypeVar +from typing import Final, Protocol, TypeVar from urllib.parse import urlparse import httpx +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.integrations.batch_utils import ( @@ -52,6 +53,19 @@ MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType( ) +class _OAuthTokenBody(TypedDict, total=False): + access_token: ReadOnly[str] + expires_in: ReadOnly[int] + + +class _OAuthTokenResponse(Protocol): + def json(self) -> _OAuthTokenBody: ... + + +def _oauth_token_body(response: _OAuthTokenResponse) -> _OAuthTokenBody: + return response.json() + + class AzureSentinelLogger(CustomBatchLogger): """ Logger that sends LiteLLM logs to Azure Sentinel via Azure Monitor Logs Ingestion API @@ -226,7 +240,7 @@ class AzureSentinelLogger(CustomBatchLogger): if response.status_code != 200: raise Exception(f"Failed to get OAuth2 token: {response.status_code} - {response.text}") - token_response: Final = response.json() + token_response: Final = _oauth_token_body(response) self.oauth_token = token_response.get("access_token") expires_in: Final = token_response.get("expires_in", 3600) diff --git a/litellm/integrations/bitbucket/bitbucket_client.py b/litellm/integrations/bitbucket/bitbucket_client.py index e4f5a50550e..5ebde5d4fb6 100644 --- a/litellm/integrations/bitbucket/bitbucket_client.py +++ b/litellm/integrations/bitbucket/bitbucket_client.py @@ -5,7 +5,7 @@ BitBucket API client for fetching .prompt files from BitBucket repositories. import base64 import urllib.parse from collections.abc import Mapping -from typing import Final, TypedDict +from typing import Final, Protocol, TypedDict from typing_extensions import NotRequired, ReadOnly @@ -36,6 +36,30 @@ class BitBucketFileMetadata(TypedDict): last_modified: ReadOnly[str | None] +class _SrcListingResponse(Protocol): + def json(self) -> BitBucketSrcListing: ... + + +class _RepositoryInfoResponse(Protocol): + def json(self) -> Mapping[str, object]: ... + + +class _BranchListingResponse(Protocol): + def json(self) -> BitBucketBranchListing: ... + + +def _src_listing(response: _SrcListingResponse) -> BitBucketSrcListing: + return response.json() + + +def _repository_info(response: _RepositoryInfoResponse) -> Mapping[str, object]: + return response.json() + + +def _branch_listing(response: _BranchListingResponse) -> BitBucketBranchListing: + return response.json() + + def _sanitize_file_path(file_path: str) -> str: """Reject path traversal and URL-encode each path segment.""" if "#" in file_path or "?" in file_path: @@ -166,7 +190,7 @@ class BitBucketClient: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data: Final[BitBucketSrcListing] = response.json() + data: Final = _src_listing(response) return [ file_path for item in data.get("values", []) @@ -201,7 +225,7 @@ class BitBucketClient: try: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - return response.json() + return _repository_info(response) except Exception as e: raise Exception(f"Failed to get repository info: {e}") @@ -231,7 +255,7 @@ class BitBucketClient: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data: Final[BitBucketBranchListing] = response.json() + data: Final = _branch_listing(response) return data.get("values", []) except Exception as e: raise Exception(f"Failed to get branches: {e}") diff --git a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py index fc885b7ee2c..5296f4b710b 100644 --- a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py +++ b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py @@ -4,12 +4,13 @@ Fetches .prompt files from BitBucket repositories and provides team-based access """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment from litellm.integrations.custom_prompt_management import CustomPromptManagement +from litellm.litellm_core_utils.core_helpers import require_attributes if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -26,6 +27,10 @@ from litellm.types.utils import StandardCallbackDynamicParams from .bitbucket_client import BitBucketClient +class _SupportsGetWithDefault(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + class BitBucketPromptTemplate: """ Represents a prompt template loaded from BitBucket. @@ -35,7 +40,7 @@ class BitBucketPromptTemplate: self, template_id: str, content: str, - metadata: dict[str, Any], + metadata: Mapping[str, object], model: str | None = None, ): self.template_id = template_id @@ -44,7 +49,9 @@ class BitBucketPromptTemplate: self.model = model or metadata.get("model") self.temperature = metadata.get("temperature") self.max_tokens = metadata.get("max_tokens") - self.input_schema = metadata.get("input", {}).get("schema", {}) + self.input_schema = require_attributes(metadata.get("input", {}), _SupportsGetWithDefault, "get").get( + "schema", {} + ) self.optional_params = {k: v for k, v in metadata.items() if k not in ["model", "input", "content"]} def __repr__(self): diff --git a/litellm/integrations/clickhouse/clickhouse_spend_logger.py b/litellm/integrations/clickhouse/clickhouse_spend_logger.py index c8453456052..cd1250b17f1 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -7,10 +7,10 @@ so `response_id` is always the raw provider response id (cache-hit suffix stripp import json import re -from collections.abc import Iterator, Mapping +from collections.abc import Iterable, Iterator, Mapping from math import isfinite from types import MappingProxyType -from typing import Any, Final +from typing import Any, Final, Protocol from pydantic import JsonValue, TypeAdapter, ValidationError @@ -19,6 +19,7 @@ from litellm._logging import verbose_logger from litellm.integrations.clickhouse.clickhouse_batch_logger import ClickHouseBatchLogger from litellm.integrations.clickhouse.context import is_lens_analysis from litellm.integrations.clickhouse.schema import SPEND_LOGS_TABLE +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload @@ -40,6 +41,14 @@ _INTERNAL_METADATA_KEYS: Final = frozenset( ) +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + +class _SupportsItems(Protocol): + def items(self) -> Iterable[tuple[object, object]]: ... + + def _request_metadata_fields(value: object) -> Iterator[tuple[str, JsonValue]]: if value is None: return @@ -85,22 +94,22 @@ def _json_mapping(value: Mapping[str, object]) -> str: return _json(dict(value)) -def _find_traceparent(metadata: Mapping[str, Any], kwargs: Mapping[str, Any]) -> tuple[str, str]: +def _find_traceparent(metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> tuple[str, str]: custom_headers = metadata.get("requester_custom_headers") or MappingProxyType({}) - proxy_request = (kwargs.get("litellm_params") or MappingProxyType({})).get( + proxy_request = require_attributes(kwargs.get("litellm_params") or MappingProxyType({}), _SupportsGet, "get").get( "proxy_server_request" ) or MappingProxyType({}) - request_headers = proxy_request.get("headers") or MappingProxyType({}) + request_headers = require_attributes(proxy_request, _SupportsGet, "get").get("headers") or MappingProxyType({}) for headers in (custom_headers, request_headers): - for name, value in headers.items(): + for name, value in require_attributes(headers, _SupportsItems, "items").items(): if str(name).lower() == "traceparent": return parse_traceparent(value) return "", "" -def _cache_tokens(usage: Mapping[str, Any]) -> tuple[int, int]: +def _cache_tokens(usage: Mapping[str, object]) -> tuple[int, int]: """(cache_read, cache_write) from a Usage dict: OpenAI prompt_tokens_details first, Anthropic fields as fallback.""" - details = usage.get("prompt_tokens_details") or MappingProxyType({}) + details = require_attributes(usage.get("prompt_tokens_details") or MappingProxyType({}), _SupportsGet, "get") cache_read = _int(details.get("cached_tokens")) or _int(usage.get("cache_read_input_tokens")) cache_write = ( _int(details.get("cache_write_tokens")) @@ -116,10 +125,17 @@ def _request_tags(value: object) -> list[str]: return [str(tag) for tag in value] -def _session_id(payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, Any]) -> str: +def _session_id(payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, object]) -> str: """Mirrors proxy `_get_session_id_for_spend_log`: explicit session id, else the payload trace id.""" - request_metadata = (kwargs.get("litellm_params") or MappingProxyType({})).get("metadata") or MappingProxyType({}) - return str(payload.get("session_id") or request_metadata.get("session_id") or payload.get("trace_id") or "") + request_metadata = require_attributes( + kwargs.get("litellm_params") or MappingProxyType({}), _SupportsGet, "get" + ).get("metadata") or MappingProxyType({}) + return str( + payload.get("session_id") + or require_attributes(request_metadata, _SupportsGet, "get").get("session_id") + or payload.get("trace_id") + or "" + ) def _is_trace_ingest(payload: StandardLoggingPayload) -> bool: @@ -128,7 +144,7 @@ def _is_trace_ingest(payload: StandardLoggingPayload) -> bool: def spend_log_row_from_payload( - payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, Any] + payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, object] ) -> SpendLogRecord: metadata: Mapping[str, Any] = payload.get("metadata") or MappingProxyType({}) hidden_params: Mapping[str, Any] = payload.get("hidden_params") or MappingProxyType({}) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 68876812c3e..ae469a104f8 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,16 +1,17 @@ import contextvars import copy import hashlib +import json import os import secrets from collections.abc import Iterable, Mapping, Sequence from datetime import datetime from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, cast, get_args import httpx -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter from litellm._internal_context import with_service_target from litellm._logging import verbose_logger @@ -20,6 +21,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, redact_nested_match_and_regex_keys, + require_attributes, ) from litellm.secret_managers.main import str_to_bool from litellm.types.guardrails import ( @@ -73,6 +75,7 @@ GUARDRAIL_SESSIONS_TARGET: Final = "guardrail_sessions" _PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16) _PRE_CALL_RAISED_TOKEN: Final = secrets.token_hex(16) _DEPLOYMENT_HOOK_RAN_TOKEN: Final = secrets.token_hex(16) +_POST_CALL_CHECKED_TOKEN: Final = secrets.token_hex(16) _GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422}) @@ -167,18 +170,22 @@ def _strict_guardrail_modes_enabled() -> bool: return True if parsed is None else parsed -def get_session_id_from_request_data(request_data: dict[str, Any]) -> str | None: +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + +def get_session_id_from_request_data(request_data: Mapping[str, object]) -> str | None: """Extract session_id from request data (litellm_session_id or metadata).""" session_id = request_data.get("litellm_session_id") if session_id: return str(session_id) - metadata: Final = request_data.get("metadata") or {} + metadata: Final = require_attributes(request_data.get("metadata") or {}, _SupportsGet, "get") session_id = metadata.get("session_id") if session_id: return str(session_id) - litellm_metadata: Final = request_data.get("litellm_metadata") or {} + litellm_metadata: Final = require_attributes(request_data.get("litellm_metadata") or {}, _SupportsGet, "get") session_id = litellm_metadata.get("session_id") if session_id: return str(session_id) @@ -281,6 +288,95 @@ def _request_carries_marker(data: Mapping[str, object], marker: str | None) -> b return False +def _record_guardrail_marker( + data: dict[str, object], # mutable-ok: the request payload the marker is recorded on + marker: str | None, +) -> None: + if marker is None: + return + for meta_key in ("metadata", "litellm_metadata"): + meta = data.get(meta_key) + if isinstance(meta, dict): + executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) + if isinstance(executed, list): + if marker not in executed: + executed.append(marker) + else: + meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker] + return + data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]} + + +_RESPONSE_METADATA_KEYS: Final = frozenset( + { + "id", + "created", + "created_at", + "model", + "usage", + "usage_info", + "system_fingerprint", + "_hidden_params", + "_response_ms", + } +) + + +def _response_fingerprint(response: object) -> str | None: + """Digest of everything but metadata set between layers, so an unknown field counts as content.""" + try: + fields: Final[Mapping[str, object]] = ( + response.model_dump(mode="json") + if isinstance(response, BaseModel) + else _REQUEST_FIELDS.validate_python(response) + ) + content: Final = {key: value for key, value in fields.items() if key not in _RESPONSE_METADATA_KEYS} + if not content: + return None + return hashlib.sha256( + json.dumps(content, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + ).hexdigest() + except (TypeError, ValueError): + return None + + +_RESPONSE_GUARDRAIL_CHECKS_KEY: Final = "guardrail_checks" + + +def _str_keyed_dict(value: object) -> Mapping[str, object] | None: + if not isinstance(value, dict): + return None + return cast("Mapping[str, object]", value) # cast-ok: response bodies and _hidden_params are keyed by field name + + +def _response_hidden_params(response: object) -> Mapping[str, object] | None: + as_dict: Final = _str_keyed_dict(response) + return _str_keyed_dict( + as_dict.get("_hidden_params") if as_dict is not None else getattr(response, "_hidden_params", None) + ) + + +def _write_response_marker(response: object, marker: str | None) -> None: + """Record a check on the response itself, so it survives the router copying request metadata per hop.""" + if marker is None: + return + hidden: Final = _response_hidden_params(response) + if not isinstance(hidden, dict): + return + hidden[_RESPONSE_GUARDRAIL_CHECKS_KEY] = tuple( + _str_collection(hidden.get(_RESPONSE_GUARDRAIL_CHECKS_KEY)) | {marker} + ) + + +def _response_carries_marker(response: object, marker: str | None) -> bool: + hidden: Final = _response_hidden_params(response) + return ( + marker is not None + and hidden is not None + and marker in _str_collection(hidden.get(_RESPONSE_GUARDRAIL_CHECKS_KEY)) + ) + + _REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object]) @@ -291,6 +387,7 @@ def _guardrails_listed_in(metadata: object) -> list[str | dict[str, DynamicGuard class CustomGuardrail(CustomLogger): + mcp_block_behavior: Literal["tool_result", "raise"] | None = None # If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path. use_native_during_call_hook: ClassVar[bool] = False @@ -912,20 +1009,7 @@ class CustomGuardrail(CustomLogger): With ``raised``, records instead that the hook raised as a policy pipeline step the pipeline then moved past. The deployment-level hook still runs it. """ - marker: Final = self._pre_call_marker(raised=raised) - if marker is None: - return - for meta_key in ("metadata", "litellm_metadata"): - meta = data.get(meta_key) - if isinstance(meta, dict): - executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) - if isinstance(executed, list): - if marker not in executed: - executed.append(marker) - else: - meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker] - return - data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]} + _record_guardrail_marker(data, self._pre_call_marker(raised=raised)) def pre_call_hook_ran_or_raised(self, data: Mapping[str, object]) -> bool: """Whether this guardrail's ``async_pre_call_hook`` ran for this request, whatever it decided.""" @@ -1123,10 +1207,26 @@ class CustomGuardrail(CustomLogger): cast("dict[str, object]", request_data) # cast-ok: proxy request payloads are str-keyed ) - if not self._is_valid_response_type(result): - return None + returned: Final = result if self._is_valid_response_type(result) else None + if target is self: + return returned + checked: Final = returned if returned is not None else response + marker: Final = self._post_call_check_marker(checked) + _record_guardrail_marker( + cast("dict[str, object]", request_data), # cast-ok: proxy request payloads are str-keyed + marker, + ) + _write_response_marker(checked, marker) + return returned - return result + def _post_call_check_marker(self, response: object) -> str | None: + fingerprint: Final = _response_fingerprint(response) + return self._marker(f"{_POST_CALL_CHECKED_TOKEN}:{fingerprint}") if fingerprint is not None else None + + def post_call_already_checked(self, data: Mapping[str, object], response: object) -> bool: + """Whether a deployment-level hook of this request already ran this guardrail on this exact response.""" + marker: Final = self._post_call_check_marker(response) + return _request_carries_marker(data, marker) or _response_carries_marker(response, marker) async def async_logging_hook( self, diff --git a/litellm/integrations/email_alerting.py b/litellm/integrations/email_alerting.py index 27ee6bb60b4..d7228628d35 100644 --- a/litellm/integrations/email_alerting.py +++ b/litellm/integrations/email_alerting.py @@ -3,17 +3,29 @@ Functions for sending Email Alerts """ import os -from typing import Final +from collections.abc import Awaitable, Mapping, Sequence +from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.repositories.team_repository import TeamRepository from litellm.types.integrations.slack_alerting import WebhookEvent +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + # we use this for the email header, please send a test email if you change this. verify it looks good on email LITELLM_LOGO_URL: Final = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png" LITELLM_SUPPORT_CONTACT: Final = "support@berri.ai" +def _query_raw_rows( + client: "PrismaClient", query: str, *args: object +) -> Awaitable[Sequence[Mapping[str, object]] | None]: + from litellm.proxy.utils import query_raw_rows + + return query_raw_rows(client.db, query, *args) + + async def get_all_team_member_emails(team_id: str | None = None) -> list: verbose_logger.debug("Email Alerting: Getting all team members for team_id=%s", team_id) if team_id is None: @@ -51,17 +63,18 @@ async def get_all_team_member_emails(team_id: str | None = None) -> list: WHERE user_id = ANY($1::TEXT[]); """ - _result: Final = await prisma_client.db.query_raw(sql_query, _team_member_user_ids) + _result: Final = await _query_raw_rows(prisma_client, sql_query, _team_member_user_ids) verbose_logger.debug("Email Alerting: Got all Emails for team, emails=%s", _result) if _result is None: return [] - emails: Final = [] - for user in _result: - if user and isinstance(user, dict) and user.get("user_email", None) is not None: - emails.append(user.get("user_email")) + emails: Final = [ + user.get("user_email") + for user in _result + if user and isinstance(user, dict) and user.get("user_email", None) is not None + ] return emails diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index 891318f1c54..044439aeec2 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -123,7 +123,9 @@ class FocusLiteLLMDatabase: ORDER BY ordinal_position; """ try: - columns_response: Final = await client.db.query_raw(info_query) + from litellm.proxy.utils import query_raw_rows + + columns_response: Final = await query_raw_rows(client.db, info_query) return {"columns": columns_response, "table_name": "LiteLLM_DailyUserSpend"} except Exception as exc: raise RuntimeError(f"Error getting table info: {exc}") from exc diff --git a/litellm/integrations/galileo.py b/litellm/integrations/galileo.py index 56147cd720d..0f7dcdbcc81 100644 --- a/litellm/integrations/galileo.py +++ b/litellm/integrations/galileo.py @@ -15,6 +15,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import require_attributes, subscript from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, get_content_from_model_response, @@ -60,6 +61,10 @@ class _JsonResponse(Protocol): def json(self) -> object: ... +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + def _login_access_token(response: _GalileoLoginResponse) -> str: """Read the bearer token out of a Galileo login response body.""" return response.json()["access_token"] @@ -463,13 +468,13 @@ class GalileoObserve(CustomLogger): pass @staticmethod - def _build_prompt(kwargs: Mapping[str, Any]) -> dict[str, object]: - optional_params: Final[Mapping[str, object]] = kwargs.get("optional_params", {}) or {} + def _build_prompt(kwargs: Mapping[str, object]) -> dict[str, object]: + optional_params: Final = require_attributes(kwargs.get("optional_params", {}) or {}, _SupportsGet, "get") prompt: Final[dict[str, object]] = {"messages": kwargs.get("messages")} if optional_params.get("functions") is not None: - prompt["functions"] = optional_params["functions"] + prompt["functions"] = subscript(optional_params, "functions") if optional_params.get("tools") is not None: - prompt["tools"] = optional_params["tools"] + prompt["tools"] = subscript(optional_params, "tools") return prompt @staticmethod diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 7057e719764..99a8a5979b4 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -5,7 +5,7 @@ import os import time from collections.abc import Mapping from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -16,7 +16,7 @@ from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase from litellm.litellm_core_utils.cloud_storage_security import ( sanitize_cloud_object_component, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, require_attributes from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus from litellm.types.integrations.gcs_bucket import * @@ -28,6 +28,10 @@ else: VertexBase = Any +class _SupportsGetWithDefault(Protocol): + def get(self, key: str, default: str, /) -> object: ... + + class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): def __init__(self, bucket_name: str | None = None) -> None: from litellm.proxy.proxy_server import premium_user @@ -298,7 +302,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): await self._send_individual_logs(items_to_process) def _get_object_name( - self, kwargs: Mapping[str, object], logging_payload: StandardLoggingPayload, response_obj: Any + self, kwargs: Mapping[str, object], logging_payload: StandardLoggingPayload, response_obj: object ) -> str: """ Get the object name to use for the current payload @@ -313,7 +317,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): return f"{current_date}/custom-{uuid.uuid4().hex}-{safe_log_id}" if logging_payload.get("error_str", None) is not None: return self._generate_failure_object_name(request_date_str=current_date) - raw_response_id: Final[object] = response_obj.get("id", "") + raw_response_id: Final[object] = require_attributes(response_obj, _SupportsGetWithDefault, "get").get("id", "") response_id: Final = raw_response_id if isinstance(raw_response_id, str) else "" return self._generate_success_object_name( request_date_str=current_date, @@ -360,7 +364,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): def _generate_success_object_name( self, request_date_str: str, - response_id: str, + response_id: object, ) -> str: return f"{request_date_str}/{response_id}" diff --git a/litellm/integrations/generic_prompt_management/generic_prompt_manager.py b/litellm/integrations/generic_prompt_management/generic_prompt_manager.py index 2601f44e0bf..fda537c55c6 100644 --- a/litellm/integrations/generic_prompt_management/generic_prompt_manager.py +++ b/litellm/integrations/generic_prompt_management/generic_prompt_manager.py @@ -132,7 +132,7 @@ class GenericPromptManager(CustomPromptManagement): async def async_fetch_prompt_from_api( self, prompt_id: str | None, prompt_spec: PromptSpec | None - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Fetch a prompt from the API asynchronously. """ diff --git a/litellm/integrations/humanloop.py b/litellm/integrations/humanloop.py index 093a68c9284..394f087c132 100644 --- a/litellm/integrations/humanloop.py +++ b/litellm/integrations/humanloop.py @@ -7,10 +7,12 @@ https://humanloop.com/ from typing import Final, cast import httpx +from pydantic import TypeAdapter from typing_extensions import TypedDict import litellm from litellm.caching import DualCache +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, decoded_json_object, subscript from litellm.llms.custom_httpx.http_handler import get_httpx_client from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -19,6 +21,8 @@ from litellm.types.utils import StandardCallbackDynamicParams from .custom_logger import CustomLogger +_DECODED_JSON: Final = TypeAdapter(object) + class PromptManagementClient(TypedDict): prompt_id: str @@ -79,18 +83,20 @@ class HumanLoopPromptManager(DualCache): raise Exception(f"Error getting prompt from Humanloop: {e.response.text}") json_response: Final = response.json() - template_message: Final = json_response["template"] - if isinstance(template_message, dict): - template_messages = [template_message] - elif isinstance(template_message, list): - template_messages = template_message + decoded_response: Final = _DECODED_JSON.validate_python(json_response) + template_message: Final = subscript(decoded_response, "template") + template_mapping: Final = as_str_mapping(template_message) + template_list: Final = as_object_list(template_message) + if template_mapping is not None: + template_messages = [template_mapping] + elif template_list is not None: + template_messages = template_list else: raise ValueError(f"Invalid template message type: {type(template_message)}") template_model: Final = json_response["model"] - optional_params: Final = {} - for k, v in json_response.items(): - if k in litellm.OPENAI_CHAT_COMPLETION_PARAMS: - optional_params[k] = v + optional_params: Final = { + k: v for k, v in decoded_json_object(decoded_response).items() if k in litellm.OPENAI_CHAT_COMPLETION_PARAMS + } return PromptManagementClient( prompt_id=humanloop_prompt_id, prompt_template=cast(list[AllMessageValues], template_messages), diff --git a/litellm/integrations/lago.py b/litellm/integrations/lago.py index ec4388bbb8d..4b79484ee84 100644 --- a/litellm/integrations/lago.py +++ b/litellm/integrations/lago.py @@ -3,7 +3,7 @@ import json import os -from typing import Final, Literal +from typing import Final, Literal, Protocol import httpx @@ -11,6 +11,7 @@ import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, get_async_httpx_client, @@ -18,6 +19,10 @@ from litellm.llms.custom_httpx.http_handler import ( ) +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + def get_utc_datetime(): import datetime as dt from datetime import datetime @@ -70,10 +75,11 @@ class LagoLogger(CustomLogger): if ( isinstance(response_obj, litellm.ModelResponse) or isinstance(response_obj, litellm.EmbeddingResponse) ) and hasattr(response_obj, "usage"): + response_usage: Final = require_attributes(response_obj["usage"], _SupportsGet, "get") usage = { - "prompt_tokens": response_obj["usage"].get("prompt_tokens", 0), - "completion_tokens": response_obj["usage"].get("completion_tokens", 0), - "total_tokens": response_obj["usage"].get("total_tokens"), + "prompt_tokens": response_usage.get("prompt_tokens", 0), + "completion_tokens": response_usage.get("completion_tokens", 0), + "total_tokens": response_usage.get("total_tokens"), } litellm_params: Final = kwargs.get("litellm_params", {}) or {} diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index 0117348da66..977fa816eed 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -7,6 +7,7 @@ from itertools import chain from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, runtime_checkable from opentelemetry import trace +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger @@ -33,6 +34,8 @@ from litellm.types.integrations.langfuse_otel import ( ) from litellm.types.utils import StandardCallbackDynamicParams +_DECODED_JSON: Final = TypeAdapter(object) + if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -643,7 +646,7 @@ def _tool_call_arguments(arguments: object) -> object: if not isinstance(arguments, str): return arguments try: - return json.loads(arguments) + return _DECODED_JSON.validate_python(json.loads(arguments)) except json.JSONDecodeError: return {} diff --git a/litellm/integrations/langfuse/langfuse_otel_attributes.py b/litellm/integrations/langfuse/langfuse_otel_attributes.py index 2873dba93ed..28530076c57 100644 --- a/litellm/integrations/langfuse/langfuse_otel_attributes.py +++ b/litellm/integrations/langfuse/langfuse_otel_attributes.py @@ -7,7 +7,7 @@ Relevant Issue: https://github.com/BerriAI/litellm/issues/13764 import json from collections.abc import Mapping from datetime import datetime -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, Protocol from pydantic import BaseModel from typing_extensions import override @@ -16,6 +16,7 @@ from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes impor BaseLLMObsOTELAttributes, safe_set_attribute, ) +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.utils import ( EmbeddingResponse, @@ -30,6 +31,10 @@ if TYPE_CHECKING: from opentelemetry.trace import Span +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + def serialize_langfuse_datetime(value: datetime) -> str: return json.dumps(value.astimezone().isoformat()) @@ -87,11 +92,11 @@ def get_output_content_by_type( class LangfuseLLMObsOTELAttributes(BaseLLMObsOTELAttributes): @staticmethod @override - def set_messages(span: "Span", kwargs: dict[str, Any]): + def set_messages(span: "Span", kwargs: Mapping[str, object]): prompt: Final = {"messages": kwargs.get("messages")} optional_params: Final = kwargs.get("optional_params", {}) - functions: Final = optional_params.get("functions") - tools: Final = optional_params.get("tools") + functions: Final = require_attributes(optional_params, _SupportsGet, "get").get("functions") + tools: Final = require_attributes(optional_params, _SupportsGet, "get").get("tools") if functions is not None: prompt["functions"] = functions if tools is not None: diff --git a/litellm/integrations/mlflow.py b/litellm/integrations/mlflow.py index 8731e96440f..9e5b1109a77 100644 --- a/litellm/integrations/mlflow.py +++ b/litellm/integrations/mlflow.py @@ -1,9 +1,11 @@ import json import threading -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, Final, Protocol from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import attribute_or_none, iter_items, require_attributes if TYPE_CHECKING: from litellm.types.utils import StandardLoggingPayload @@ -11,6 +13,25 @@ else: StandardLoggingPayload = Any +class _ChoiceWithMessage(Protocol): + @property + def message(self) -> object: ... + + +class _ExcludeNoneDumpable(Protocol): + def model_dump(self, *, exclude_none: bool) -> object: ... + + +def _dumped_choice_message(choice: object) -> object: + message: Final = require_attributes(choice, _ChoiceWithMessage, "message").message + return require_attributes(message, _ExcludeNoneDumpable, "model_dump").model_dump(exclude_none=True) + + +def _output_messages(response_obj: object) -> Sequence[object]: + choices: Final = attribute_or_none(response_obj, "choices") if hasattr(response_obj, "choices") else () + return [_dumped_choice_message(c) for c in iter_items(choices)] + + class MlflowLogger(CustomLogger): def __init__(self): from mlflow.tracking import MlflowClient @@ -62,7 +83,7 @@ class MlflowLogger(CustomLogger): inputs: Final = self._construct_input(kwargs) input_messages: Final = inputs.get("messages", []) - output_messages: Final = [c.message.model_dump(exclude_none=True) for c in getattr(response_obj, "choices", [])] + output_messages: Final = _output_messages(response_obj) if messages := [*input_messages, *output_messages]: set_span_chat_messages(span, messages) if tools := inputs.get("tools"): diff --git a/litellm/integrations/openmeter.py b/litellm/integrations/openmeter.py index 918e849a086..9668fe8fa89 100644 --- a/litellm/integrations/openmeter.py +++ b/litellm/integrations/openmeter.py @@ -4,13 +4,13 @@ import json import os from collections.abc import Mapping -from typing import Final +from typing import Final, Protocol import httpx import litellm from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, require_attributes from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, @@ -19,6 +19,10 @@ from litellm.llms.custom_httpx.http_handler import ( ) +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + def get_utc_datetime(): import datetime as dt from datetime import datetime @@ -60,10 +64,11 @@ class OpenMeterLogger(CustomLogger): if ( isinstance(response_obj, litellm.ModelResponse) or isinstance(response_obj, litellm.EmbeddingResponse) ) and hasattr(response_obj, "usage"): + response_usage: Final = require_attributes(response_obj["usage"], _SupportsGet, "get") usage = { - "prompt_tokens": response_obj["usage"].get("prompt_tokens", 0), - "completion_tokens": response_obj["usage"].get("completion_tokens", 0), - "total_tokens": response_obj["usage"].get("total_tokens"), + "prompt_tokens": response_usage.get("prompt_tokens", 0), + "completion_tokens": response_usage.get("completion_tokens", 0), + "total_tokens": response_usage.get("total_tokens"), } # OPENMETER_TRUST_REQUEST_USER (default "true"): when set to "false", diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 39c485cb31f..5cd201a2f47 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -6,7 +6,7 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -35,6 +35,7 @@ from litellm.integrations.otel.model.metadata import flatten_metadata from litellm.integrations.otel.model.semconv import LiteLLM, Metric from litellm.integrations.otel.plumbing.otlp_tls import resolve_otlp_http_tls from litellm.integrations.otel.routing import routing_decision_attributes +from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.secret_redaction import redact_string @@ -47,7 +48,6 @@ from litellm.secret_managers.main import get_secret_bool, str_to_bool from litellm.types.integrations.otel_span_attributes import SpanAttributes from litellm.types.services import ServiceLoggerPayload from litellm.types.utils import ( - ChatCompletionMessageToolCall, CostBreakdown, Function, LLMResponseTypes, @@ -130,6 +130,22 @@ def _text_as_stored_on_span(span: Span, key: str, text: str) -> str: _MAX_DYNAMIC_TRACER_PROVIDERS: Final = 256 +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + +def _response_field(response_obj: object, key: str) -> object: + if not response_obj: + return None + return require_attributes(response_obj, _SupportsGet, "get").get(key) + + +def _billed_usage(response_obj: object, call_type: str | None, litellm_params: Mapping[str, object] | None) -> object: + if not response_obj or is_unbilled_non_inference_call_from_params(call_type, litellm_params, response_obj): + return None + return _response_field(response_obj, "usage") + + class _CacheHitView(TypedDict, total=False): cache_hit: ReadOnly[object] saved_cache_cost: ReadOnly[object] @@ -2407,17 +2423,17 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): @staticmethod def _tool_calls_kv_pair( - tool_calls: list[ChatCompletionMessageToolCall], + tool_calls: Iterable[object], ) -> dict[str, object]: kv_pairs: Final[dict[str, object]] = {} for idx, tool_call in enumerate(tool_calls): - _function = tool_call.get("function") + _function: object = require_attributes(tool_call, _SupportsGet, "get").get("function") if not _function: continue keys = Function.__annotations__.keys() for key in keys: - _value = _function.get(key) + _value = require_attributes(_function, _SupportsGet, "get").get(key) if _value: kv_pairs[f"{SpanAttributes.LLM_COMPLETIONS.value}.{idx}.function_call.{key}"] = _value @@ -2584,7 +2600,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # but Embeddings and Image-gen responses do not. Fall back to # the litellm call ID so every call type can be correlated # across LiteLLM UI, Phoenix traces, and provider logs (Issue #8). - response_id = (response_obj.get("id") if response_obj else None) or standard_logging_payload.get("id") + response_id = _response_field(response_obj, "id") or standard_logging_payload.get("id") if response_id: self.safe_set_attribute( span=span, @@ -2601,42 +2617,37 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) # The model used to generate the response. - if response_obj and response_obj.get("model"): + response_model: Final = _response_field(response_obj, "model") + if response_model: self.safe_set_attribute( span=span, key=SpanAttributes.LLM_RESPONSE_MODEL.value, - value=response_obj.get("model"), + value=response_model, ) self._set_service_tier_attributes(span=span, standard_logging_payload=standard_logging_payload) - usage: Final = ( - response_obj.get("usage") - if response_obj - and not is_unbilled_non_inference_call_from_params( - kwargs.get("call_type"), litellm_params, response_obj - ) - else None - ) + usage: Final = _billed_usage(response_obj, kwargs.get("call_type"), litellm_params) if usage: + usage_fields: Final = require_attributes(usage, _SupportsGet, "get") self.safe_set_attribute( span=span, key=SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS.value, - value=usage.get("total_tokens"), + value=usage_fields.get("total_tokens"), ) # The number of tokens used in the LLM response (completion). self.safe_set_attribute( span=span, key=SpanAttributes.GEN_AI_USAGE_OUTPUT_TOKENS.value, - value=usage.get("completion_tokens"), + value=usage_fields.get("completion_tokens"), ) # The number of tokens used in the LLM prompt. self.safe_set_attribute( span=span, key=SpanAttributes.GEN_AI_USAGE_INPUT_TOKENS.value, - value=usage.get("prompt_tokens"), + value=usage_fields.get("prompt_tokens"), ) ######################################################################## @@ -2712,7 +2723,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ########## LLM Response Attributes ########## ############################################# if response_obj is not None: - if response_obj.get("choices"): + response_view: Final = require_attributes(response_obj, _SupportsGet, "get") + if response_view.get("choices"): transformed_choices: Final = self._transform_choices_to_otel_semantic_conventions( response_obj.get("choices") ) @@ -2722,10 +2734,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): value=safe_dumps(transformed_choices), ) - finish_reasons: Final = [] - for idx, choice in enumerate(response_obj.get("choices")): - if choice.get("finish_reason"): - finish_reasons.append(choice.get("finish_reason")) + choice_views: Final = [ + require_attributes(choice, _SupportsGet, "get") + for choice in iter_items(response_view.get("choices")) + ] + finish_reasons: Final = [ + choice.get("finish_reason") for choice in choice_views if choice.get("finish_reason") + ] if finish_reasons: self.safe_set_attribute( @@ -2734,11 +2749,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): value=safe_dumps(finish_reasons), ) - for idx, choice in enumerate(response_obj.get("choices")): + for choice in choice_views: if choice.get("finish_reason"): - message = choice.get("message") - if tool_calls := tool_calls_with_server_tool_calls(message): - kv_pairs = OpenTelemetry._tool_calls_kv_pair(list(tool_calls)) + tool_calls = tool_calls_with_server_tool_calls(choice.get("message")) + if tool_calls: + kv_pairs = OpenTelemetry._tool_calls_kv_pair(tool_calls) for key, value in kv_pairs.items(): self.safe_set_attribute( span=span, @@ -2746,13 +2761,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): value=value, ) - elif response_obj.get("output"): + elif response_view.get("output"): # Responses API: ResponsesAPIResponse has an "output" # list instead of "choices". Each item with # type="message" contains a "content" list of # OutputText objects (type="output_text"). - output_items: Final = response_obj.get("output") - output_messages: Final = self._transform_responses_api_output_to_otel(output_items) + output_items: Final = response_view.get("output") + output_messages: Final = self._transform_responses_api_output_to_otel(iter_items(output_items)) if output_messages: self.safe_set_attribute( span=span, @@ -2766,7 +2781,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # ChatCompletionMessageToolCall format expected by # _tool_calls_kv_pair. tool_calls = [] - for out_item in output_items: + for out_item in iter_items(output_items): item_d = self._to_dict(out_item) if item_d and item_d.get("type") == "function_call": tool_calls.append( @@ -2787,7 +2802,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) # Extract finish reason from ResponsesAPIResponse.status - status: Final = response_obj.get("status") + status: Final = response_view.get("status") if status: self.safe_set_attribute( span=span, @@ -2907,7 +2922,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): return obj.model_dump() return None - def _transform_responses_api_output_to_otel(self, output: list) -> list[dict]: + def _transform_responses_api_output_to_otel(self, output: Iterable[object]) -> list[dict]: """ Transform Responses API output items into OTEL GenAI 1.38 format. diff --git a/litellm/integrations/opentelemetry_utils/base_otel_llm_obs_attributes.py b/litellm/integrations/opentelemetry_utils/base_otel_llm_obs_attributes.py index 97a970c8966..350c61d22f6 100644 --- a/litellm/integrations/opentelemetry_utils/base_otel_llm_obs_attributes.py +++ b/litellm/integrations/opentelemetry_utils/base_otel_llm_obs_attributes.py @@ -1,4 +1,5 @@ from abc import ABC +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final if TYPE_CHECKING: @@ -7,7 +8,7 @@ if TYPE_CHECKING: class BaseLLMObsOTELAttributes(ABC): @staticmethod - def set_messages(span: "Span", kwargs: dict[str, Any]): + def set_messages(span: "Span", kwargs: Mapping[str, object]): pass @staticmethod diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index b5b65d299e0..41811f836c7 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -1034,10 +1034,12 @@ def emit_guardrail_span(entry: "StandardLoggingGuardrailInformation") -> None: pass -def seed_request_identity(user_api_key_dict: object, model: str | None = None) -> None: +def seed_request_identity(user_api_key_dict: object, model: str | None = None, request_data: object = None) -> None: logger: Final = _registered_v2_logger() if logger is not None: - logger.seed_request_identity(user_api_key_dict, model=model) + logger.seed_request_identity( + user_api_key_dict, model=model, request_metadata=metadata_from_request_data(request_data) + ) def fan_out_provider() -> ApiTracerProvider: diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index 1197ae055d0..27ee20362a2 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -11,6 +11,7 @@ identical metrics. The attribute cardinality filter is reused from v1 by import from collections.abc import Mapping from dataclasses import dataclass from datetime import datetime +from types import MappingProxyType from typing import Any, Final, Literal, Protocol, TypeAlias from opentelemetry.metrics import Histogram, Meter @@ -34,6 +35,7 @@ from litellm.integrations.otel.model.semconv import ( resolve_provider, ) from litellm.integrations.otel.model.utils import to_seconds +from litellm.litellm_core_utils.core_helpers import attribute_or_none, require_attributes from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -177,7 +179,14 @@ class _MetricKwargs(TypedDict, total=False): end_time: ReadOnly[datetime | float | str | None] -def resolve_error_type(kwargs: Mapping[str, Any]) -> str: +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def resolve_error_type(kwargs: Mapping[str, object]) -> str: """The ``error.type`` value for a failed request. Bounded by construction: the mapped provider exception's class name (the same @@ -188,8 +197,13 @@ def resolve_error_type(kwargs: Mapping[str, Any]) -> str: span and its exception event, where high cardinality is free. """ std_log: Final = kwargs.get("standard_logging_object") - info: Final = getattr(std_log, "error_information", None) or (std_log or {}).get("error_information") or {} - error_class: Final = info.get("error_class") or info.get("error_code") + info: Final = ( + attribute_or_none(std_log, "error_information") + or require_attributes(std_log or _NO_FIELDS, _SupportsGet, "get").get("error_information") + or _NO_FIELDS + ) + info_fields: Final = require_attributes(info, _SupportsGet, "get") + error_class: Final = info_fields.get("error_class") or info_fields.get("error_code") if error_class: return str(error_class) exception: Final = kwargs.get("exception") diff --git a/litellm/integrations/otel/runtime.py b/litellm/integrations/otel/runtime.py index 89227770448..dda79f71bb1 100644 --- a/litellm/integrations/otel/runtime.py +++ b/litellm/integrations/otel/runtime.py @@ -67,12 +67,12 @@ def phase_event(name: str, attributes: PhaseEventAttributes | None = None) -> No runtime[2](name, attributes) -def seed_request_identity(user_api_key_dict: object, model: object = None) -> None: - """Seed request-identity Baggage at the auth boundary (no-op without V2).""" +def seed_request_identity(user_api_key_dict: object, model: object = None, request_data: object = None) -> None: + """Seed request-identity Baggage (no-op without V2).""" runtime: Final = _otel_runtime() if runtime is None: return - runtime[1](user_api_key_dict, model=model) + runtime[1](user_api_key_dict, model=model, request_data=request_data) def phase_attributes(attributes: PhaseAttributes) -> None: diff --git a/litellm/integrations/prompt_layer.py b/litellm/integrations/prompt_layer.py index 1efd7560a1d..5bdc56fedad 100644 --- a/litellm/integrations/prompt_layer.py +++ b/litellm/integrations/prompt_layer.py @@ -4,9 +4,12 @@ import os import traceback from typing import Final -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter import litellm +from litellm.litellm_core_utils.core_helpers import decoded_json_object + +_DECODED_JSON: Final = TypeAdapter(object) class PromptLayerLogger: @@ -61,8 +64,8 @@ class PromptLayerLogger: }, ) - response_json: Final = request_response.json() - if not request_response.json().get("success", False): + response_json: Final = decoded_json_object(_DECODED_JSON.validate_python(request_response.json())) + if not response_json.get("success", False): raise Exception("Promptlayer did not successfully log the response!") print_verbose(f"Prompt Layer Logging: success - final response object: {request_response.text}") diff --git a/litellm/integrations/weave/weave_otel.py b/litellm/integrations/weave/weave_otel.py index 2bff3682de7..68a9dc3e162 100644 --- a/litellm/integrations/weave/weave_otel.py +++ b/litellm/integrations/weave/weave_otel.py @@ -4,7 +4,8 @@ import base64 import json import os from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Protocol from opentelemetry.trace import Status, StatusCode from typing_extensions import override @@ -19,6 +20,7 @@ from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes impor BaseLLMObsOTELAttributes, safe_set_attribute, ) +from litellm.litellm_core_utils.core_helpers import has_attributes, require_attributes from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.types.integrations.weave_otel import WeaveOtelConfig, WeaveSpanAttributes from litellm.types.utils import StandardCallbackDynamicParams @@ -34,6 +36,21 @@ WEAVE_BASE_URL: Final = "https://trace.wandb.ai" WEAVE_OTEL_ENDPOINT: Final = "/otel/v1/traces" +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + +class _SupportsReplace(Protocol): + def replace(self, old: str, new: str, /) -> object: ... + + +class _ModelDumpable(Protocol): + def model_dump(self) -> object: ... + + class WeaveLLMObsOTELAttributes(BaseLLMObsOTELAttributes): """ Weave-specific LLM observability OTEL attributes. @@ -44,15 +61,15 @@ class WeaveLLMObsOTELAttributes(BaseLLMObsOTELAttributes): @staticmethod @override - def set_messages(span: Span, kwargs: dict[str, Any]): + def set_messages(span: Span, kwargs: Mapping[str, object]): """Set input messages as span attributes using OpenInference conventions.""" messages: Final = kwargs.get("messages") or [] - optional_params: Final = kwargs.get("optional_params") or {} + optional_params: Final = kwargs.get("optional_params") or _NO_FIELDS prompt: Final = {"messages": messages} - functions: Final = optional_params.get("functions") - tools: Final = optional_params.get("tools") + functions: Final = require_attributes(optional_params, _SupportsGet, "get").get("functions") + tools: Final = require_attributes(optional_params, _SupportsGet, "get").get("tools") if functions is not None: prompt["functions"] = functions if tools is not None: @@ -60,7 +77,7 @@ class WeaveLLMObsOTELAttributes(BaseLLMObsOTELAttributes): safe_set_attribute(span, OpenInferenceSpanAttributes.INPUT_VALUE, json.dumps(prompt)) -def _set_weave_specific_attributes(span: Span, kwargs: Mapping[str, Any], response_obj: Any): +def _set_weave_specific_attributes(span: Span, kwargs: Mapping[str, object], response_obj: object): """ Sets Weave-specific metadata attributes onto the OTEL span. @@ -69,25 +86,27 @@ def _set_weave_specific_attributes(span: Span, kwargs: Mapping[str, Any], respon """ # Extract all needed data upfront - litellm_params: Final = kwargs.get("litellm_params") or {} + litellm_params: Final = kwargs.get("litellm_params") or _NO_FIELDS # optional_params = kwargs.get("optional_params") or {} - metadata: Final = kwargs.get("metadata") or {} + metadata: Final = kwargs.get("metadata") or _NO_FIELDS model: Final = kwargs.get("model") or "" - custom_llm_provider: Final = litellm_params.get("custom_llm_provider") or "" + custom_llm_provider: Final = ( + require_attributes(litellm_params, _SupportsGet, "get").get("custom_llm_provider") or "" + ) # Weave supports a custom display name and will default to the model name if not provided. - display_name = metadata.get("display_name") + display_name = require_attributes(metadata, _SupportsGet, "get").get("display_name") if not display_name and model: if custom_llm_provider: display_name = f"{custom_llm_provider}/{model}" else: display_name = model if display_name: - display_name = display_name.replace("/", "__") + display_name = require_attributes(display_name, _SupportsReplace, "replace").replace("/", "__") safe_set_attribute(span, WeaveSpanAttributes.DISPLAY_NAME.value, display_name) # Weave threads are OpenInference sessions. - if (session_id := metadata.get("session_id")) is not None: + if (session_id := require_attributes(metadata, _SupportsGet, "get").get("session_id")) is not None: if isinstance(session_id, (list, dict)): session_id = safe_dumps(session_id) safe_set_attribute(span, WeaveSpanAttributes.THREAD_ID.value, session_id) @@ -97,7 +116,7 @@ def _set_weave_specific_attributes(span: Span, kwargs: Mapping[str, Any], respon # but we override them here to better match Weave's expectations if response_obj: output_dict = None - if hasattr(response_obj, "model_dump"): + if has_attributes(response_obj, _ModelDumpable, "model_dump"): output_dict = response_obj.model_dump() elif hasattr(response_obj, "get"): output_dict = response_obj diff --git a/litellm/integrations/websearch_interception/tools.py b/litellm/integrations/websearch_interception/tools.py index 36b02c48dc5..227f9f33baa 100644 --- a/litellm/integrations/websearch_interception/tools.py +++ b/litellm/integrations/websearch_interception/tools.py @@ -7,9 +7,19 @@ to this format for consistent interception and execution. """ from collections.abc import Mapping -from typing import Any, Final +from typing import Final, Protocol from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME +from litellm.litellm_core_utils.core_helpers import require_attributes + + +class _SupportsGetWithDefault(Protocol): + def get(self, key: str, default: str, /) -> object: ... + + +class _SupportsStartswith(Protocol): + def startswith(self, prefix: str, /) -> object: ... + _WEB_SEARCH_TOOL_DESCRIPTION: Final = ( "Search the web for information. Use this when you need current " @@ -164,7 +174,7 @@ def is_chat_completion_request_web_search_tool(tool: Mapping[str, object]) -> bo return tool.get("type") == "web_search" or is_web_search_tool(dict(tool)) -def is_web_search_tool_chat_completion(tool: dict[str, Any]) -> bool: +def is_web_search_tool_chat_completion(tool: Mapping[str, object]) -> bool: """ Check if a tool is a web search tool for Chat Completions API (strict check). @@ -197,7 +207,7 @@ def is_web_search_tool_chat_completion(tool: dict[str, Any]) -> bool: # Check for OpenAI format: {"type": "function", "function": {"name": "litellm_web_search"}} if tool_type == "function" and "function" in tool: function_def: Final = tool.get("function", {}) - function_name: Final = function_def.get("name", "") + function_name: Final = require_attributes(function_def, _SupportsGetWithDefault, "get").get("name", "") if function_name == LITELLM_WEB_SEARCH_TOOL_NAME: return True @@ -229,7 +239,7 @@ def is_anthropic_native_web_search_tool(tool: Mapping[str, object]) -> bool: return tool_type.startswith("web_search_") and tool_type != "function" -def is_web_search_tool(tool: dict[str, Any]) -> bool: +def is_web_search_tool(tool: Mapping[str, object]) -> bool: """ Check if a tool is a web search tool (native or LiteLLM standard). @@ -281,7 +291,7 @@ def is_web_search_tool(tool: dict[str, Any]) -> bool: # Check for OpenAI format: {"type": "function", "function": {"name": "..."}} if tool_type == "function" and "function" in tool: function_def: Final = tool.get("function", {}) - function_name: Final = function_def.get("name", "") + function_name: Final = require_attributes(function_def, _SupportsGetWithDefault, "get").get("name", "") if function_name == LITELLM_WEB_SEARCH_TOOL_NAME: return True @@ -290,7 +300,7 @@ def is_web_search_tool(tool: dict[str, Any]) -> bool: return True # Check for native Anthropic web_search_* types - if tool_type.startswith("web_search_"): + if require_attributes(tool_type, _SupportsStartswith, "startswith").startswith("web_search_"): return True # Check for Claude Code's web_search with a type field diff --git a/litellm/integrations/weights_biases.py b/litellm/integrations/weights_biases.py index 9bc070a1f9a..01670501f22 100644 --- a/litellm/integrations/weights_biases.py +++ b/litellm/integrations/weights_biases.py @@ -1,11 +1,13 @@ from typing import Final +from litellm.litellm_core_utils.core_helpers import iter_items, subscript + imported_openAIResponse = True try: import io import logging from collections.abc import Mapping - from typing import Any, Literal, Protocol, TypeVar + from typing import Literal, Protocol, TypeVar from wandb.sdk.data_types import trace_tree @@ -25,7 +27,7 @@ try: class OpenAIRequestResponseResolver: def __call__( self, - request: dict[str, Any], + request: Mapping[str, object], response: OpenAIResponse, time_elapsed: float, ) -> trace_tree.WBTraceTree | None: @@ -110,14 +112,14 @@ try: def _resolve_chat_completion( self, - request: dict[str, Any], + request: Mapping[str, object], response: OpenAIResponse, time_elapsed: float, ) -> trace_tree.WBTraceTree: """Resolves the request and response objects for `openai.Completion`.""" prompt: Final = io.StringIO() - for message in request["messages"]: - prompt.write(f"\n\n**{message['role']}**: {message['content']}\n") + for message in iter_items(request["messages"]): + prompt.write(f"\n\n**{subscript(message, 'role')}**: {subscript(message, 'content')}\n") request_str: Final = prompt.getvalue() choices: Final = [ diff --git a/litellm/interactions/litellm_responses_transformation/transformation.py b/litellm/interactions/litellm_responses_transformation/transformation.py index 2e9b62dcfe2..6bc7413dfa4 100644 --- a/litellm/interactions/litellm_responses_transformation/transformation.py +++ b/litellm/interactions/litellm_responses_transformation/transformation.py @@ -8,11 +8,11 @@ This module handles transforming between: from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Final, cast +from typing import Final, Protocol, cast from pydantic import BaseModel -from litellm.litellm_core_utils.core_helpers import billed_usage_cost +from litellm.litellm_core_utils.core_helpers import attribute_or_none, billed_usage_cost, has_attributes from litellm.types.interactions import ( InteractionInput, InteractionsAPIOptionalRequestParams, @@ -27,6 +27,10 @@ from litellm.types.llms.openai import ( _STEP_TYPE_ROLES: Final = MappingProxyType({"user_input": "user", "model_output": "assistant"}) +class _IsoFormattable(Protocol): + def isoformat(self) -> object: ... + + class LiteLLMResponsesInteractionsConfig: """Configuration class for transforming between Interactions API and Responses API.""" @@ -232,12 +236,12 @@ class LiteLLMResponsesInteractionsConfig: ) # Convert created_at to ISO string - created_at: Final = getattr(responses_response, "created_at", None) + created_at: Final = attribute_or_none(responses_response, "created_at") if isinstance(created_at, int): from datetime import datetime created = datetime.fromtimestamp(created_at).isoformat() - elif created_at is not None and hasattr(created_at, "isoformat"): + elif created_at is not None and has_attributes(created_at, _IsoFormattable, "isoformat"): created = created_at.isoformat() else: created = None diff --git a/litellm/interactions/streaming_iterator.py b/litellm/interactions/streaming_iterator.py index c7faa3fc510..674c05a7545 100644 --- a/litellm/interactions/streaming_iterator.py +++ b/litellm/interactions/streaming_iterator.py @@ -7,8 +7,9 @@ from the Google Interactions API, similar to the responses API streaming iterato import asyncio import json +from collections.abc import Mapping from datetime import datetime -from typing import Any, Final +from typing import Final, Protocol import httpx from pydantic import ConfigDict, TypeAdapter @@ -16,7 +17,7 @@ from pydantic import ConfigDict, TypeAdapter from litellm._logging import verbose_logger from litellm.constants import STREAM_SSE_DONE_STRING from litellm.litellm_core_utils.asyncify import run_async_function -from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.core_helpers import process_response_headers, require_attributes from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.thread_pool_executor import executor @@ -32,6 +33,10 @@ _INTERACTIONS_STREAMING_CHUNK_ADAPTER: Final = TypeAdapter( ) +class _SupportsGet(Protocol): + def get(self, key: str, default: object, /) -> object: ... + + class BaseInteractionsAPIStreamingIterator: """ Base class for streaming iterators that process responses from the Interactions API. @@ -45,7 +50,7 @@ class BaseInteractionsAPIStreamingIterator: model: str | None, interactions_api_config: BaseInteractionsAPIConfig, logging_obj: LiteLLMLoggingObj, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, object] | None = None, custom_llm_provider: str | None = None, ): self.response = response @@ -65,9 +70,9 @@ class BaseInteractionsAPIStreamingIterator: model=model or "", optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), ) - _model_info: Final[dict] = litellm_metadata.get("model_info", {}) if litellm_metadata else {} + _model_info: Final[object] = litellm_metadata.get("model_info", {}) if litellm_metadata else {} self._hidden_params = { - "model_id": _model_info.get("id", None), + "model_id": require_attributes(_model_info, _SupportsGet, "get").get("id", None), "api_base": _api_base, } self._hidden_params["additional_headers"] = process_response_headers(self.response.headers or {}) @@ -135,7 +140,7 @@ class InteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator): model: str | None, interactions_api_config: BaseInteractionsAPIConfig, logging_obj: LiteLLMLoggingObj, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, object] | None = None, custom_llm_provider: str | None = None, ): super().__init__( @@ -202,7 +207,7 @@ class SyncInteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator) model: str | None, interactions_api_config: BaseInteractionsAPIConfig, logging_obj: LiteLLMLoggingObj, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, object] | None = None, custom_llm_provider: str | None = None, ): super().__init__( diff --git a/litellm/litellm_core_utils/audio_utils/utils.py b/litellm/litellm_core_utils/audio_utils/utils.py index 29b10ac3c83..808a424815e 100644 --- a/litellm/litellm_core_utils/audio_utils/utils.py +++ b/litellm/litellm_core_utils/audio_utils/utils.py @@ -313,7 +313,7 @@ def calculate_request_duration(file: FileTypes) -> float | None: file_content = content elif hasattr(content, "read") and not isinstance(content, (str, os.PathLike)): # File-like object in tuple - current_pos: Final = getattr(content, "tell", lambda: None)() + current_pos: Final = content.tell() if hasattr(content, "tell") else None # Seek to start to ensure we read the entire content if hasattr(content, "seek"): content.seek(0) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index b287e564666..39198c8f2e8 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -4,6 +4,7 @@ import copy import logging import math import re +import sys from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable @@ -133,7 +134,7 @@ class _SupportsIter(Protocol): @runtime_checkable class _SupportsGetItem(Protocol): - def __getitem__(self, index: int, /) -> object: ... + def __getitem__(self, key: str | int, /) -> object: ... def attribute_or_none(source: object, name: str) -> object: @@ -166,6 +167,28 @@ def iter_items(value: object) -> Iterable[object]: raise TypeError(f"'{type(value).__name__}' object is not iterable") +@runtime_checkable +class _SupportsContains(Protocol): + def __contains__(self, key: object, /) -> bool: ... + + +_NOT_A_CONTAINER: Final = "is not a container or iterable" if sys.version_info >= (3, 14) else "is not iterable" + + +def json_contains(value: object, key: object) -> bool: + if isinstance(value, _SupportsContains): + return key in value + if isinstance(value, (_SupportsIter, _SupportsGetItem)): + return key in tuple(iter_items(value)) + raise TypeError(f"argument of type '{type(value).__name__}' {_NOT_A_CONTAINER}") + + +def subscript(value: object, key: str | int) -> object: + if isinstance(value, _SupportsGetItem): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + + def attribute_items(source: object, name: str) -> Iterable[object]: value: Final = attribute_or_none(source, name) return iter_items(value) if value else () diff --git a/litellm/litellm_core_utils/dd_tracing.py b/litellm/litellm_core_utils/dd_tracing.py index a4db3bcd5c2..e72c7d4980b 100644 --- a/litellm/litellm_core_utils/dd_tracing.py +++ b/litellm/litellm_core_utils/dd_tracing.py @@ -5,8 +5,9 @@ If the ddtrace package is not installed, the tracer will be a no-op. """ from contextlib import contextmanager -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.secret_managers.main import get_secret_bool if TYPE_CHECKING: @@ -15,6 +16,14 @@ else: DD_TRACER = Any +class _TagStrSpan(Protocol): + def set_tag_str(self, key: str, value: str, /) -> object: ... + + +class _TagSpan(Protocol): + def set_tag(self, key: str, value: str, /) -> object: ... + + class NullSpan: """A no-op span implementation.""" @@ -81,7 +90,7 @@ else: tracer = NullTracer() -def get_active_span() -> Any | None: +def get_active_span() -> object: """ Return the active Datadog span, checking current span first and then root span. """ @@ -115,10 +124,10 @@ def set_active_span_tag(tag_key: str, tag_value: str) -> bool: return False try: - if hasattr(span, "set_tag_str"): + if has_attributes(span, _TagStrSpan, "set_tag_str"): span.set_tag_str(tag_key, str(tag_value)) return True - if hasattr(span, "set_tag"): + if has_attributes(span, _TagSpan, "set_tag"): span.set_tag(tag_key, str(tag_value)) return True except Exception: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ed6db5946af..18d29603c49 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -19,7 +19,7 @@ from datetime import datetime as dt_object from enum import Enum from functools import lru_cache from types import MappingProxyType, TracebackType -from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, TypeVar, cast from httpx import Response, Timeout from openai import AsyncStream as OpenAIAsyncStream @@ -90,6 +90,7 @@ from litellm.litellm_core_utils.classifier_logging import ( is_classifier_call, ) from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, finite_cost, get_provider_response_headers_from_hidden_params, is_expected_client_error, @@ -99,6 +100,7 @@ from litellm.litellm_core_utils.core_helpers import ( proxy_stamped_used_client_oauth_token, reconstruct_model_name, reports_native_messages_cost, + require_attributes, set_response_cost_in_hidden_params, single_log_line, ) @@ -666,6 +668,31 @@ class UsageObservation(Enum): ESTIMATED = "estimated" +class _A2AUsageSlot(Protocol): + usage: Mapping[str, object] | None + + +class _A2AResponseCopier(Protocol): + model_copy: Callable[..., _A2AUsageSlot] + + +class _MetadataLookup(Protocol): + get: Callable[[str], object] + + +class _ParamsLookup(Protocol): + def get(self, key: str, default: object, /) -> object: ... + + +def _callback_control_params(litellm_params: object) -> dict[str, object]: + mapping: Final = as_str_mapping(litellm_params) + if isinstance(mapping, dict): + return mapping + if mapping is not None: + return dict(mapping) + return litellm_params.model_dump() if isinstance(litellm_params, BaseModel) else {} + + class Logging(LiteLLMLoggingBaseClass): global \ supabaseClient, \ @@ -929,7 +956,9 @@ class Logging(LiteLLMLoggingBaseClass): new_call.start_time = start_time new_call.litellm_params = litellm_params new_call.model_call_details = { - **{key: value for key, value in self.model_call_details.items() if not key.startswith("_accounting_")}, + **{ + key: value for key, value in self._model_call_details_view.items() if not key.startswith("_accounting_") + }, "litellm_call_id": litellm_call_id, "start_time": start_time, "first_api_call_start_time": None, @@ -2467,7 +2496,7 @@ class Logging(LiteLLMLoggingBaseClass): return self.response_cost_calculator(result=result, cache_hit=cache_hit) @staticmethod - def is_sync_litellm_request(litellm_params: Mapping[str, object]) -> bool: + def is_sync_litellm_request(litellm_params: _ParamsLookup) -> bool: """True for sync SDK entrypoints (``completion``), false for async (``acompletion``, etc.).""" return ( litellm_params.get(CallTypes.acompletion.value, False) is not True @@ -2918,11 +2947,12 @@ class Logging(LiteLLMLoggingBaseClass): """Deprecated alias of mark_logging_complete, kept for callers of the old name""" self.mark_logging_complete(event_type=event_type) - def should_run_callback(self, callback: litellm.CALLBACK_TYPES, litellm_params: dict, event_hook: str) -> bool: + def should_run_callback(self, callback: litellm.CALLBACK_TYPES, litellm_params: object, event_hook: str) -> bool: if litellm.global_disable_no_log_param: return True - if litellm_params.get("no-log", False) is True: + request_params: Final = require_attributes(litellm_params, _ParamsLookup, "get") + if request_params.get("no-log", False) is True: # proxy cost tracking cal backs should run if not (isinstance(callback, CustomLogger) and "_PROXY_" in callback.__class__.__name__): @@ -2932,7 +2962,7 @@ class Logging(LiteLLMLoggingBaseClass): # Check for dynamically disabled callbacks via headers if EnterpriseCallbackControls is not None and EnterpriseCallbackControls.is_callback_disabled_dynamically( callback=callback, - litellm_params=litellm_params, + litellm_params=_callback_control_params(litellm_params), standard_callback_dynamic_params=self.standard_callback_dynamic_params, ): verbose_logger.debug( @@ -3643,8 +3673,8 @@ class Logging(LiteLLMLoggingBaseClass): build_logging_payload=True, publish_logging_payload=not for_async_request, ) - litellm_params = self.model_call_details.get("litellm_params", {}) - is_sync_request: Final = self.is_sync_litellm_request(litellm_params) + litellm_params = self._model_call_details_view.get("litellm_params", {}) + is_sync_request: Final = self.is_sync_litellm_request(require_attributes(litellm_params, _ParamsLookup, "get")) try: ## BUILD COMPLETE STREAMED RESPONSE complete_streaming_response: ( @@ -3800,18 +3830,23 @@ class Logging(LiteLLMLoggingBaseClass): model = call_details_view["model"] messages = call_details_view["messages"] - optional_params = kwargs.get("optional_params", {}) - litellm_params = kwargs.get("litellm_params", {}) + optional_params = call_details_view.get("optional_params", {}) + litellm_params = call_details_view.get("litellm_params", {}) supabaseClient.log_event( model=model, messages=messages, - end_user=optional_params.get("user", "default"), + end_user=require_attributes(optional_params, _ParamsLookup, "get").get("user", "default"), response_obj=result, start_time=start_time, end_time=end_time, litellm_call_id=( current_call_id - if (current_call_id := litellm_params.get("litellm_call_id")) is not None + if ( + current_call_id := require_attributes(litellm_params, _ParamsLookup, "get").get( + "litellm_call_id", None + ) + ) + is not None else str(uuid.uuid4()) ), print_verbose=print_verbose, @@ -4369,7 +4404,7 @@ class Logging(LiteLLMLoggingBaseClass): for callback in delivery_callbacks: # check if callback can run for this request - litellm_params = self.model_call_details.get("litellm_params", {}) + litellm_params = self._model_call_details_view.get("litellm_params", {}) should_run = self.should_run_callback( callback=callback, litellm_params=litellm_params, @@ -4658,8 +4693,8 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug("Logging Details LiteLLM-Failure Call: %s", litellm.failure_callback) if not self.should_run_logging(event_type="sync_failure"): # prevent double logging return - litellm_params: Final = self.model_call_details.get("litellm_params", {}) - is_sync_request: Final = self.is_sync_litellm_request(litellm_params) + litellm_params: Final = self._model_call_details_view.get("litellm_params", {}) + is_sync_request: Final = self.is_sync_litellm_request(require_attributes(litellm_params, _ParamsLookup, "get")) try: start_time, end_time = self._failure_handler_helper_fn( @@ -4879,7 +4914,7 @@ class Logging(LiteLLMLoggingBaseClass): self.mark_logging_complete(event_type="async_failure") for callback in callbacks: try: - litellm_params = self.model_call_details.get("litellm_params", {}) + litellm_params = self._model_call_details_view.get("litellm_params", {}) should_run = self.should_run_callback( callback=callback, litellm_params=litellm_params, @@ -5441,7 +5476,7 @@ class Logging(LiteLLMLoggingBaseClass): httpx_response: Final[Response | None] = self.model_call_details.get("httpx_response", None) if httpx_response is None: raise ValueError("Google GenAI Generate Content: httpx_response is None") - dict_result: Final = httpx_response.json() + dict_result: Final[GenerateContentResponseBody] = httpx_response.json() result = litellm.VertexGeminiConfig().transform_google_generate_content_to_openai_model_response( completion_response=dict_result, model_response=litellm.ModelResponse(), @@ -5454,7 +5489,7 @@ class Logging(LiteLLMLoggingBaseClass): ) return result - def _handle_a2a_response_logging(self, result: Any) -> object: + def _handle_a2a_response_logging(self, result: object) -> object: """ Handles logging for A2A (Agent-to-Agent) responses. @@ -5473,7 +5508,7 @@ class Logging(LiteLLMLoggingBaseClass): return result # Deep copy result and add usage - result_copy: Final = result.model_copy(deep=True) + result_copy: Final = require_attributes(result, _A2AResponseCopier, "model_copy").model_copy(deep=True) result_copy.usage = usage.model_dump() if hasattr(usage, "model_dump") else dict(usage) return result_copy @@ -7475,7 +7510,7 @@ class StandardLoggingPayloadSetup: @staticmethod def get_standard_logging_payload_trace_id( logging_obj: Logging, - litellm_params: Mapping[str, Any], + litellm_params: Mapping[str, object], ) -> str: """ Returns the `litellm_trace_id` for this request @@ -7491,9 +7526,13 @@ class StandardLoggingPayloadSetup: """ dynamic_litellm_session_id: Final = litellm_params.get("litellm_session_id") dynamic_litellm_trace_id: Final = litellm_params.get("litellm_trace_id") - metadata: Final[Mapping[str, object] | None] = litellm_params.get("metadata") - metadata_session_id: Final = metadata.get("session_id") if metadata else None - metadata_trace_id: Final = metadata.get("trace_id") if metadata else None + metadata: Final = litellm_params.get("metadata") + metadata_session_id: Final = ( + require_attributes(metadata, _MetadataLookup, "get").get("session_id") if metadata else None + ) + metadata_trace_id: Final = ( + require_attributes(metadata, _MetadataLookup, "get").get("trace_id") if metadata else None + ) ordered_candidates: Final[tuple[object, object, object, object]] = ( (dynamic_litellm_trace_id, dynamic_litellm_session_id, metadata_trace_id, metadata_session_id) @@ -7508,7 +7547,7 @@ class StandardLoggingPayloadSetup: @staticmethod def get_standard_logging_payload_session_id( logging_obj: Logging, - litellm_params: Mapping[str, Any], + litellm_params: Mapping[str, object], ) -> str: """ Returns the end-user/conversation `litellm_session_id` for this request, independent of trace_id. @@ -7523,8 +7562,10 @@ class StandardLoggingPayloadSetup: dynamic_litellm_session_id: Final[object] = litellm_params.get("litellm_session_id") if dynamic_litellm_session_id: return str(dynamic_litellm_session_id) - metadata: Final[Mapping[str, object] | None] = litellm_params.get("metadata") - metadata_session_id: Final = metadata.get("session_id") if metadata else None + metadata: Final = litellm_params.get("metadata") + metadata_session_id: Final = ( + require_attributes(metadata, _MetadataLookup, "get").get("session_id") if metadata else None + ) if metadata_session_id: return str(metadata_session_id) return logging_obj.litellm_session_id diff --git a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py index 78a8cb51bac..9a3197cc044 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py +++ b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py @@ -1,7 +1,8 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Any, Final +from typing import Any, Final, Protocol +from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes from litellm.types.utils import ( CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, @@ -55,19 +56,29 @@ _INTERACTIONS_MODALITY_FIELDS: Final[Mapping[str, str]] = MappingProxyType( ) -def _modality_field(entry: Mapping[str, object]) -> str | None: - return _INTERACTIONS_MODALITY_FIELDS.get(str(entry.get("modality", "")).lower()) +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + +def _modality_field(entry: object) -> str | None: + return _INTERACTIONS_MODALITY_FIELDS.get( + str(require_attributes(entry, _SupportsGet, "get").get("modality", "")).lower() + ) def _token_count(value: object) -> int: return value if isinstance(value, int) else 0 -def _modality_token_sums(entries: Sequence[Mapping[str, object]]) -> Mapping[str, int]: +def _modality_token_sums(entries: Sequence[object]) -> Mapping[str, int]: fields: Final = frozenset(field for entry in entries if (field := _modality_field(entry)) is not None) return MappingProxyType( { - field: sum(_token_count(entry.get("tokens")) for entry in entries if _modality_field(entry) == field) + field: sum( + _token_count(require_attributes(entry, _SupportsGet, "get").get("tokens")) + for entry in entries + if _modality_field(entry) == field + ) for field in fields } ) @@ -119,12 +130,16 @@ class InteractionsUsageObjectTransformation: return "total_input_tokens" in usage_object or "total_output_tokens" in usage_object @staticmethod - def transform_interactions_usage_object(usage_object: Mapping[str, Any]) -> Usage: - input_entries: Final = tuple(usage_object.get("input_tokens_by_modality") or ()) + tuple( - usage_object.get("tool_use_tokens_by_modality") or () + def transform_interactions_usage_object(usage_object: Mapping[str, object]) -> Usage: + input_entries: Final = tuple(iter_items(usage_object.get("input_tokens_by_modality") or ())) + tuple( + iter_items(usage_object.get("tool_use_tokens_by_modality") or ()) + ) + cached_sums: Final = _modality_token_sums( + tuple(iter_items(usage_object.get("cached_tokens_by_modality") or ())) + ) + output_sums: Final = _modality_token_sums( + tuple(iter_items(usage_object.get("output_tokens_by_modality") or ())) ) - cached_sums: Final = _modality_token_sums(tuple(usage_object.get("cached_tokens_by_modality") or ())) - output_sums: Final = _modality_token_sums(tuple(usage_object.get("output_tokens_by_modality") or ())) total_cached_tokens: Final = _token_count(usage_object.get("total_cached_tokens")) input_sums: Final = _subtract_cached_from_input( diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 9ac4f9d5ee8..ae7c4b47ece 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -6,10 +6,11 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass, field, replace from datetime import datetime, timezone, tzinfo from types import MappingProxyType -from typing import Final, Literal, Protocol, TypedDict, cast +from typing import Final, Literal, Protocol, SupportsFloat, SupportsIndex, TypedDict, cast from zoneinfo import ZoneInfo, ZoneInfoNotFoundError -from typing_extensions import ReadOnly +from pydantic import TypeAdapter +from typing_extensions import Buffer, ReadOnly import litellm from litellm._internal_context import current_billing_time @@ -48,10 +49,10 @@ _IMAGE_RESPONSE_CALL_TYPES: Final = frozenset( } ) -# Pre-resolved DataResidency enum values for fast membership checks _VALID_DATA_RESIDENCIES: Final = frozenset(r.value for r in DataResidency) _DEPLOYMENT_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields) +_IMAGE_REQUEST_PARAMS: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) _BASE_CACHE_WRITE_RATE_KEY: Final = "cache_creation_input_token_cost" @@ -1428,6 +1429,13 @@ def get_regional_uplift_multiplier(model_info: ModelInfo, data_residency: str | multiplier: Final = model_info.get(f"regional_processing_uplift_multiplier_{residency}") if multiplier is None: return 1.0 + if not isinstance(multiplier, (str, Buffer, SupportsFloat, SupportsIndex)): + verbose_logger.error( + "Invalid regional_processing_uplift_multiplier_%s for model (got %s); defaulting to 1.0", + residency, + type(multiplier).__name__, + ) + return 1.0 try: return float(multiplier) except (TypeError, ValueError): @@ -1464,7 +1472,7 @@ def get_fedramp_uplift_multiplier(model_info: ModelInfo, data_residency: str | N return 1.0 -def get_vertex_regional_endpoint_uplift(model_info: ModelInfo, vertex_location: str | None) -> float: +def get_vertex_regional_endpoint_uplift(model_info: Mapping[str, object], vertex_location: str | None) -> float: """ Resolve the per-model uplift multiplier for Vertex AI non-global (regional and multi-region) endpoints. @@ -2318,6 +2326,20 @@ class CostCalculatorUtils: ) resolved_n: Final = n if n is not None else (len(completion_response.data) if completion_response.data else 0) pricing: Final = deployment_pricing(model_info) + if custom_llm_provider == litellm.LlmProviders.XAI.value: + from litellm.llms.xai.image_generation.cost_calculator import cost_calculator as xai_image_cost_calculator + + xai_cost: Final = xai_image_cost_calculator( + model=model, + image_response=completion_response, + quality=quality, + n=resolved_n, + optional_params=_IMAGE_REQUEST_PARAMS.validate_python(optional_params), + call_type=call_type, + model_info=pricing, + ) + if xai_cost is not None: + return xai_cost model_priced_calculator: Final = _model_priced_image_cost_calculator(custom_llm_provider, model, pricing) if model_priced_calculator is not None: diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index ec857218c79..0572ca66fc9 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -6,7 +6,7 @@ import json from collections.abc import Hashable, Iterable, Mapping, Sequence from dataclasses import dataclass from itertools import chain, groupby -from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast +from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast from pydantic import TypeAdapter @@ -1143,7 +1143,7 @@ def strip_prefill_prefix_field(message: AllMessageValues) -> AllMessageValues: return _message_without_fields(message, frozenset({"prefix"})) -def filter_value_from_dict(dictionary: dict, key: str, depth: int = 0) -> Any: +def filter_value_from_dict(dictionary: dict, key: str, depth: int = 0) -> Mapping[object, object]: """ Filters a value from a dictionary diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 280abad77c5..87a7abefec8 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -10,7 +10,7 @@ from collections.abc import Callable, Container, Iterable, Iterator, Mapping, Se from dataclasses import dataclass, replace from enum import Enum from types import MappingProxyType -from typing import Annotated, Any, Final, TypeAlias, TypedDict, TypeVar, cast, overload +from typing import Annotated, Any, Final, Protocol, TypeAlias, TypedDict, TypeVar, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment from pydantic import BaseModel, BeforeValidator, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator @@ -22,7 +22,7 @@ import litellm.types.llms from litellm import verbose_logger from litellm._uuid import uuid from litellm.constants import REDACTED_BY_LITELLM -from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, require_attributes from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( anthropic_system_messages, find_anthropic_server_tool_result, @@ -4150,7 +4150,7 @@ class BedrockImageProcessor: _validate_format = validate_format @staticmethod - def _get_document_format(mime_type: str, supported_doc_formats: list[str]) -> str: + def _get_document_format(mime_type: str, supported_doc_formats: Sequence[str]) -> str: """ Get the document format from the mime type @@ -6571,6 +6571,20 @@ def has_tool_with_name(tools: object, tool_name: str) -> bool: return False +class _SupportsModelDump(Protocol): + def model_dump(self) -> dict[str, object]: ... + + +class _StructuredMessagesSource(Protocol): + def get_structured_messages(self, data: dict[str, object], /) -> Sequence[object] | None: ... + + +def _instantiate_structured_message_sources( + source_classes: Iterable[Callable[[], _StructuredMessagesSource]], +) -> Sequence[_StructuredMessagesSource]: + return [source_cls() for source_cls in source_classes] + + def resolve_structured_messages( messages: list[dict[str, object]] | None, request_kwargs: dict[str, Any], @@ -6608,16 +6622,17 @@ def resolve_structured_messages( break # 2. Fallback: try each mapped handler until one produces messages - handlers_to_try: Final[list[Any]] = [] - if call_type is not None and call_type in mappings: - handlers_to_try.append(mappings[call_type]()) - else: - handlers_to_try.extend(handler_cls() for handler_cls in mappings.values()) + handlers_to_try: Final = _instantiate_structured_message_sources( + (mappings[call_type],) if call_type is not None and call_type in mappings else mappings.values() + ) for handler in handlers_to_try: structured = handler.get_structured_messages(request_kwargs) if structured: - return [msg if isinstance(msg, dict) else msg.model_dump() for msg in structured] + return [ + msg if isinstance(msg, dict) else require_attributes(msg, _SupportsModelDump, "model_dump").model_dump() + for msg in structured + ] return None diff --git a/litellm/litellm_core_utils/prompt_templates/tpm_media_utils.py b/litellm/litellm_core_utils/prompt_templates/tpm_media_utils.py index a4cb52d689c..a3e2d2e15d9 100644 --- a/litellm/litellm_core_utils/prompt_templates/tpm_media_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/tpm_media_utils.py @@ -202,7 +202,7 @@ def _tpm_embedding_leaf(value: object) -> _TPMMediaProjection: part: Final = as_str_mapping(value) if part is not None and part.get("type") == "file": return _tpm_embedding_file_block(part) - if not isinstance(value, str) or not value.startswith("data:"): + if not isinstance(value, str) or is_plain_embedding_text(value): return _tpm_media_leaf(value) if not _is_base64_data_uri(value): return _TPMMediaProjection(value) @@ -210,6 +210,10 @@ def _tpm_embedding_leaf(value: object) -> _TPMMediaProjection: return _TPMMediaProjection("" if media else value, media) +def is_plain_embedding_text(value: object) -> bool: + return isinstance(value, str) and not value.startswith("data:") + + def _is_json_array(value: object) -> bool: return isinstance(value, (list, tuple)) diff --git a/litellm/litellm_core_utils/realtime_guardrails.py b/litellm/litellm_core_utils/realtime_guardrails.py new file mode 100644 index 00000000000..382eecc5a8c --- /dev/null +++ b/litellm/litellm_core_utils/realtime_guardrails.py @@ -0,0 +1,316 @@ +import json +from collections.abc import AsyncIterator, Iterable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_logger +from litellm.exceptions import ModifyResponseException + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + +_DEFAULT_VIOLATION_MESSAGE: Final = "I'm sorry, that request was blocked by the content filter." +BLOCKED_TOOL_OUTPUT: Final = json.dumps({"error": "Tool output blocked by content policy"}) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_JSON_ARRAY: Final = TypeAdapter(tuple[object, ...]) +_TEXTS: Final = TypeAdapter(tuple[str, ...]) +_TYPED_TEXT_PART_TYPES: Final = frozenset({"input_text", "output_text", "text"}) + + +@dataclass(frozen=True, slots=True) +class RealtimeGuardrailBlock: + message: str + ends_session_on_violation: bool + end_session_after_n_fails: int | None + + def ends_session(self, violation_count: int) -> bool: + return self.ends_session_on_violation or ( + self.end_session_after_n_fails is not None and violation_count >= self.end_session_after_n_fails + ) + + +@dataclass(frozen=True, slots=True) +class RealtimeGuardrailRewrite: + texts: tuple[str, ...] + + +RealtimeGuardrailVerdict: TypeAlias = RealtimeGuardrailBlock | RealtimeGuardrailRewrite | None + + +@dataclass(frozen=True, slots=True) +class RealtimeMessageText: + role: str + texts: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class RealtimeToolOutput: + texts: tuple[str, ...] + sanitized_event: str + + +RealtimeGuardedItem: TypeAlias = RealtimeMessageText | RealtimeToolOutput + + +def _json_object(value: object) -> Mapping[str, object] | None: + try: + return _JSON_OBJECT.validate_python(value) + except ValidationError: + return None + + +def _turn_copy(value: object) -> object: + mapping: Final = _json_object(value) + return value if mapping is None else dict(mapping) + + +def _typed_text(part: object) -> str | None: + part_object: Final = _json_object(part) + if part_object is None or part_object.get("type") not in _TYPED_TEXT_PART_TYPES: + return None + text: Final[object] = part_object.get("text", "") + return text if isinstance(text, str) else None + + +def _content_parts(content: object) -> Sequence[object]: + return _JSON_ARRAY.validate_python(content) if isinstance(content, list) else () + + +def _typed_texts(content: object) -> tuple[str, ...]: + return tuple(text for text in map(_typed_text, _content_parts(content)) if text is not None) + + +def _with_text(part: object, text: str | None) -> object: + part_object: Final = _json_object(part) + return part if text is None or part_object is None else {**part_object, "text": text} + + +def _tool_output(client_event: Mapping[str, object], item: Mapping[str, object]) -> RealtimeToolOutput | None: + output: Final[object] = item.get("output", "") + text: Final = output if isinstance(output, str) else json.dumps(output) + if not text: + return None + sanitized_event: Final = json.dumps({**client_event, "item": {**item, "output": BLOCKED_TOOL_OUTPUT}}) + return RealtimeToolOutput(texts=(text,), sanitized_event=sanitized_event) + + +def realtime_guarded_item(client_event: Mapping[str, object]) -> RealtimeGuardedItem | None: + if client_event.get("type") != "conversation.item.create": + return None + item: Final = _json_object(client_event.get("item")) + if item is None: + return None + if item.get("type") == "function_call_output": + return _tool_output(client_event, item) + texts: Final = _typed_texts(item.get("content")) + role: Final[object] = item.get("role") + return RealtimeMessageText(role=role if isinstance(role, str) else "user", texts=texts) if any(texts) else None + + +def rewritten_item(client_event: Mapping[str, object], texts: tuple[str, ...]) -> Mapping[str, object]: + item: Final = _json_object(client_event.get("item")) or {} + if item.get("type") == "function_call_output": + return {**item, "output": " ".join(texts)} + parts: Final = _content_parts(item.get("content")) + text_positions: Final = tuple(index for index, part in enumerate(parts) if _typed_text(part) is not None) + replacements: Final = dict(zip(text_positions, texts)) + merged_away: Final = frozenset(text_positions[len(texts) :]) + return { + **item, + "content": [ + _with_text(part, replacements.get(index)) for index, part in enumerate(parts) if index not in merged_away + ], + } + + +def _violation_text(error: Exception) -> str: + detail: Final[object] = getattr(error, "detail", None) + detail_object: Final = _json_object(detail) + if detail_object is not None: + reason: Final[object] = detail_object.get("error") + return reason if isinstance(reason, str) and reason else str(error) + if detail is not None: + return str(detail) + return str(error) or _DEFAULT_VIOLATION_MESSAGE + + +def _is_guardrail_block(error: Exception) -> bool: + return hasattr(error, "status_code") or isinstance(error, (ValueError, ModifyResponseException)) + + +def _screened_texts(screened: object, texts: tuple[str, ...]) -> tuple[str, ...]: + screened_object: Final = _json_object(screened) + try: + returned: Final = _TEXTS.validate_python(screened_object.get("texts") if screened_object else None) + except ValidationError: + return texts + return returned if len(returned) == len(texts) else texts + + +async def _guardrail_outcome( + guardrail: "CustomGuardrail", texts: tuple[str, ...], request_data: Mapping[str, object] +) -> RealtimeGuardrailBlock | tuple[str, ...]: + try: + screened: Final = await guardrail.apply_guardrail( + inputs={"texts": list(texts), "images": []}, + request_data=dict(request_data), + input_type="request", + ) + except Exception as error: + if not _is_guardrail_block(error): + verbose_logger.exception("[realtime guardrail] unexpected error in apply_guardrail: %s", error) + raise + return RealtimeGuardrailBlock( + message=guardrail.realtime_violation_message or _violation_text(error), + ends_session_on_violation=guardrail.on_violation == "end_session", + end_session_after_n_fails=guardrail.end_session_after_n_fails, + ) + return _screened_texts(screened, texts) + + +async def _screen( + guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] +) -> RealtimeGuardrailBlock | tuple[str, ...]: + if not guardrails: + return texts + outcome: Final = await _guardrail_outcome(guardrails[0], texts, request_data) + if isinstance(outcome, RealtimeGuardrailBlock): + return outcome + return await _screen(guardrails[1:], outcome, request_data) + + +async def _screened_until_block( + guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] +) -> AsyncIterator[RealtimeGuardrailBlock | str]: + for text in texts: + if isinstance(outcome := await _screen(guardrails, (text,), request_data), RealtimeGuardrailBlock): + yield outcome + return + for screened in outcome: + yield screened + + +async def _screen_each( + guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] +) -> RealtimeGuardrailBlock | tuple[str, ...]: + outcomes: Final = tuple([outcome async for outcome in _screened_until_block(guardrails, texts, request_data)]) + block: Final = next((outcome for outcome in outcomes if isinstance(outcome, RealtimeGuardrailBlock)), None) + return block if block is not None else tuple(outcome for outcome in outcomes if isinstance(outcome, str)) + + +def _joins(texts: tuple[str, ...]) -> tuple[str, ...]: + unspaced: Final = "".join(texts) + boundaries: Final = zip(texts, texts[1:]) + word_break_unknown: Final = any(not left[-1:].isspace() and not right[:1].isspace() for left, right in boundaries) + return (unspaced, " ".join(texts)) if word_break_unknown else (unspaced,) + + +async def _rewritten_parts( + guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] +) -> RealtimeGuardrailBlock | tuple[str, ...]: + parts: Final = await _screen_each(guardrails, texts, request_data) + if isinstance(parts, RealtimeGuardrailBlock): + return parts + rejoined: Final = _joins(parts) + screened_rejoined: Final = await _screen_each(guardrails, rejoined, request_data) + if isinstance(screened_rejoined, RealtimeGuardrailBlock): + return screened_rejoined + masked_join: Final = next((text for text, join in zip(screened_rejoined, rejoined) if text != join), None) + return parts if masked_join is None else (masked_join,) + + +def _in_role_scope(guardrail: "CustomGuardrail", role: str) -> bool: + from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, + effective_skip_system_message_for_guardrail, + effective_skip_tool_message_for_guardrail, + role_out_of_guardrail_scope, + ) + + return not role_out_of_guardrail_scope( + role, + skip_system_message=effective_skip_system_message_for_guardrail(guardrail), + skip_tool_message=effective_skip_tool_message_for_guardrail(guardrail), + scan_only_tool_results=effective_scan_only_tool_results_for_guardrail(guardrail), + ) + + +def realtime_runner_guardrail( + callback: object, + check_data: Mapping[str, object], + event_hooks: Sequence["GuardrailEventHooks"], +) -> "CustomGuardrail | None": + from litellm.integrations.custom_guardrail import CustomGuardrail + + if not isinstance(callback, CustomGuardrail) or callback.use_native_lifecycle_hooks: + return None + runs: Final = any( + callback.should_run_guardrail(data=check_data, event_type=event_hook) for event_hook in event_hooks + ) + return callback if runs else None + + +def realtime_screening_guardrails( + role: str, + event_hooks: Sequence["GuardrailEventHooks"], + check_data: Mapping[str, object], + callbacks: Iterable[object], +) -> tuple["CustomGuardrail", ...]: + runner_guardrails: Final = (realtime_runner_guardrail(callback, check_data, event_hooks) for callback in callbacks) + guardrails: Final = { + id(guardrail): guardrail + for guardrail in runner_guardrails + if guardrail is not None and _in_role_scope(guardrail, role) + } + return tuple(guardrails.values()) + + +async def realtime_guardrail_verdict( + texts: tuple[str, ...], + role: str, + event_hooks: Sequence["GuardrailEventHooks"], + request_data: Mapping[str, object], + user_api_key_dict: object | None, + callbacks: Iterable[object], +) -> RealtimeGuardrailVerdict: + check_data: Final = {**request_data, "transcript": " ".join(texts), "stream": True} + in_scope: Final = realtime_screening_guardrails(role, event_hooks, check_data, callbacks) + guardrail_request_data: Final = { + **{key: _turn_copy(value) for key, value in request_data.items()}, + "user_api_key_dict": user_api_key_dict, + } + joined: Final = _joins(texts) + screened: Final = await _screen_each(in_scope, joined, guardrail_request_data) + if isinstance(screened, RealtimeGuardrailBlock): + return screened + if screened == joined: + return None + rewritten: Final = screened if len(texts) < 2 else await _rewritten_parts(in_scope, texts, guardrail_request_data) + if isinstance(rewritten, RealtimeGuardrailBlock): + return rewritten + return None if rewritten == texts else RealtimeGuardrailRewrite(texts=rewritten) + + +def guardrail_violation_error_event(message: str) -> str: + return json.dumps( + { + "type": "error", + "error": {"type": "guardrail_violation", "message": message, "code": "content_policy_violation"}, + } + ) + + +def guardrail_voice_prompt_events(message: str) -> tuple[str, str]: + prompt: Final = f"Say exactly the following message to the user, word for word, do not add anything else: {message}" + return ( + json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": prompt}]}, + } + ), + json.dumps({"type": "response.create"}), + ) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index b6af181980e..01aa7296b0f 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -47,13 +47,27 @@ from litellm.types.realtime import ALL_DELTA_TYPES from .litellm_logging import Logging as LiteLLMLogging from .realtime_errors import client_close_code, realtime_error_event, websocket_close_reason +from .realtime_guardrails import ( + RealtimeGuardedItem, + RealtimeGuardrailBlock, + RealtimeGuardrailRewrite, + RealtimeGuardrailVerdict, + RealtimeMessageText, + RealtimeToolOutput, + guardrail_violation_error_event, + guardrail_voice_prompt_events, + realtime_guarded_item, + realtime_guardrail_verdict, + realtime_runner_guardrail, + realtime_screening_guardrails, + rewritten_item, +) if TYPE_CHECKING: from openai.types.realtime.realtime_audio_formats import RealtimeAudioFormats from websockets.asyncio.client import ClientConnection from websockets.exceptions import ConnectionClosed - from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.types.guardrails import GuardrailEventHooks CLIENT_CONNECTION_CLASS = ClientConnection @@ -437,21 +451,6 @@ def _opted_out_of_always_on_guardrail(callback: object, request_data: Mapping[st ) -def _realtime_runner_guardrail( - callback: object, - check_data: Mapping[str, object], - event_hooks: Sequence["GuardrailEventHooks"], -) -> "CustomGuardrail | None": - from litellm.integrations.custom_guardrail import CustomGuardrail - - if not isinstance(callback, CustomGuardrail) or callback.use_native_lifecycle_hooks: - return None - runs: Final = any( - callback.should_run_guardrail(data=check_data, event_type=event_hook) for event_hook in event_hooks - ) - return callback if runs else None - - def translation_guardrails_apply(request_data: Mapping[str, object]) -> bool: from litellm.types.guardrails import GuardrailEventHooks @@ -479,7 +478,7 @@ def realtime_guardrails_apply(request_data: Mapping[str, object], callbacks: Seq check_data: Final = {**request_data, "stream": True} runner_event_hooks: Final = (GuardrailEventHooks.realtime_input_transcription, GuardrailEventHooks.pre_call) return any( - _realtime_runner_guardrail(callback, check_data, runner_event_hooks) is not None for callback in callbacks + realtime_runner_guardrail(callback, check_data, runner_event_hooks) is not None for callback in callbacks ) @@ -1663,16 +1662,60 @@ class RealTimeStreaming: async def run_realtime_guardrails( self, - transcript: str, + transcript: object, item_id: str | None = None, pre_block_backend_message: str | None = None, event_hooks: Sequence["GuardrailEventHooks"] | None = None, ) -> bool: """ - Run registered guardrails on realtime text (transcript, user message, tool output). - Returns True if blocked (synthetic warning already sent to client). Returns False if clean (caller should send response.create to the backend). + """ + from litellm.types.guardrails import GuardrailEventHooks + + if isinstance(transcript, str): + verdict: Final = await self.apply_realtime_guardrails( + (transcript,), "user", item_id, pre_block_backend_message, event_hooks + ) + return isinstance(verdict, RealtimeGuardrailBlock) + if realtime_screening_guardrails( + "user", + event_hooks or (GuardrailEventHooks.realtime_input_transcription,), + {**self.request_data, "stream": True}, + litellm.callbacks, + ): + raise TypeError(f"Realtime guardrails cannot screen a {type(transcript).__name__} transcript") + return False + + async def _screen_client_item(self, guarded_item: RealtimeGuardedItem | None) -> RealtimeGuardrailVerdict: + from litellm.types.guardrails import GuardrailEventHooks + + match guarded_item: + case RealtimeToolOutput(texts=texts, sanitized_event=sanitized_event): + return await self.apply_realtime_guardrails( + texts, + "tool", + pre_block_backend_message=sanitized_event, + event_hooks=[GuardrailEventHooks.pre_call], + ) + case RealtimeMessageText(role=role, texts=texts): + return await self.apply_realtime_guardrails(texts, role, event_hooks=[GuardrailEventHooks.pre_call]) + case None: + return None + + async def apply_realtime_guardrails( + self, + texts: tuple[str, ...], + role: str, + item_id: str | None = None, + pre_block_backend_message: str | None = None, + event_hooks: Sequence["GuardrailEventHooks"] | None = None, + ) -> RealtimeGuardrailVerdict: + """ + Run registered guardrails on realtime text (transcript, user message, tool output). + + A block has already been reported to the client and the backend when this returns. + A rewrite carries the masked texts the caller forwards in place of the originals. ``pre_block_backend_message`` (if provided) is sent to the backend BEFORE any of the guardrail's own backend messages when a block is @@ -1687,102 +1730,33 @@ class RealTimeStreaming: """ from litellm.types.guardrails import GuardrailEventHooks - if event_hooks is None: - event_hooks = [GuardrailEventHooks.realtime_input_transcription] - _realtime_event_types: Final = event_hooks - _check_data: Final = {**self.request_data, "transcript": transcript, "stream": True} - _already_run: Final[set] = set() - - runner_guardrails: Final = ( - _realtime_runner_guardrail(callback, _check_data, _realtime_event_types) for callback in litellm.callbacks + block: Final = await realtime_guardrail_verdict( + texts, + role, + event_hooks if event_hooks is not None else (GuardrailEventHooks.realtime_input_transcription,), + self.request_data, + self.user_api_key_dict, + litellm.callbacks, ) - for callback in runner_guardrails: - if callback is None or id(callback) in _already_run: - continue - _already_run.add(id(callback)) - try: - await callback.apply_guardrail( - inputs={"texts": [transcript], "images": []}, - request_data={"user_api_key_dict": self.user_api_key_dict}, - input_type="request", - ) - except Exception as e: - # Re-raise unexpected errors (no status_code/detail = programming bug, not a block). - # HTTPException and guardrail-raised exceptions have a status_code or detail attr. - is_guardrail_block = hasattr(e, "status_code") or isinstance(e, ValueError) - if not is_guardrail_block: - verbose_logger.exception( - "[realtime guardrail] unexpected error in apply_guardrail: %s", - e, - ) - raise - # Extract the human-readable error from the detail dict (HTTPException) - # or fall back to str(e) for plain ValueError. - detail: object | None = getattr(e, "detail", None) - if isinstance(detail, dict): - safe_msg = detail.get("error") or str(e) - elif detail is not None: - safe_msg = str(detail) - else: - safe_msg = str(e) or "I'm sorry, that request was blocked by the content filter." + if not isinstance(block, RealtimeGuardrailBlock): + return block + await self._send_blocked_turn_to_backend(pre_block_backend_message, item_id) + await self.websocket.send_text(guardrail_violation_error_event(block.message)) + if not self._is_transcription_session: + for voice_prompt_event in guardrail_voice_prompt_events(block.message): + await self._send_to_backend(voice_prompt_event) - # Use realtime_violation_message if configured; fall back to guardrail error text. - error_msg = getattr(callback, "realtime_violation_message", None) or safe_msg + self._violation_count += 1 + if block.ends_session(self._violation_count): + verbose_logger.warning("[realtime guardrail] ending session after violation %d", self._violation_count) + await self.backend_ws.close() - await self._send_blocked_turn_to_backend(pre_block_backend_message, item_id) - # Send the policy violation hint (shows as small gray status text in UI). - await self.websocket.send_text( - json.dumps( - { - "type": "error", - "error": { - "type": "guardrail_violation", - "message": error_msg, - "code": "content_policy_violation", - }, - } - ) - ) - if not self._is_transcription_session: - # Ask the LLM to voice the exact guardrail message so the - # user hears it as audio in voice sessions (not just text). - guardrail_prompt = ( - f"Say exactly the following message to the user, word for word, " - f"do not add anything else: {error_msg}" - ) - await self._send_to_backend( - json.dumps( - { - "type": "conversation.item.create", - "item": { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": guardrail_prompt}], - }, - } - ) - ) - await self._send_to_backend(json.dumps({"type": "response.create"})) - - self._violation_count += 1 - end_session_after: int | None = getattr(callback, "end_session_after_n_fails", None) - should_end = getattr(callback, "on_violation", None) == "end_session" or ( - end_session_after is not None and self._violation_count >= end_session_after - ) - if should_end: - verbose_logger.warning( - "[realtime guardrail] ending session after violation %d", - self._violation_count, - ) - await self.backend_ws.close() - - verbose_logger.warning( - "[realtime guardrail] BLOCKED transcript (violation %d): %r", - self._violation_count, - transcript[:80], - ) - return True - return False + verbose_logger.warning( + "[realtime guardrail] BLOCKED transcript (violation %d): %r", + self._violation_count, + " ".join(texts)[:80], + ) + return block async def _handle_provider_config_message(self, raw_response: str) -> None: """Process a backend message when a provider_config is set (transformed path).""" @@ -1866,7 +1840,7 @@ class RealTimeStreaming: self.store_message(event_str) self._capture_transcription_usage(event) await self._send_event_to_client(event, event_str) - blocked = await self.run_realtime_guardrails(cast(str, transcript)) + blocked = await self.run_realtime_guardrails(transcript) if not blocked and self._should_disable_vad_auto_response(): await self._send_to_backend(json.dumps({"type": "response.create"})) continue @@ -2267,92 +2241,22 @@ class RealTimeStreaming: guardrail_turn_detection_injected = False msg_type: str | None = None try: - from litellm.types.guardrails import GuardrailEventHooks - msg_obj = json.loads(message) client_event = msg_obj msg_type = client_event.get("type") - if msg_type == "conversation.item.create": - # Check user text messages for prompt injection - item = client_event.get("item", {}) - # Check function_call_output first so a client cannot - # bypass the tool-result guardrail by also setting - # role="user" on a function_call_output item. - if item.get("type") == "function_call_output": - # Tool results are client-controlled and fed to the - # model; check them with the same guardrail used for - # user text so an attacker cannot smuggle blocked - # content into a function_call_output. - output: object = item.get("output", "") - output_text = output if isinstance(output, str) else json.dumps(output) - if output_text: - # Build the sanitized function_call_output up - # front so we can hand it to the guardrail - # runner as the pre-block message. Providers - # that pair every toolCall with a toolResponse - # (e.g. Gemini/Vertex Live) require the - # toolResponse to arrive BEFORE any other - # client message — otherwise the guardrail's - # own clientContent would violate the - # pending-tool-call protocol contract and the - # backend could close the connection before - # the sanitized response ever lands. Dropping - # the blocked item outright would similarly - # leave such providers waiting indefinitely. - # The sanitized payload carries no blocked - # content — only a generic policy marker. - sanitized_msg = json.dumps( - { - **msg_obj, - "item": { - **item, - "output": json.dumps( - { - "error": "Tool output blocked by content policy", - } - ), - }, - } - ) - blocked = await self.run_realtime_guardrails( - output_text, - pre_block_backend_message=sanitized_msg, - event_hooks=[GuardrailEventHooks.pre_call], - ) - if blocked: - # ``_pending_guardrail_message`` is - # intentionally NOT set here. That flag - # exists to swallow the reflexive - # ``response.create`` an OpenAI client - # sends immediately after a user text - # message. In a tool-calling flow the - # client may not send a ``response.create`` - # at all (e.g. Gemini SDKs auto-respond), - # so leaving the flag set would - # incorrectly drop an unrelated - # ``response.create`` from a later - # interaction turn. - continue - elif item.get("role") == "user": - content_list: Sequence[object] = item.get("content", []) - texts = [ - c.get("text", "") - for c in content_list - if isinstance(c, dict) and c.get("type") == "input_text" - ] - combined_text = " ".join(texts) - if combined_text: - blocked = await self.run_realtime_guardrails( - combined_text, - event_hooks=[GuardrailEventHooks.pre_call], - ) - if blocked: - # Store the guardrail reason so the next response.create - # (sent automatically by the client) is rewritten to - # include it as response instructions. - self._pending_guardrail_message = combined_text - continue # don't forward the original blocked message + self._pending_guardrail_message = None + + match await self._screen_client_item(guarded_item := realtime_guarded_item(client_event)): + case RealtimeGuardrailBlock(): + if isinstance(guarded_item, RealtimeMessageText): + self._pending_guardrail_message = " ".join(guarded_item.texts) + continue + case RealtimeGuardrailRewrite(texts=masked_texts): + msg_obj["item"] = rewritten_item(client_event, masked_texts) + message = json.dumps(msg_obj) + case None: + pass if msg_type == "response.create" and self._pending_guardrail_message: # The guardrail already sent the synthetic AI bubble — drop this diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index a5e79d7b2b4..c13c9bc31fb 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -12,7 +12,7 @@ import copy import inspect from collections.abc import Iterable, Mapping, Sequence from dataclasses import replace -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError @@ -21,7 +21,9 @@ from litellm.constants import REDACTED_BY_LITELLM from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.classifier_logging import without_classifier_audit from litellm.litellm_core_utils.core_helpers import ( + attribute_or_none, get_metadata_variable_name_from_kwargs, + has_attributes, ) from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items @@ -42,6 +44,11 @@ if TYPE_CHECKING: else: LiteLLMLoggingObject = Any + +class _HasArguments(Protocol): + arguments: str + + MESSAGE_REDACTION_ENABLE_HEADERS: Final = ("litellm-enable-message-redaction", "x-litellm-enable-message-redaction") @@ -163,7 +170,7 @@ def _redact_bedrock_guardrail_trace(response: object) -> None: } response.update(traces) return - for field, value in ((name, getattr(response, name, None)) for name in BEDROCK_GUARDRAIL_TRACE_FIELDS): + for field, value in ((name, attribute_or_none(response, name)) for name in BEDROCK_GUARDRAIL_TRACE_FIELDS): if value is not None: setattr(response, field, _redacted_bedrock_trace(value)) @@ -191,8 +198,8 @@ def _redact_tool_calls(tool_calls) -> None: if not tool_calls: return for tool_call in tool_calls: - function = getattr(tool_call, "function", None) - if function is not None and hasattr(function, "arguments"): + function = attribute_or_none(tool_call, "function") + if function is not None and has_attributes(function, _HasArguments, "arguments"): function.arguments = REDACTED_BY_LITELLM diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 8c402fe9b44..f6bad7f611b 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -117,7 +117,9 @@ class SensitiveDataMasker: elif key_is_sensitive and isinstance(item, str): masked_items.append(self.mask_value(item)) else: - masked_items.append(item if isinstance(item, (int, float, bool, str, list)) else str(item)) + masked_items.append( + item if item is None or isinstance(item, (int, float, bool, str, list)) else str(item) + ) return masked_items def mask_dict( @@ -144,7 +146,7 @@ class SensitiveDataMasker: str_value = str(v) if v is not None else "" masked_data[k] = self.mask_value(str_value) else: - masked_data[k] = v if isinstance(v, (int, float, bool, str, list)) else str(v) + masked_data[k] = v if v is None or isinstance(v, (int, float, bool, str, list)) else str(v) except Exception: masked_data[k] = "" diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index aa6ce6ea99e..6839e11f8c7 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -7,7 +7,7 @@ import logging import threading import time import traceback -from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence, Sized from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeVar, cast @@ -17,7 +17,7 @@ import httpx from openai.types.chat import ChatCompletionChunk from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice from openai.types.chat.chat_completion_chunk import ChoiceDelta -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import NotRequired, TypedDict import litellm @@ -27,7 +27,9 @@ from litellm.constants import STREAM_CHUNK_CHOICE_FIELDS, STREAM_FINISH_CHOICE_F from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import ( UPSTREAM_USAGE_COST, + attribute_or_none, chunk_usage, + decoded_json_object, provider_reported_usage_cost, stamp_litellm_bill_on_usage, ) @@ -129,6 +131,13 @@ def _apply_provider_fields(response: ModelResponseStream, fields: Mapping[str, o setattr(response, key, value) +def _entry_len(entries: Mapping[str, object], key: str) -> int: + value: Final = entries[key] + if isinstance(value, Sized): + return len(value) + raise TypeError(f"object of type '{type(value).__name__}' has no len()") + + def _prepend_upstream_choice_extras( choice: StreamingChoices, upstream_choice: ChunkChoice, extra_keys: frozenset[str] ) -> None: @@ -346,6 +355,9 @@ class _ParsedChunkHiddenParams(LiteLLMBaseModel): upstream_provider: str | None = None +_DECODED_JSON: Final = TypeAdapter(object) + + def _provider_response_model(chunk: object) -> str | None: model: Final[object] = chunk.get("model") if isinstance(chunk, Mapping) else getattr(chunk, "model", None) return model if isinstance(model, str) and model else None @@ -432,9 +444,8 @@ class CustomStreamWrapper: self._closed = False self._stream_created_time: float = time.time() - litellm_params: Final[GenericLiteLLMParams] = GenericLiteLLMParams.model_validate( - dict(**self.logging_obj.model_call_details.get("litellm_params", {})) - ) + raw_litellm_params: Final = dict[str, object](**self.logging_obj.model_call_details.get("litellm_params", {})) + litellm_params: Final[GenericLiteLLMParams] = GenericLiteLLMParams.model_validate(raw_litellm_params) self.merge_reasoning_content_in_choices: bool = litellm_params.merge_reasoning_content_in_choices or False self.sent_first_thinking_block = False self.sent_last_thinking_block = False @@ -465,7 +476,7 @@ class CustomStreamWrapper: _api_base: Final = get_api_base( model=model or "", - optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), + optional_params=raw_litellm_params, ) self._hidden_params: dict[str, object] = { @@ -953,7 +964,7 @@ class CustomStreamWrapper: if "text_output" in chunk: response = CustomStreamWrapper.strip_sse_data_from_chunk(chunk) or "" response = response.strip() - parsed_response = json.loads(response) + parsed_response = decoded_json_object(_DECODED_JSON.validate_python(json.loads(response))) else: return { "text": "", @@ -1087,7 +1098,7 @@ class CustomStreamWrapper: def is_chunk_non_empty( self, - completion_obj: dict[str, Any], + completion_obj: Mapping[str, object], model_response: ModelResponseStream, response_obj: Mapping[str, object], ) -> bool: @@ -1097,7 +1108,7 @@ class CustomStreamWrapper: or ( "tool_calls" in completion_obj and completion_obj["tool_calls"] is not None - and len(completion_obj["tool_calls"]) > 0 + and _entry_len(completion_obj, "tool_calls") > 0 ) or ("function_call" in completion_obj and completion_obj["function_call"] is not None) or ( @@ -1850,7 +1861,9 @@ class CustomStreamWrapper: elif delta.tool_calls is not None: if isinstance(delta.tool_calls, list): for t in delta.tool_calls: - if hasattr(t, "functions") and hasattr(getattr(t, "functions"), "arguments"): + if hasattr(t, "functions") and hasattr( + attribute_or_none(t, "functions"), "arguments" + ): if t.function.arguments is None: t.function.arguments = "" _json_delta: Final[_DeltaDumpDict] = delta.model_dump() @@ -1888,7 +1901,7 @@ class CustomStreamWrapper: return ## CHECK FOR TOOL USE - if "tool_calls" in completion_obj and len(completion_obj["tool_calls"]) > 0: + if "tool_calls" in completion_obj and _entry_len(completion_obj, "tool_calls") > 0: if self.is_function_call is True: # user passed in 'functions' param completion_obj["function_call"] = completion_obj["tool_calls"][0]["function"] completion_obj["tool_calls"] = None @@ -2817,8 +2830,9 @@ class CustomStreamWrapper: def _record_and_strip_usage_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream | None: self.chunks.append(chunk.model_copy()) + dumped_chunk: Final = MappingProxyType[str, object](chunk.model_dump()) usage_free_chunk: Final = self.model_response_creator( - chunk={key: value for key, value in chunk.model_dump().items() if key != "usage"}, + chunk={key: value for key, value in dumped_chunk.items() if key != "usage"}, hidden_params=cast( # cast-ok: preserve mapping operations on dynamic storage Mapping[str, object], getattr(chunk, HIDDEN_PARAMS_ATTR) ), diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index 5dc9679fe78..48a947c31ca 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -11,13 +11,15 @@ A2A Protocol Format: """ import json -from collections.abc import Sequence -from typing import TYPE_CHECKING, Any, Final, Optional +from collections.abc import Mapping, Sequence +from functools import reduce +from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, runtime_checkable from pydantic import ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import has_attributes, require_attributes, subscript from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, @@ -30,6 +32,38 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) +_UNSUPPORTED_RESPONSE: Final = object() + + +class _ModelDumpable(Protocol): + def model_dump(self) -> object: ... + + +class _ResultLookup(Protocol): + def get(self, key: str, default: object, /) -> object: ... + + +@runtime_checkable +class _SupportsSetItem(Protocol): + def __setitem__(self, key: str, value: object, /) -> None: ... + + +def _item_assignment_target(value: object) -> _SupportsSetItem: + if isinstance(value, _SupportsSetItem): + return value + raise TypeError(f"'{type(value).__name__}' object does not support item assignment") + + +def _response_payload(response: object) -> object: + if has_attributes(response, _ModelDumpable, "model_dump"): + return response.model_dump() + if isinstance(response, dict): + return response + return _UNSUPPORTED_RESPONSE + + +def _path_step(current: object, key: str) -> object: + return subscript(current, int(key) if key.isdigit() else key) class _A2ATextPart(TypedDict, total=False): @@ -123,7 +157,7 @@ class A2AGuardrailHandler(BaseTranslation): async def process_output_response( self, - response: Any, + response: object, guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, @@ -147,17 +181,12 @@ class A2AGuardrailHandler(BaseTranslation): Returns: Modified response with guardrails applied to text content """ - # Handle both dict and Pydantic model responses - if hasattr(response, "model_dump"): - response_dict = response.model_dump() - is_pydantic = True - elif isinstance(response, dict): - response_dict = response - is_pydantic = False - else: + payload: Final = _response_payload(response) + if payload is _UNSUPPORTED_RESPONSE: verbose_proxy_logger.warning("A2A: Unknown response type %s, skipping guardrail", type(response)) return response + response_dict: Final = require_attributes(payload, _ResultLookup, "get") result: Final = response_dict.get("result", {}) if not result or not isinstance(result, dict): verbose_proxy_logger.debug("A2A: No result in response, skipping guardrail") @@ -218,15 +247,8 @@ class A2AGuardrailHandler(BaseTranslation): verbose_proxy_logger.debug("A2A: Processed output response") - # Update the original response - if is_pydantic: - # For Pydantic models, we need to update the underlying dict - # and the model will reflect the changes - response_dict["result"] = result - return response - else: - response["result"] = result - return response + _item_assignment_target(response_dict)["result"] = result + return response async def process_output_streaming_response( self, @@ -448,20 +470,11 @@ class A2AGuardrailHandler(BaseTranslation): def _apply_text_to_path( self, - result: dict[str | int, Any], + result: Mapping[str, object], path: tuple[str, ...], part_idx: int, text: str, ) -> None: """Apply guardrailed text back to the specified path in the result.""" - # Navigate to the parts list - current = result - for key in path: - if key.isdigit(): - # Array index - current = current[int(key)] - else: - current = current[key] - - # Update the text in the part - current[part_idx]["text"] = text + parts: Final = reduce(_path_step, path, result) + _item_assignment_target(subscript(parts, part_idx))["text"] = text diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 8f1a08d2964..ca9529ba246 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -795,18 +795,21 @@ class AnthropicMessagesHandler(BaseTranslation): data: dict, # mutable-ok: the request body the adapter reads *, keep_cache_control: bool = False, + keep_tool_errors: bool = False, ) -> ChatCompletionRequest: """Translate Anthropic request to OpenAI chat completion format.""" ( chat_completion_compatible_request, _tool_name_mapping, - ) = LiteLLMAnthropicMessagesAdapter(keep_cache_control=keep_cache_control).translate_anthropic_to_openai( + ) = LiteLLMAnthropicMessagesAdapter( + keep_cache_control=keep_cache_control, keep_tool_errors=keep_tool_errors + ).translate_anthropic_to_openai( anthropic_message_request=cast(AnthropicMessagesRequest, scannable_tool_result_request(data)), preserve_midturn_system=True, ) return chat_completion_compatible_request - def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: + def get_structured_messages(self, data: dict, *, keep_tool_errors: bool = False) -> list[AllMessageValues] | None: """ Convert Anthropic messages request data to OpenAI-spec structured messages. @@ -815,7 +818,7 @@ class AnthropicMessagesHandler(BaseTranslation): messages: Final = data.get("messages") if messages is None: return None - chat_completion_compatible_request: Final = self._translate_to_openai(data) + chat_completion_compatible_request: Final = self._translate_to_openai(data, keep_tool_errors=keep_tool_errors) result: Final = cast( list[AllMessageValues], chat_completion_compatible_request.get("messages", []), diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 4062a957e00..ed4ed295dd3 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -1081,7 +1081,7 @@ class ModelResponseIterator: if "caller" in content_block_start["content_block"]: caller_data: Final = content_block_start["content_block"]["caller"] if caller_data and tool_use is not None: - tool_use["caller"] = cast(dict[str, Any], caller_data) + tool_use["caller"] = caller_data elif content_block_start["content_block"]["type"] == "redacted_thinking": ( thinking_blocks, diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 5d4957fa166..1c42dd15853 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -2107,7 +2107,7 @@ def _strip_empty_content_blocks_from_message(message: object) -> object: def strip_empty_content_blocks_from_anthropic_messages( messages: Sequence[object], -) -> list[Any]: +) -> Sequence[object]: """ Return a new message list with empty or whitespace-only ``{"type": "text"}`` content blocks and empty *unsigned* ``{"type": "thinking"}`` blocks removed. @@ -2220,7 +2220,7 @@ def _sanitize_tool_use_id_content_block(block: object) -> object: return block -def sanitize_tool_use_ids_in_anthropic_messages(messages: Sequence[object]) -> list[Any]: +def sanitize_tool_use_ids_in_anthropic_messages(messages: Sequence[object]) -> Sequence[object]: """ Return a new message list with ``tool_use`` / ``server_tool_use`` ``id`` and ``tool_result`` ``tool_use_id`` values rewritten to satisfy Anthropic's diff --git a/litellm/llms/anthropic/messages/adapters/transformation.py b/litellm/llms/anthropic/messages/adapters/transformation.py index f09bee7cc4e..323743a8bb4 100644 --- a/litellm/llms/anthropic/messages/adapters/transformation.py +++ b/litellm/llms/anthropic/messages/adapters/transformation.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast from pydantic import JsonValue, TypeAdapter import litellm -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, normalize_drop_params from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.llms.anthropic.messages.convertible_blocks import tool_result_text_part from litellm.llms.anthropic.messages.utils import ( @@ -434,8 +434,9 @@ _CLAUDE_OVER_REASONING_EFFORT_PROVIDERS: Final = frozenset({LlmProviders.KYMA.va class LiteLLMAnthropicMessagesAdapter: - def __init__(self, *, keep_cache_control: bool = False) -> None: + def __init__(self, *, keep_cache_control: bool = False, keep_tool_errors: bool = False) -> None: self.keep_cache_control: Final = keep_cache_control + self.keep_tool_errors: Final = keep_tool_errors ### FOR [BETA] `/v1/messages` endpoint support @@ -615,11 +616,7 @@ class LiteLLMAnthropicMessagesAdapter: self._add_cache_control_if_applicable(content, doc_obj, model) new_user_content_list.append(doc_obj) elif content.get("type") == "tool_result": - tool_result = ChatCompletionToolMessage( - role="tool", - tool_call_id=content.get("tool_use_id", ""), - content=self._tool_result_content(content.get("content")), - ) + tool_result = self._tool_message(content) self._add_cache_control_if_applicable(content, tool_result, model) tool_message_list.append(tool_result) @@ -1432,6 +1429,15 @@ class LiteLLMAnthropicMessagesAdapter: return None return anthropic_image_source_to_openai_url(image_source) + def _tool_message(self, block: object) -> ChatCompletionToolMessage: + fields: Final = as_str_mapping(block) or {} + is_error: Final = fields.get("is_error") + tool_call_id: Final = tool_use_id if isinstance(tool_use_id := fields.get("tool_use_id"), str) else "" + content: Final = self._tool_result_content(fields.get("content")) + if self.keep_tool_errors and isinstance(is_error, bool): + return ChatCompletionToolMessage(role="tool", tool_call_id=tool_call_id, content=content, is_error=is_error) + return ChatCompletionToolMessage(role="tool", tool_call_id=tool_call_id, content=content) + def _tool_result_content(self, raw_content: object) -> ToolResultContent: if isinstance(raw_content, str): return raw_content diff --git a/litellm/llms/anthropic/messages/native/agentic_streaming_iterator.py b/litellm/llms/anthropic/messages/native/agentic_streaming_iterator.py index 0aa30998557..dfda86339e7 100644 --- a/litellm/llms/anthropic/messages/native/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/messages/native/agentic_streaming_iterator.py @@ -17,11 +17,14 @@ execute it. import asyncio import contextlib import json -from collections.abc import AsyncIterator +from collections.abc import AsyncIterable, AsyncIterator from typing import TYPE_CHECKING, Any, Final, cast +from pydantic import TypeAdapter + from litellm._logging import verbose_logger from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES +from litellm.litellm_core_utils.core_helpers import decoded_json_object, has_attributes from litellm.types.integrations.custom_logger import AgenticLoopLimitError, AgenticLoopSafetyError if TYPE_CHECKING: @@ -42,6 +45,7 @@ AGENTIC_LOOP_LIMIT_ERROR_SSE_BYTES: Final = ( b'"The agentic loop reached its limit before this turn finished. ' b'Raise max_agentic_loops to allow more loops."}}\n\n' ) +_DECODED_JSON: Final = TypeAdapter(object) def is_server_fulfilled_tool_leak_error(chunk: object) -> bool: @@ -76,10 +80,10 @@ def _parse_sse_events(raw: bytes) -> list[tuple]: continue data_str = stripped[len("data:") :].strip() try: - data = json.loads(data_str) + data = _DECODED_JSON.validate_python(json.loads(data_str)) except (json.JSONDecodeError, ValueError): continue - event_type = current_event_type or data.get("type", "") + event_type = current_event_type or decoded_json_object(data).get("type", "") current_event_type = None events.append((event_type, data)) return events @@ -371,7 +375,7 @@ class AgenticAnthropicStreamingIterator: verbose_logger.debug("AgenticStreamingIterator: Could not rebuild response from SSE bytes") return - result: Final = await self._http_handler._call_agentic_completion_hooks( + result: Final[object] = await self._http_handler._call_agentic_completion_hooks( response=rebuilt, model=self._model, messages=self._messages, @@ -386,7 +390,7 @@ class AgenticAnthropicStreamingIterator: if result is None: return - if hasattr(result, "__aiter__"): + if has_attributes(result, AsyncIterable[object], "__aiter__"): self._follow_up_iterator = result.__aiter__() elif isinstance(result, dict): from litellm.llms.anthropic.messages.native.fake_stream_iterator import ( diff --git a/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py index 5c899e8e8bf..20dd13a329e 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py @@ -2,8 +2,8 @@ ## Translates OpenAI call to Anthropic `/v1/messages` format import asyncio from collections import deque -from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Iterable, Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Final import anyio from pydantic import ConfigDict, field_validator @@ -12,6 +12,7 @@ from litellm import verbose_logger from litellm._logging import redact_internal_details_from_client_message from litellm._uuid import uuid from litellm.exceptions import APIError, MidStreamFallbackError +from litellm.litellm_core_utils.core_helpers import has_attributes, iter_items from litellm.litellm_core_utils.prompt_templates.reasoning_utils import ( encrypted_reasoning_signature, ) @@ -126,7 +127,7 @@ class AnthropicResponsesStreamWrapper: def __init__( self, - responses_stream: Any, + responses_stream: object, model: str, litellm_logging_obj: "LiteLLMLoggingObject | None" = None, safeguards_evaluator: "SafeguardsEvaluator | None" = None, @@ -414,7 +415,9 @@ class AnthropicResponsesStreamWrapper: response_obj: Final = getattr(event, "response", None) or ( event.get("response") if isinstance(event, dict) else None ) - output: Final = (getattr(response_obj, "output", None) or ()) if response_obj is not None else () + output: Final[Iterable[object]] = ( + (getattr(response_obj, "output", None) or ()) if response_obj is not None else () + ) refusal_text: Final = responses_output_refusal_text(output) or (self._refusal_text or None) status: Final = getattr(response_obj, "status", None) if response_obj is not None else None has_tool_call: Final = any( @@ -508,14 +511,14 @@ class AnthropicResponsesStreamWrapper: return self._chunk_queue.popleft() try: - if hasattr(self.responses_stream, "__aiter__"): + if has_attributes(self.responses_stream, AsyncIterable[object], "__aiter__"): async for event in self.responses_stream: self._process_event(event) if self._chunk_queue: return self._chunk_queue.popleft() else: if self._sync_responses_iterator is None: - self._sync_responses_iterator = iter(self.responses_stream) + self._sync_responses_iterator = iter(iter_items(self.responses_stream)) sync_iterator: Final = self._sync_responses_iterator missing: Final = object() while (event := await asyncio.to_thread(next, sync_iterator, missing)) is not missing: diff --git a/litellm/llms/anthropic/messages/responses_adapters/transformation.py b/litellm/llms/anthropic/messages/responses_adapters/transformation.py index b4455f1a4e2..d7b8592f45c 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/messages/responses_adapters/transformation.py @@ -10,6 +10,8 @@ from collections.abc import Iterable, Mapping, Sequence from itertools import groupby from typing import Any, Final, cast +from pydantic import TypeAdapter + from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, @@ -58,6 +60,7 @@ from litellm.types.llms.openai import ( REASONING_SUMMARY_PART_SEPARATOR: Final = "\n\n" RESPONSES_INCLUDE_ENCRYPTED_REASONING: Final = "reasoning.encrypted_content" +_DECODED_JSON: Final = TypeAdapter(object) class LiteLLMAnthropicToResponsesAPIAdapter: @@ -805,15 +808,17 @@ class LiteLLMAnthropicToResponsesAPIAdapter: content.append(reasoning_block) elif item_type == "function_call": try: - input_data = json.loads(item.get("arguments", "{}")) + input_data = _DECODED_JSON.validate_python(json.loads(item.get("arguments", "{}"))) except (json.JSONDecodeError, TypeError): input_data = {} content.append( - AnthropicResponseContentBlockToolUse( - type="tool_use", - id=item.get("call_id") or item.get("id", ""), - name=item.get("name", ""), - input=input_data, + AnthropicResponseContentBlockToolUse.model_validate( + { + "type": "tool_use", + "id": item.get("call_id") or item.get("id", ""), + "name": item.get("name", ""), + "input": input_data, + } ).model_dump(exclude_none=True) ) stop_reason = "tool_use" diff --git a/litellm/llms/anthropic/messages/stream_assembly.py b/litellm/llms/anthropic/messages/stream_assembly.py index 843e8a6e495..9a31d171bc8 100644 --- a/litellm/llms/anthropic/messages/stream_assembly.py +++ b/litellm/llms/anthropic/messages/stream_assembly.py @@ -4,10 +4,13 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import Final +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.core_helpers import as_str_mapping, map_finish_reason +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object, map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import is_str_keyed_dict from litellm.llms.anthropic.chat.handler import ( ModelResponseIterator as AnthropicModelResponseIterator, ) @@ -21,6 +24,8 @@ from litellm.types.utils import ( TextCompletionResponse, ) +_DECODED_JSON: Final = TypeAdapter(object) + def split_sse_chunk_into_events(chunk: str | bytes) -> list[str]: """ @@ -98,6 +103,14 @@ _FAST_PATH_DISALLOWED_DELTA_TYPES: Final = frozenset( ) +def _joined_text(parts: Sequence[object]) -> str: + strings: Final = [part for part in parts if isinstance(part, str)] + if len(strings) == len(parts): + return "".join(strings) + first_bad: Final = next(position for position, part in enumerate(parts) if not isinstance(part, str)) + raise TypeError(f"sequence item {first_bad}: expected str instance, {type(parts[first_bad]).__name__} found") + + def collapse_pure_text_chunks( all_chunks: Sequence[str | bytes], ) -> list[str] | None: @@ -120,10 +133,10 @@ def collapse_pure_text_chunks( if ev: normalized.append(ev) - text_block_indexes: Final[set] = set() + text_block_indexes: Final[set[object]] = set() out: Final[list[str]] = [] - pending_text: list[str] = [] - pending_index: int | None = None + pending_text: list[object] = [] + pending_index: object = None saw_any_text_delta = False def flush() -> None: @@ -132,7 +145,7 @@ def collapse_pure_text_chunks( merged: Final = { "type": "content_block_delta", "index": pending_index if pending_index is not None else 0, - "delta": {"type": "text_delta", "text": "".join(pending_text)}, + "delta": {"type": "text_delta", "text": _joined_text(pending_text)}, } out.append("data: " + json.dumps(merged)) pending_text = [] @@ -158,20 +171,20 @@ def collapse_pure_text_chunks( if not json_str.startswith("{"): continue try: - data = json.loads(json_str) + data = decoded_json_object(_DECODED_JSON.validate_python(json.loads(json_str))) except (json.JSONDecodeError, ValueError): return None etype = data.get("type") if etype == "content_block_start": - block = data.get("content_block") or {} + block = decoded_json_object(data.get("content_block") or {}) if block.get("type") != "text": return None text_block_indexes.add(data.get("index")) flush() out.append(ev) elif etype == "content_block_delta": - delta = data.get("delta") or {} + delta = decoded_json_object(data.get("delta") or {}) dtype = delta.get("type") if dtype in _FAST_PATH_DISALLOWED_DELTA_TYPES: return None @@ -291,9 +304,10 @@ def extract_sse_data(event_str: str) -> dict | None: if payload is None or not payload.startswith("{"): return None try: - return json.loads(payload) + decoded: Final = _DECODED_JSON.validate_python(json.loads(payload)) except (ValueError, TypeError): return None + return decoded if is_str_keyed_dict(decoded) else None def build_usage_only_response_from_chunks( @@ -492,7 +506,7 @@ def _carries_no_data(fields: Sequence[tuple[str, str]]) -> bool: def _json_mapping(payload: str) -> Mapping[str, object] | None: try: - value: Final[object] = json.loads(payload) + value: Final = _DECODED_JSON.validate_python(json.loads(payload)) except ValueError: return None return as_str_mapping(value) @@ -555,7 +569,7 @@ def is_anthropic_sse_stream(all_chunks: Sequence[object]) -> bool: def _is_json_no_reader_can_take(payload: str) -> bool: try: - value: Final[object] = json.loads(payload) + value: Final = _DECODED_JSON.validate_python(json.loads(payload)) except ValueError: return False if value is None: diff --git a/litellm/llms/apiserpent/search/transformation.py b/litellm/llms/apiserpent/search/transformation.py index faf1903ffe8..2e5fd2032c1 100644 --- a/litellm/llms/apiserpent/search/transformation.py +++ b/litellm/llms/apiserpent/search/transformation.py @@ -12,7 +12,9 @@ from typing import Final, Literal, cast from urllib.parse import urlencode import httpx +from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object, iter_items from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.apiserpent.search.defaults import ( @@ -33,6 +35,7 @@ from litellm.secret_managers.main import get_secret_str DEEP_SEARCH_PARAM: Final = "deep" APISERPENT_BASE: Final = "https://apiserpent.com" APISERPENT_PARAMS_KEY: Final = "_apiserpent_params" +_DECODED_JSON: Final = TypeAdapter(object) class APISerpentSearchConfig(BaseSearchConfig): @@ -150,22 +153,24 @@ class APISerpentSearchConfig(BaseSearchConfig): Full format nests results under ``results.organic[]``; simple format returns a flat ``results[]`` array. Both expose title/url/snippet. """ - response_json: Final = raw_response.json() + response_json: Final = decoded_json_object(_DECODED_JSON.validate_python(raw_response.json())) raw_results: Final = response_json.get("results") or {} - organic: Final = raw_results.get("organic", []) if isinstance(raw_results, dict) else raw_results + nested_results: Final = as_str_mapping(raw_results) + organic: Final = nested_results.get("organic", []) if nested_results is not None else raw_results - results: Final[list[SearchResult]] = [] - for result in organic: - results.append( - SearchResult( - title=result.get("title", ""), - url=result.get("url", ""), - snippet=result.get("snippet", ""), - date=result.get("date"), - last_updated=None, - ) + results: Final = [ + SearchResult.model_validate( + { + "title": result.get("title", ""), + "url": result.get("url", ""), + "snippet": result.get("snippet", ""), + "date": result.get("date"), + "last_updated": None, + } ) + for result in map(decoded_json_object, iter_items(organic)) + ] return SearchResponse( results=results, diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 69ffac5d6f9..ef74042156e 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -2,7 +2,7 @@ import asyncio import json import time from collections.abc import Callable, Coroutine -from typing import Final +from typing import Final, Protocol import httpx from openai import ( @@ -17,6 +17,7 @@ from pydantic import BaseModel import litellm from litellm.constants import DEFAULT_MAX_RETRIES +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing from litellm.llms.custom_httpx.http_handler import ( @@ -60,6 +61,11 @@ from .image_generation import ( from .image_generation.http_utils import azure_deployment_image_generation_json_body +class _HasResponseText(Protocol): + @property + def text(self) -> str: ... + + class AzureOpenAIAssistantsAPIConfig: """ Reference: https://learn.microsoft.com/en-us/azure/ai-services/openai/assistants-reference-messages?tabs=python#create-message @@ -853,11 +859,11 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): except Exception as e: status_code: Final = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) error_text = str(e) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - error_text = error_response.text + error_text = require_attributes(error_response, _HasResponseText, "text").text raise AzureOpenAIError( status_code=status_code, message=error_text, diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index a814a9fbac0..6a0740b2f15 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -251,9 +251,8 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): def compact_url(self, api_base: str) -> str: params: Final = httpx.URL(api_base).params - v1_params: Final = ( - params if BaseAzureLLM.is_azure_v1_api_version(params.get("api-version")) else params.remove("api-version") - ) + api_version: Final = next(iter(params.get_list("api-version")), None) + v1_params: Final = params if BaseAzureLLM.is_azure_v1_api_version(api_version) else params.remove("api-version") return super().compact_url(str(_v1_responses_url(api_base).copy_with(params=v1_params))) ######################################################### diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index 423a09c5c92..5dae6cc9fc3 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -308,7 +308,7 @@ class AzureAIAgentsHandler: logging_obj: LiteLLMLoggingObj, optional_params: dict, litellm_params: dict, - timeout: float, + timeout: float | str | httpx.Timeout | None, client: HTTPHandler | None = None, headers: dict | None = None, ) -> ModelResponse: @@ -432,7 +432,7 @@ class AzureAIAgentsHandler: logging_obj: LiteLLMLoggingObj, optional_params: dict, litellm_params: dict, - timeout: float, + timeout: float | str | httpx.Timeout | None, client: AsyncHTTPHandler | None = None, headers: dict | None = None, ) -> ModelResponse: @@ -559,7 +559,7 @@ class AzureAIAgentsHandler: logging_obj: LiteLLMLoggingObj, optional_params: dict, litellm_params: dict, - timeout: float, + timeout: float | str | httpx.Timeout | None, headers: dict | None = None, ) -> AsyncIterator[ModelResponseStream]: """Execute async streaming completion using Azure Agent Service with native SSE.""" diff --git a/litellm/llms/azure_ai/agents/transformation.py b/litellm/llms/azure_ai/agents/transformation.py index 18f8f7916b7..f6e70496247 100644 --- a/litellm/llms/azure_ai/agents/transformation.py +++ b/litellm/llms/azure_ai/agents/transformation.py @@ -323,7 +323,7 @@ class AzureAIAgentsConfig(BaseConfig): logging_obj: LiteLLMLoggingObj, optional_params: dict, litellm_params: dict, - timeout: float | Any, + timeout: float | str | httpx.Timeout | None, acompletion: bool, stream: bool | None = False, headers: dict | None = None, diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 3609871fdf8..5fcb84593d2 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -8,7 +8,7 @@ from collections.abc import Mapping from typing import Final import httpx -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter import litellm from litellm._logging import verbose_logger @@ -22,6 +22,8 @@ from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +_DECODED_JSON: Final = TypeAdapter(object) + class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): """ @@ -42,7 +44,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): system: object = None, optional_params: Mapping[str, JsonValue] | None = None, context_management: JsonValue = None, - ) -> dict[str, object]: + ) -> object: """ Handle a CountTokens request using httpx with Azure authentication. @@ -126,7 +128,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): message=error_text, ) - azure_response: Final = response.json() + azure_response: Final = _DECODED_JSON.validate_python(response.json()) verbose_logger.debug("Azure AI Anthropic response: %s", azure_response) diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index 14f7d4bfd2b..8657fcddab4 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -11,6 +11,7 @@ from pydantic import JsonValue import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountFields, ChatCountTranslationError, @@ -139,13 +140,14 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): ) if result is not None: + response_fields: Final = decoded_json_object(result) return TokenCountResponse.model_validate( { - "total_tokens": result.get("input_tokens", 0), + "total_tokens": response_fields.get("input_tokens", 0), "request_model": request_model, "model_used": model_to_use, "tokenizer_type": "azure_ai_anthropic_api", - "original_response": result, + "original_response": response_fields, } ) except AnthropicError as e: diff --git a/litellm/llms/azure_ai/anthropic/messages_transformation.py b/litellm/llms/azure_ai/anthropic/messages_transformation.py index 1f9ded827b6..63f9f726c71 100644 --- a/litellm/llms/azure_ai/anthropic/messages_transformation.py +++ b/litellm/llms/azure_ai/anthropic/messages_transformation.py @@ -3,7 +3,7 @@ Azure Anthropic messages transformation config - extends AnthropicMessagesConfig """ from collections.abc import Sequence -from typing import Any, Final +from typing import Final from litellm.llms.anthropic.messages.native.transformation import ( AnthropicMessagesConfig, @@ -109,14 +109,10 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): Processes both `system` and `messages` content blocks. """ - def _sanitize(cache_control: Any) -> None: - if isinstance(cache_control, dict): - cache_control.pop("scope", None) - def _process_content_list(content: list) -> None: for item in content: - if isinstance(item, dict) and "cache_control" in item: - _sanitize(item["cache_control"]) + if isinstance(item, dict) and "cache_control" in item and isinstance(item["cache_control"], dict): + item["cache_control"].pop("scope", None) if "system" in anthropic_messages_request: system: Final = anthropic_messages_request["system"] diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index 9bd476efdff..00d3ad63fda 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -2,11 +2,12 @@ import base64 from collections.abc import Mapping, Sequence from io import BufferedReader from types import MappingProxyType -from typing import Any, Final +from typing import Final, Protocol from httpx._types import RequestFiles import litellm +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.llms.azure_ai.common_utils import ( AzureFoundryModelInfo, get_azure_ai_auth_headers, @@ -21,6 +22,10 @@ from litellm.types.llms.openai import FileTypes from litellm.types.router import GenericLiteLLMParams +class _BinaryReadable(Protocol): + def read(self) -> bytes: ... + + class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): """ Azure AI Foundry FLUX 2 image edit config @@ -100,7 +105,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): if image is None: raise ValueError("FLUX 2 image edit requires an image.") - images: Final = tuple(image) if isinstance(image, list) else (image,) + images: Final[Sequence[object]] = [*image] if isinstance(image, list) else [image] if not images: raise ValueError("FLUX 2 image edit requires at least one image.") max_reference_images: Final = 10 if "flex" in model.lower() else 8 @@ -121,14 +126,14 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): } return request_body, [] - def _convert_image_to_base64(self, image: Any) -> str: + def _convert_image_to_base64(self, image: object) -> str: """Convert image file to base64 string""" if isinstance(image, BufferedReader): image_bytes = image.read() image.seek(0) # Reset file pointer for potential reuse elif isinstance(image, bytes): image_bytes = image - elif hasattr(image, "read"): + elif has_attributes(image, _BinaryReadable, "read"): image_bytes = image.read() else: raise ValueError(f"Unsupported image type: {type(image)}") diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 6399ded9555..6794739be24 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,7 +1,9 @@ from abc import ABC, abstractmethod from collections.abc import Sequence from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional +from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional, Protocol + +from litellm.litellm_core_utils.core_helpers import has_attributes if TYPE_CHECKING: from fastapi import HTTPException @@ -15,6 +17,10 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +class _ModelDumpable(Protocol): + def model_dump(self) -> object: ... + + @dataclass(slots=True) class StreamTransformSink: """Out-parameter used by ``process_output_streaming_response`` to hand the @@ -81,7 +87,7 @@ class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( - user_api_key_dict: Any | None, + user_api_key_dict: object, ) -> dict[str, object]: """ Transform user_api_key_dict to a metadata dict with prefixed keys. @@ -99,7 +105,11 @@ class BaseTranslation(ABC): return {} # Convert to dict if it's a Pydantic object - user_dict = user_api_key_dict.model_dump() if hasattr(user_api_key_dict, "model_dump") else user_api_key_dict + user_dict: Final = ( + user_api_key_dict.model_dump() + if has_attributes(user_api_key_dict, _ModelDumpable, "model_dump") + else user_api_key_dict + ) if not isinstance(user_dict, dict): return {} @@ -257,11 +267,12 @@ class BaseTranslation(ABC): """ return None - def get_structured_messages(self, data: dict) -> list["AllMessageValues"] | None: + def get_structured_messages(self, data: dict, *, keep_tool_errors: bool = False) -> list["AllMessageValues"] | None: """ Convert request data to OpenAI-spec structured messages. - Override in subclasses for format-specific conversion. + Override in subclasses for format-specific conversion. ``keep_tool_errors`` copies a provider's + tool-failure marker onto tool messages as ``is_error``; only for copies that are never sent to a model. Returns None if no convertible content is found. """ diff --git a/litellm/llms/base_llm/managed_resources/isolation.py b/litellm/llms/base_llm/managed_resources/isolation.py index 09ca23a8ded..135986a6563 100644 --- a/litellm/llms/base_llm/managed_resources/isolation.py +++ b/litellm/llms/base_llm/managed_resources/isolation.py @@ -11,8 +11,9 @@ are denied so an empty user_id can never select an unscoped query. """ from collections.abc import Mapping, Sequence -from typing import Any, Final +from typing import Final, Protocol +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.types.proxy.auth.user_api_key_auth import ( UserAPIKeyAuth, ) @@ -47,15 +48,20 @@ def resolve_resource_owner_id( return None -def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, object]: +class _Identified(Protocol): + @property + def id(self) -> object: ... + + +def build_list_page(items: Sequence[object], has_more: bool = False) -> dict[str, object]: """Build the OpenAI-style paginated list response shape used by managed file/batch/vector-store listings. ``first_id`` and ``last_id`` are sourced from each item's ``.id`` attribute.""" return { "object": "list", "data": items, - "first_id": items[0].id if items else None, - "last_id": items[-1].id if items else None, + "first_id": require_attributes(items[0], _Identified, "id").id if items else None, + "last_id": require_attributes(items[-1], _Identified, "id").id if items else None, "has_more": has_more, } diff --git a/litellm/llms/base_llm/responses/codex_compat.py b/litellm/llms/base_llm/responses/codex_compat.py index fd769832217..62036e0a830 100644 --- a/litellm/llms/base_llm/responses/codex_compat.py +++ b/litellm/llms/base_llm/responses/codex_compat.py @@ -10,6 +10,11 @@ variant``. Both Amazon Bedrock endpoints reject them: They are *history* items, so they only appear from the second turn of a session onward -- a first-turn request succeeds and hides the problem entirely. +An ``agent_message`` content block's ``encrypted_content`` is ciphertext that only the +backend that issued it can decrypt for the recipient model. A backend that declares +``supports_encrypted_agent_messages`` gets those items intact, since turning them into +text would hand the model ciphertext it believes it wrote. + Codex also sends a ``web_search`` tool on every turn. api.openai.com runs that tool itself; a backend with no server-side tools rejects the whole request over it, so the same providers drop the tool types their backend does not accept. @@ -24,6 +29,7 @@ from typing import Final from typing_extensions import ReadOnly, TypedDict +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping from litellm.types.llms.openai import ResponseInputParam AGENT_MESSAGE_INPUT_ITEM_TYPE: Final = "agent_message" @@ -63,6 +69,15 @@ def _agent_message_text(item: "Mapping[str, object]") -> str: ) +def _is_encrypted_block(block: object) -> bool: + mapping: Final = as_str_mapping(block) + return mapping is not None and bool(mapping.get("encrypted_content")) + + +def _carries_encrypted_content(item: "Mapping[str, object]") -> bool: + return any(_is_encrypted_block(block) for block in as_object_list(item.get("content")) or ()) + + def _normalize_agent_message_item(item: "Mapping[str, object]") -> "_RewrittenAssistantMessageItem | None": text: Final = _agent_message_text(item) if not text: @@ -97,28 +112,35 @@ def _normalize_local_shell_call_item(item: "Mapping[str, object]") -> "_Rewritte return rewritten -def _normalize_input_item(item: object) -> "tuple[object, str | None]": +def _normalize_input_item(item: object, keep_encrypted_agent_messages: bool) -> "tuple[object, str | None]": """Returns (normalized item, or None to drop it; original type when rewritten).""" - if not isinstance(item, dict): + mapping: Final = as_str_mapping(item) + if mapping is None: return item, None - item_type: Final = item.get("type") + item_type: Final = mapping.get("type") if item_type == AGENT_MESSAGE_INPUT_ITEM_TYPE: - return _normalize_agent_message_item(item), item_type + if keep_encrypted_agent_messages and _carries_encrypted_content(mapping): + return item, None + return _normalize_agent_message_item(mapping), AGENT_MESSAGE_INPUT_ITEM_TYPE if item_type == CONTEXT_COMPACTION_INPUT_ITEM_TYPE: - return _normalize_context_compaction_item(item), item_type + return _normalize_context_compaction_item(mapping), CONTEXT_COMPACTION_INPUT_ITEM_TYPE if item_type == LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: - return _normalize_local_shell_call_item(item), item_type + return _normalize_local_shell_call_item(mapping), LOCAL_SHELL_CALL_INPUT_ITEM_TYPE return item, None def normalize_codex_input_items( input: "str | ResponseInputParam", + *, + keep_encrypted_agent_messages: bool, ) -> "tuple[str | ResponseInputParam, tuple[str, ...]]": """Rewrite the Codex history item types a Responses backend rejects. - ``agent_message`` (Codex multi-agent traffic; its ``encrypted_content`` slot - carries the plaintext payload when the model never issued encrypted args) - becomes an assistant message, ``context_compaction`` becomes the ``compaction`` + ``agent_message`` (Codex multi-agent traffic) carrying an ``encrypted_content`` + block is forwarded intact when ``keep_encrypted_agent_messages`` is set, because + only the backend can read it. Otherwise it becomes an assistant message: a backend + that never encrypts gets the plaintext payload Codex puts in that slot when the + model issued no encrypted args. ``context_compaction`` becomes the ``compaction`` spelling these backends accept, and ``local_shell_call`` becomes the ``function_call`` its recorded ``function_call_output`` already pairs with. @@ -127,7 +149,7 @@ def normalize_codex_input_items( """ if not isinstance(input, list): return input, () - normalized: Final = tuple(_normalize_input_item(item) for item in input) + normalized: Final = tuple(_normalize_input_item(item, keep_encrypted_agent_messages) for item in input) rewritten_types: Final = tuple(sorted(frozenset(item_type for _, item_type in normalized if item_type is not None))) kept: Final = [i for i, _ in normalized if i is not None] # Codex passthrough items sit outside the OpenAI input union. diff --git a/litellm/llms/base_llm/videos/async_task_utils.py b/litellm/llms/base_llm/videos/async_task_utils.py index 507aefea6c6..ffc847127ea 100644 --- a/litellm/llms/base_llm/videos/async_task_utils.py +++ b/litellm/llms/base_llm/videos/async_task_utils.py @@ -7,11 +7,104 @@ from typing import Final from pydantic import TypeAdapter, ValidationError +from litellm.exceptions import UnsupportedParamsError from litellm.images.utils import ImageEditRequestUtils +from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoFrameImage, VideoInputReference from litellm.types.videos.utils import decode_video_id_with_provider, id_safe_name _FIELDS_ADAPTER: Final = TypeAdapter(dict[str, object]) _FILE_TUPLE_ADAPTER: Final = TypeAdapter(tuple[object, ...]) +_FRAME_IMAGES_ADAPTER: Final = TypeAdapter(tuple[VideoFrameImage, ...]) +_INPUT_REFERENCES_ADAPTER: Final = TypeAdapter(tuple[VideoInputReference, ...]) +OPENROUTER_VIDEO_FIELDS: Final = frozenset( + {"duration", "aspect_ratio", "frame_images", "input_references", "generate_audio", "seed"} +) + + +def unsupported_video_fields( + params: VideoCreateOptionalRequestParams, + model: str, + provider: str, + drop_params: bool, + supported_native: frozenset[str] = frozenset(), +) -> frozenset[str]: + unsupported: Final = tuple( + key for key, value in params.items() if key in OPENROUTER_VIDEO_FIELDS - supported_native and value is not None + ) + if unsupported and not drop_params: + raise UnsupportedParamsError( + message=f"{provider} video generation does not support {', '.join(unsupported)}, set drop_params=True to ignore them", + model=model, + llm_provider=provider, + ) + return frozenset(unsupported) + + +def video_image_params( + params: Mapping[str, object], model: str, provider: str, drop_params: bool +) -> Mapping[str, object]: + frames: Final = params.get("frame_images") + references: Final = params.get("input_references") + conflict: Final = ( + "content cannot be combined with frame_images or input_references" + if params.get("content") is not None and (frames is not None or references is not None) + else "frame_images cannot be combined with input_references or input_reference" + if frames is not None and (references is not None or params.get("input_reference") is not None) + else "input_references cannot be combined with input_reference" + if references is not None and params.get("input_reference") is not None + else None + ) + if conflict is not None and not drop_params: + raise UnsupportedParamsError( + message=f"{conflict}, set drop_params=True to ignore it", model=model, llm_provider=provider + ) + try: + parsed_frames: Final = _FRAME_IMAGES_ADAPTER.validate_python(() if frames is None else frames) + parsed_references: Final = _INPUT_REFERENCES_ADAPTER.validate_python( + () if references is None or frames is not None else references + ) + except ValidationError as e: + if not drop_params: + raise UnsupportedParamsError( + message=f"Invalid video image items, set drop_params=True to ignore them. {e}", + model=model, + llm_provider=provider, + ) from e + return _mapped_video_image_params( + {k: v for k, v in params.items() if k not in {"frame_images", "input_references"}}, + (), + ) + roles: Final = tuple(item["frame_type"] for item in parsed_frames) + first_indices: Final = {role: roles.index(role) for role in frozenset(roles)} + if len(roles) != len(frozenset(roles)) and not drop_params: + raise UnsupportedParamsError( + message="Duplicate frame_type, set drop_params=True to ignore it", model=model, llm_provider=provider + ) + frame_items: Final = tuple( + {"type": item["type"], "image_url": item["image_url"], "role": item["frame_type"]} + for index, item in enumerate(parsed_frames) + if index == first_indices[item["frame_type"]] + ) + reference_items: Final = tuple( + {"type": item["type"], "image_url": item["image_url"], "role": "reference_image"} for item in parsed_references + ) + media: Final = frame_items if frames is not None else reference_items + return _mapped_video_image_params(params, media) + + +def _mapped_video_image_params( + params: Mapping[str, object], media: tuple[Mapping[str, object], ...] +) -> Mapping[str, object]: + frames: Final = params.get("frame_images") + references: Final = params.get("input_references") + excluded: Final[frozenset[str]] = frozenset({"frame_images", "input_references", "aspect_ratio"}) | ( + frozenset({"content", "input_reference"}) if frames is not None or references is not None else frozenset[str]() + ) + return { + **{key: value for key, value in params.items() if key not in excluded}, + **({"ratio": params["aspect_ratio"]} if params.get("aspect_ratio") is not None else {}), + **({"content": list(media)} if frames is not None or references is not None else {}), + } def _fields(value: object) -> Mapping[str, object] | None: diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index b93d40738ea..12c5031628b 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -960,7 +960,8 @@ class AmazonConverseConfig(BaseConfig): ) -> bool: """Check if the Bedrock model takes native structured outputs (outputConfig.textFormat) for this request. - ``supports_native_structured_output`` in ``litellm.model_cost`` turns it on for every request, and + ``supports_native_structured_output`` or ``supports_bedrock_native_structured_output_by_default`` in + ``litellm.model_cost`` turns it on for every request, and ``supports_bedrock_native_structured_output`` (the deployment's ``model_info`` first) also when Claude will think, which the json tool path would drop. When the request does not ask for thinking, a model that cannot force the json tool gets it unless that key is false. diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index bc2ef2441d2..57cc39a03f4 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -11,7 +11,7 @@ from litellm import verbose_logger from litellm._uuid import uuid from litellm.caching.caching import InMemoryCache from litellm.constants import RESPONSE_FORMAT_TOOL_NAME -from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.litellm_core_utils.core_helpers import decoded_json_object, map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers @@ -50,6 +50,7 @@ from litellm.types.utils import GenericStreamingChunk as GChunk from ..common_utils import ( BedrockError, BedrockEventStreamResponseDict, + bedrock_event_chunk_text, bedrock_event_stream_header, bedrock_event_stream_response, bedrock_stream_event_error_status, @@ -59,6 +60,7 @@ from ..common_utils import ( get_bedrock_response_stream_shape, get_bedrock_stream_event_statuses, get_bedrock_tool_name, + parse_bedrock_event_stream, ) bedrock_tool_name_mappings: Final[InMemoryCache] = InMemoryCache(max_size_in_memory=50, default_ttl=600) @@ -922,7 +924,7 @@ class AWSEventStreamDecoder: response_stream_shape: "Shape", request_id_headers: httpx.Headers | None = None, ) -> str | None: - parsed_response: Final = self.parser.parse(response_dict, response_stream_shape) + parsed_response: Final = parse_bedrock_event_stream(self.parser, response_dict, response_stream_shape) if response_dict["status_code"] != 200: raise build_bedrock_stream_error(response_dict, response_stream_shape, headers=request_id_headers) @@ -932,11 +934,12 @@ class AWSEventStreamDecoder: raise build_bedrock_stream_event_error( event_type, event_error_status, response_dict["body"], headers=request_id_headers ) - if "chunk" in parsed_response: - chunk = parsed_response.get("chunk") - if not chunk: + parsed_fields: Final = decoded_json_object(parsed_response) + if "chunk" in parsed_fields: + event_chunk: Final = parsed_fields.get("chunk") + if not event_chunk: return None - return chunk.get("bytes").decode() + return bedrock_event_chunk_text(event_chunk) else: chunk = response_dict.get("body") if not chunk: diff --git a/litellm/llms/bedrock/chat/native_output_fallback.py b/litellm/llms/bedrock/chat/native_output_fallback.py index 79742d0810c..7d5662e2024 100644 --- a/litellm/llms/bedrock/chat/native_output_fallback.py +++ b/litellm/llms/bedrock/chat/native_output_fallback.py @@ -9,7 +9,9 @@ from litellm.llms.bedrock.common_utils import ( NativeOutputFallback, bedrock_strict_tools_fallback, is_bedrock_native_output_rejection, + is_bedrock_stream_permission_denial, warn_bedrock_native_output_rejected, + warn_bedrock_native_output_stream_denied, warn_bedrock_strict_tools_dropped, ) from litellm.types.utils import ModelResponse @@ -25,22 +27,55 @@ NativeOutputFallbackFactory: TypeAlias = Callable[[], NativeOutputFallback | Non class _SelectedFallback: fallback: NativeOutputFallback native_output_retry: bool + buffered_native_output: bool = False _RetryFor: TypeAlias = Callable[[BaseLLMException], _SelectedFallback | None] +_NextRetryFor: TypeAlias = Callable[[BaseLLMException, _SelectedFallback], _SelectedFallback | None] -def _native_output_retry( - error: BaseLLMException, fallback: NativeOutputFallbackFactory | None, model: str +def _native_output_rejection_retry( + error: BaseLLMException, fallback: NativeOutputFallbackFactory, model: str ) -> NativeOutputFallback | None: - 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: warn_bedrock_native_output_rejected(model, error.message) return retry +def _buffered(retry: NativeOutputFallback) -> NativeOutputFallback: + return NativeOutputFallback( + optional_params={**retry.optional_params, "fake_stream": True}, litellm_params=retry.litellm_params + ) + + +def _buffered_stream_retry(fallback: NativeOutputFallbackFactory, model: str) -> NativeOutputFallback | None: + retry: Final = fallback() + if retry is None or retry.optional_params.get("fake_stream") is not True: + return None + warn_bedrock_native_output_stream_denied(model) + return retry + + +def _native_output_retry( + error: BaseLLMException, + fallback: NativeOutputFallbackFactory | None, + buffered_stream_fallback: NativeOutputFallbackFactory | None, + model: str, +) -> _SelectedFallback | None: + if fallback is not None and is_bedrock_native_output_rejection(error.status_code, error.message): + rejection_retry: Final = _native_output_rejection_retry(error, fallback, model) + return None if rejection_retry is None else _SelectedFallback(rejection_retry, native_output_retry=True) + if buffered_stream_fallback is not None and is_bedrock_stream_permission_denial(error.status_code, error.message): + stream_retry: Final = _buffered_stream_retry(buffered_stream_fallback, model) + return ( + None + if stream_retry is None + else _SelectedFallback(stream_retry, native_output_retry=True, buffered_native_output=True) + ) + return None + + def _strict_tools_retry( error: BaseLLMException, fallback: NativeOutputFallbackFactory | None, model: str ) -> NativeOutputFallback | None: @@ -55,12 +90,13 @@ def _strict_tools_retry( def _fallback_for( error: BaseLLMException, fallback: NativeOutputFallbackFactory | None, + buffered_stream_fallback: NativeOutputFallbackFactory | None, strict_tools_fallback: NativeOutputFallbackFactory | None, model: str, ) -> _SelectedFallback | None: - native_output_retry: Final = _native_output_retry(error, fallback, model) + native_output_retry: Final = _native_output_retry(error, fallback, buffered_stream_fallback, model) if native_output_retry is not None: - return _SelectedFallback(native_output_retry, native_output_retry=True) + return native_output_retry strict_tools_retry: Final = _strict_tools_retry(error, strict_tools_fallback, model) if strict_tools_retry is None: return None @@ -87,12 +123,81 @@ def _strict_tools_retry_after_native_output_retry( return composed_retry +def _next_retry( + error: BaseLLMException, + retry: _SelectedFallback, + fallback: NativeOutputFallbackFactory | None, + strict_tools_fallback: NativeOutputFallbackFactory | None, + model: str, +) -> _SelectedFallback | None: + if ( + retry.buffered_native_output + and fallback is not None + and is_bedrock_native_output_rejection(error.status_code, error.message) + ): + rejection_retry: Final = _native_output_rejection_retry(error, fallback, model) + return ( + None if rejection_retry is None else _SelectedFallback(_buffered(rejection_retry), native_output_retry=True) + ) + composed_retry: Final = _strict_tools_retry_after_native_output_retry(error, retry, strict_tools_fallback, model) + return None if composed_retry is None else _SelectedFallback(composed_retry, native_output_retry=False) + + +def _attempt_or_next_retry( + attempt: NativeOutputAttempt, retry: _SelectedFallback, next_retry_for: _NextRetryFor +) -> NativeOutputAttemptResult | _SelectedFallback: + try: + return attempt(retry.fallback) + except BaseLLMException as error: + next_retry: Final = next_retry_for(error, retry) + if next_retry is None: + raise + return next_retry + + +def _attempt_retries( + attempt: NativeOutputAttempt, retry: _SelectedFallback, next_retry_for: _NextRetryFor +) -> NativeOutputAttemptResult: + after_first_retry: Final = _attempt_or_next_retry(attempt, retry, next_retry_for) + if not isinstance(after_first_retry, _SelectedFallback): + return after_first_retry + after_second_retry: Final = _attempt_or_next_retry(attempt, after_first_retry, next_retry_for) + if not isinstance(after_second_retry, _SelectedFallback): + return after_second_retry + return attempt(after_second_retry.fallback) + + +async def _await_or_next_retry( + attempt: NativeOutputAttempt, retry: _SelectedFallback, next_retry_for: _NextRetryFor +) -> NativeOutputResponse | _SelectedFallback: + try: + result: Final = attempt(retry.fallback) + return await result if isinstance(result, Coroutine) else result + except BaseLLMException as error: + next_retry: Final = next_retry_for(error, retry) + if next_retry is None: + raise + return next_retry + + +async def _await_retries( + attempt: NativeOutputAttempt, retry: _SelectedFallback, next_retry_for: _NextRetryFor +) -> NativeOutputResponse: + after_first_retry: Final = await _await_or_next_retry(attempt, retry, next_retry_for) + if not isinstance(after_first_retry, _SelectedFallback): + return after_first_retry + after_second_retry: Final = await _await_or_next_retry(attempt, after_first_retry, next_retry_for) + if not isinstance(after_second_retry, _SelectedFallback): + return after_second_retry + last_attempt: Final = attempt(after_second_retry.fallback) + return await last_attempt if isinstance(last_attempt, Coroutine) else last_attempt + + async def _await_with_native_output_fallback( first_attempt: Coroutine[object, object, NativeOutputResponse], attempt: NativeOutputAttempt, retry_for: _RetryFor, - strict_tools_fallback: NativeOutputFallbackFactory | None, - model: str, + next_retry_for: _NextRetryFor, ) -> NativeOutputResponse: try: return await first_attempt @@ -100,17 +205,7 @@ async def _await_with_native_output_fallback( retry: Final = retry_for(error) if retry is None: raise - try: - second_attempt: Final = attempt(retry.fallback) - return await second_attempt if isinstance(second_attempt, Coroutine) else second_attempt - except BaseLLMException as error: - composed_retry: Final = _strict_tools_retry_after_native_output_retry( - error, retry, strict_tools_fallback, model - ) - if composed_retry is None: - raise - third_attempt: Final = attempt(composed_retry) - return await third_attempt if isinstance(third_attempt, Coroutine) else third_attempt + return await _await_retries(attempt, retry, next_retry_for) def call_with_native_output_fallback( @@ -118,11 +213,19 @@ def call_with_native_output_fallback( fallback: NativeOutputFallbackFactory | None, model: str, strict_tools_fallback: NativeOutputFallbackFactory | None = None, + buffered_stream_fallback: NativeOutputFallbackFactory | None = None, ) -> NativeOutputAttemptResult: - if fallback is None and strict_tools_fallback is None: + if fallback is None and strict_tools_fallback is None and buffered_stream_fallback is None: return attempt(None) retry_for: Final = partial( - _fallback_for, fallback=fallback, strict_tools_fallback=strict_tools_fallback, model=model + _fallback_for, + fallback=fallback, + buffered_stream_fallback=buffered_stream_fallback, + strict_tools_fallback=strict_tools_fallback, + model=model, + ) + next_retry_for: Final = partial( + _next_retry, fallback=fallback, strict_tools_fallback=strict_tools_fallback, model=model ) try: first_attempt: Final = attempt(None) @@ -130,15 +233,7 @@ def call_with_native_output_fallback( retry: Final = retry_for(error) if retry is None: raise - try: - return attempt(retry.fallback) - except BaseLLMException as retry_error: - composed_retry: Final = _strict_tools_retry_after_native_output_retry( - retry_error, retry, strict_tools_fallback, model - ) - if composed_retry is None: - raise - return attempt(composed_retry) + return _attempt_retries(attempt, retry, next_retry_for) if isinstance(first_attempt, Coroutine): - return _await_with_native_output_fallback(first_attempt, attempt, retry_for, strict_tools_fallback, model) + return _await_with_native_output_fallback(first_attempt, attempt, retry_for, next_retry_for) return first_attempt diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index f0627c7362c..85e2c5fe1a1 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -13,7 +13,7 @@ import warnings from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias from typing_extensions import ( ReadOnly, @@ -24,6 +24,7 @@ from typing_extensions import ( if TYPE_CHECKING: from botocore.eventstream import EventStreamMessage from botocore.model import ServiceModel, Shape + from botocore.parsers import EventStreamJSONParser from litellm.types.llms.bedrock import BedrockCreateBatchRequest @@ -33,7 +34,7 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError import litellm from litellm import verbose_logger from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix -from litellm.litellm_core_utils.core_helpers import lax_bool +from litellm.litellm_core_utils.core_helpers import decoded_json_object, lax_bool, require_attributes from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.litellm_core_utils.prompt_templates.json_schema_utils import ( schema_property_count, @@ -135,7 +136,7 @@ _BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = (*AWS_AUTH_PARAM_KEYS def merge_bedrock_aws_request_params( litellm_params: Mapping[str, object], optional_params: Mapping[str, object], -) -> dict[str, Any]: +) -> dict[str, object]: """Merge deployment and request parameters without allowing auth escalation. Deployment configuration is authoritative for AWS authentication. When a @@ -297,9 +298,16 @@ def _bedrock_model_supports(model: str, key: str, custom_llm_provider: str = "be def _bedrock_native_output_enabled( model: str, custom_llm_provider: str, model_info: Mapping[str, object] | None ) -> bool: - deployment_flag: Final = (model_info or {}).get("supports_native_structured_output") + deployment_info: Final = model_info or {} + deployment_flag: Final = deployment_info.get("supports_native_structured_output") if isinstance(deployment_flag, bool): return deployment_flag + return deployment_info.get("supports_bedrock_native_structured_output") is not False and ( + bedrock_model_always_takes_native_output(model, custom_llm_provider) + ) + + +def bedrock_model_always_takes_native_output(model: str, custom_llm_provider: str = "bedrock") -> bool: return _bedrock_model_supports(model, "supports_native_structured_output", custom_llm_provider) @@ -443,6 +451,18 @@ def warn_bedrock_native_output_rejected(model: str, error_text: str) -> None: ) +def is_bedrock_stream_permission_denial(status_code: int, error_text: str) -> bool: + return status_code == 403 and "bedrock:InvokeModelWithResponseStream" in error_text + + +def warn_bedrock_native_output_stream_denied(model: str) -> None: + verbose_logger.warning( + "Bedrock denied bedrock:InvokeModelWithResponseStream for model=%s, so litellm sent the structured output " + "request again as a buffered stream. Allow that action to stream it.", + model, + ) + + @dataclass(frozen=True, slots=True) class NativeOutputFallback: optional_params: Mapping[str, object] @@ -543,6 +563,22 @@ def json_object_or_none(value: object) -> Mapping[str, object] | None: BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES: Final = 64 +def _response_format_json_schema(response_format: object) -> Mapping[str, object] | None: + json_schema: Final = json_object_or_none((json_object_or_none(response_format) or {}).get("json_schema")) + return json_object_or_none((json_schema or {}).get("schema")) + + +def bedrock_batch_record_takes_native_output( + model: str, response_format: object, model_info: Mapping[str, object] | None = None +) -> bool: + schema: Final = _response_format_json_schema(response_format) + return _bedrock_native_output_enabled(model, "bedrock", model_info) and ( + schema is None + or schema_property_count(schema) <= BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES + or not _bedrock_model_supports(model, "supports_bedrock_native_structured_output_by_default") + ) + + def bedrock_cleaned_output_schema( model: str, schema: Mapping[str, object], custom_llm_provider: str = "bedrock" ) -> Mapping[str, object] | None: @@ -2481,15 +2517,16 @@ class BedrockEventStreamDecoderBase: ), ) response_dict: Final = event.to_response_dict() - parsed_response: Final = self.parser.parse(response_dict, response_stream_shape) + parsed_response: Final = parse_bedrock_event_stream(self.parser, response_dict, response_stream_shape) if response_dict["status_code"] != 200: raise build_bedrock_stream_error(response_dict, response_stream_shape) - if "chunk" in parsed_response: - chunk = parsed_response.get("chunk") - if not chunk: + parsed_fields: Final = decoded_json_object(parsed_response) + if "chunk" in parsed_fields: + event_chunk: Final = parsed_fields.get("chunk") + if not event_chunk: return None - return chunk.get("bytes").decode() + return bedrock_event_chunk_text(event_chunk) else: chunk = response_dict.get("body") if not chunk: @@ -2506,6 +2543,20 @@ def _decoded_json_value(raw: str) -> object: return _JSON_VALUE.validate_python(json.loads(raw)) +class _DecodablePayload(Protocol): + def decode(self) -> str: ... + + +def parse_bedrock_event_stream( + parser: EventStreamJSONParser, response_dict: Mapping[str, object], response_stream_shape: Shape +) -> object: + return _JSON_VALUE.validate_python(parser.parse(response_dict, response_stream_shape)) + + +def bedrock_event_chunk_text(chunk: object) -> str: + return require_attributes(decoded_json_object(chunk).get("bytes"), _DecodablePayload, "decode").decode() + + def get_anthropic_beta_from_headers(headers: dict) -> list[str]: """ Extract anthropic-beta header values and convert them to a list. @@ -2785,3 +2836,13 @@ class CommonBatchFilesUtils: Get Bedrock-specific error class. """ return BedrockError(status_code=status_code, message=error_message, headers=headers) + + +def bedrock_model_encrypts_agent_messages(model: str) -> bool: + """Only OpenAI's GPT models on Bedrock lock Codex worker messages. + + Open-weight models (gpt-oss, Gemma, Grok) never lock them, so Codex puts plain text in the locked + slot, which Bedrock rejects with ``invalid_encrypted_content`` unless it is rewritten to text first. + """ + name: Final = f".{model.rsplit('/', 1)[-1]}" + return ".openai.gpt-" in name and ".openai.gpt-oss" not in name diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index dec099f37c0..6f0a18b0088 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -45,7 +45,7 @@ class _CountTokensReply(LiteLLMPydanticObjectBase): @dataclass(frozen=True, slots=True) class CountedTokens: input_tokens: int - original_response: Mapping[str, JsonValue] + original_response: Mapping[str, object] tokenizer_type: str diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 5f8f478e4cd..84067923788 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -7,14 +7,18 @@ Simplified handler leveraging existing LiteLLM Bedrock infrastructure. from typing import Any, Final import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.llms.bedrock.base_aws_llm import run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client +_DECODED_JSON: Final = TypeAdapter(object) + class BedrockCountTokensHandler(BedrockCountTokensConfig): """ @@ -29,7 +33,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): litellm_params: dict[str, Any], resolved_model: str, client: AsyncHTTPHandler | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Handle a CountTokens request using existing LiteLLM patterns. @@ -111,12 +115,12 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): response=response, ) - bedrock_response: Final = response.json() + bedrock_response: Final = _DECODED_JSON.validate_python(response.json()) verbose_logger.debug("Bedrock response: %s", bedrock_response) # Transform response back to expected format - final_response: Final = self.transform_bedrock_response_to_anthropic(bedrock_response) + final_response: Final = self.transform_bedrock_response_to_anthropic(decoded_json_object(bedrock_response)) verbose_logger.debug("Final response: %s", final_response) diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index df5a3827c1f..cd59e714ba0 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -247,9 +247,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): ) return join_api_base(base_url, f"model/{encoded_model_id}/count-tokens") - def transform_bedrock_response_to_anthropic( - self, bedrock_response: Mapping[str, JsonValue] - ) -> dict[str, JsonValue]: + def transform_bedrock_response_to_anthropic(self, bedrock_response: Mapping[str, object]) -> dict[str, object]: """ Transform Bedrock CountTokens response to Anthropic format. diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 4d137876890..e5f470fb87b 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -11,7 +11,7 @@ from datetime import datetime from email.utils import parsedate_to_datetime from itertools import chain from types import MappingProxyType -from typing import Any, Final, Literal, TypeAlias, TypedDict +from typing import Any, Final, Literal, Protocol, TypeAlias, TypedDict from urllib.parse import quote, unquote, urlencode import httpx @@ -19,6 +19,7 @@ from httpx import Headers, Response from openai.types.file_deleted import FileDeleted from pydantic import ConfigDict, Field, TypeAdapter from typing_extensions import ReadOnly +from typing_extensions import TypedDict as ExtraItemsTypedDict from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -38,6 +39,7 @@ from litellm.litellm_core_utils.cloud_storage_security import ( split_configured_cloud_bucket_name, validate_managed_cloud_file_id, ) +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.prompt_templates.common_utils import ( text_completion_prompt_to_messages, ) @@ -71,6 +73,7 @@ from litellm.utils import get_llm_provider, get_optional_params from ..base_aws_llm import BaseAWSLLM from ..common_utils import ( BedrockError, + bedrock_batch_record_takes_native_output, deployment_model_info, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id, @@ -150,6 +153,10 @@ def _retrieved_s3_file_size(raw_response: Response) -> int: ) +class _SupportsRstrip(Protocol): + def rstrip(self, chars: str, /) -> str: ... + + class _S3DeleteContext(LiteLLMBaseModel): file_id: str = Field(min_length=1) @@ -167,6 +174,24 @@ _LITELLM_PARAMS_THE_MAPPER_TAKES: Final = frozenset({"allowed_openai_params"}) _MAPPED_PARAMS_THE_REQUEST_HANDLER_STRIPS: Final = frozenset({"json_mode"}) +class _ParamsTheMapperOwns(ExtraItemsTypedDict, total=False, extra_items=object): + pass + + +_UPLOAD_PARAMS: Final = TypeAdapter(dict[str, object]) + + +def _upload_params(params: object) -> Mapping[str, object]: + return _UPLOAD_PARAMS.validate_python(params) if isinstance(params, Mapping) else {} + + +def _upload_deployment_model_info(litellm_params: object) -> Mapping[str, object]: + upload_params: Final = _upload_params(litellm_params) + return deployment_model_info(upload_params) or deployment_model_info( + _upload_params(upload_params.get("litellm_metadata")) + ) + + def _invoke_route_model(model: str) -> str: return f"invoke/{_strip_llm_routing_prefix(model).removeprefix('invoke/')}" @@ -695,9 +720,11 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): encoded_object_name: Final = encode_s3_object_key_for_url(object_name) # S3 endpoint URL format - s3_endpoint_url: Final = ( + s3_endpoint_url: Final = require_attributes( request_params.get("s3_endpoint_url") - or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" + or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}", + _SupportsRstrip, + "rstrip", ).rstrip("/") return f"{s3_endpoint_url}/{bucket_name}/{encoded_object_name}" @@ -1034,6 +1061,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): openai_request_body: Mapping[str, Any], model: str, provider: str | None = None, + model_info: Mapping[str, object] | None = None, ) -> dict[str, object]: """ Transform OpenAI request body to Bedrock-compatible modelInput @@ -1046,7 +1074,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): from litellm.types.utils import LlmProviders messages: Final = map_developer_role_to_system_role(openai_request_body.get("messages", [])) - optional_params: Final = {k: v for k, v in openai_request_body.items() if k not in ["model", "messages"]} + request_params: Final[Mapping[str, object]] = openai_request_body + optional_params: Final = {k: v for k, v in request_params.items() if k not in ["model", "messages"]} # --- Anthropic: use existing AmazonAnthropicClaudeConfig --- if provider == LlmProviders.ANTHROPIC: @@ -1055,14 +1084,22 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) config: Final = AmazonAnthropicClaudeConfig() - litellm_params: Final = without_bedrock_native_output_opt_in({}) + litellm_params: Final = without_bedrock_native_output_opt_in( + { + "model_info": { + "supports_native_structured_output": bedrock_batch_record_takes_native_output( + _invoke_route_model(model), optional_params.get("response_format"), model_info + ) + } + } + ) mapped_params = get_optional_params( model=_invoke_route_model(model), custom_llm_provider="bedrock", messages=messages, model_info=deployment_model_info(litellm_params), - **MappingProxyType( - { + **_ParamsTheMapperOwns( + **{ k: v for k, v in optional_params.items() if not is_litellm_owned_kwarg(k) or k in _LITELLM_PARAMS_THE_MAPPER_TAKES @@ -1150,25 +1187,32 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ] def transform_batch_input_lines( - self, batch_input_lines: Sequence[_BatchInputLine], target_model: str = "" + self, + batch_input_lines: Sequence[_BatchInputLine], + target_model: str = "", + model_info: Mapping[str, object] | None = None, ) -> tuple[_BedrockBatchRecord, ...]: return tuple( - self._transform_batch_input_line(idx, batch_input_line, target_model) + self._transform_batch_input_line(idx, batch_input_line, target_model, model_info) for idx, batch_input_line in enumerate(batch_input_lines) ) def _transform_batch_input_line( - self, idx: int, batch_input_line: _BatchInputLine, target_model: str + self, idx: int, batch_input_line: _BatchInputLine, target_model: str, model_info: Mapping[str, object] | None ) -> _BedrockBatchRecord: try: - return self._transform_openai_batch_record(idx, batch_input_line.record, target_model) + return self._transform_openai_batch_record(idx, batch_input_line.record, target_model, model_info) except _BATCH_RECORD_CONVERSION_ERRORS as e: raise _batch_line_error( batch_input_line.line_number, f"could not be converted to a Bedrock batch record: {e}" ) from e def _transform_openai_batch_record( - self, idx: int, openai_record: _OpenAIBatchRecord, target_model: str + self, + idx: int, + openai_record: _OpenAIBatchRecord, + target_model: str, + model_info: Mapping[str, object] | None = None, ) -> _BedrockBatchRecord: import litellm @@ -1192,6 +1236,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): openai_request_body=self._transform_batch_body_to_chat_body(openai_body, record_kind), model=model_for_transform, provider=provider, + model_info=model_info, ) ) return _BedrockBatchRecord( @@ -1211,7 +1256,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): model: str, create_file_data: CreateFileRequest, optional_params: dict, - litellm_params: dict, + litellm_params: dict[str, object], ) -> bytes | str | dict: """ Transform file request and return a pre-signed request for S3. @@ -1237,7 +1282,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): litellm_params_model: Final = litellm_params.get("model") target_model: Final = model or (litellm_params_model if isinstance(litellm_params_model, str) else "") bedrock_jsonl_content: Final = self.transform_batch_input_lines( - batch_input_lines, target_model=target_model + batch_input_lines, target_model=target_model, model_info=_upload_deployment_model_info(litellm_params) ) file_content = "\n".join(json.dumps(item) for item in bedrock_jsonl_content) elif isinstance(extracted_file_data_content, bytes): diff --git a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py index 1cece8743de..def53da4493 100644 --- a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py +++ b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py @@ -18,8 +18,10 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles +from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.images.main import ImageEditOptionalRequestParams @@ -38,6 +40,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + def _nova_canvas_task_body( *, @@ -401,7 +405,7 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): json_mode: bool | None = None, ) -> ImageResponse: try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error parsing Nova Canvas image edit response: {e}", @@ -416,18 +420,19 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): headers=raw_response.headers, ) - images: Final[list[str]] = response_data.get("images") or [] + response_object: Final = decoded_json_object(response_data) + images: Final[object] = response_object.get("images") or [] - if "errors" in response_data and not images: + if "errors" in response_object and not images: raise self.get_error_class( - error_message=f"Nova Canvas image edit error: {response_data['errors']}", + error_message=f"Nova Canvas image edit error: {response_object['errors']}", status_code=raw_response.status_code, headers=raw_response.headers, ) # Nova Canvas InvokeModel success body uses "images" and optional "error" (AWS docs); # it does not use Stability-style "finish_reasons". - error_msg = response_data.get("message") or response_data.get("error") + error_msg = response_object.get("message") or response_object.get("error") if error_msg and not images: if not isinstance(error_msg, str): error_msg = str(error_msg) @@ -439,7 +444,7 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): model_response: Final = ImageResponse() model_response.data = [] - for image_b64 in images: + for image_b64 in iter_items(images): if image_b64: model_response.data.append( ImageObject( diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py index d0ce5895970..8605e7e29e9 100644 --- a/litellm/llms/bedrock/image_edit/stability_transformation.py +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -22,11 +22,13 @@ API Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parame """ import base64 -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, runtime_checkable import httpx from httpx._types import RequestFiles +from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items, json_contains from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.images.main import ImageEditOptionalRequestParams @@ -44,6 +46,19 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + + +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, key: str | int, /) -> object: ... + + +def _subscript(value: object, key: str | int) -> object: + if isinstance(value, _SupportsGetItem): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + class BedrockStabilityImageEditConfig(BaseImageEditConfig): """ @@ -299,7 +314,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): OpenAI expects: {"data": [{"b64_json": "base64..."}], "created": timestamp} """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error parsing Bedrock Stability response: {e}", @@ -308,18 +323,19 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): ) # Check for errors in response - if "errors" in response_data: + if json_contains(response_data, "errors"): raise self.get_error_class( - error_message=f"Bedrock Stability error: {response_data['errors']}", + error_message=f"Bedrock Stability error: {_subscript(response_data, 'errors')}", status_code=raw_response.status_code, headers=raw_response.headers, ) # Check finish_reasons - finish_reasons: Final = response_data.get("finish_reasons", []) - if finish_reasons and finish_reasons[0]: + response_fields: Final = decoded_json_object(response_data) + finish_reasons: Final = response_fields.get("finish_reasons", []) + if finish_reasons and _subscript(finish_reasons, 0): raise self.get_error_class( - error_message=f"Bedrock Stability error: {finish_reasons[0]}", + error_message=f"Bedrock Stability error: {_subscript(finish_reasons, 0)}", status_code=400, headers=raw_response.headers, ) @@ -329,9 +345,9 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): model_response.data = [] # Extract images from response - images: Final = response_data.get("images", []) + images: Final = response_fields.get("images", []) if images: - for image_b64 in images: + for image_b64 in iter_items(images): if image_b64: model_response.data.append( ImageObject( diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index b587528cd2f..85c2384f9a8 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -14,7 +14,7 @@ from pydantic import InstanceOf, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, @@ -849,6 +849,7 @@ class InvokeStreamFrame: _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) _Frame = TypeVar("_Frame", ConverseStreamFrame, InvokeStreamFrame) @@ -1298,7 +1299,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): texts: Sequence[tuple[_StreamGroupKey, str]] = () if event_type == "contentBlockDelta": try: - payload_dict: dict[str, object] = json.loads(payload_bytes) + payload_dict = decoded_json_object(_DECODED_JSON.validate_python(json.loads(payload_bytes))) texts = _delta_frame_texts(payload_dict) except Exception as e: verbose_proxy_logger.debug( diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 49939399ec8..9f4fb8cfecc 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -6,12 +6,14 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic. import asyncio import contextlib +import functools import importlib.metadata import json -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mapping, MutableMapping, Sequence from dataclasses import dataclass, field +from enum import Enum, auto from types import MappingProxyType -from typing import Final, NoReturn, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Final, NoReturn, Protocol, runtime_checkable from pydantic import JsonValue, TypeAdapter @@ -24,11 +26,25 @@ from litellm.constants import ( BEDROCK_REALTIME_SDK_DISTRIBUTION, BEDROCK_REALTIME_SDK_SUPPORTED_RANGE, BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY, + REALTIME_SESSION_ENDED_ABNORMALLY_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, ) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.realtime_errors import realtime_error_event +from litellm.litellm_core_utils.realtime_guardrails import ( + RealtimeGuardedItem, + RealtimeGuardrailBlock, + RealtimeGuardrailRewrite, + RealtimeGuardrailVerdict, + RealtimeToolOutput, + guardrail_violation_error_event, + guardrail_voice_prompt_events, + realtime_guarded_item, + realtime_guardrail_verdict, + rewritten_item, +) from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeEventTypes from litellm.llms.base_llm.realtime.transformation import RealtimeSessionUpdateRejection from litellm.types.llms.bedrock import AwsAuthParams @@ -40,9 +56,13 @@ from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig, requested_modalities +if TYPE_CHECKING: + from litellm.types.guardrails import GuardrailEventHooks + _CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _REQUEST_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) _SPEND_LOG_METADATA_KEYS: Final = ("litellm_metadata", "metadata") +GUARDRAIL_UNAVAILABLE_MESSAGE: Final = "A guardrail could not check this request, so the session was ended" _EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) @@ -259,11 +279,11 @@ def _raise_provider_failure(scope: MutableMapping[str, object], failure: BaseExc raise error from failure -def _parse_client_message(message: str) -> Mapping[str, JsonValue]: +def _parse_client_message(message: str) -> Mapping[str, JsonValue] | None: try: return _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message)) except ValueError: - return _EMPTY_JSON_OBJECT + return None async def _send_opening_message( @@ -282,6 +302,27 @@ async def _send_opening_message( client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt +async def _forward_client_message( + client_ws: RealtimeClientWebSocket, + bedrock_stream: BedrockBidirectionalStream, + transformation_config: BedrockRealtimeConfig, + model: str, + session_state: RealtimeResponseTransformInput, + send_to_bedrock: Callable[[str], Awaitable[None]], + client_message: str, +) -> None: + transformed_messages: Final = transformation_config.transform_realtime_request( + message=client_message, + model=model, + session_configuration_request=session_state.get("session_configuration_request"), + ) + if transformed_messages and not _is_session_committed(client_ws.scope): + await _send_opening_message(client_ws, bedrock_stream, client_message, transformed_messages, send_to_bedrock) + return + for bedrock_message in transformed_messages: + await send_to_bedrock(bedrock_message) + + async def _ack_session_update( client_ws: RealtimeClientWebSocket, transformation_config: BedrockRealtimeConfig, @@ -317,6 +358,161 @@ async def _reject_session_update( return True +async def _send_guardrail_violation(client_ws: RealtimeClientWebSocket, block: RealtimeGuardrailBlock) -> bool: + try: + await client_ws.send_text(guardrail_violation_error_event(block.message)) + except Exception as e: # noqa: BLE001 # any send failure means the client is gone + verbose_proxy_logger.debug("Bedrock Realtime: guardrail error not delivered: %s", e, exc_info=True) + return False + return True + + +async def _end_guarded_session(client_ws: RealtimeClientWebSocket) -> None: + verbose_proxy_logger.warning("Bedrock Realtime: ending session after a guardrail violation") + with contextlib.suppress(Exception): + await client_ws.close() + + +async def _end_unscreened_session(client_ws: RealtimeClientWebSocket) -> None: + with contextlib.suppress(Exception): + await client_ws.send_text(realtime_error_event(GUARDRAIL_UNAVAILABLE_MESSAGE, error_type="guardrail_error")) + with contextlib.suppress(Exception): + await client_ws.close() + + +@dataclass(frozen=True, slots=True) +class _GuardrailUnavailable: + pass + + +def _registered_callbacks() -> Iterable[object]: + return litellm.callbacks + + +class _ClientEventRoute(Enum): + FORWARD = auto() + DROP = auto() + END_SESSION = auto() + + +class BedrockRealtimeGuardrails: + """One Bedrock realtime session's guardrail checks, shared by both bridge directions. + + Nova Sonic has no way to cancel a response or remove a turn, so a blocked user transcript ends + the session before that turn's answer reaches the client. + """ + + def __init__( + self, + request_data: Mapping[str, object] = _EMPTY_JSON_OBJECT, + user_api_key_dict: object | None = None, + callbacks: Callable[[], Iterable[object]] = _registered_callbacks, + ) -> None: + self._request_data: Final = request_data + self._user_api_key_dict: Final = user_api_key_dict + self._callbacks: Final = callbacks + self._violation_count = 0 + self._drop_next_response_create = False + self._failed = False + + @property + def failed(self) -> bool: + return self._failed + + async def _verdict( + self, texts: tuple[str, ...], role: str, hook: "GuardrailEventHooks" + ) -> RealtimeGuardrailVerdict | _GuardrailUnavailable: + try: + return await realtime_guardrail_verdict( + texts, role, (hook,), self._request_data, self._user_api_key_dict, self._callbacks() + ) + except Exception as e: # noqa: BLE001 # a guardrail that cannot decide fails closed, never as a Bedrock failure + verbose_proxy_logger.warning("Bedrock Realtime: a guardrail failed, ending the session: %s", e) + self._failed = True + return _GuardrailUnavailable() + + def _ends_session(self, block: RealtimeGuardrailBlock) -> bool: + self._violation_count += 1 + return block.ends_session(self._violation_count) + + async def screen_client_event( + self, + client_event: Mapping[str, JsonValue], + client_ws: RealtimeClientWebSocket, + forward_to_bedrock: Callable[[str], Awaitable[None]], + ) -> _ClientEventRoute: + from litellm.types.guardrails import GuardrailEventHooks + + event_type: Final = client_event.get("type") + if event_type == "response.create" and self._drop_next_response_create: + self._drop_next_response_create = False + return _ClientEventRoute.DROP + if event_type == "conversation.item.create": + self._drop_next_response_create = False + guarded_item: Final = realtime_guarded_item(client_event) + if guarded_item is None: + return _ClientEventRoute.FORWARD + held: Final = json.dumps(client_event) + client_ws.scope[BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY] = held # rebind-ok: replayed on fallback + role: Final = "tool" if isinstance(guarded_item, RealtimeToolOutput) else "user" + match await self._verdict(guarded_item.texts, role, GuardrailEventHooks.pre_call): + case None: + return _ClientEventRoute.FORWARD + case RealtimeGuardrailRewrite(texts=texts): + await forward_to_bedrock(json.dumps({**client_event, "item": rewritten_item(client_event, texts)})) + return _ClientEventRoute.DROP + case RealtimeGuardrailBlock() as block: + return await self._route_blocked_item(block, guarded_item, client_ws, forward_to_bedrock) + case _GuardrailUnavailable(): + await _end_unscreened_session(client_ws) + return _ClientEventRoute.END_SESSION + + async def _route_blocked_item( + self, + block: RealtimeGuardrailBlock, + guarded_item: RealtimeGuardedItem, + client_ws: RealtimeClientWebSocket, + forward_to_bedrock: Callable[[str], Awaitable[None]], + ) -> _ClientEventRoute: + if isinstance(guarded_item, RealtimeToolOutput): + await forward_to_bedrock(guarded_item.sanitized_event) + if not await _send_guardrail_violation(client_ws, block): + return _ClientEventRoute.END_SESSION + if self._ends_session(block): + await _end_guarded_session(client_ws) + return _ClientEventRoute.END_SESSION + if isinstance(guarded_item, RealtimeToolOutput): + return _ClientEventRoute.DROP + for voice_prompt_event in guardrail_voice_prompt_events(f'"{block.message}"'): + await forward_to_bedrock(voice_prompt_event) + self._drop_next_response_create = True + return _ClientEventRoute.DROP + + async def continues_after( + self, provider_event: Mapping[str, object], client_ws: RealtimeClientWebSocket, client_audio_streamed: bool + ) -> bool: + from litellm.types.guardrails import GuardrailEventHooks + + transcript: Final = provider_event.get("transcript") + if ( + not client_audio_streamed + or provider_event.get("type") != "conversation.item.input_audio_transcription.completed" + or not isinstance(transcript, str) + or not transcript + ): + return True + match await self._verdict((transcript,), "user", GuardrailEventHooks.realtime_input_transcription): + case RealtimeGuardrailBlock() as block: + if await _send_guardrail_violation(client_ws, block): + await _end_guarded_session(client_ws) + return False + case _GuardrailUnavailable(): + await _end_unscreened_session(client_ws) + return False + case None | RealtimeGuardrailRewrite(): + return True + + async def _close_bedrock_input( bedrock_stream: BedrockBidirectionalStream, close_messages: Sequence[str], @@ -381,6 +577,8 @@ class BedrockRealtime(BaseAWSLLM): aws_external_id: str | None = None, aws_session_tags: object = None, drop_params: bool = False, + user_api_key_dict: object | None = None, + litellm_metadata: Mapping[str, object] | None = None, **kwargs: object, ): """ @@ -461,7 +659,15 @@ class BedrockRealtime(BaseAWSLLM): try: await self._run_session( - websocket, open_bidirectional_stream, model, logging_obj, pending_opening_message, drop_params + websocket, + open_bidirectional_stream, + model, + logging_obj, + pending_opening_message, + drop_params, + BedrockRealtimeGuardrails( + request_data={"litellm_metadata": litellm_metadata or {}}, user_api_key_dict=user_api_key_dict + ), ) finally: await _close_bedrock_client(bedrock_client) @@ -474,6 +680,7 @@ class BedrockRealtime(BaseAWSLLM): logging_obj: LiteLLMLogging, pending_opening_message: str | None, drop_params: bool = False, + guardrails: BedrockRealtimeGuardrails | None = None, ) -> None: transformation_config: Final = BedrockRealtimeConfig(drop_params=drop_params) @@ -496,6 +703,7 @@ class BedrockRealtime(BaseAWSLLM): "session_configuration_request": None, } + session_guardrails: Final = guardrails if guardrails is not None else BedrockRealtimeGuardrails() outcome: Final = await self._bridge( websocket, bedrock_stream, @@ -504,7 +712,10 @@ class BedrockRealtime(BaseAWSLLM): session_state, logging_obj, initial_message=pending_opening_message, + guardrails=session_guardrails, ) + if session_guardrails.failed: + logging_obj.model_call_details[REALTIME_SESSION_ENDED_ABNORMALLY_KEY] = True logged_events: Final = ( *outcome.logged_events, @@ -544,10 +755,12 @@ class BedrockRealtime(BaseAWSLLM): session_state: RealtimeResponseTransformInput, logging_obj: LiteLLMLogging, initial_message: str | None, + guardrails: BedrockRealtimeGuardrails | None = None, ) -> _BridgeOutcome: """Run both forwarding directions until the client leaves or either side fails.""" logged: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: events forwarded before a failure are still spend signals: Final = SessionEndSignals() + session_guardrails: Final = guardrails if guardrails is not None else BedrockRealtimeGuardrails() async def collect_logged_events() -> None: try: @@ -559,6 +772,7 @@ class BedrockRealtime(BaseAWSLLM): logging_obj, session_state, signals=signals, + guardrails=session_guardrails, ): logged.append(event) except Exception: @@ -575,6 +789,7 @@ class BedrockRealtime(BaseAWSLLM): logging_obj, initial_message, signals=signals, + guardrails=session_guardrails, ) ) bedrock_task: Final = asyncio.create_task(collect_logged_events()) @@ -610,6 +825,7 @@ class BedrockRealtime(BaseAWSLLM): initial_message: str | None = None, *, signals: SessionEndSignals, + guardrails: BedrockRealtimeGuardrails | None = None, ) -> None: """Forward messages from client WebSocket to Bedrock stream. @@ -631,10 +847,23 @@ class BedrockRealtime(BaseAWSLLM): raise _stream_ended_before_session_closed() from e verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200]) + forward_to_bedrock: Final = functools.partial( + _forward_client_message, + client_ws, + bedrock_stream, + transformation_config, + model, + session_state, + send_to_bedrock, + ) + session_guardrails: Final = guardrails if guardrails is not None else BedrockRealtimeGuardrails() try: async for message in _client_messages(client_ws, initial_message): verbose_proxy_logger.debug("Bedrock Realtime: Received from client: %s", message[:200]) parsed_client_message = _parse_client_message(message) + if parsed_client_message is None: + verbose_proxy_logger.warning("Bedrock Realtime: dropped a client message that is not valid JSON") + continue is_session_update = _json_str(parsed_client_message.get("type")) == "session.update" if is_session_update: rejection = transformation_config.session_update_rejection(parsed_client_message) @@ -643,18 +872,15 @@ class BedrockRealtime(BaseAWSLLM): break continue - transformed_messages = transformation_config.transform_realtime_request( - message=message, - model=model, - session_configuration_request=session_state.get("session_configuration_request"), - ) - if transformed_messages and not _is_session_committed(client_ws.scope): - await _send_opening_message( - client_ws, bedrock_stream, message, transformed_messages, send_to_bedrock - ) - else: - for bedrock_message in transformed_messages: - await send_to_bedrock(bedrock_message) + match await session_guardrails.screen_client_event( + parsed_client_message, client_ws, forward_to_bedrock + ): + case _ClientEventRoute.END_SESSION: + break + case _ClientEventRoute.DROP: + continue + case _ClientEventRoute.FORWARD: + await forward_to_bedrock(message) if is_session_update and not await _ack_session_update( client_ws, transformation_config, model, logging_obj, parsed_client_message @@ -679,6 +905,7 @@ class BedrockRealtime(BaseAWSLLM): session_state: RealtimeResponseTransformInput, *, signals: SessionEndSignals, + guardrails: BedrockRealtimeGuardrails | None = None, ) -> AsyncIterator[OpenAIRealtimeEvents]: """Forward messages from Bedrock to the client, yielding the ones to record for spend logging. @@ -694,6 +921,7 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200]) return True + session_guardrails: Final = guardrails if guardrails is not None else BedrockRealtimeGuardrails() output: Final = await bedrock_stream.await_output() while True: result = await output[1].receive() @@ -750,3 +978,7 @@ class BedrockRealtime(BaseAWSLLM): return if _should_log_event(openai_message): yield openai_message + if not await session_guardrails.continues_after( + openai_message, client_ws, transformation_config.client_audio_streamed + ): + return diff --git a/litellm/llms/bedrock/responses/transformation.py b/litellm/llms/bedrock/responses/transformation.py index ff44d973d97..96854f7e601 100644 --- a/litellm/llms/bedrock/responses/transformation.py +++ b/litellm/llms/bedrock/responses/transformation.py @@ -54,6 +54,7 @@ from litellm.llms.bedrock.common_utils import ( BEDROCK_CHAT_COMPLETIONS_ROUTE_PREFIX, DROP_DISABLED_REASONING_EFFORT_WARNING, BedrockError, + bedrock_model_encrypts_agent_messages, bedrock_reasoning_effort_disabled, bedrock_supports_openai_responses, ) @@ -187,10 +188,12 @@ class BedrockOpenAIResponsesConfig(BaseAWSLLM, OpenAIResponsesAPIConfig): self, fetch_image: "Callable[[str], str]" = convert_url_to_base64, async_fetch_image: "Callable[[str], Awaitable[str]]" = async_convert_url_to_base64, + encrypted_agent_messages: bool = False, ) -> None: super().__init__() self.fetch_image = fetch_image self.async_fetch_image = async_fetch_image + self.encrypted_agent_messages = encrypted_agent_messages @classmethod def for_model(cls, model: str | None) -> "BedrockOpenAIResponsesConfig | None": @@ -207,7 +210,7 @@ class BedrockOpenAIResponsesConfig(BaseAWSLLM, OpenAIResponsesAPIConfig): _without_chat_completions_route(model), litellm.model_cost ): return None - return cls() + return cls(encrypted_agent_messages=bedrock_model_encrypts_agent_messages(model)) @property def custom_llm_provider(self) -> LlmProviders: @@ -245,6 +248,9 @@ class BedrockOpenAIResponsesConfig(BaseAWSLLM, OpenAIResponsesAPIConfig): def supports_native_file_search(self) -> bool: return False + def supports_encrypted_agent_messages(self) -> bool: + return self.encrypted_agent_messages + def validate_environment( self, headers: dict, # mutable-ok: signature fixed by the BaseResponsesAPIConfig override contract @@ -358,7 +364,9 @@ class BedrockOpenAIResponsesConfig(BaseAWSLLM, OpenAIResponsesAPIConfig): litellm_params: GenericLiteLLMParams, headers: dict, # mutable-ok: same ) -> dict: # mutable-ok: same - normalized_input, rewritten_types = normalize_codex_input_items(input) + normalized_input, rewritten_types = normalize_codex_input_items( + input, keep_encrypted_agent_messages=self.supports_encrypted_agent_messages() + ) if rewritten_types: verbose_logger.warning( "Bedrock Runtime Responses API: rewrote Codex input item type(s) %s that the endpoint rejects.", diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 96d1da19bcc..b2b3a9b0faf 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -1,11 +1,12 @@ from collections.abc import Mapping from copy import deepcopy -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, runtime_checkable from urllib.parse import urlparse import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -34,6 +35,32 @@ else: LiteLLMLoggingObj = Any +@runtime_checkable +class _SupportsSetDefault(Protocol): + def setdefault(self, key: str, default: object, /) -> object: ... + + +@runtime_checkable +class _SupportsSetItem(Protocol): + def __setitem__(self, key: str, value: object, /) -> None: ... + + +def _setdefault_target(value: object) -> _SupportsSetDefault: + if isinstance(value, _SupportsSetDefault): + return value + raise AttributeError(f"{type(value).__name__!r} object has no attribute 'setdefault'", name="setdefault", obj=value) + + +def _item_assignment_target(value: object) -> _SupportsSetItem: + if isinstance(value, _SupportsSetItem): + return value + raise TypeError(f"'{type(value).__name__}' object does not support item assignment") + + +def _vector_search_configuration(retrieval_config: object) -> Mapping[str, object]: + return decoded_json_object(decoded_json_object(retrieval_config).get("vectorSearchConfiguration", {})) + + class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): """Vector store configuration for AWS Bedrock Knowledge Bases.""" @@ -206,7 +233,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, - extra_body: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: if isinstance(query, list): query = " ".join(query) @@ -218,30 +245,33 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): "retrievalQuery": BedrockKBRetrievalQuery(text=query), } - retrieval_config: dict[str, Any] = {} - - if isinstance(extra_body, dict): - retrieval_config = deepcopy( - extra_body.get("retrievalConfiguration") or extra_body.get("retrieval_configuration") or {} - ) + retrieval_config: Final[object] = ( + deepcopy(extra_body.get("retrievalConfiguration") or extra_body.get("retrieval_configuration") or {}) + if isinstance(extra_body, dict) + else {} + ) max_results: Final = vector_store_search_optional_params.get("max_num_results") if max_results is not None: - existing_number_of_results = retrieval_config.get("vectorSearchConfiguration", {}).get("numberOfResults") + existing_number_of_results: Final = _vector_search_configuration(retrieval_config).get("numberOfResults") if existing_number_of_results is not None and existing_number_of_results != max_results: verbose_logger.debug( "Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.numberOfResults (%s) with max_num_results=%s", existing_number_of_results, max_results, ) - retrieval_config.setdefault("vectorSearchConfiguration", {})["numberOfResults"] = max_results + _item_assignment_target(_setdefault_target(retrieval_config).setdefault("vectorSearchConfiguration", {}))[ + "numberOfResults" + ] = max_results filters: Final = vector_store_search_optional_params.get("filters") if filters is not None: - existing_filter: Final = retrieval_config.get("vectorSearchConfiguration", {}).get("filter") + existing_filter: Final = _vector_search_configuration(retrieval_config).get("filter") if existing_filter is not None and existing_filter != filters: verbose_logger.debug( "Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.filter with filters from vector_store_search_optional_params" ) - retrieval_config.setdefault("vectorSearchConfiguration", {})["filter"] = filters + _item_assignment_target(_setdefault_target(retrieval_config).setdefault("vectorSearchConfiguration", {}))[ + "filter" + ] = filters if retrieval_config: request_body["retrievalConfiguration"] = cast(BedrockKBRetrievalConfiguration, retrieval_config) user_context: Final = self._user_context(extra_body=extra_body, litellm_params=litellm_params) diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 7190a4f28c6..46419d7a5b1 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -75,10 +75,12 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI self, aws_signer: BaseAWSLLM | None = None, use_openai_path: bool = True, + encrypted_agent_messages: bool = False, ): super().__init__() self._aws_signer = aws_signer or BaseAWSLLM() self.use_openai_path = use_openai_path + self.encrypted_agent_messages = encrypted_agent_messages @property def custom_llm_provider(self) -> LlmProviders: @@ -87,6 +89,9 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI def auto_reasoning_summary(self) -> ReasoningSummaryLevel: return "auto" if self.use_openai_path else super().auto_reasoning_summary() + def supports_encrypted_agent_messages(self) -> bool: + return self.use_openai_path and self.encrypted_agent_messages + def get_error_class( self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers ) -> BaseLLMException: @@ -215,7 +220,9 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI hoisted: Final = hoist_additional_tools( input, params.get("tools"), drop_params=drops_unsupported_params(litellm_params.drop_params) ) - normalized_input, rewritten_types = normalize_codex_input_items(hoisted.input) + normalized_input, rewritten_types = normalize_codex_input_items( + hoisted.input, keep_encrypted_agent_messages=self.supports_encrypted_agent_messages() + ) if rewritten_types: verbose_logger.warning( "Bedrock Mantle Responses API: rewrote Codex input item type(s) %s that Mantle rejects.", diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 38bcbb0d905..8eb02def465 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -10,7 +10,7 @@ API Reference: https://docs.bfl.ai/ import base64 import time from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol import httpx from httpx._types import RequestFiles @@ -19,6 +19,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.litellm_core_utils.core_helpers import as_object_list, has_attributes from litellm.litellm_core_utils.url_utils import async_safe_get, join_api_base, safe_get from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str @@ -52,6 +53,19 @@ class _BflPolledResult(TypedDict): _BFL_POLLED_RESULT_ADAPTER: Final = TypeAdapter(_BflPolledResult) _DECODED_JSON: Final = TypeAdapter(object) +_UNTYPED_VALUE: Final = TypeAdapter(object) + + +class _BinaryReadable(Protocol): + def read(self) -> bytes: ... + + +class _Tellable(Protocol): + def tell(self) -> int: ... + + +class _Seekable(Protocol): + def seek(self, offset: int, /) -> object: ... _BFL_REQUEST_PARAMS: Final = ( @@ -197,7 +211,7 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): def _read_image_bytes( self, - image: Any, + image: object, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, ) -> bytes: @@ -208,9 +222,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): ) if isinstance(image, bytes): return image - elif isinstance(image, list): - # If it's a list, take the first image - return self._read_image_bytes(image[0], depth=depth + 1, max_depth=max_depth) + elif (image_list := as_object_list(image)) is not None: + return self._read_image_bytes(image_list[0], depth=depth + 1, max_depth=max_depth) elif isinstance(image, str): if image.startswith(("http://", "https://")): response: Final = safe_get(litellm.module_level_client, image, timeout=60.0) @@ -221,13 +234,12 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): "Unsupported image input: plain string values that are not URLs are not accepted. " "Provide image bytes or a file-like object." ) - elif hasattr(image, "read"): - # File-like object - pos: Final = getattr(image, "tell", lambda: 0)() - if hasattr(image, "seek"): + elif has_attributes(image, _BinaryReadable, "read"): + pos: Final = image.tell() if has_attributes(image, _Tellable, "tell") else 0 + if has_attributes(image, _Seekable, "seek"): image.seek(0) data: Final = image.read() - if hasattr(image, "seek"): + if has_attributes(image, _Seekable, "seek"): image.seek(pos) return data else: @@ -265,7 +277,7 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): # Handle mask if provided (for inpainting) if "mask" in image_edit_optional_request_params: - mask: Final = image_edit_optional_request_params["mask"] + mask: Final = _UNTYPED_VALUE.validate_python(image_edit_optional_request_params["mask"]) mask_bytes: Final = self._read_image_bytes(mask) request_body["mask"] = base64.b64encode(mask_bytes).decode("utf-8") diff --git a/litellm/llms/bytez/chat/transformation.py b/litellm/llms/bytez/chat/transformation.py index c297eb13db1..465f505c729 100644 --- a/litellm/llms/bytez/chat/transformation.py +++ b/litellm/llms/bytez/chat/transformation.py @@ -6,7 +6,9 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.url_utils import encode_url_path_segments @@ -36,6 +38,8 @@ else: # 5 minute timeout (models may need to load) STREAMING_TIMEOUT: Final = 60 * 5 +_DECODED_JSON: Final = TypeAdapter(object) + class BytezChatConfig(BaseConfig): """ @@ -194,11 +198,11 @@ class BytezChatConfig(BaseConfig): ) -> ModelResponse: json: Final = raw_response.json() - error: Final = json.get("error") + error: Final = decoded_json_object(_DECODED_JSON.validate_python(json)).get("error") if error is not None: raise BytezError( - message=str(json["error"]), + message=str(error), status_code=raw_response.status_code, ) diff --git a/litellm/llms/cohere/common_utils.py b/litellm/llms/cohere/common_utils.py index 65531451d52..4ba5f9bd57b 100644 --- a/litellm/llms/cohere/common_utils.py +++ b/litellm/llms/cohere/common_utils.py @@ -1,6 +1,8 @@ import json from typing import Final, Literal +from pydantic import TypeAdapter + from litellm.llms.base_llm.base_utils import BaseLLMModelInfo from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.cohere import CohereV2StreamedToolCall @@ -12,6 +14,8 @@ from litellm.types.utils import ( ProviderSpecificModelInfo, ) +_DECODED_JSON: Final = TypeAdapter(object) + class CohereError(BaseLLMException): def __init__(self, status_code, message): @@ -190,8 +194,10 @@ class ModelResponseIterator: if index != -1: str_line = str_line[index:] - data_json: Final = json.loads(str_line) - return self.chunk_parser(chunk=data_json) + data_json: Final = _DECODED_JSON.validate_python(json.loads(str_line)) + if isinstance(data_json, dict): + return self.chunk_parser(chunk=data_json) + raise AttributeError(f"{type(data_json).__name__!r} object has no attribute 'get'", name="get", obj=data_json) # Async iterator def __aiter__(self): diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 169c21cf635..8c9de272c0d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -153,7 +153,7 @@ from litellm.types.containers.main import ( ) from litellm.types.decisions import DecisionsIRRequest, DecisionsIRResponse from litellm.types.extraction import ExtractionResponse -from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig +from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig, TwoStepFileUploadRequest from litellm.types.integrations.custom_logger import ( NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, AgenticLoopLimitError, @@ -188,6 +188,7 @@ from litellm.types.responses.main import DeleteResponseResult from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( + AgenticLoopParams, CallTypes, EmbeddingResponse, FileTypes, @@ -239,6 +240,10 @@ def _client_request_headers( return _merge_headers_case_insensitive(optional_headers or {}, extra_headers) +def _requested_stream(response_api_optional_request_params: Mapping[str, object]) -> object: + return response_api_optional_request_params.get("stream", False) + + def _rust_responses_websocket_enabled( custom_llm_provider: str | None, ) -> bool: @@ -386,12 +391,19 @@ class _SignedBodyKwargs(TypedDict, total=False): json: ReadOnly[dict[str, object]] +class _AgenticLoopParamsView(TypedDict): + agentic_loop_params: ReadOnly[AgenticLoopParams] + + _GOAWAY_RETRY_MESSAGE: Final = "Vector store search connection was closed by an HTTP/2 GOAWAY, retrying once" def _closed_by_http2_goaway(error: httpx.RemoteProtocolError) -> bool: cause: Final = error.__cause__ - return cause is not None and any(isinstance(arg, ConnectionTerminated) for arg in cause.args) + if cause is None: + return False + cause_args: Final[tuple[object, ...]] = cause.args + return any(isinstance(arg, ConnectionTerminated) for arg in cause_args) async def _asend_with_goaway_retry(send: Callable[[], Awaitable[httpx.Response]]) -> httpx.Response: @@ -433,8 +445,12 @@ def _signed_body_kwargs(*, signed_body: bytes | None, data: dict[str, object]) - return {"json": data} +def _response_content_type(headers: Mapping[str, str]) -> str: + return headers.get("content-type", "") + + def _answered_json(response: httpx.Response) -> bool: - return "application/json" in response.headers.get("content-type", "").lower() + return "application/json" in _response_content_type(response.headers).lower() def _stream_query_params(stream: bool, starting_after: int | None) -> Mapping[str, object]: @@ -897,7 +913,7 @@ class BaseLLMHTTPHandler: fake_stream = ( fake_stream - or optional_params.pop("fake_stream", False) + or optional_params.pop("fake_stream", False) is True or provider_config.should_fake_stream(model=model, custom_llm_provider=custom_llm_provider, stream=stream) ) @@ -920,7 +936,7 @@ class BaseLLMHTTPHandler: api_key=api_key, model=model, optional_params=optional_params, - stream=stream, + stream=False if fake_stream else stream, litellm_params=litellm_params, ) @@ -1056,7 +1072,7 @@ class BaseLLMHTTPHandler: provider_config=provider_config, fake_stream=fake_stream, ) - if provider_config.has_custom_stream_wrapper is True: + if provider_config.has_custom_stream_wrapper is True and not fake_stream: return provider_config.get_sync_custom_stream_wrapper( model=model, custom_llm_provider=custom_llm_provider, @@ -1225,7 +1241,7 @@ class BaseLLMHTTPHandler: json_mode: bool | None = None, signed_json_body: bytes | None = None, ): - if provider_config.has_custom_stream_wrapper is True: + if provider_config.has_custom_stream_wrapper is True and not fake_stream: return await provider_config.get_async_custom_stream_wrapper( model=model, custom_llm_provider=custom_llm_provider, @@ -3064,7 +3080,7 @@ class BaseLLMHTTPHandler: } # Check if streaming is requested - stream = response_api_optional_request_params.get("stream", False) + stream = _requested_stream(response_api_optional_request_params) caller_requested_stream: Final = bool(stream or (extra_body or {}).get("stream")) api_base: Final = responses_api_provider_config.get_complete_url( @@ -4136,7 +4152,7 @@ class BaseLLMHTTPHandler: ): # Handle pre-signed requests (e.g., from Bedrock S3 uploads) # Type narrowing: this is a plain dict, not TwoStepFileUploadConfig - presigned_request: Final = cast(dict[str, Any], transformed_request) + presigned_request: Final = cast(TwoStepFileUploadRequest, transformed_request) upload_response = getattr(sync_httpx_client, presigned_request["method"].lower())( url=presigned_request["url"], headers=presigned_request["headers"], @@ -4350,7 +4366,7 @@ class BaseLLMHTTPHandler: ): # Handle pre-signed requests (e.g., from Bedrock S3 uploads) # Type narrowing: this is a plain dict, not TwoStepFileUploadConfig - presigned_request: Final = cast(dict[str, Any], transformed_request) + presigned_request: Final = cast(TwoStepFileUploadRequest, transformed_request) upload_response = await getattr(async_httpx_client, presigned_request["method"].lower())( url=presigned_request["url"], headers=presigned_request["headers"], @@ -6129,8 +6145,10 @@ class BaseLLMHTTPHandler: full_model_name = model if logging_obj is not None: - agentic_params: Final[Mapping[str, object]] = logging_obj.model_call_details.get("agentic_loop_params", {}) - full_model_name = cast(str, agentic_params.get("model", model)) + agentic_view: Final[_AgenticLoopParamsView] = { + "agentic_loop_params": logging_obj.model_call_details.get("agentic_loop_params", {}) + } + full_model_name = agentic_view["agentic_loop_params"].get("model", model) optional_params: Final = dict(anthropic_messages_optional_request_params) optional_params.update(patch.optional_params) @@ -11116,7 +11134,7 @@ class BaseLLMHTTPHandler: extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, - ): + ) -> Mapping[str, object]: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( llm_provider=litellm.LlmProviders(custom_llm_provider), @@ -11407,7 +11425,7 @@ class BaseLLMHTTPHandler: extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, - ): + ) -> Mapping[str, object]: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( llm_provider=litellm.LlmProviders(custom_llm_provider), diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py index 0f01a4b3172..e1efa075dc8 100644 --- a/litellm/llms/dashscope/embed/transformation.py +++ b/litellm/llms/dashscope/embed/transformation.py @@ -11,12 +11,13 @@ Endpoint Docs - https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api """ -from typing import Final +from typing import Final, Protocol, runtime_checkable import httpx from pydantic import ConfigDict, TypeAdapter, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict +from litellm.litellm_core_utils.core_helpers import json_contains from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -29,6 +30,19 @@ from ..common_utils import DashScopeError, resolve_dashscope_openai_api_base DEFAULT_API_BASE: Final = "https://dashscope.aliyuncs.com/compatible-mode/v1" +_DECODED_JSON: Final = TypeAdapter(object) + + +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, key: str, /) -> object: ... + + +def _subscript(value: object, key: str) -> object: + if isinstance(value, _SupportsGetItem): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + @with_config(ConfigDict(extra="allow", strict=True)) class _DashScopeUsage(TypedDict, total=False): @@ -149,7 +163,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, ) -> EmbeddingResponse: try: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise DashScopeError( status_code=raw_response.status_code, @@ -163,8 +177,8 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): original_response=response_json, ) - if "error" in response_json: - error: Final[object] = response_json["error"] + if json_contains(response_json, "error"): + error: Final = _subscript(response_json, "error") message: Final = error.get("message", str(error)) if isinstance(error, dict) else str(error) raise DashScopeError( status_code=raw_response.status_code, diff --git a/litellm/llms/dashscope/videos/transformation.py b/litellm/llms/dashscope/videos/transformation.py index 59e42e7b926..c815e3790cf 100644 --- a/litellm/llms/dashscope/videos/transformation.py +++ b/litellm/llms/dashscope/videos/transformation.py @@ -22,6 +22,7 @@ from litellm.llms.base_llm.videos.async_task_utils import ( input_reference_url, polled_video_models, routed_video_models, + unsupported_video_fields, ) from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_httpx.http_handler import ( @@ -558,6 +559,13 @@ class DashScopeVideoConfig(BaseVideoConfig): model: str, drop_params: bool, ) -> dict[str, object]: # mutable-ok: BaseVideoConfig contract + excluded_openrouter: Final = unsupported_video_fields( + video_create_optional_params, + model, + "dashscope", + drop_params, + frozenset({"duration", "seed", "aspect_ratio", "generate_audio"}), + ) unsupported: Final = next( ( key @@ -578,7 +586,7 @@ class DashScopeVideoConfig(BaseVideoConfig): { key: value for key, value in video_create_optional_params.items() - if value is not None and key not in _MAPPED_ELSEWHERE + if value is not None and key not in _MAPPED_ELSEWHERE and key not in excluded_openrouter } ) mapped_params: Final[dict[str, object]] = { # mutable-ok: BaseVideoConfig contract; extra_body merges into it diff --git a/litellm/llms/dataforseo/search/transformation.py b/litellm/llms/dataforseo/search/transformation.py index 5cf7735325d..b88f7df0f91 100644 --- a/litellm/llms/dataforseo/search/transformation.py +++ b/litellm/llms/dataforseo/search/transformation.py @@ -4,13 +4,14 @@ Calls DataForSEO SERP API to search the web. DataForSEO API Reference: https://docs.dataforseo.com/v3/serp/google/organic/live/advanced/?bash """ -from collections.abc import Iterable, Mapping -from typing import Final, Literal +from collections.abc import Iterable, Mapping, Sized +from typing import Final, Literal, Protocol, runtime_checkable import httpx from pydantic import ConfigDict, TypeAdapter from litellm.constants import DEFAULT_DATAFORSEO_LOCATION_CODE +from litellm.litellm_core_utils.core_helpers import decoded_json_object, json_contains from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, @@ -21,6 +22,24 @@ from litellm.secret_managers.main import get_secret_str _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) + + +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, key: str | int, /) -> object: ... + + +def _subscript(value: object, key: str | int) -> object: + if isinstance(value, _SupportsGetItem): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + + +def _json_len(value: object) -> int: + if isinstance(value, Sized): + return len(value) + raise TypeError(f"object of type '{type(value).__name__}' has no len()") class DataForSEOSearchConfig(BaseSearchConfig): @@ -193,20 +212,20 @@ class DataForSEOSearchConfig(BaseSearchConfig): Returns: SearchResponse with standardized format """ - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) # Transform results to SearchResult objects results: Final = [] # DataForSEO wraps results in tasks array - if "tasks" in response_json and len(response_json["tasks"]) > 0: - task: Final = response_json["tasks"][0] + if json_contains(response_json, "tasks") and _json_len(_subscript(response_json, "tasks")) > 0: + task: Final = decoded_json_object(_subscript(_subscript(response_json, "tasks"), 0)) # Check if task was successful if task.get("status_code") == 20000 and "result" in task: # Result is an array, take first element - if len(task["result"]) > 0: - result: Final = _JSON_OBJECT.validate_python(task["result"][0]) + if _json_len(task["result"]) > 0: + result: Final = _JSON_OBJECT.validate_python(_subscript(task["result"], 0)) # Items contain the actual search results for raw_item in _JSON_VALUES.validate_python(result.get("items", [])): diff --git a/litellm/llms/edenai/videos/transformation.py b/litellm/llms/edenai/videos/transformation.py index c9f53523acf..0cb9c5afaec 100644 --- a/litellm/llms/edenai/videos/transformation.py +++ b/litellm/llms/edenai/videos/transformation.py @@ -15,7 +15,7 @@ from httpx._types import RequestFiles from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.router import GenericLiteLLMParams -from litellm.types.videos.main import VideoObject +from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoObject from ..common_utils import EdenAIException, authorized_headers, endpoint_url, reported_cost @@ -36,6 +36,14 @@ def _usage_with_reported_cost( class EdenAIVideoConfig(OpenAIVideoConfig): + def map_openai_params( + self, + video_create_optional_params: VideoCreateOptionalRequestParams, + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: inherited BaseVideoConfig return contract + return dict(video_create_optional_params) + def validate_environment( self, headers: dict[str, object], # mutable-ok: inherited contract diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index 86aa6388ff0..c658dab4d3f 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -12,6 +12,7 @@ from pydantic import TypeAdapter from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.videos.async_task_utils import unsupported_video_fields from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -304,6 +305,13 @@ class FalAIVideoConfig(BaseVideoConfig): model: str, drop_params: bool, ) -> _VideoParams: + unsupported: Final = unsupported_video_fields( + video_create_optional_params, + model, + "fal_ai", + drop_params, + frozenset({"generate_audio", "aspect_ratio", "seed", "duration"}), + ) supported_params: Final[frozenset[str]] = frozenset(self.get_supported_openai_params(model)) input_reference: Final[object] = video_create_optional_params.get("input_reference") if "input_reference" in video_create_optional_params and not isinstance(input_reference, str): @@ -338,7 +346,11 @@ class FalAIVideoConfig(BaseVideoConfig): **duration_params, **size_params, **user_params, - **{key: value for key, value in video_create_optional_params.items() if key not in supported_params}, + **{ + key: value + for key, value in video_create_optional_params.items() + if key not in supported_params and key not in unsupported + }, } return mapped_params diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py index d599f0b65f7..a45c25b27f3 100644 --- a/litellm/llms/fireworks_ai/responses/transformation.py +++ b/litellm/llms/fireworks_ai/responses/transformation.py @@ -1,10 +1,11 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # the base signature leaves the params bag untyped from urllib.parse import unquote import httpx +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.responses.transformation import ( STANDARD_MODE_WITHOUT_CONTEXT_REASONING_FIELD_VALUES, @@ -20,8 +21,9 @@ from litellm.llms.fireworks_ai.common_utils import ( ) from litellm.llms.openai.responses.instruction_items import developer_items_as_system, instruction_items_folded from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.responses.additional_tools import HoistedAdditionalTools, hoist_additional_tools from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import ResponseInputParam +from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS, ResponseInputParam from litellm.types.responses.main import DeleteResponseResult from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders @@ -98,6 +100,23 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): headers: dict[str, str], # mutable-ok: overrides the base class signature ) -> dict[str, object]: # mutable-ok: overrides the base class signature validated_input: Final = self._validate_input_param(input) + tools_param: Final = response_api_optional_request_params.get("tools") + hoisted: Final = ( + hoist_additional_tools( + validated_input, + cast( # cast-ok: the base signature leaves the params bag untyped + "Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None", tools_param + ), + drop_params=drops_unsupported_params(litellm_params.drop_params), + ) + if tools_param is None or isinstance(tools_param, (list, tuple)) + else HoistedAdditionalTools(input=validated_input, tools=(), hoisted=()) + ) + hoisted_params: Final = ( + {**response_api_optional_request_params, "tools": list(hoisted.tools)} + if hoisted.hoisted + else response_api_optional_request_params + ) model_info: Final = deployment_model_info(litellm_params) declared: Final = declared_mid_conversation_system_support(model, self.custom_llm_provider.value, model_info) keeps_developer_items_in_place: Final = declared is True and ( @@ -107,14 +126,9 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): folded: Final = ( None if keeps_developer_items_in_place - else instruction_items_folded( - validated_input, response_api_optional_request_params, "user" if declared is False else "system" - ) - ) - placed_input, placed_params = folded or ( - developer_items_as_system(validated_input), - response_api_optional_request_params, + else instruction_items_folded(hoisted.input, hoisted_params, "user" if declared is False else "system") ) + placed_input, placed_params = folded or (developer_items_as_system(hoisted.input), hoisted_params) request: Final = super().transform_responses_api_request( model=resolve_fireworks_resource_name(model), input=placed_input, diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 74cd6f22704..b95219348e4 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -4,10 +4,11 @@ import json import math from collections.abc import Iterable, Mapping, Sequence from types import MappingProxyType -from typing import Any, Final, cast +from typing import Final, cast import httpx from pydantic import JsonValue, TypeAdapter +from typing_extensions import TypedDict as ExtraItemsTypedDict import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -510,6 +511,13 @@ def get_api_key_from_env() -> str | None: return get_secret_str("GOOGLE_API_KEY") or get_secret_str("GEMINI_API_KEY") +_COUNT_TOKENS_REQUEST_FIELDS: Final = frozenset({"model", "contents", "system_instruction", "tools"}) + + +class _ForwardedLitellmParams(ExtraItemsTypedDict, total=False, extra_items=object): + pass + + class GoogleAIStudioTokenCounter(BaseTokenCounter): """Token counter implementation for Google AI Studio provider.""" @@ -526,7 +534,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): model_to_use: str, messages: list[dict[str, object]] | None, contents: list[dict[str, object]] | None, - deployment: dict[str, Any] | None = None, + deployment: dict[str, object] | None = None, request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, @@ -537,6 +545,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): ) -> TokenCountResponse | None: import copy + from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.llms.anthropic.messages.context_management.inference_view import ( InferenceRequestViewError, count_context_edits, @@ -553,11 +562,12 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): vertex_gemini_count_context, ) + litellm_params: Final = decoded_json_object((deployment or {}).get("litellm_params", {})) request_view: Final = await inference_request_view( messages or [], system, messages_format, - count_context_edits(context_management, tools, (deployment or {}).get("litellm_params", {}), model_to_use), + count_context_edits(context_management, tools, litellm_params, model_to_use), ) if isinstance(request_view, InferenceRequestViewError): return count_context_error_response(request_view, request_model, model_to_use) @@ -570,8 +580,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): if blocks_error is not None: return count_context_error_response(blocks_error, request_model, model_to_use) - deployment = deployment or {} - count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) + count_tokens_params_request: Final = copy.deepcopy(litellm_params) resolved_contents: Final = ( contents if contents is not None @@ -581,16 +590,15 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): ) if isinstance(resolved_contents, VertexGeminiCountContextError): return count_context_error_response(resolved_contents, request_model, model_to_use) - count_tokens_params: Final = { - "model": model_to_use, - "contents": resolved_contents, - "system_instruction": count_context.system_instruction, - "tools": count_context.tools, - } - count_tokens_params_request.update(count_tokens_params) try: result: Final = await GoogleAIStudioTokenCounter().acount_tokens( - **count_tokens_params_request, + **_ForwardedLitellmParams( + **{k: v for k, v in count_tokens_params_request.items() if k not in _COUNT_TOKENS_REQUEST_FIELDS} + ), + model=model_to_use, + contents=resolved_contents, + system_instruction=count_context.system_instruction, + tools=count_context.tools, ) except litellm.APIError as error: return provider_error_response(error, request_model, model_to_use) diff --git a/litellm/llms/gemini/count_tokens/handler.py b/litellm/llms/gemini/count_tokens/handler.py index 4ed47323d4f..2c78d04eea5 100644 --- a/litellm/llms/gemini/count_tokens/handler.py +++ b/litellm/llms/gemini/count_tokens/handler.py @@ -1,11 +1,12 @@ -from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Protocol, runtime_checkable import httpx from pydantic import ConfigDict, RootModel, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm +from litellm.litellm_core_utils.core_helpers import has_attributes, iter_items, json_contains from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.types.utils import LlmProviders @@ -26,8 +27,44 @@ class _GeminiCountTokensResponseRoot(RootModel[_GeminiCountTokensResponse]): model_config = ConfigDict(frozen=True, hide_input_in_errors=True) +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, key: str, /) -> object: ... + + +@runtime_checkable +class _SupportsSetItem(Protocol): + def __setitem__(self, key: str, value: object, /) -> None: ... + + +class _KeywordArguments(Protocol): + def keys(self) -> Iterable[str]: ... + + def __getitem__(self, key: str, /) -> object: ... + + +def _subscript(value: object, key: str) -> object: + if isinstance(value, _SupportsGetItem): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + + +def _item_assignment_target(value: object) -> _SupportsSetItem: + if isinstance(value, _SupportsSetItem): + return value + raise TypeError(f"'{type(value).__name__}' object does not support item assignment") + + +def _function_response_kwargs(value: object) -> _KeywordArguments: + if has_attributes(value, _KeywordArguments, "keys"): + return value + raise TypeError( + f"google.genai.types.FunctionResponse() argument after ** must be a mapping, not {type(value).__name__}" + ) + + class GoogleAIStudioTokenCounter: - def _clean_contents_for_gemini_api(self, contents: Any) -> object: + def _clean_contents_for_gemini_api(self, contents: Sequence[Mapping[str, object]] | None) -> object: """ Clean up contents to remove unsupported fields for the Gemini API. @@ -52,12 +89,14 @@ class GoogleAIStudioTokenCounter: for content in cleaned_contents: parts = content["parts"] - for part in parts: - if "functionResponse" in part: - function_response_data = part["functionResponse"] - function_response_part = FunctionResponse(**function_response_data) + for part in iter_items(parts): + if json_contains(part, "functionResponse"): + function_response_data = _subscript(part, "functionResponse") + function_response_part = FunctionResponse(**_function_response_kwargs(function_response_data)) function_response_part.id = None - part["functionResponse"] = function_response_part.model_dump(exclude_none=True) + _item_assignment_target(part)["functionResponse"] = function_response_part.model_dump( + exclude_none=True + ) return cleaned_contents diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index 50da73cc8b9..a02afa3a13f 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -9,13 +9,14 @@ Per OpenAPI spec (https://ai.google.dev/static/api/interactions.openapi.json): Requests use Api-Revision 2026-05-20 (`steps` schema). """ +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias import httpx from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, process_response_headers +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, has_attributes, process_response_headers from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.interactions.transformation import BaseInteractionsAPIConfig from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo @@ -83,6 +84,18 @@ class _CancelPayloadSource(Protocol): def json(self) -> _CancelPayload: ... +class _ImageConfigItems(Protocol): + def keys(self) -> Iterable[str]: ... + + def __getitem__(self, key: str, /) -> object: ... + + +def _image_config_items(image_config: object) -> _ImageConfigItems: + if has_attributes(image_config, _ImageConfigItems, "keys"): + return image_config + raise TypeError(f"'{type(image_config).__name__}' object is not a mapping") + + def _interaction_body(response: _InteractionPayloadSource) -> _InteractionPayload: """Decode the body of an Interactions API interaction response.""" return response.json() @@ -241,7 +254,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): request_body["response_format"] = response_format # image_config moves out of generation_config into response_format. - generation_config: dict[str, Any] | None = optional_params.get("generation_config") + generation_config: Mapping[str, object] | None = optional_params.get("generation_config") if generation_config is not None: image_config = None if isinstance(generation_config, dict): @@ -255,7 +268,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): if image_config is not None: # Move image_config to response_format with type=image. - image_rf: Final[_JsonObject] = {"type": "image", **image_config} + image_rf: Final[_JsonObject] = {"type": "image", **_image_config_items(image_config)} existing_rf: Final = request_body.get("response_format") if existing_rf is None: request_body["response_format"] = image_rf diff --git a/litellm/llms/gemini/videos/transformation.py b/litellm/llms/gemini/videos/transformation.py index 3f6b900fd01..65a81aeee0f 100644 --- a/litellm/llms/gemini/videos/transformation.py +++ b/litellm/llms/gemini/videos/transformation.py @@ -10,6 +10,7 @@ import litellm from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS from litellm.images.utils import ImageEditRequestUtils from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.base_llm.videos.async_task_utils import unsupported_video_fields from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.vertex_ai.videos.transformation import veo_video_count_from_parameters from litellm.secret_managers.main import get_secret_str @@ -123,6 +124,13 @@ class GeminiVideoConfig(BaseVideoConfig): All other params are passed through as-is to support Gemini-specific parameters. """ + unsupported: Final = unsupported_video_fields( + video_create_optional_params, + model, + "gemini", + drop_params, + frozenset({"seed", "duration", "aspect_ratio", "generate_audio"}), + ) mapped_params: Final[dict[str, object]] = {} # Get supported OpenAI params (exclude "model" and "prompt" which are handled separately) @@ -158,7 +166,7 @@ class GeminiVideoConfig(BaseVideoConfig): # Pass through any other params that weren't mapped (Gemini-specific params) for key, value in video_create_optional_params.items(): - if key not in openai_params_to_map and key not in mapped_params: + if key not in openai_params_to_map and key not in mapped_params and key not in unsupported: mapped_params[key] = value return mapped_params diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 5324aa6f94e..08b5de10d41 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -4,9 +4,12 @@ GigaChat Streaming Response Handler import json import uuid -from collections.abc import Mapping, Sequence -from typing import Any, Final +from collections.abc import Awaitable, Mapping, Sequence +from typing import Final, Protocol +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.llms.gigachat.utils import convert_usage from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, @@ -14,13 +17,27 @@ from litellm.types.llms.openai import ( ) from litellm.types.utils import ChatCompletionUsageBlock, GenericStreamingChunk +_DECODED_JSON: Final = TypeAdapter(object) + + +class _SyncChunkSource(Protocol): + def __next__(self) -> object: ... + + +class _AsyncChunkSource(Protocol): + def __anext__(self) -> Awaitable[object]: ... + + +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + class GigaChatModelResponseIterator: """Iterator for GigaChat streaming responses.""" def __init__( self, - streaming_response: Any, + streaming_response: object, sync_stream: bool, json_mode: bool | None = False, ): @@ -28,7 +45,7 @@ class GigaChatModelResponseIterator: self.response_iterator = self.streaming_response self.json_mode = json_mode - def chunk_parser(self, chunk: Mapping[str, object]) -> GenericStreamingChunk: + def chunk_parser(self, chunk: _SupportsGet) -> GenericStreamingChunk: """Parse a single streaming chunk from GigaChat.""" choices: Sequence = chunk.get("choices") or () if not choices: @@ -106,14 +123,14 @@ class GigaChatModelResponseIterator: def __next__(self) -> GenericStreamingChunk: try: - chunk = self.response_iterator.__next__() + chunk: Final = require_attributes(self.response_iterator, _SyncChunkSource, "__next__").__next__() if isinstance(chunk, str): # Parse SSE format: data: {...} - chunk = chunk.removeprefix("data: ") - if chunk.strip() == "[DONE]": + payload: Final = chunk.removeprefix("data: ") + if payload.strip() == "[DONE]": raise StopIteration try: - chunk = json.loads(chunk) + decoded: Final = _DECODED_JSON.validate_python(json.loads(payload)) except json.JSONDecodeError: return GenericStreamingChunk( text="", @@ -123,7 +140,8 @@ class GigaChatModelResponseIterator: usage=None, index=0, ) - return self.chunk_parser(chunk) + return self.chunk_parser(require_attributes(decoded, _SupportsGet, "get")) + return self.chunk_parser(require_attributes(chunk, _SupportsGet, "get")) except StopIteration: raise @@ -132,14 +150,14 @@ class GigaChatModelResponseIterator: async def __anext__(self) -> GenericStreamingChunk: try: - chunk = await self.response_iterator.__anext__() + chunk: Final = await require_attributes(self.response_iterator, _AsyncChunkSource, "__anext__").__anext__() if isinstance(chunk, str): # Parse SSE format - chunk = chunk.removeprefix("data: ") - if chunk.strip() == "[DONE]": + payload: Final = chunk.removeprefix("data: ") + if payload.strip() == "[DONE]": raise StopAsyncIteration try: - chunk = json.loads(chunk) + decoded: Final = _DECODED_JSON.validate_python(json.loads(payload)) except json.JSONDecodeError: return GenericStreamingChunk( text="", @@ -149,6 +167,7 @@ class GigaChatModelResponseIterator: usage=None, index=0, ) - return self.chunk_parser(chunk) + return self.chunk_parser(require_attributes(decoded, _SupportsGet, "get")) + return self.chunk_parser(require_attributes(chunk, _SupportsGet, "get")) except StopAsyncIteration: raise diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 5d997babc6e..9174b77a9f0 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -1,11 +1,13 @@ import json import os -from typing import Final, Literal, Protocol, get_args +from collections.abc import Iterator +from typing import Final, Literal, Protocol, get_args, runtime_checkable import httpx from pydantic import TypeAdapter import litellm +from litellm.litellm_core_utils.core_helpers import as_object_list, iter_items, json_contains from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -38,6 +40,28 @@ class _SupportsTokenEncode(Protocol): def encode(self, text: str) -> list[int]: ... +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, key: str | int, /) -> object: ... + + +def _subscript(value: object, key: str | int) -> object: + if isinstance(value, _SupportsGetItem): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + + +def _unpack_pair(value: object) -> tuple[object, object]: + if not hasattr(value, "__iter__") and not hasattr(value, "__getitem__"): + raise TypeError(f"cannot unpack non-iterable {type(value).__name__} object") + first, second = iter_items(value) + return first, second + + +def _pairs(value: object) -> Iterator[tuple[object, object]]: + return (_unpack_pair(item) for item in iter_items(value)) + + def get_hf_task_embedding_for_model( model: str, task_type: str | None, api_base: str, timeout: float | httpx.Timeout | None = None ) -> object: @@ -183,15 +207,15 @@ class HuggingFaceEmbedding(BaseLLM): def _process_embedding_response( self, - embeddings: dict, + embeddings: object, model_response: EmbeddingResponse, model: str, input: list, encoding: _SupportsTokenEncode, ) -> EmbeddingResponse: output_data: Final = [] - if "similarities" in embeddings: - for idx, embedding in embeddings["similarities"]: + if json_contains(embeddings, "similarities"): + for idx, embedding in _pairs(_subscript(embeddings, "similarities")): output_data.append( { "object": "embedding", @@ -200,8 +224,9 @@ class HuggingFaceEmbedding(BaseLLM): } ) else: - for idx, embedding in enumerate(embeddings): - if isinstance(embedding, float) or isinstance(embedding, list) and isinstance(embedding[0], float): + for idx, embedding in enumerate(iter_items(embeddings)): + embedding_list = as_object_list(embedding) + if isinstance(embedding, float) or embedding_list is not None and isinstance(embedding_list[0], float): output_data.append( { "object": "embedding", @@ -214,7 +239,7 @@ class HuggingFaceEmbedding(BaseLLM): { "object": "embedding", "index": idx, - "embedding": embedding[0][0], # flatten list returned from hf + "embedding": _subscript(_subscript(embedding, 0), 0), # flatten list returned from hf } ) model_response.object = "list" @@ -288,10 +313,10 @@ class HuggingFaceEmbedding(BaseLLM): original_response=response, ) - embeddings: Final = response.json() + embeddings: Final = _DECODED_JSON.validate_python(response.json()) - if "error" in embeddings: - raise HuggingFaceError(status_code=500, message=embeddings["error"]) + if json_contains(embeddings, "error"): + raise HuggingFaceError(status_code=500, message=_subscript(embeddings, "error")) ## PROCESS RESPONSE ## return self._process_embedding_response( @@ -394,10 +419,10 @@ class HuggingFaceEmbedding(BaseLLM): original_response=response, ) - embeddings: Final = response.json() + embeddings: Final = _DECODED_JSON.validate_python(response.json()) - if "error" in embeddings: - raise HuggingFaceError(status_code=500, message=embeddings["error"]) + if json_contains(embeddings, "error"): + raise HuggingFaceError(status_code=500, message=_subscript(embeddings, "error")) ## PROCESS RESPONSE ## return self._process_embedding_response( diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index 6e6a117227e..bb836eb1b35 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -365,7 +365,10 @@ class HuggingFaceEmbeddingConfig(BaseConfig): ) content: str = "" for chunk in streamed_response: - content += chunk["choices"][0]["delta"]["content"] + delta_content = chunk.choices[0].delta.content + if delta_content is None: + raise TypeError('can only concatenate str (not "NoneType") to str') + content += delta_content completion_response: Final[list[dict[str, str]]] = [{"generated_text": content}] ## LOGGING logging_obj.post_call( diff --git a/litellm/llms/litellm_proxy/skills/code_execution.py b/litellm/llms/litellm_proxy/skills/code_execution.py index fbc287589b3..af26b4aa095 100644 --- a/litellm/llms/litellm_proxy/skills/code_execution.py +++ b/litellm/llms/litellm_proxy/skills/code_execution.py @@ -15,11 +15,12 @@ import base64 import json from collections.abc import Mapping, Sequence from enum import Enum -from typing import Any, Final, Protocol, TypedDict +from typing import Final, Protocol, TypedDict from typing_extensions import ReadOnly from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import require_attributes class _ToolParameterSchema(TypedDict, total=False): @@ -237,17 +238,21 @@ class CodeExecutionHandler: execution_results: Final[list[_ExecutionResult]] = [] executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) - response: Any = None # Initialize to avoid possibly unbound error + response: _CompletionResponse | None = None # Initialize to avoid possibly unbound error for iteration in range(self.max_iterations): verbose_logger.debug("CodeExecutionHandler: Iteration %s/%s", iteration + 1, self.max_iterations) # Make LLM call - response = await litellm.acompletion( - model=model, - messages=current_messages, - tools=tools, - **kwargs, + response = require_attributes( + await litellm.acompletion( + model=model, + messages=current_messages, + tools=tools, + **kwargs, + ), + _CompletionResponse, + "choices", ) choice: _ResponseChoice = response.choices[0] diff --git a/litellm/llms/litellm_proxy/skills/transformation.py b/litellm/llms/litellm_proxy/skills/transformation.py index 73e26606497..3b488784443 100644 --- a/litellm/llms/litellm_proxy/skills/transformation.py +++ b/litellm/llms/litellm_proxy/skills/transformation.py @@ -10,8 +10,6 @@ Pattern follows litellm/llms/litellm_proxy/responses/transformation.py from collections.abc import Coroutine, Sequence from typing import TYPE_CHECKING, Final, Optional -from pydantic import JsonValue - from litellm.types.llms.anthropic_skills import ( DeleteSkillResponse, ListSkillsResponse, @@ -47,7 +45,7 @@ class LiteLLMSkillsTransformationHandler: file_content: bytes | None = None, file_name: str | None = None, file_type: str | None = None, - metadata: dict[str, JsonValue] | None = None, + metadata: dict[str, object] | None = None, user_id: str | None = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, _is_async: bool = False, @@ -130,7 +128,7 @@ class LiteLLMSkillsTransformationHandler: file_content: bytes | None = None, file_name: str | None = None, file_type: str | None = None, - metadata: dict[str, JsonValue] | None = None, + metadata: dict[str, object] | None = None, user_id: str | None = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, ) -> Skill: diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 14470922e57..fb707a39362 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -252,10 +252,11 @@ class MilvusVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): # Build attributes with all available metadata # Exclude system fields and already-processed fields - attributes = {} - for key, value in result.items(): - if key not in ["id", "content", "distance", text_field]: - attributes[key] = value + attributes = { + key: value + for key, value in decoded_json_object(_DECODED_JSON.validate_python(result)).items() + if key not in ["id", "content", "distance", text_field] + } result_obj = VectorStoreSearchResult( score=score, diff --git a/litellm/llms/minimax/videos/transformation.py b/litellm/llms/minimax/videos/transformation.py index a1505d2a71e..39ee783a4ba 100644 --- a/litellm/llms/minimax/videos/transformation.py +++ b/litellm/llms/minimax/videos/transformation.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.url_utils import ( safe_get, ) from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.videos.async_task_utils import video_image_params from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -488,6 +489,9 @@ class MinimaxVideoConfig(BaseVideoConfig): "aigc_watermark", "extra", "parameters", + "frame_images", + "input_references", + "aspect_ratio", ] def map_openai_params( @@ -496,12 +500,20 @@ class MinimaxVideoConfig(BaseVideoConfig): model: str, drop_params: bool, ) -> dict[str, object]: # mutable-ok: BaseVideoConfig contract + requested: Final = video_image_params(video_create_optional_params, model, "minimax", drop_params) mapped_params: Final[dict[str, object]] = {} # mutable-ok: BaseVideoConfig contract; extra_body merges into it dropped_keys: Final = _DROP_FROM_CREATE_BODY | _additional_drop_keys(video_create_optional_params) - for key, value in video_create_optional_params.items(): + for key, value in requested.items(): if value is None or key in dropped_keys: continue - if key == "seconds": + if key == "seconds" and requested.get("duration") is not None: + if not drop_params and _duration_param(value, model) != _duration_param(requested["duration"], model): + raise UnsupportedParamsError( + message="duration and seconds must agree, set drop_params=True to ignore seconds", + model=model, + llm_provider="minimax", + ) + elif key == "seconds": mapped_params["duration"] = _duration_param(value, model) elif key == "size": mapped_params.setdefault("ratio", _ratio_from_size(value, model)) diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 61d237de2d8..d13e46658ae 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -21,6 +21,7 @@ from litellm._logging import verbose_logger from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import is_str_keyed_dict from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, strip_litellm_internal_message_fields, @@ -82,6 +83,7 @@ _MISTRAL_REQUEST_ROUTING: Final = TypeAdapter(MistralRequestRouting) _MISTRAL_CACHE_KEY_PARAMS: Final = TypeAdapter(MistralCacheKeyParams) _MISTRAL_REASONING_EFFORT_PARAMS: Final = TypeAdapter(MistralReasoningEffortParams) _MISTRAL_FILE_PART: Final = TypeAdapter(MistralFilePart) +_DECODED_JSON: Final = TypeAdapter(object) _REASONING_PROMPT_REASON: Final = "litellm sends Magistral a reasoning system prompt instead of the requested level." _NON_LEVEL_EFFORT_REASON: Final = ( "the Mistral API takes it only as a level such as 'high', so Mistral applies its default effort." @@ -1005,9 +1007,13 @@ class MistralConfig(OpenAIGPTConfig): logging_obj.model_call_details["response_headers"] = raw_response.headers # Handle Mistral-specific response transformations - response_data = raw_response.json() - response_data = self._handle_empty_content_response(response_data) - response_data = self._handle_content_list_to_str_conversion(response_data) + decoded_response: Final = _DECODED_JSON.validate_python(raw_response.json()) + if not is_str_keyed_dict(decoded_response): + missing_get: Final = f"{type(decoded_response).__name__!r} object has no attribute 'get'" + raise AttributeError(missing_get, name="get", obj=decoded_response) + response_data: Final = self._handle_content_list_to_str_conversion( + self._handle_empty_content_response(decoded_response) + ) final_response_obj: Final = cast( ModelResponse, diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 9163ddfa5de..846f5f27e0d 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -286,7 +286,7 @@ def handle_cohere_response( def handle_cohere_stream_chunk( - dict_chunk: Mapping[str, JsonValue], + dict_chunk: Mapping[str, object], prior_tool_calls_emitted: bool = False, prior_text_emitted: bool = False, ) -> ModelResponseStream: diff --git a/litellm/llms/oci/chat/generic.py b/litellm/llms/oci/chat/generic.py index 13cea58495d..5cf4c7aa535 100644 --- a/litellm/llms/oci/chat/generic.py +++ b/litellm/llms/oci/chat/generic.py @@ -8,6 +8,7 @@ parsing, and streaming chunk parsing for models served with import datetime import hashlib +from collections.abc import Mapping from typing import Final import httpx @@ -363,7 +364,7 @@ def handle_generic_response( return model_response -def handle_generic_stream_chunk(dict_chunk: dict) -> ModelResponseStream: +def handle_generic_stream_chunk(dict_chunk: Mapping) -> ModelResponseStream: """Parse a single GENERIC SSE chunk into a LiteLLM ModelResponseStream.""" # OCI streams tool calls progressively — early chunks may omit required fields. if dict_chunk.get("message") and dict_chunk["message"].get("toolCalls"): diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index f97f8f4ca59..256d0a74623 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -19,6 +19,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger from litellm.constants import DEFAULT_OCI_CHAT_MAX_TOKENS +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.prompt_templates.common_utils import ( merge_consecutive_system_messages, @@ -891,13 +892,14 @@ class OCIStreamWrapper(CustomStreamWrapper): if payload == "[DONE]": return None try: - dict_chunk: Final = json.loads(payload) + decoded_chunk: Final = _DECODED_JSON.validate_python(json.loads(payload)) except json.JSONDecodeError as e: raise OCIError( status_code=500, message=f"Chunk cannot be parsed as JSON: {e}", ) + dict_chunk: Final = decoded_json_object(decoded_chunk) if dict_chunk.get("apiFormat") == "COHERE": result: Final = handle_cohere_stream_chunk( dict_chunk, diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index fcf6a396a13..d6b9ff46410 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -1,6 +1,7 @@ import base64 import io -from typing import Final, overload +from collections.abc import Iterable +from typing import Final, Protocol, overload import httpx from openai.types.chat import ChatCompletionChunk @@ -8,6 +9,7 @@ from pydantic import BaseModel, TypeAdapter from litellm import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -19,6 +21,14 @@ _THINK_OFF_REASON: Final = "Ollama takes think only as on or off, and litellm ma _DECODED_JSON: Final = TypeAdapter(object) +class _JsonLookup(Protocol): + def get(self, key: str, default: object, /) -> object: ... + + +class _JsonItems(Protocol): + def items(self) -> Iterable[tuple[str, object]]: ... + + class OllamaResponseIdNormalizer: def __init__(self) -> None: self.response_id: str | None = None @@ -206,15 +216,15 @@ class OllamaModelInfo(BaseLLMModelInfo): return any(name.lower() in model_cost_keys for name in potential_model_names) @staticmethod - def _supports_function_calling(ollama_model_info: dict) -> bool: - _template: Final[str] = str(ollama_model_info.get("template", "") or "") + def _supports_function_calling(ollama_model_info: object) -> bool: + _template: Final[str] = str(require_attributes(ollama_model_info, _JsonLookup, "get").get("template", "") or "") return "tools" in _template.lower() @staticmethod - def _get_max_tokens(ollama_model_info: dict) -> int | None: - _model_info: Final[dict] = ollama_model_info.get("model_info", {}) + def _get_max_tokens(ollama_model_info: object) -> object: + _model_info: Final = require_attributes(ollama_model_info, _JsonLookup, "get").get("model_info", {}) - for key, value in _model_info.items(): + for key, value in require_attributes(_model_info, _JsonItems, "items").items(): if "context_length" in key: return value return None @@ -253,7 +263,7 @@ class OllamaModelInfo(BaseLLMModelInfo): "max_output_tokens": None, } - model_info: Final = response.json() + model_info: Final = _DECODED_JSON.validate_python(response.json()) max_tokens: Final = self._get_max_tokens(model_info) return { diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 1e70141f939..e1fc9a8d262 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -6,12 +6,13 @@ Ollama /chat/completion calls handled in llm_http_handler.py from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Any, Final, Protocol, TypedDict +from typing import Final, Protocol, TypedDict import httpx from typing_extensions import NotRequired, ReadOnly import litellm +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.types.mcp import has_header from litellm.types.utils import EmbeddingResponse @@ -23,6 +24,14 @@ class TokenEncoder(Protocol): def encode(self, text: str, /) -> Sequence[int]: ... +class _DebugLogger(Protocol): + def debug(self, message: str, /) -> object: ... + + +class _WarningLogger(Protocol): + def warning(self, message: str, /) -> object: ... + + class OllamaEmbeddingResponse(TypedDict): """Body of an Ollama ``/api/embed`` response.""" @@ -51,7 +60,7 @@ def _process_ollama_embedding_response( prompts: list[str], model: str, model_response: EmbeddingResponse, - logging_obj: Any, + logging_obj: object, encoding: TokenEncoder | None, ) -> EmbeddingResponse: output_data: Final = [] @@ -66,11 +75,15 @@ def _process_ollama_embedding_response( if encoding is not None: input_tokens = len(encoding.encode("".join(prompts))) if logging_obj: - logging_obj.debug("Ollama response missing prompt_eval_count; estimated with encoding.") + require_attributes(logging_obj, _DebugLogger, "debug").debug( + "Ollama response missing prompt_eval_count; estimated with encoding." + ) else: input_tokens = 0 if logging_obj: - logging_obj.warning("Missing prompt_eval_count and no encoding provided; defaulted to 0.") + require_attributes(logging_obj, _WarningLogger, "warning").warning( + "Missing prompt_eval_count and no encoding provided; defaulted to 0." + ) model_response.object = "list" model_response.data = output_data diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index c5488bc8a98..26246cfd7b6 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -103,7 +103,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): delivers_ended_stream_rewrites = True assembles_streamed_response = True - def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: + def get_structured_messages(self, data: dict, *, keep_tool_errors: bool = False) -> list[AllMessageValues] | None: """ Convert chat completions request data to OpenAI-spec structured messages. diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 739966c95d1..2cf48a75524 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -9,6 +9,7 @@ from collections.abc import Mapping, Sequence from typing import Final import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger @@ -18,6 +19,8 @@ from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, ) +_DECODED_JSON: Final = TypeAdapter(object) + class OpenAICountTokensHandler(OpenAICountTokensConfig): """ @@ -34,7 +37,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): tools: Sequence[Mapping[str, object]] | None = None, instructions: str | None = None, text: Mapping[str, object] | None = None, - ) -> Mapping[str, object]: + ) -> object: """ Handle a token counting request to OpenAI's Responses API. @@ -67,7 +70,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): api_key: str, api_base: str | None = None, timeout: float | httpx.Timeout | None = None, - ) -> Mapping[str, object]: + ) -> object: try: endpoint_url: Final = self.get_openai_count_tokens_endpoint(api_base) @@ -96,7 +99,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): message=error_text, ) - openai_response: Final = response.json() + openai_response: Final = _DECODED_JSON.validate_python(response.json()) verbose_logger.debug("OpenAI response: %s", openai_response) return openai_response diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index 87a75dae95e..c82a2982a54 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -272,7 +272,7 @@ class OpenAITokenCounter(BaseTokenCounter): @staticmethod async def _counted( resolved: _CountCredential, - send: Callable[[str, str | None], Awaitable[Mapping[str, object]]], + send: Callable[[str, str | None], Awaitable[object]], model_to_use: str, request_model: str, ) -> TokenCountResponse | None: diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 45f44c5e72f..ba357219bca 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -78,6 +78,7 @@ from litellm.llms.openai.responses.guardrail_translation.tool_merge import ( append_returned_tools, merge_guardrailed_tools, ) +from litellm.responses.additional_tools import additional_tools_in_input from litellm.responses.litellm_completion_transformation.compaction_items import ( compaction_summary, decode_compaction_item, @@ -776,7 +777,7 @@ class OpenAIResponsesHandler(BaseTranslation): delivers_ended_stream_rewrites = True assembles_streamed_response = True - def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: + def get_structured_messages(self, data: dict, *, keep_tool_errors: bool = False) -> list[AllMessageValues] | None: """ Convert Responses API request data to OpenAI-spec structured messages. @@ -989,9 +990,14 @@ class OpenAIResponsesHandler(BaseTranslation): def extract_request_tool_names(self, data: Mapping[str, object]) -> list[str]: """Extract tool names from Responses API request (tools[].name for function - and custom, tools[].server_label for mcp).""" + and custom, tools[].server_label for mcp), including tools carried by + ``additional_tools`` input items.""" names: Final[list[str]] = [] - for tool in list_field(data, "tools"): + request_tools: Final[tuple[object, ...]] = ( + *list_field(data, "tools"), + *additional_tools_in_input(data.get("input")), + ) + for tool in request_tools: if not isinstance(tool, dict): continue if tool.get("type") in ("function", "custom") and tool.get("name"): diff --git a/litellm/llms/openai/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index 0560e1c2e3a..abd157a9a08 100644 --- a/litellm/llms/openai/videos/transformation.py +++ b/litellm/llms/openai/videos/transformation.py @@ -9,6 +9,7 @@ import httpx from httpx._types import FileContent, FileTypes, RequestFiles from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base +from litellm.llms.base_llm.videos.async_task_utils import unsupported_video_fields from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.openai.image_edit.transformation import ImageEditRequestUtils from litellm.llms.openai.workload_identity import resolve_openai_bearer_token @@ -66,9 +67,16 @@ class OpenAIVideoConfig(BaseVideoConfig): video_create_optional_params: VideoCreateOptionalRequestParams, model: str, drop_params: bool, - ) -> dict: + ) -> dict[str, object]: """No mapping applied since inputs are in OpenAI spec already""" - return dict(video_create_optional_params) + unsupported: Final = unsupported_video_fields( + video_create_optional_params, + model, + "openai", + drop_params, + frozenset({"seed", "duration", "aspect_ratio", "generate_audio"}), + ) + return {key: value for key, value in video_create_optional_params.items() if key not in unsupported} def validate_environment( self, diff --git a/litellm/llms/recraft/image_edit/transformation.py b/litellm/llms/recraft/image_edit/transformation.py index 74bfa52c5e0..45b02a84379 100644 --- a/litellm/llms/recraft/image_edit/transformation.py +++ b/litellm/llms/recraft/image_edit/transformation.py @@ -1,10 +1,12 @@ from io import BufferedReader -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, runtime_checkable import httpx from httpx._types import RequestFiles +from pydantic import TypeAdapter from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str @@ -21,6 +23,24 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + + +@runtime_checkable +class _SupportsStrKeyLookup(Protocol): + def __getitem__(self, key: str, /) -> object: ... + + +def _json_field(value: object, key: str) -> object: + if isinstance(value, _SupportsStrKeyLookup): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + + +def _image_object(image_data: object) -> ImageObject: + fields: Final = decoded_json_object(image_data) + return ImageObject(url=fields.get("url", None), b64_json=fields.get("b64_json", None)) + class RecraftImageEditConfig(BaseImageEditConfig): DEFAULT_BASE_URL: str = "https://external.api.recraft.ai" @@ -152,7 +172,7 @@ class RecraftImageEditConfig(BaseImageEditConfig): ) -> ImageResponse: model_response: Final = ImageResponse() try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image edit response: {e}", @@ -162,12 +182,7 @@ class RecraftImageEditConfig(BaseImageEditConfig): if not model_response.data: model_response.data = [] - for image_data in response_data["data"]: - model_response.data.append( - ImageObject( - url=image_data.get("url", None), - b64_json=image_data.get("b64_json", None), - ) - ) + for image_data in iter_items(_json_field(response_data, "data")): + model_response.data.append(_image_object(image_data)) return model_response diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index ad8a24c1535..7ad76e62a96 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -11,6 +11,7 @@ import litellm from litellm.constants import RUNWAYML_DEFAULT_API_VERSION from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.videos.async_task_utils import unsupported_video_fields from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -155,12 +156,23 @@ class RunwayMLVideoConfig(BaseVideoConfig): - size -> ratio (convert "WIDTHxHEIGHT" to "WIDTH:HEIGHT") - seconds -> duration (convert to integer) """ + unsupported: Final = unsupported_video_fields( + video_create_optional_params, + model, + "runwayml", + drop_params, + frozenset({"seed", "duration", "aspect_ratio", "generate_audio"}), + ) supported_openai_params: Final = self.get_supported_openai_params(model) return { **self._prompt_image_param(video_create_optional_params), **self._ratio_param(video_create_optional_params), **self._duration_param(video_create_optional_params), - **{key: value for key, value in video_create_optional_params.items() if key not in supported_openai_params}, + **{ + key: value + for key, value in video_create_optional_params.items() + if key not in supported_openai_params and key not in unsupported + }, } @staticmethod diff --git a/litellm/llms/sap/credentials.py b/litellm/llms/sap/credentials.py index 74d065956cc..38a7cfedd34 100644 --- a/litellm/llms/sap/credentials.py +++ b/litellm/llms/sap/credentials.py @@ -3,7 +3,7 @@ from __future__ import annotations import json import os import tempfile -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone from pathlib import Path @@ -16,6 +16,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object, iter_items from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_httpx_client @@ -74,12 +75,12 @@ def _get_nested(d: object, path: Sequence[str]) -> object: return cur -def _load_json_env(var_name: str) -> Mapping[str, Sequence[Mapping[str, object]]] | None: +def _load_json_env(var_name: str) -> object: raw: Final = os.environ.get(var_name) if not raw: return None try: - return json.loads(raw) + return _DECODED_JSON.validate_python(json.loads(raw)) except json.JSONDecodeError: return None @@ -91,15 +92,23 @@ def _str_or_none(value: object) -> str | None: return None -def _load_vcap() -> Mapping[str, Sequence[Mapping[str, object]]]: +def _load_vcap() -> object: return _load_json_env(VCAP_SERVICES_ENV_VAR) or {} +def _vcap_service_groups(vcap: object) -> Iterable[object]: + services_by_name: Final = as_str_mapping(vcap) + if services_by_name is not None: + return services_by_name.values() + raise AttributeError(f"'{type(vcap).__name__}' object has no attribute 'values'", name="values", obj=vcap) + + def _get_vcap_service(label: str) -> Mapping[str, object] | None: - for services in _load_vcap().values(): - for svc in services: - if svc.get("label") == label: - return svc + for services in _vcap_service_groups(_load_vcap()): + for svc in iter_items(services): + service = decoded_json_object(svc) + if service.get("label") == label: + return service return None diff --git a/litellm/llms/snowflake/embedding/transformation.py b/litellm/llms/snowflake/embedding/transformation.py index 45bec1283a3..c5e61ad5573 100644 --- a/litellm/llms/snowflake/embedding/transformation.py +++ b/litellm/llms/snowflake/embedding/transformation.py @@ -1,9 +1,10 @@ from collections.abc import Mapping -from typing import Final +from typing import Final, Protocol, runtime_checkable import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -14,6 +15,29 @@ from litellm.types.utils import EmbeddingResponse from ..utils import SnowflakeBaseConfig, SnowflakeException _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) + + +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, key: str | int, /) -> object: ... + + +@runtime_checkable +class _SupportsSetItem(Protocol): + def __setitem__(self, key: str, value: object, /) -> None: ... + + +def _subscript(value: object, key: str | int) -> object: + if isinstance(value, _SupportsGetItem): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + + +def _item_assignment_target(value: object) -> _SupportsSetItem: + if isinstance(value, _SupportsSetItem): + return value + raise TypeError(f"'{type(value).__name__}' object does not support item assignment") class SnowflakeEmbeddingConfig(SnowflakeBaseConfig, BaseEmbeddingConfig): @@ -54,10 +78,10 @@ class SnowflakeEmbeddingConfig(SnowflakeBaseConfig, BaseEmbeddingConfig): optional_params: dict, litellm_params: dict, ) -> EmbeddingResponse: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) # convert embeddings to 1d array - for item in response_json["data"]: - item["embedding"] = item["embedding"][0] + for item in iter_items(_subscript(response_json, "data")): + _item_assignment_target(item)["embedding"] = _subscript(_subscript(item, "embedding"), 0) returned_response: Final = EmbeddingResponse.model_validate(_JSON_OBJECT.validate_python(response_json)) returned_response.model = "snowflake/" + (returned_response.model or "") diff --git a/litellm/llms/soniox/realtime/transformation.py b/litellm/llms/soniox/realtime/transformation.py index ed7590b6f39..1bd0176b468 100644 --- a/litellm/llms/soniox/realtime/transformation.py +++ b/litellm/llms/soniox/realtime/transformation.py @@ -164,16 +164,19 @@ def _ignored_settings(update: TranscriptionSessionUpdate) -> tuple[str, ...]: def build_start_request( model: str, update: TranscriptionSessionUpdate | None, options: SonioxRealtimeOptions, drop_params: bool = False ) -> SonioxStartRequest: - return SonioxStartRequest( - **options.model_dump(exclude={"language_hints"}), - language_hints=_language_hints( - None if update is None else update.language, options.language_hints, options.language_hints_strict - ), - model=model, - audio_format="pcm_s16le", - sample_rate=_sample_rate(update), - num_channels=1, - enable_endpoint_detection=update is None or _uses_endpoint_detection(update, drop_params), + configured_options: Final[Mapping[str, object]] = options.model_dump(exclude={"language_hints"}) + return SonioxStartRequest.model_validate( + { + **configured_options, + "language_hints": _language_hints( + None if update is None else update.language, options.language_hints, options.language_hints_strict + ), + "model": model, + "audio_format": "pcm_s16le", + "sample_rate": _sample_rate(update), + "num_channels": 1, + "enable_endpoint_detection": update is None or _uses_endpoint_detection(update, drop_params), + } ) diff --git a/litellm/llms/valkey/vector_stores/transformation.py b/litellm/llms/valkey/vector_stores/transformation.py index cd95d9c95b7..bb9b2da3dbf 100644 --- a/litellm/llms/valkey/vector_stores/transformation.py +++ b/litellm/llms/valkey/vector_stores/transformation.py @@ -15,6 +15,7 @@ import httpx from pydantic import ConfigDict import litellm +from litellm.litellm_core_utils.core_helpers import attribute_items from litellm.llms.base_llm.vector_store.transformation import ( BaseDirectVectorStoreConfig, VectorStoreEmbeddingExecutor, @@ -33,7 +34,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from redis import Redis from redis.asyncio import Redis as AsyncRedis - from redis.commands.search.document import Document from redis.commands.search.query import Query from redis.commands.search.result import Result @@ -185,7 +185,7 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig): return query_cls(knn_expr).return_fields(text_field, DISTANCE_FIELD_NAME).paging(0, k).dialect(2) @staticmethod - def _to_result(doc: "Document", text_field: str) -> VectorStoreSearchResult: + def _to_result(doc: object, text_field: str) -> VectorStoreSearchResult: content: Final = [VectorStoreResultContent(text=str(getattr(doc, text_field, "")), type="text")] return VectorStoreSearchResult( score=1.0 - float(getattr(doc, DISTANCE_FIELD_NAME)), @@ -196,7 +196,7 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig): @classmethod def _to_response(cls, search_result: "Result", query_text: str, text_field: str) -> VectorStoreSearchResponse: - docs: Final = getattr(search_result, "docs", None) or () + docs: Final = attribute_items(search_result, "docs") data: Final = sorted( (cls._to_result(doc, text_field) for doc in docs), key=lambda result: result.get("score") or 0.0, diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index e04d741f9ee..6701900cf6b 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1444,7 +1444,7 @@ class VertexAITokenCounter(BaseTokenCounter): async def count_tokens( self, model_to_use: str, - messages: list[dict[str, Any]] | None, + messages: list[dict[str, object]] | None, contents: list[dict[str, object]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 158ed979bad..1f03de4625c 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -1299,7 +1299,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _map_openai_to_vertex_params( self, - openai_request_body: dict[str, Any], + openai_request_body: dict[str, object], deployment_drop_params: object = None, ) -> dict[str, object]: """ diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 9999c75a91c..6e801ed5756 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -8,7 +8,7 @@ import os import re from collections.abc import Callable, Mapping, Sequence from functools import partial -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast from urllib.parse import quote import httpx @@ -18,7 +18,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import asyncify -from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params, has_attributes from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.prompt_templates.factory import ( GeminiMediaSource, @@ -271,8 +271,21 @@ def _image_url_payload_may_need_sync_gcs_metadata_fetch( return _gs_uri_requires_content_type_metadata(url) +class _SupportsModelDump(Protocol): + def model_dump(self, *, exclude_none: bool) -> object: ... + + +def _message_fields(raw: object) -> Mapping[object, object] | None: + if isinstance(raw, dict): + return raw + if not has_attributes(raw, _SupportsModelDump, "model_dump"): + return None + dumped: Final = raw.model_dump(exclude_none=False) + return dumped if isinstance(dumped, dict) else None + + def _openai_messages_may_need_sync_media_fetch( - messages: list[AllMessageValues], + messages: Sequence[object], custom_llm_provider: str = "vertex_ai", model: str | None = None, ) -> bool: @@ -283,10 +296,7 @@ def _openai_messages_may_need_sync_media_fetch( model is None or VertexGeminiConfig.supports_multimodal_function_responses(model) ) for raw in messages: - msg: Any = raw - if not isinstance(msg, dict) and hasattr(msg, "model_dump"): - msg = msg.model_dump(exclude_none=False) - if not isinstance(msg, dict): + if (msg := _message_fields(raw)) is None: continue images_field = msg.get("images") if isinstance(images_field, list): diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index 87b0dfdbcbc..894c33b5fbc 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -4,7 +4,7 @@ import os from collections.abc import Mapping, Sequence from io import BufferedRandom, BufferedReader, BytesIO from pathlib import Path -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import httpx from httpx._types import RequestFiles @@ -13,6 +13,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, map_vertex_image_size @@ -39,6 +40,18 @@ class _VertexImagenPredictResponse(TypedDict, total=False): predictions: ReadOnly[Sequence[object]] +class _SupportsRead(Protocol): + def read(self) -> bytes | str: ... + + +def _image_items(image: object) -> Sequence[object] | None: + return image if isinstance(image, (list, tuple)) else None + + +def _image_fields(image: object) -> Mapping[object, object] | None: + return image if isinstance(image, dict) else None + + class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): """ Vertex AI Imagen Image Edit Configuration @@ -293,30 +306,30 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): return reference_images - def _read_all_bytes(self, image: Any, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> bytes: + def _read_all_bytes(self, image: object, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> bytes: if depth > max_depth: raise ValueError( f"Max recursion depth {max_depth} reached while reading image bytes for Vertex AI Imagen image edit." ) - if isinstance(image, (list, tuple)): - for item in image: + if (items := _image_items(image)) is not None: + for item in items: if item is not None: return self._read_all_bytes(item, depth=depth + 1, max_depth=max_depth) raise ValueError("Unsupported image type for Vertex AI Imagen image edit.") - if isinstance(image, dict): + if (fields := _image_fields(image)) is not None: for key in ("data", "bytes", "content"): - if key in image and image[key] is not None: - value = image[key] + if key in fields and fields[key] is not None: + value = fields[key] if isinstance(value, str): try: return base64.b64decode(value) except Exception: continue return self._read_all_bytes(value, depth=depth + 1, max_depth=max_depth) - if "path" in image: - return self._read_all_bytes(image["path"], depth=depth + 1, max_depth=max_depth) + if "path" in fields: + return self._read_all_bytes(fields["path"], depth=depth + 1, max_depth=max_depth) if isinstance(image, bytes): return image @@ -350,7 +363,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): "Unsupported image input: filesystem paths are not accepted for " "Vertex AI Imagen image edit. Provide image bytes or a file-like object." ) - if hasattr(image, "read"): + if has_attributes(image, _SupportsRead, "read"): data = image.read() if isinstance(data, str): data = data.encode("utf-8") diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index 8d39ae76456..28d6321453f 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -12,7 +12,7 @@ from typing import Any, Final from pydantic import JsonValue, TypeAdapter from litellm.anthropic_beta_headers_manager import aload_beta_headers_config, filter_and_transform_beta_headers -from litellm.litellm_core_utils.core_helpers import drops_unsupported_params +from litellm.litellm_core_utils.core_helpers import decoded_json_object, drops_unsupported_params from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.common_utils import ( OBJECT_LIST_ADAPTER, @@ -34,6 +34,7 @@ _COUNT_TOKENS_PUBLISHERS: Final = ( _REQUEST_MESSAGES: Final = TypeAdapter(Sequence[object]) _CONTEXT_MANAGEMENT: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_DECODED_JSON: Final = TypeAdapter(object) def count_tokens_request_body( @@ -226,7 +227,7 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): raise ValueError(f"Token counting request failed with status {response.status_code}: {error_text}") # Parse response - result: Final = response.json() + result: Final = decoded_json_object(_DECODED_JSON.validate_python(response.json())) # Return token count # Vertex AI Anthropic returns: {"input_tokens": 123} diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index 080899d08b0..2b88861b79b 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -9,16 +9,18 @@ import base64 import time from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, ClassVar, Final, TypedDict, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, SupportsFloat, SupportsIndex, TypedDict, cast import httpx from httpx._types import FileContent, RequestFiles -from typing_extensions import ReadOnly +from typing_extensions import Buffer, ReadOnly import litellm from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.base_llm.videos.async_task_utils import unsupported_video_fields from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.vertex_ai.common_utils import ( convert_vertex_datetime_to_openai_datetime, @@ -80,16 +82,16 @@ def veo_video_count_from_parameters(parameters: Mapping[str, object]) -> int | N def _build_vertex_video_usage_from_request_data( - request_data: dict[str, Any] | None, + request_data: Mapping[str, object] | None, ) -> dict[str, float | str]: """Build usage metadata (duration, resolution, video count) for video cost calculation.""" usage_data: Final[dict[str, float | str]] = {} if not request_data: return usage_data - parameters: Final = request_data.get("parameters", {}) + parameters: Final = decoded_json_object(request_data.get("parameters", {})) duration: Final = parameters.get("durationSeconds") or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS - if duration is not None: + if isinstance(duration, (str, Buffer, SupportsFloat, SupportsIndex)): try: usage_data["duration_seconds"] = float(duration) except (ValueError, TypeError): @@ -197,6 +199,13 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): skipped if ``resolution`` is already set - seconds → durationSeconds (defaults to 4 seconds if not provided) """ + _ = unsupported_video_fields( + video_create_optional_params, + model, + "vertex_ai", + drop_params, + frozenset({"duration", "seed", "aspect_ratio", "generate_audio"}), + ) mapped_params: Final[dict[str, object]] = {} # Map input_reference to image (will be processed in transform_video_create_request) diff --git a/litellm/llms/volcengine/videos/transformation.py b/litellm/llms/volcengine/videos/transformation.py index 9d6dba15826..17627c963ec 100644 --- a/litellm/llms/volcengine/videos/transformation.py +++ b/litellm/llms/volcengine/videos/transformation.py @@ -16,6 +16,7 @@ from litellm.llms.base_llm.videos.async_task_utils import ( input_reference_url, polled_video_models, routed_video_models, + video_image_params, ) from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_httpx.http_handler import ( @@ -546,6 +547,9 @@ class VolcEngineVideoConfig(BaseVideoConfig): "extra_headers", "content", "parameters", + "frame_images", + "input_references", + "aspect_ratio", *_CREATE_PASSTHROUGH_KEYS, ] @@ -555,7 +559,19 @@ class VolcEngineVideoConfig(BaseVideoConfig): model: str, drop_params: bool, ) -> dict[str, object]: # mutable-ok: BaseVideoConfig contract - requested: Final = _without(video_create_optional_params, _additional_drop_keys(video_create_optional_params)) + requested: Final = video_image_params( + _without(video_create_optional_params, _additional_drop_keys(video_create_optional_params)), + model, + "volcengine", + drop_params, + ) + if ( + not drop_params + and requested.get("duration") is not None + and requested.get("seconds") is not None + and _duration_param(requested["duration"]) != _duration_param(requested["seconds"]) + ): + raise _droppable(model, "duration and seconds must agree") unsupported: Final = next( (key for key, value in requested.items() if key in _UNSUPPORTED_BY_VOLCENGINE and value is not None), None, @@ -563,13 +579,17 @@ class VolcEngineVideoConfig(BaseVideoConfig): if unsupported is not None and not drop_params: raise _droppable(model, f"'{unsupported}' is not supported by Volcengine video generation") passthrough: Final = MappingProxyType( - {key: value for key, value in requested.items() if value is not None and key not in _MAPPED_ELSEWHERE} + { + key: value + for key, value in requested.items() + if value is not None and key not in _MAPPED_ELSEWHERE and key != "duration" + } ) return { **_user_fields(requested.get("user"), model, drop_params), **_size_fields(requested.get("size"), model, drop_params), - **_duration_fields(requested.get("seconds"), model, drop_params), **passthrough, + **_duration_fields(requested.get("duration", requested.get("seconds")), model, drop_params), **_parameters_fields(requested.get("parameters"), model, drop_params), } diff --git a/litellm/llms/xai/image_generation/cost_calculator.py b/litellm/llms/xai/image_generation/cost_calculator.py new file mode 100644 index 00000000000..5cf575ae743 --- /dev/null +++ b/litellm/llms/xai/image_generation/cost_calculator.py @@ -0,0 +1,70 @@ +from collections.abc import Mapping +from typing import Final + +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.llm_cost_calc.utils import get_cost_per_unit, resolve_image_model_info +from litellm.llms.xai.common_utils import xai_reported_cost_in_usd +from litellm.types.utils import ImageResponse, ModelInfo + +_PARAMS: Final = TypeAdapter(Mapping[str, object]) + + +def _request_param(optional_params: Mapping[str, object] | None, key: str) -> object: + if optional_params is None: + return None + if key in optional_params: + return optional_params[key] + extra_body: Final = optional_params.get("extra_body") + return _PARAMS.validate_python(extra_body).get(key) if isinstance(extra_body, Mapping) else None + + +def cost_calculator( + model: str, + image_response: ImageResponse, + n: int, + quality: str | None = None, + optional_params: Mapping[str, object] | None = None, + call_type: str | None = None, + model_info: ModelInfo | None = None, +) -> float | None: + if model_info is not None and model_info.get("input_cost_per_pixel") is not None: + return None + + caller_set_prices: Final = model_info is not None and any( + model_info.get(key) is not None + for key in ( + "input_cost_per_image", + "output_cost_per_image", + "output_cost_per_image_1K", + "output_cost_per_image_2K", + ) + ) + billed_cost: Final = xai_reported_cost_in_usd(getattr(image_response.usage, "cost_in_usd_ticks", None)) + if billed_cost is not None and not caller_set_prices: + return billed_cost + + requested_quality: Final = quality or image_response.quality or _request_param(optional_params, "quality") or "auto" + resolution: Final = _request_param(optional_params, "resolution") or "1k" + is_edit: Final = call_type in ("image_edit", "aimage_edit") + low_quality: Final = requested_quality == "low" or (requested_quality == "auto" and not is_edit) + if low_quality and resolution in ("1k", "2k"): + try: + pricing: Final = resolve_image_model_info(model, "xai", model_info) + except Exception: # noqa: BLE001 # model lookup raises bare Exception for unmapped models + return None + resolution_price: Final = get_cost_per_unit(pricing, f"output_cost_per_image_{str(resolution).upper()}", None) + deployment_flat_price: Final = next( + ( + get_cost_per_unit(model_info, key, None) + for key in ("input_cost_per_image", "output_cost_per_image") + if model_info is not None and model_info.get(key) is not None + ), + None, + ) + if deployment_flat_price is not None: + return deployment_flat_price * n + if resolution_price is not None: + return resolution_price * n + + return None diff --git a/litellm/main.py b/litellm/main.py index 6b9964a450f..1d18100d1d5 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -40,6 +40,7 @@ import openai from openai import AsyncStream, Stream from openai.types.audio import TranscriptionStreamEvent from pydantic import BaseModel, TypeAdapter, ValidationError +from typing_extensions import TypedDict as ExtraItemsTypedDict from typing_extensions import assert_never, overload import litellm @@ -407,6 +408,10 @@ def _resolve_control_options(kwargs: Mapping[str, object], model: str) -> Contro return assert_never(control) +class _ForwardedCompletionParams(ExtraItemsTypedDict, total=False, extra_items=object): + pass + + @dataclass(frozen=True, slots=True) class _SupplementalProviderParams: params: GenericLiteLLMParams @@ -4393,6 +4398,14 @@ def _converse_native_output_fallback( return None if AmazonConverseConfig.sends_native_text_format(fallback.optional_params) else fallback +def _buffered_stream_fallback( + first_attempt_params: Mapping[str, object], first_attempt_litellm_params: Mapping[str, object] +) -> NativeOutputFallback: + return NativeOutputFallback( + optional_params={**first_attempt_params, "fake_stream": True}, litellm_params=first_attempt_litellm_params + ) + + def _native_output_attempt_timeout( ctx: CompletionDispatchContext, fallback: NativeOutputFallback | None ) -> float | str | httpx.Timeout | None: @@ -4496,11 +4509,14 @@ def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResu ) converse_params: Final = _OPTIONAL_PARAMS_ADAPTER.validate_python(optional_params) + converse_uses_native_output: Final = AmazonConverseConfig().uses_native_output( + converse_model, converse_params, litellm_params + ) response = call_with_native_output_fallback( converse_attempt, ( partial(_converse_native_output_fallback, ctx, converse_params, dict(litellm_params)) - if AmazonConverseConfig().uses_native_output(converse_model, converse_params, litellm_params) + if converse_uses_native_output else None ), converse_model, @@ -4509,6 +4525,11 @@ def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResu if drops_unsupported_params(litellm_params.get("drop_params")) else None ), + buffered_stream_fallback=( + partial(_buffered_stream_fallback, converse_params, dict(litellm_params)) + if converse_uses_native_output + else None + ), ) elif bedrock_route == "converse_like": model = model.replace("converse_like/", "") @@ -4551,15 +4572,20 @@ def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResu ) invoke_params: Final = _OPTIONAL_PARAMS_ADAPTER.validate_python(optional_params) + invoke_uses_native_output: Final = isinstance( + provider_config, AmazonAnthropicClaudeConfig + ) and provider_config.uses_native_output(model, invoke_params, litellm_params) response = call_with_native_output_fallback( invoke_attempt, - ( - partial(_native_output_fallback, ctx, invoke_params, dict(litellm_params)) - if isinstance(provider_config, AmazonAnthropicClaudeConfig) - and provider_config.uses_native_output(model, invoke_params, litellm_params) + partial(_native_output_fallback, ctx, invoke_params, dict(litellm_params)) + if invoke_uses_native_output + else None, + model, + buffered_stream_fallback=( + partial(_buffered_stream_fallback, invoke_params, dict(litellm_params)) + if invoke_uses_native_output else None ), - model, ) return response @@ -8060,7 +8086,7 @@ def text_completion( text_completion_response = TextCompletionResponse() - optional_params: Final[dict[str, Any]] = {} + optional_params: Final[dict[str, object]] = {} # default values for all optional params are none, litellm only passes them to the llm when they are set to non None values if best_of is not None: optional_params["best_of"] = best_of @@ -8127,7 +8153,7 @@ def text_completion( def process_prompt(i, individual_prompt): decoded_prompt: Final = tokenizer.decode(individual_prompt) - all_params: Final = {**kwargs, **optional_params} + all_params: Final = _ForwardedCompletionParams(**{**kwargs, **optional_params}) response: Final[TextCompletionResponse] = text_completion( model=model, prompt=decoded_prompt, @@ -8201,7 +8227,7 @@ def text_completion( messages=messages, *args, **kwargs, - **optional_params, + **_ForwardedCompletionParams(**optional_params), ) if kwargs.get("acompletion", False) is True: return response @@ -8545,7 +8571,7 @@ def transcription( ) -> ( TranscriptionResponse | Stream[TranscriptionStreamEvent] - | Coroutine[Any, Any, TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]] + | Coroutine[object, object, TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]] ): """ Calls openai + azure whisper endpoints. @@ -8634,7 +8660,7 @@ def transcription( TranscriptionResponse | Stream[TranscriptionStreamEvent] | AsyncStream[TranscriptionStreamEvent] - | Coroutine[Any, Any, TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]] + | Coroutine[object, object, TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]] ) provider_config: Final = ProviderConfigManager.get_provider_audio_transcription_config( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index aa358f4f6f3..77244a32f6f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1961,6 +1961,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2041,6 +2042,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2116,6 +2118,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2242,7 +2245,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2287,6 +2291,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2331,7 +2336,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2376,6 +2382,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2420,7 +2427,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2465,6 +2473,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2504,7 +2513,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2549,6 +2559,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2593,7 +2604,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2633,6 +2645,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2715,7 +2728,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -2759,7 +2773,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -2803,7 +2818,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -2847,7 +2863,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -2891,7 +2908,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -3023,6 +3041,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -3068,6 +3087,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -3113,6 +3133,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -3158,6 +3179,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -3203,6 +3225,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -12097,16 +12120,19 @@ }, "azure/tts-1": { "deprecation_date": "2026-12-15", + "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "input_cost_per_character": 1.5e-05, "litellm_provider": "azure", - "mode": "audio_speech" + "mode": "audio_speech", + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule" }, "azure/tts-1-hd": { "deprecation_date": "2026-12-15", + "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "input_cost_per_character": 3.3e-05, "litellm_provider": "azure", "mode": "audio_speech", - "source": "https://azure.microsoft.com/en-us/pricing/details/azure-openai/" + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule" }, "azure/us/gpt-4.1-2025-04-14": { "deprecation_date": "2027-04-14", @@ -12859,11 +12885,13 @@ "supports_vision": true }, "azure/whisper-1": { - "deprecation_date": "2026-12-15", + "deprecation_date": "2027-02-26", + "deprecation_source": "https://management.azure.com/subscriptions/a88257a5-3323-4ef9-95bf-5681a46c9624/providers/Microsoft.CognitiveServices/locations/eastus/models?api-version=2024-10-01", "input_cost_per_second": 0.0001, "litellm_provider": "azure", "mode": "audio_transcription", - "output_cost_per_second": 0.0001 + "output_cost_per_second": 0.0001, + "source": "https://management.azure.com/subscriptions/a88257a5-3323-4ef9-95bf-5681a46c9624/providers/Microsoft.CognitiveServices/locations/eastus/models?api-version=2024-10-01" }, "azure_ai/Cohere-embed-v3-english": { "deprecation_date": "2027-02-15", @@ -13837,12 +13865,12 @@ "azure_ai/jais-30b-chat": { "deprecation_date": "2026-01-30", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", - "input_cost_per_token": 0.0032, + "input_cost_per_token": 3.2e-06, "litellm_provider": "azure_ai", "max_input_tokens": 8192, "max_output_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.00971, + "output_cost_per_token": 9.71e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", "supports_pdf_input": false }, @@ -31630,7 +31658,7 @@ "input_cost_per_audio_token_batches": 2.5e-07 }, "gemini/gemini-embedding-2-preview": { - "deprecation_date": "2026-08-10", + "deprecation_date": "2026-11-06", "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", "input_cost_per_audio_token": 6.5e-06, "input_cost_per_audio_token_batches": 3.25e-06, @@ -31961,7 +31989,8 @@ "rpm": 1000.0, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", + "successor_model": "gemini/gemini-nano-banana-2.1", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -32516,7 +32545,7 @@ "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, "rpm": 2000, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -32562,6 +32591,7 @@ "input_cost_per_token_flex": 2.5e-07, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, + "successor_model": "gemini/gemini-3.6-flash", "supports_audio_input": true, "supports_computer_use": true }, @@ -32578,7 +32608,7 @@ "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, "rpm": 2000, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -32625,7 +32655,8 @@ "output_cost_per_token_batches": 4.5e-06, "input_cost_per_token_flex": 7.5e-07, "output_cost_per_token_flex": 4.5e-06, - "cache_read_input_token_cost_flex": 8e-08 + "cache_read_input_token_cost_flex": 8e-08, + "successor_model": "gemini/gemini-3.6-flash" }, "gemini/gemini-3.6-flash": { "prompt_cache_min_tokens": 4096, @@ -32705,7 +32736,7 @@ "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, "rpm": 2000, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -32748,7 +32779,8 @@ }, "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014 + "google_maps_grounding_cost_per_query": 0.014, + "successor_model": "gemini/gemini-3.8-flash" }, "gemini/gemini-3.8-flash": { "prompt_cache_min_tokens": 4096, @@ -33610,7 +33642,7 @@ "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 0, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_modalities": [ "text" ], @@ -33625,7 +33657,8 @@ "supports_system_messages": false, "supports_vision": false, "supports_web_search": false, - "output_cost_per_image": 0.08 + "output_cost_per_image": 0.08, + "successor_model": "gemini/lyria-3.5" }, "gemini/veo-3.1-fast-generate-preview": { "deprecation_date": "2026-10-22", @@ -44089,6 +44122,7 @@ "supports_vision": true }, "mistral/mistral-large-2512": { + "deprecation_date": "2027-04-30", "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 5e-07, "litellm_provider": "mistral", @@ -44164,6 +44198,7 @@ "supports_vision": true }, "mistral/mistral-medium-3-5": { + "deprecation_date": "2027-04-30", "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 1.5e-06, "litellm_provider": "mistral", @@ -55051,14 +55086,16 @@ "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference" }, "wandb/deepseek-ai/DeepSeek-V3.1": { + "cache_read_input_token_cost": 2e-07, "supports_reasoning": true, "max_input_tokens": 128000, "max_output_tokens": 128000, - "input_cost_per_token": 5.5e-07, - "output_cost_per_token": 1.65e-06, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 1.5e-06, "litellm_provider": "wandb", "mode": "chat", - "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference" + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", + "supports_prompt_caching": true }, "wandb/meta-llama/Llama-3.3-70B-Instruct": { "max_input_tokens": 128000, @@ -55130,8 +55167,8 @@ "watsonx/core42/jais-13b-chat": { "max_input_tokens": 8192, "max_output_tokens": 8192, - "input_cost_per_token": 0.0005, - "output_cost_per_token": 0.002, + "input_cost_per_token": 1.8e-06, + "output_cost_per_token": 1.8e-06, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -64183,6 +64220,39 @@ "supports_tool_choice": true, "supports_vision": false }, + "tencent/hy4": { + "cache_read_input_token_cost": 4.2e-08, + "comment": "TokenHub bills Hy4 at 50% of list until 2026-12-31 23:59:59 Beijing time; these are the list rates.", + "input_cost_per_token": 8.34e-07, + "input_cost_per_token_cache_hit": 4.2e-08, + "litellm_provider": "tencent", + "max_input_tokens": 960000, + "max_output_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.501e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true + }, + "tencent/hy4-preview": { + "cache_read_input_token_cost": 4.2e-08, + "comment": "TokenHub bills Hy4 preview at 50% of list until 2026-12-31 23:59:59 Beijing time; these are the list rates.", + "input_cost_per_token": 8.34e-07, + "input_cost_per_token_cache_hit": 4.2e-08, + "litellm_provider": "tencent", + "max_input_tokens": 960000, + "max_output_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.501e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true + }, "tencent/minimax-m3": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 6e-08, @@ -64726,6 +64796,7 @@ "source": "https://docs.mistral.ai/models/model-cards/mistral-moderation-26-03" }, "mistral/voxtral-mini-2602": { + "deprecation_date": "2027-04-30", "input_cost_per_second": 5e-05, "litellm_provider": "mistral", "mode": "audio_transcription", @@ -64742,6 +64813,7 @@ "supports_audio_input": true }, "mistral/voxtral-mini-transcribe-realtime-2602": { + "deprecation_date": "2027-04-30", "input_cost_per_second": 0.0001, "litellm_provider": "mistral", "mode": "audio_transcription", @@ -64758,6 +64830,7 @@ "supports_audio_input": true }, "mistral/voxtral-mini-tts-2603": { + "deprecation_date": "2027-04-30", "input_cost_per_character": 1.6e-05, "litellm_provider": "mistral", "mode": "audio_speech", @@ -66431,13 +66504,13 @@ "max_output_tokens": 262144, "max_input_tokens": 262000, "input_cost_per_token": 1.3e-07, - "output_cost_per_token": 2.8e-07, - "cache_read_input_token_cost": 7e-08, + "output_cost_per_token": 2.7e-07, + "cache_read_input_token_cost": 5e-08, "supports_prompt_caching": true, "litellm_provider": "wandb", "mode": "chat", "supports_vision": false, - "source": "https://wandb.ai/site/pricing/tokens/" + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference" }, "wandb/deepseek-ai/DeepSeek-V4-Pro": { "supports_reasoning": true, @@ -66514,14 +66587,14 @@ "supports_reasoning": true, "max_output_tokens": 262144, "max_input_tokens": 262000, - "input_cost_per_token": 7.1e-07, - "output_cost_per_token": 3.5e-06, - "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 3e-06, + "cache_read_input_token_cost": 1.2e-07, "supports_prompt_caching": true, "litellm_provider": "wandb", "mode": "chat", "supports_vision": true, - "source": "https://wandb.ai/site/pricing/tokens/" + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference" }, "wandb/moonshotai/Kimi-K2.6": { "supports_reasoning": true, @@ -66587,15 +66660,17 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/Qwen/Qwen3.6-35B-A3B": { + "cache_read_input_token_cost": 3.5e-08, "supports_reasoning": true, "max_output_tokens": 262144, "max_input_tokens": 262000, - "input_cost_per_token": 2.5e-07, - "output_cost_per_token": 1.25e-06, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 8.9e-07, "litellm_provider": "wandb", "mode": "chat", "supports_vision": true, - "source": "https://wandb.ai/site/pricing/tokens/" + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", + "supports_prompt_caching": true }, "wandb/Qwen/Qwen3.6-27B": { "supports_reasoning": true, @@ -66649,18 +66724,18 @@ "litellm_provider": "wandb", "mode": "chat", "supports_reasoning": true, - "input_cost_per_token": 1e-07, + "input_cost_per_token": 6e-08, "output_cost_per_token": 1.5e-07, - "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost": 1.5e-08, "max_input_tokens": 131000, "supports_prompt_caching": true, - "source": "https://wandb.ai/site/pricing/tokens/" + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference" }, "wandb/zai-org/GLM-5.2": { "supports_reasoning": true, "max_output_tokens": 262144, "max_input_tokens": 262000, - "input_cost_per_token": 7.6e-07, + "input_cost_per_token": 6.2e-07, "output_cost_per_token": 2.42e-06, "cache_read_input_token_cost": 1.4e-07, "supports_prompt_caching": true, @@ -67955,7 +68030,9 @@ "input_cost_per_image": 0.06, "litellm_provider": "xai", "mode": "image_generation", - "source": "https://docs.x.ai/docs/models", + "output_cost_per_image_1K": 0.04, + "output_cost_per_image_2K": 0.06, + "source": "https://docs.x.ai/developers/migration/imagine-image-quality-nov-2", "supported_endpoints": [ "/v1/images/generations" ], @@ -70849,6 +70926,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token_above_272k_tokens": 1e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, "source": "https://docs.perplexity.ai/docs/agent-api/models" }, "perplexity/openai/gpt-5.4": { @@ -74215,14 +74293,14 @@ ] }, "wandb/deepseek-ai/DeepSeek-V4.1-Flash": { - "cache_read_input_token_cost": 3e-08, - "input_cost_per_token": 2e-07, + "cache_read_input_token_cost": 2e-09, + "input_cost_per_token": 1.3e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "max_output_tokens": 1048576, "mode": "chat", "output_cost_per_token": 6.5e-07, - "source": "https://wandb.ai/site/pricing/tokens/", + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_prompt_caching": true, "supports_reasoning": true, "supports_vision": true @@ -74243,13 +74321,13 @@ "supports_vision": true }, "wandb/zai-org/GLM-5.3-Flash": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 1.5e-07, + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 1.2e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "mode": "chat", - "output_cost_per_token": 5e-07, - "source": "https://docs.wandb.ai/inference/models.md", + "output_cost_per_token": 3.8e-07, + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -75544,7 +75622,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -75584,7 +75663,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -75624,6 +75704,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -75697,6 +75778,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -76995,13 +77077,13 @@ "supports_vision": true }, "azure/whisper": { - "deprecation_date": "2026-12-15", - "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", + "deprecation_date": "2027-02-26", + "deprecation_source": "https://management.azure.com/subscriptions/a88257a5-3323-4ef9-95bf-5681a46c9624/providers/Microsoft.CognitiveServices/locations/eastus/models?api-version=2024-10-01", "input_cost_per_second": 0.0001, "litellm_provider": "azure", "mode": "audio_transcription", "output_cost_per_second": 0.0001, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://management.azure.com/subscriptions/a88257a5-3323-4ef9-95bf-5681a46c9624/providers/Microsoft.CognitiveServices/locations/eastus/models?api-version=2024-10-01" }, "gemini/gemini-nano-banana-2.1": { "input_cost_per_token": 1.5e-06, @@ -77412,6 +77494,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -77469,6 +77552,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -77646,6 +77730,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -77698,6 +77783,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -77745,6 +77831,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -77858,6 +77945,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -78605,13 +78693,13 @@ "supports_vision": true }, "wandb/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2.5e-07, + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 1.5e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "max_output_tokens": 943718, "mode": "chat", - "output_cost_per_token": 1.5e-06, + "output_cost_per_token": 5e-07, "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_function_calling": true, "supports_prompt_caching": true, @@ -78907,12 +78995,12 @@ "supports_vision": true }, "wandb/reflection/Beam-1.0-501B-A23B": { - "cache_read_input_token_cost": 1e-07, - "input_cost_per_token": 1e-06, + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 3e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "mode": "chat", - "output_cost_per_token": 3e-06, + "output_cost_per_token": 1e-06, "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_function_calling": true, "supports_prompt_caching": true, @@ -78923,11 +79011,11 @@ }, "wandb/zai-org/GLM-5.3": { "cache_read_input_token_cost": 1.9e-07, - "input_cost_per_token": 1.2e-06, + "input_cost_per_token": 7.5e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "mode": "chat", - "output_cost_per_token": 4e-06, + "output_cost_per_token": 3e-06, "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_function_calling": true, "supports_prompt_caching": true, @@ -83064,13 +83152,16 @@ "mode": "chat", "max_input_tokens": 1000000, "max_output_tokens": 393216, - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 8e-07, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-1-flash-1017" + "source": "https://empiriolabs.ai/models/deepseek-v4-1-flash-1017", + "cache_read_input_token_cost": 2e-08, + "supports_prompt_caching": true, + "supports_function_calling": true }, "empiriolabs/deepseek-v4-flash-0731": { "litellm_provider": "empiriolabs", @@ -83082,7 +83173,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731" + "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-flash-0731:variant1": { "litellm_provider": "empiriolabs", @@ -83094,7 +83186,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731" + "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro": { "litellm_provider": "empiriolabs", @@ -83106,7 +83199,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro-0813": { "litellm_provider": "empiriolabs", @@ -83118,7 +83212,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro-0813" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro-0813", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro:variant1": { "litellm_provider": "empiriolabs", @@ -83136,7 +83231,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro:variant2": { "litellm_provider": "empiriolabs", @@ -83156,7 +83252,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro:variant3": { "litellm_provider": "empiriolabs", @@ -83168,7 +83265,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/fugu-max": { "litellm_provider": "empiriolabs", @@ -83366,7 +83464,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/glm-5-2" + "source": "https://empiriolabs.ai/models/glm-5-2", + "supports_function_calling": true }, "empiriolabs/glm-5-2:variant2": { "litellm_provider": "empiriolabs", @@ -83384,7 +83483,8 @@ "search_context_size_high": 0.013 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/glm-5-2" + "source": "https://empiriolabs.ai/models/glm-5-2", + "supports_function_calling": true }, "empiriolabs/glm-5-3": { "litellm_provider": "empiriolabs", @@ -83438,7 +83538,8 @@ "supports_reasoning": true, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-6" + "source": "https://empiriolabs.ai/models/kimi-k2-6", + "supports_function_calling": true }, "empiriolabs/kimi-k2-7-code": { "litellm_provider": "empiriolabs", @@ -83457,7 +83558,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-7-code" + "source": "https://empiriolabs.ai/models/kimi-k2-7-code", + "supports_function_calling": true }, "empiriolabs/kimi-k2-7-code-highspeed": { "litellm_provider": "empiriolabs", @@ -83476,7 +83578,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-7-code-highspeed" + "source": "https://empiriolabs.ai/models/kimi-k2-7-code-highspeed", + "supports_function_calling": true }, "empiriolabs/kimi-k2-7-code:variant1": { "litellm_provider": "empiriolabs", @@ -83497,7 +83600,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-7-code" + "source": "https://empiriolabs.ai/models/kimi-k2-7-code", + "supports_function_calling": true }, "empiriolabs/kimi-k3": { "litellm_provider": "empiriolabs", @@ -83516,7 +83620,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k3" + "source": "https://empiriolabs.ai/models/kimi-k3", + "supports_function_calling": true }, "empiriolabs/kimi-k3-1-pro": { "litellm_provider": "empiriolabs", @@ -83558,7 +83663,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/mimo-v2-6-flash" + "source": "https://empiriolabs.ai/models/mimo-v2-6-flash", + "supports_function_calling": true }, "empiriolabs/mimo-v2-6-pro": { "litellm_provider": "empiriolabs", @@ -83578,7 +83684,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/mimo-v2-6-pro" + "source": "https://empiriolabs.ai/models/mimo-v2-6-pro", + "supports_function_calling": true }, "empiriolabs/mimo-v2-6-pro-ultraspeed": { "litellm_provider": "empiriolabs", @@ -83592,7 +83699,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/mimo-v2-6-pro-ultraspeed" + "source": "https://empiriolabs.ai/models/mimo-v2-6-pro-ultraspeed", + "supports_function_calling": true }, "empiriolabs/minimax-m2-7": { "litellm_provider": "empiriolabs", @@ -83707,7 +83815,8 @@ "search_context_size_high": 0.14 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/mistral-small-4" + "source": "https://empiriolabs.ai/models/mistral-small-4", + "supports_function_calling": true }, "empiriolabs/muse-glimmer-30b": { "litellm_provider": "empiriolabs", @@ -83722,7 +83831,8 @@ "supports_reasoning": true, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/muse-glimmer-30b" + "source": "https://empiriolabs.ai/models/muse-glimmer-30b", + "supports_function_calling": true }, "empiriolabs/muse-spark-1-1": { "litellm_provider": "empiriolabs", @@ -83788,7 +83898,8 @@ "supports_system_messages": true, "supports_vision": true, "supports_prompt_caching": true, - "source": "https://empiriolabs.ai/models/nova-lite-1-0" + "source": "https://empiriolabs.ai/models/nova-lite-1-0", + "supports_function_calling": true }, "empiriolabs/nova-lite-2": { "litellm_provider": "empiriolabs", @@ -83802,7 +83913,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_prompt_caching": true, - "source": "https://empiriolabs.ai/models/nova-lite-2" + "source": "https://empiriolabs.ai/models/nova-lite-2", + "supports_function_calling": true }, "empiriolabs/nova-micro-1-0": { "litellm_provider": "empiriolabs", @@ -83814,7 +83926,8 @@ "cache_read_input_token_cost": 2.24e-08, "supports_system_messages": true, "supports_prompt_caching": true, - "source": "https://empiriolabs.ai/models/nova-micro-1-0" + "source": "https://empiriolabs.ai/models/nova-micro-1-0", + "supports_function_calling": true }, "empiriolabs/nova-pro-1-0": { "litellm_provider": "empiriolabs", @@ -83825,7 +83938,8 @@ "output_cost_per_token": 9.6e-06, "supports_system_messages": true, "supports_vision": true, - "source": "https://empiriolabs.ai/models/nova-pro-1-0" + "source": "https://empiriolabs.ai/models/nova-pro-1-0", + "supports_function_calling": true }, "empiriolabs/perplexity-advanced-deep-research": { "litellm_provider": "empiriolabs", @@ -83925,7 +84039,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b" + "source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-27b": { "litellm_provider": "empiriolabs", @@ -83946,7 +84061,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-27b" + "source": "https://empiriolabs.ai/models/qwen3-5-27b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-35b-a3b": { "litellm_provider": "empiriolabs", @@ -83967,7 +84083,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b" + "source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-397b-a17b": { "litellm_provider": "empiriolabs", @@ -83988,7 +84105,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b" + "source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-4b": { "litellm_provider": "empiriolabs", @@ -84039,7 +84157,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-flash" + "source": "https://empiriolabs.ai/models/qwen3-5-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-5-flash:variant1": { "litellm_provider": "empiriolabs", @@ -84062,7 +84181,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-flash" + "source": "https://empiriolabs.ai/models/qwen3-5-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-5-omni-flash": { "litellm_provider": "empiriolabs", @@ -84080,7 +84200,8 @@ "search_context_size_high": 0.015 }, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/qwen3-5-omni-flash" + "source": "https://empiriolabs.ai/models/qwen3-5-omni-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-5-omni-plus": { "litellm_provider": "empiriolabs", @@ -84099,7 +84220,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/qwen3-5-omni-plus" + "source": "https://empiriolabs.ai/models/qwen3-5-omni-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-5-plus": { "litellm_provider": "empiriolabs", @@ -84108,8 +84230,8 @@ "max_output_tokens": 65536, "input_cost_per_token": 3.6e-07, "output_cost_per_token": 2.21e-06, - "input_cost_per_token_above_256k_tokens": 1.08e-06, - "output_cost_per_token_above_256k_tokens": 6.62e-06, + "input_cost_per_token_above_256k_tokens": 4.5e-07, + "output_cost_per_token_above_256k_tokens": 2.76e-06, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, @@ -84120,7 +84242,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-plus" + "source": "https://empiriolabs.ai/models/qwen3-5-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-5-plus:variant1": { "litellm_provider": "empiriolabs", @@ -84143,7 +84266,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-plus" + "source": "https://empiriolabs.ai/models/qwen3-5-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-6-27b": { "litellm_provider": "empiriolabs", @@ -84162,7 +84286,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-27b" + "source": "https://empiriolabs.ai/models/qwen3-6-27b", + "supports_function_calling": true }, "empiriolabs/qwen3-6-35b-a3b": { "litellm_provider": "empiriolabs", @@ -84176,7 +84301,8 @@ "supports_reasoning": true, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-35b-a3b" + "source": "https://empiriolabs.ai/models/qwen3-6-35b-a3b", + "supports_function_calling": true }, "empiriolabs/qwen3-6-flash": { "litellm_provider": "empiriolabs", @@ -84197,7 +84323,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-flash" + "source": "https://empiriolabs.ai/models/qwen3-6-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-6-flash:variant1": { "litellm_provider": "empiriolabs", @@ -84218,7 +84345,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-flash" + "source": "https://empiriolabs.ai/models/qwen3-6-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-6-plus": { "litellm_provider": "empiriolabs", @@ -84239,7 +84367,8 @@ "search_context_size_high": 0.026 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-plus" + "source": "https://empiriolabs.ai/models/qwen3-6-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-6-plus:variant1": { "litellm_provider": "empiriolabs", @@ -84260,7 +84389,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-plus" + "source": "https://empiriolabs.ai/models/qwen3-6-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-7-flash": { "litellm_provider": "empiriolabs", @@ -84287,7 +84417,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-flash" + "source": "https://empiriolabs.ai/models/qwen3-7-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-7-flash:variant1": { "litellm_provider": "empiriolabs", @@ -84314,7 +84445,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-flash" + "source": "https://empiriolabs.ai/models/qwen3-7-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-7-max": { "litellm_provider": "empiriolabs", @@ -84332,7 +84464,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-max" + "source": "https://empiriolabs.ai/models/qwen3-7-max", + "supports_function_calling": true }, "empiriolabs/qwen3-7-max:variant1": { "litellm_provider": "empiriolabs", @@ -84350,7 +84483,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-max" + "source": "https://empiriolabs.ai/models/qwen3-7-max", + "supports_function_calling": true }, "empiriolabs/qwen3-7-plus": { "litellm_provider": "empiriolabs", @@ -84371,7 +84505,8 @@ "search_context_size_high": 0.03 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-plus" + "source": "https://empiriolabs.ai/models/qwen3-7-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-7-plus:variant1": { "litellm_provider": "empiriolabs", @@ -84395,7 +84530,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-plus" + "source": "https://empiriolabs.ai/models/qwen3-7-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-8-27b": { "litellm_provider": "empiriolabs", @@ -84430,7 +84566,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-flash" + "source": "https://empiriolabs.ai/models/qwen3-8-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-8-max": { "litellm_provider": "empiriolabs", @@ -84449,7 +84586,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-max" + "source": "https://empiriolabs.ai/models/qwen3-8-max", + "supports_function_calling": true }, "empiriolabs/qwen3-8-max-0902": { "litellm_provider": "empiriolabs", @@ -84468,7 +84606,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-max-0902" + "source": "https://empiriolabs.ai/models/qwen3-8-max-0902", + "supports_function_calling": true }, "empiriolabs/qwen3-8-max:variant1": { "litellm_provider": "empiriolabs", @@ -84487,7 +84626,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-max" + "source": "https://empiriolabs.ai/models/qwen3-8-max", + "supports_function_calling": true }, "empiriolabs/qwen3-8-omni-flash": { "litellm_provider": "empiriolabs", @@ -84507,7 +84647,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/qwen3-8-omni-flash" + "source": "https://empiriolabs.ai/models/qwen3-8-omni-flash", + "supports_function_calling": true }, "empiriolabs/seed-2-0-code": { "litellm_provider": "empiriolabs", @@ -84521,7 +84662,8 @@ "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, - "source": "https://empiriolabs.ai/models/seed-2-0-code" + "source": "https://empiriolabs.ai/models/seed-2-0-code", + "supports_function_calling": true }, "empiriolabs/seed-2-0-lite": { "litellm_provider": "empiriolabs", @@ -84536,7 +84678,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-0-lite" + "source": "https://empiriolabs.ai/models/seed-2-0-lite", + "supports_function_calling": true }, "empiriolabs/seed-2-0-mini": { "litellm_provider": "empiriolabs", @@ -84551,7 +84694,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-0-mini" + "source": "https://empiriolabs.ai/models/seed-2-0-mini", + "supports_function_calling": true }, "empiriolabs/seed-2-0-pro": { "litellm_provider": "empiriolabs", @@ -84566,7 +84710,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-0-pro" + "source": "https://empiriolabs.ai/models/seed-2-0-pro", + "supports_function_calling": true }, "empiriolabs/seed-2-1-turbo": { "litellm_provider": "empiriolabs", @@ -84579,7 +84724,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-1-turbo" + "source": "https://empiriolabs.ai/models/seed-2-1-turbo", + "supports_function_calling": true }, "empiriolabs/skylark-embedding-vision": { "litellm_provider": "empiriolabs", @@ -86235,7 +86381,7 @@ "input_cost_per_token": 2e-06, "litellm_provider": "baseten", "max_input_tokens": 1000000, - "max_output_tokens": 262144, + "max_output_tokens": 1048576, "mode": "chat", "output_cost_per_token": 1e-05, "source": "https://inference.baseten.co/v1/models", @@ -86570,5 +86716,170 @@ "supports_tool_choice": true, "supports_vision": true, "token_count_framing": 100.0 + }, + "moonshotai.kimi-k3.1": { + "cache_creation_input_token_cost": 5.5e-06, + "cache_creation_input_token_cost_flex": 2.75e-06, + "cache_creation_input_token_cost_priority": 9.625e-06, + "cache_read_input_token_cost": 2.2e-07, + "cache_read_input_token_cost_flex": 1.1e-07, + "cache_read_input_token_cost_priority": 3.85e-07, + "input_cost_per_token": 4.4e-06, + "input_cost_per_token_flex": 2.2e-06, + "input_cost_per_token_priority": 7.7e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 868928, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.2e-05, + "output_cost_per_token_flex": 1.1e-05, + "output_cost_per_token_priority": 3.85e-05, + "reasoning_effort_levels": [ + "none", + "low", + "medium", + "high", + "xhigh", + "max" + ], + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k3-1-highspeed.md", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_none_reasoning_effort": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_regex_lookaround": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "us.moonshotai.kimi-k3.1": { + "cache_creation_input_token_cost": 5.5e-06, + "cache_creation_input_token_cost_flex": 2.75e-06, + "cache_creation_input_token_cost_priority": 9.625e-06, + "cache_read_input_token_cost": 2.2e-07, + "cache_read_input_token_cost_flex": 1.1e-07, + "cache_read_input_token_cost_priority": 3.85e-07, + "input_cost_per_token": 4.4e-06, + "input_cost_per_token_flex": 2.2e-06, + "input_cost_per_token_priority": 7.7e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 868928, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.2e-05, + "output_cost_per_token_flex": 1.1e-05, + "output_cost_per_token_priority": 3.85e-05, + "reasoning_effort_levels": [ + "none", + "low", + "medium", + "high", + "xhigh", + "max" + ], + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k3-1-highspeed.md", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_none_reasoning_effort": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_regex_lookaround": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.moonshotai.kimi-k3.1": { + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_flex": 2.5e-06, + "cache_creation_input_token_cost_priority": 8.75e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_flex": 1e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_flex": 2e-06, + "input_cost_per_token_priority": 7e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 868928, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-05, + "output_cost_per_token_flex": 1e-05, + "output_cost_per_token_priority": 3.5e-05, + "reasoning_effort_levels": [ + "none", + "low", + "medium", + "high", + "xhigh", + "max" + ], + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k3-1-highspeed.md", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_none_reasoning_effort": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_regex_lookaround": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "in.moonshotai.kimi-k3.1": { + "cache_creation_input_token_cost": 5.5e-06, + "cache_creation_input_token_cost_flex": 2.75e-06, + "cache_creation_input_token_cost_priority": 9.625e-06, + "cache_read_input_token_cost": 2.2e-07, + "cache_read_input_token_cost_flex": 1.1e-07, + "cache_read_input_token_cost_priority": 3.85e-07, + "input_cost_per_token": 4.4e-06, + "input_cost_per_token_flex": 2.2e-06, + "input_cost_per_token_priority": 7.7e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 868928, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.2e-05, + "output_cost_per_token_flex": 1.1e-05, + "output_cost_per_token_priority": 3.85e-05, + "reasoning_effort_levels": [ + "none", + "low", + "medium", + "high", + "xhigh", + "max" + ], + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k3-1-highspeed.md", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_none_reasoning_effort": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_regex_lookaround": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "wandb/deepseek-ai/DeepSeek-V4.1-Flash-1017": { + "cache_read_input_token_cost": 2e-09, + "input_cost_per_token": 1.3e-07, + "litellm_provider": "wandb", + "max_input_tokens": 1049000, + "max_output_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6.5e-07, + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": true } } diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 01c22dc6725..680b0e6ede4 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -371,8 +371,8 @@ class LazyFeatureMiddleware: fastapi_app.state.lazy_locks = {} @property - def _loaded(self) -> set: - return self._fastapi_app.state.lazy_loaded + def _loaded(self) -> set[str]: + return _lazy_loaded(self._fastapi_app) async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: # Short-circuit once every feature has loaded. diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 5cb63c98d56..0182c27fac4 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -9633,7 +9633,7 @@ "paths": { "/backend-api/codex/{endpoint}": { "get": { - "description": "Forward a client's ChatGPT subscription with separate LiteLLM authentication.", + "description": "Forward a client's ChatGPT subscription with separate LiteLLM authentication, listing and serving the\nkey's other chat models through LiteLLM.", "operationId": "codex_passthrough_route_backend_api_codex__endpoint__get", "parameters": [ { @@ -9681,7 +9681,7 @@ ] }, "post": { - "description": "Forward a client's ChatGPT subscription with separate LiteLLM authentication.", + "description": "Forward a client's ChatGPT subscription with separate LiteLLM authentication, listing and serving the\nkey's other chat models through LiteLLM.", "operationId": "codex_passthrough_route_backend_api_codex__endpoint__post", "parameters": [ { @@ -9724,7 +9724,7 @@ } } }, - "description": "Unmodified Codex response or server-sent events" + "description": "Codex response, the model catalog with the key's LiteLLM models, or server-sent events" }, "422": { "content": { @@ -16540,6 +16540,22 @@ "description": "Character budget for each text sent to the guardrail's model. Implemented by guardrail='decision_model'.", "title": "Max Input Chars" }, + "mcp_block_behavior": { + "anyOf": [ + { + "enum": [ + "tool_result", + "raise" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "description": "How to handle guardrail blocks when executing MCP tools for an LLM request. 'raise' propagates the guardrail error to the client. 'tool_result' or None returns the block as a tool result and lets the LLM continue.", + "title": "Mcp Block Behavior" + }, "model": { "anyOf": [ { @@ -16597,7 +16613,7 @@ "type": "null" } ], - "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", + "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. Bedrock Nova Sonic cannot cancel a spoken turn's answer, so a blocked spoken turn always ends its session. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", "title": "On Violation" }, "only_scan_new_messages": { @@ -16673,7 +16689,7 @@ "type": "null" } ], - "description": "The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set.", + "description": "The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to the guardrail's block reason if not set.", "title": "Realtime Violation Message" }, "run_in_parallel": { @@ -19116,6 +19132,22 @@ "description": "Character budget for each text sent to the guardrail's model. Implemented by guardrail='decision_model'.", "title": "Max Input Chars" }, + "mcp_block_behavior": { + "anyOf": [ + { + "enum": [ + "tool_result", + "raise" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "description": "How to handle guardrail blocks when executing MCP tools for an LLM request. 'raise' propagates the guardrail error to the client. 'tool_result' or None returns the block as a tool result and lets the LLM continue.", + "title": "Mcp Block Behavior" + }, "metadata": { "anyOf": [ { @@ -19271,7 +19303,7 @@ "type": "null" } ], - "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", + "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. Bedrock Nova Sonic cannot cancel a spoken turn's answer, so a blocked spoken turn always ends its session. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", "title": "On Violation" }, "only_scan_new_messages": { @@ -19692,7 +19724,7 @@ "type": "null" } ], - "description": "The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set.", + "description": "The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to the guardrail's block reason if not set.", "title": "Realtime Violation Message" }, "rules": { @@ -49950,6 +49982,10 @@ ], "title": "Content" }, + "is_error": { + "title": "Is Error", + "type": "boolean" + }, "role": { "const": "tool", "title": "Role", diff --git a/litellm/proxy/admin_mcp.py b/litellm/proxy/admin_mcp.py index 32576ea8113..55b0777a27c 100644 --- a/litellm/proxy/admin_mcp.py +++ b/litellm/proxy/admin_mcp.py @@ -71,8 +71,12 @@ def _admin_route_allowed(auth: UserAPIKeyAuth, request: Request, method: str, pa return True +def _scope_app(scope: Mapping[str, object]) -> object: + return scope["app"] + + def _admin_routes(request: Request) -> Iterator[tuple[tuple[str, str], APIRoute]]: - app: Final = request.app + app: Final = _scope_app(request.scope) if not isinstance(app, FastAPI): return for route in app.routes: diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 42cf44a3c95..2bd9b4357d7 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -69,7 +69,7 @@ from litellm.proxy.agent_endpoints.kill_switch import ( from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user -from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity +from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity, unless_database_fails from litellm.proxy.utils import PrismaClient, get_custom_url from litellm.repositories.chunked_in import find_many_in from litellm.types.agents import ( @@ -394,10 +394,12 @@ async def resolve_agent_daily_activity_scope( prisma_client=prisma_client, ) - agent_records: Final = ( - await agents_table(prisma_client).find_many(where={}) + agent_records: Final = await unless_database_fails( + agents_table(prisma_client).find_many(where={}) if resolved_agent_ids is None - else await find_many_in(agents_table(prisma_client), "agent_id", resolved_agent_ids) + else find_many_in(agents_table(prisma_client), "agent_id", resolved_agent_ids), + (), + "its agent names", ) agent_metadata: Final[Mapping[str, dict[str, object]]] = MappingProxyType( {agent.agent_id: {"agent_name": agent.agent_name} for agent in agent_records} diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 11ca3d9083c..b40eb451979 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -180,6 +180,7 @@ from litellm.repositories.table_repositories import ( ) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository +from litellm.responses.additional_tools import additional_tools_in_input from litellm.router import Router from litellm.router_api import RegistryReader from litellm.types.proxy.auth.auth_checks import UserNotFoundError @@ -1037,9 +1038,6 @@ async def check_tools_allowlist( call_types: Final = get_call_types_for_route(route) if not call_types or not any(ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types): return - tool_names: Final = extract_request_tool_names(route, request_body) - if not tool_names: - return key_meta: Final = (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {} team_meta: Final = (valid_token.team_metadata or {}) if isinstance(valid_token.team_metadata, dict) else {} key_allowed: Final = key_meta.get("allowed_tools") @@ -1047,6 +1045,9 @@ async def check_tools_allowlist( effective: Final = key_allowed if (isinstance(key_allowed, list) and len(key_allowed) > 0) else team_allowed if not isinstance(effective, list) or len(effective) == 0: return + tool_names: Final = extract_request_tool_names(route, request_body) + if not tool_names: + return allowed_set: Final = {str(t) for t in effective} disallowed: Final = [n for n in tool_names if n not in allowed_set] if disallowed: @@ -7988,6 +7989,10 @@ def _tools(tools: object) -> tuple[object, ...]: return () +def _request_tools(request_body: Mapping[str, object]) -> tuple[object, ...]: + return (*_tools(request_body.get("tools")), *additional_tools_in_input(request_body.get("input"))) + + def _tool_vector_store_ids(tool: object) -> tuple[str, ...]: try: tool_fields: Final = _TOOL_ADAPTER.validate_python(tool, strict=True) @@ -7998,14 +8003,14 @@ def _tool_vector_store_ids(tool: object) -> tuple[str, ...]: def _strict_requested_vector_store_ids(request_body: Mapping[str, object]) -> tuple[str, ...]: """ - Same fields VectorStoreRegistry.get_vector_store_ids_to_run reads, but a vector_store_ids that is - not a list of strings is a 400 instead of being skipped or iterated, and tools that are not - objects name no store. + Same fields VectorStoreRegistry.get_vector_store_ids_to_run reads, plus the tools of `additional_tools` + input items, but a vector_store_ids that is not a list of strings is a 400 instead of being skipped or + iterated, and tools that are not objects name no store. """ return tuple( chain( _validated_vector_store_ids(request_body.get("vector_store_ids")), - chain.from_iterable(_tool_vector_store_ids(tool) for tool in _tools(request_body.get("tools"))), + chain.from_iterable(_tool_vector_store_ids(tool) for tool in _request_tools(request_body)), ) ) @@ -8044,7 +8049,7 @@ def _legacy_requested_vector_store_ids( return tuple( chain( top_level, - chain.from_iterable(vector_store_ids_in_tool(tool) for tool in _tools(request_body.get("tools"))), + chain.from_iterable(vector_store_ids_in_tool(tool) for tool in _request_tools(request_body)), ) ) @@ -8113,8 +8118,8 @@ async def vector_store_access_check( """ Checks whether the caller may use every vector store the request names. - Requested stores come from `vector_store_ids`, `tools[].vector_store_ids` and the RAG - `retrieval_config.vector_store_id`. Grants come from each identity's + Requested stores come from `vector_store_ids`, `tools[].vector_store_ids`, the tools of + `additional_tools` input items and the RAG `retrieval_config.vector_store_id`. Grants come from each identity's `object_permission.vector_stores`. Both modes read the same request fields, with or without a vector store registry, and answer 400 diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 2b5fc6d877e..24602dec677 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -863,7 +863,7 @@ async def pre_db_read_auth_checks( Raises: - HTTPException if request fails initial auth checks """ - from litellm.proxy.proxy_server import general_settings, llm_router, premium_user + from litellm.proxy.proxy_server import general_settings, llm_router # Check 1. request size await check_if_request_size_is_safe(request=request) @@ -891,19 +891,12 @@ async def pre_db_read_auth_checks( ) # Check 4. Check if request route is an allowed route on the proxy - if "allowed_routes" in general_settings: - _allowed_routes: Final = general_settings["allowed_routes"] - if premium_user is not True: - verbose_proxy_logger.error( - "Trying to set allowed_routes. This is an Enterprise feature. %s", - CommonProxyErrors.not_premium_user.value, - ) - if route not in _allowed_routes: - verbose_proxy_logger.error("Route %s not in allowed_routes=%s", route, _allowed_routes) - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"Access forbidden: Route {route} not allowed", - ) + if not RouteChecks.is_allowed_on_proxy(route): + verbose_proxy_logger.error("Route %s not in allowed_routes=%s", route, general_settings["allowed_routes"]) + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"Access forbidden: Route {route} not allowed", + ) def route_in_additonal_public_routes(current_route: str): diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 9bcfa18fc51..1a7356fb035 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -192,34 +192,25 @@ def get_team_models( - Empty list if no models set - If model_access_groups is provided, only return models that are in the access groups """ - all_models_set: Final[set[str]] = set() - if len(team_models) > 0: - all_models_set.update(team_models) - if SpecialModelNames.all_team_models.value in all_models_set: - all_models_set.update(team_models) - # GH#30619: expand all-team-models sentinel - # to the actual proxy model list - all_models_set.discard(SpecialModelNames.all_team_models.value) - all_models_set.update(proxy_model_list) - if include_model_access_groups: - all_models_set.update(model_access_groups.keys()) - if SpecialModelNames.all_proxy_models.value in all_models_set: - all_models_set.update(proxy_model_list) - if include_model_access_groups: - all_models_set.update(model_access_groups.keys()) - - all_models = _get_models_from_access_groups( + expands_to_proxy_models: Final = ( + SpecialModelNames.all_team_models.value in team_models + or SpecialModelNames.all_proxy_models.value in team_models + ) + expansion: Final = ( + (*proxy_model_list, *(model_access_groups if include_model_access_groups else ())) + if expands_to_proxy_models + else () + ) + named_models: Final = (model for model in team_models if model != SpecialModelNames.all_team_models.value) + all_models: Final = _get_models_from_access_groups( model_access_groups=model_access_groups, - all_models=list(all_models_set), + all_models=list(dict.fromkeys((*named_models, *expansion))), include_model_access_groups=include_model_access_groups, proxy_model_list=proxy_model_list, ) - # deduplicate while preserving order - all_models = list(dict.fromkeys(all_models)) - verbose_proxy_logger.debug("ALL TEAM MODELS - %s", len(all_models)) - return all_models + return list(dict.fromkeys(all_models)) def get_complete_model_list( diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 125bfab4edf..0aa2528d46b 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -439,6 +439,19 @@ class RouteChecks: detail=f"user not allowed to access this route. Route={route} is an admin only route", ) + @staticmethod + def is_allowed_on_proxy(route: str) -> bool: + from litellm.proxy.proxy_server import general_settings, premium_user + + if "allowed_routes" not in general_settings: + return True + if premium_user is not True: + verbose_proxy_logger.error( + "Trying to set allowed_routes. This is an Enterprise feature. %s", + CommonProxyErrors.not_premium_user.value, + ) + return route in general_settings["allowed_routes"] + @staticmethod def is_unbilled_route(route: str) -> bool: resource, _, action = route.rsplit("/", 1)[-1].partition(":") diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ba253b44ea3..a52ad14e293 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -309,6 +309,17 @@ async def _normalize_claude_model( rewrite_request_model(request_data, request, source) +_CODEX_NATIVE_RESPONSES_ROUTE: Final = "/backend-api/codex/responses" + + +def _is_codex_native_litellm_turn(route: str, model: object, llm_router: litellm.Router, team_id: str | None) -> bool: + if route != _CODEX_NATIVE_RESPONSES_ROUTE: + return False + from litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints import routes_to_litellm + + return routes_to_litellm(model, llm_router, team_id) + + async def _resolve_router_settings_model_group_alias( request_data: dict[str, object], # mutable-ok: the request body is rewritten in place for every downstream reader valid_token: UserAPIKeyAuth, @@ -325,9 +336,12 @@ async def _resolve_router_settings_model_group_alias( if request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is True: return request.scope[MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY] = True - if request_dispatched_to_pass_through_endpoint(request) or request_dispatched_to_provider_pass_through(request): - return requested: Final = request_data.get("model") + if request_dispatched_to_pass_through_endpoint(request) or ( + request_dispatched_to_provider_pass_through(request) + and not _is_codex_native_litellm_turn(route, requested, llm_router, valid_token.team_id) + ): + return if not isinstance(requested, str) or await read_raw_json_body(request=request) is None: return settings: Final = await proxy_config.get_hierarchical_router_settings( @@ -338,6 +352,11 @@ async def _resolve_router_settings_model_group_alias( target: Final = resolve_model_group_alias(settings.get("model_group_alias"), requested) if target is None or target == requested: return + if request_dispatched_to_provider_pass_through(request): + from litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints import native_turn_serves_through_litellm + + if not await native_turn_serves_through_litellm(requested, llm_router, request, valid_token): + return verbose_proxy_logger.debug( "router_settings.model_group_alias resolved %s -> %s before auth", requested.replace("\r", "").replace("\n", ""), @@ -347,6 +366,27 @@ async def _resolve_router_settings_model_group_alias( rewrite_request_model(request_data, request, target) +def _resolve_codex_native_chatgpt_model( + request_data: dict[str, object], # mutable-ok: the request body is rewritten in place for every downstream reader + valid_token: UserAPIKeyAuth, + request: Request | None, + route: str, +) -> None: + requested: Final = request_data.get("model") + if route != _CODEX_NATIVE_RESPONSES_ROUTE or request is None or not isinstance(requested, str): + return + from litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints import native_chatgpt_model + from litellm.proxy.proxy_server import llm_router + + if llm_router is None: + return + target: Final = native_chatgpt_model(requested, llm_router, valid_token) + if target is None: + return + request.scope.setdefault(CLIENT_REQUESTED_MODEL_SCOPE_KEY, requested) + rewrite_request_model(request_data, request, target) + + def _get_model_names_for_budget_checks( model: str | list[str] | None, ) -> list[str]: @@ -4172,6 +4212,7 @@ async def _authorize_authenticated_request( RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request) await _normalize_claude_model(request_data, user_api_key_auth_obj, request, route) await _resolve_router_settings_model_group_alias(request_data, user_api_key_auth_obj, request, route) + _resolve_codex_native_chatgpt_model(request_data, user_api_key_auth_obj, request, route) # Single authorization point. Builder paths MUST NOT call common_checks. # Route through the same exception handler the builder uses so @@ -4650,6 +4691,7 @@ async def enforce_key_and_fallback_model_access( """ await _normalize_claude_model(request_data, valid_token, request, route) await _resolve_router_settings_model_group_alias(request_data, valid_token, request, route) + _resolve_codex_native_chatgpt_model(request_data, valid_token, request, route) config: Final = valid_token.config if config != {}: diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 1e10d628c1e..512a9166e31 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -9,13 +9,14 @@ import os from collections.abc import Mapping, MutableMapping from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Final, Literal, Protocol, cast from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest +from litellm.litellm_core_utils.core_helpers import attribute_or_none, has_attributes, iter_items from litellm.litellm_core_utils.hidden_params import get_hidden_params, set_hidden_param from litellm.llms.openai.workload_identity import resolve_openai_workload_identity_config from litellm.proxy._types import * @@ -82,6 +83,18 @@ if TYPE_CHECKING: router: Final = APIRouter() +class _UserBatchLister(Protocol): + async def list_user_batches( + self, + user_api_key_dict: UserAPIKeyAuth, + limit: int | None = None, + after: str | None = None, + provider: str | None = None, + target_model_names: str | None = None, + llm_router: Router | None = None, + ) -> Mapping[str, object]: ... + + def _hidden_param_string(hidden_params: Mapping[str, object], key: str) -> str: value: Final = hidden_params.get(key) return value if isinstance(value, str) else "" @@ -917,9 +930,9 @@ async def list_batches( # Try to use managed objects table for listing batches (returns encoded IDs). managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") - if managed_files_obj is not None and hasattr(managed_files_obj, "list_user_batches"): + if managed_files_obj is not None and has_attributes(managed_files_obj, _UserBatchLister, "list_user_batches"): verbose_proxy_logger.debug("Using managed objects table for batch listing") - response = await cast(Any, managed_files_obj).list_user_batches( + response = await managed_files_obj.list_user_batches( user_api_key_dict=user_api_key_dict, limit=limit, after=after, @@ -949,9 +962,9 @@ async def list_batches( # Encode batch IDs in the list response so clients can use # them for retrieve/cancel/file downloads through the proxy. - response_data: Final = getattr(response, "data", None) + response_data: Final = attribute_or_none(response, "data") if response_data: - for batch in response_data: + for batch in iter_items(response_data): encode_batch_response_ids(batch, model=model_param) verbose_proxy_logger.debug("Listed batches using model: %s", model_param) diff --git a/litellm/proxy/caching_routes.py b/litellm/proxy/caching_routes.py index 775d00ed53e..dcb56aef927 100644 --- a/litellm/proxy/caching_routes.py +++ b/litellm/proxy/caching_routes.py @@ -2,10 +2,12 @@ from collections.abc import Mapping from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import RedisCache +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import ProxyErrorTypes, ProxyException @@ -13,6 +15,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.caching import CachePingResponse, HealthCheckCacheParams masker: Final = SensitiveDataMasker() +_DECODED_JSON: Final = TypeAdapter(object) router: Final = APIRouter( prefix="/cache", @@ -149,7 +152,7 @@ async def cache_delete(request: Request): if litellm.cache is None: raise HTTPException(status_code=503, detail="Cache not initialized. litellm.cache is None") - request_data: Final = await request.json() + request_data: Final = decoded_json_object(_DECODED_JSON.validate_python(await request.json())) keys: Final = request_data.get("keys", None) if litellm.cache.type == "redis": diff --git a/litellm/proxy/client/chat.py b/litellm/proxy/client/chat.py index ccdef44f9e8..bba339faf7e 100644 --- a/litellm/proxy/client/chat.py +++ b/litellm/proxy/client/chat.py @@ -7,6 +7,8 @@ from pydantic import ConfigDict, TypeAdapter from .exceptions import StreamError, UnauthorizedError +_DECODED_JSON: Final = TypeAdapter(object) + _SSE_LINE: Final[TypeAdapter[bytes | bytearray]] = TypeAdapter( bytes | bytearray, config=ConfigDict(arbitrary_types_allowed=True, strict=True, hide_input_in_errors=True) ) @@ -59,7 +61,7 @@ class ChatClient: frequency_penalty: float | None = None, user: str | None = None, return_request: bool = False, - ) -> dict[str, object] | requests.Request: + ) -> object | requests.Request: """ Create a chat completion. @@ -114,7 +116,7 @@ class ChatClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) diff --git a/litellm/proxy/client/cli/commands/http.py b/litellm/proxy/client/cli/commands/http.py index 529564c685c..27104052613 100644 --- a/litellm/proxy/client/cli/commands/http.py +++ b/litellm/proxy/client/cli/commands/http.py @@ -4,10 +4,13 @@ from typing import Final, Protocol import click import requests import rich +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict from ...http_client import HTTPClient +_DECODED_JSON: Final = TypeAdapter(object) + class _CliContext(TypedDict): base_url: ReadOnly[str] @@ -85,7 +88,7 @@ def request( json_data = None if json: try: - json_data = json_lib.loads(json) + json_data = _DECODED_JSON.validate_python(json_lib.loads(json)) except ValueError as e: raise click.BadParameter(f"Invalid JSON format: {e}") diff --git a/litellm/proxy/client/credentials.py b/litellm/proxy/client/credentials.py index 649fe529936..920aa238ace 100644 --- a/litellm/proxy/client/credentials.py +++ b/litellm/proxy/client/credentials.py @@ -3,9 +3,12 @@ from typing import Any, Final from urllib.parse import quote import requests +from pydantic import TypeAdapter from .exceptions import UnauthorizedError +_DECODED_JSON: Final = TypeAdapter(object) + class CredentialsManagementClient: def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): @@ -75,7 +78,7 @@ class CredentialsManagementClient: credential_values: Mapping[str, object], return_request: bool = False, display_name: str | None = None, - ) -> dict[str, object] | requests.Request: + ) -> object | requests.Request: """ Create a new credential. @@ -111,7 +114,7 @@ class CredentialsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -121,7 +124,7 @@ class CredentialsManagementClient: self, credential_name: str, return_request: bool = False, - ) -> dict[str, object] | requests.Request: + ) -> object | requests.Request: """ Delete a credential by name. @@ -148,7 +151,7 @@ class CredentialsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -159,7 +162,7 @@ class CredentialsManagementClient: credential_name: str, display_name: str | None, return_request: bool = False, - ) -> Mapping[str, object] | requests.Request: + ) -> object | requests.Request: url: Final = f"{self._base_url}/credentials/{quote(credential_name, safe='')}" data: Final[Mapping[str, object]] = {"display_name": display_name, "credential_info": {}} @@ -172,7 +175,7 @@ class CredentialsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -182,7 +185,7 @@ class CredentialsManagementClient: self, credential_name: str, return_request: bool = False, - ) -> dict[str, object] | requests.Request: + ) -> object | requests.Request: """ Get a credential by name. @@ -209,7 +212,7 @@ class CredentialsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) diff --git a/litellm/proxy/client/health.py b/litellm/proxy/client/health.py index 685efa61d15..d40c64b5f41 100644 --- a/litellm/proxy/client/health.py +++ b/litellm/proxy/client/health.py @@ -1,4 +1,6 @@ -from typing import Any, Final +from typing import Final + +from litellm.litellm_core_utils.core_helpers import decoded_json_object from .http_client import HTTPClient @@ -19,7 +21,7 @@ class HealthManagementClient: """ self._http = HTTPClient(base_url=base_url, api_key=api_key, timeout=timeout) - def get_readiness(self) -> dict[str, Any]: + def get_readiness(self) -> object: """ Check the readiness of the LiteLLM proxy server. @@ -32,7 +34,7 @@ class HealthManagementClient: """ return self._http.request("GET", "/health/readiness") - def get_server_version(self) -> str | None: + def get_server_version(self) -> object: """ Get the LiteLLM server version from the readiness endpoint. @@ -40,4 +42,4 @@ class HealthManagementClient: Optional[str]: The server version if available, otherwise None. """ readiness: Final = self.get_readiness() - return readiness.get("litellm_version") + return decoded_json_object(readiness).get("litellm_version") diff --git a/litellm/proxy/client/http_client.py b/litellm/proxy/client/http_client.py index 96a170c75cb..e00529b1915 100644 --- a/litellm/proxy/client/http_client.py +++ b/litellm/proxy/client/http_client.py @@ -1,11 +1,14 @@ """HTTP client for making requests to the LiteLLM proxy server.""" from collections.abc import Mapping -from typing import Any, Final +from typing import Final import requests +from pydantic import TypeAdapter from typing_extensions import TypedDict, Unpack +_DECODED_JSON: Final = TypeAdapter(object) + class _RequestOptions(TypedDict, total=False, extra_items=object): pass @@ -32,10 +35,10 @@ class HTTPClient: uri: str, *, data: Mapping[str, object] | list | bytes | None = None, - json: Mapping[str, object] | list | None = None, + json: object = None, headers: dict[str, str] | None = None, **kwargs: Unpack[_RequestOptions], - ) -> Any: + ) -> object: """Make an HTTP request to the LiteLLM proxy server. This method is used to make generic requests to the LiteLLM proxy @@ -99,4 +102,4 @@ class HTTPClient: response.raise_for_status() # Parse and return JSON response - return response.json() + return _DECODED_JSON.validate_python(response.json()) diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index 028b338f412..dcd49cf375d 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -3,11 +3,14 @@ from collections.abc import Mapping from typing import Any, Final import requests +from pydantic import TypeAdapter from litellm.litellm_core_utils.secret_redaction import redact_string from .exceptions import UnauthorizedError +_DECODED_JSON: Final = TypeAdapter(object) + class KeysManagementClient: def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): @@ -122,7 +125,7 @@ class KeysManagementClient: budget_id: str | None = None, config: Mapping[str, object] | None = None, return_request: bool = False, - ) -> dict[str, object] | requests.Request: + ) -> object | requests.Request: """ Generate an API key based on the provided data. @@ -179,7 +182,7 @@ class KeysManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -190,7 +193,7 @@ class KeysManagementClient: keys: builtins.list[str] | None = None, key_aliases: builtins.list[str] | None = None, return_request: bool = False, - ) -> dict[str, object] | requests.Request: + ) -> object | requests.Request: """ Delete existing keys @@ -223,7 +226,7 @@ class KeysManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -239,7 +242,7 @@ class KeysManagementClient: key_alias: str | None = None, team_id: str | None = None, user_id: str | None = None, - ) -> dict[str, object] | requests.Request: + ) -> object | requests.Request: """ Update an existing API key's parameters. @@ -285,11 +288,11 @@ class KeysManagementClient: response: Final = session.send(request.prepare(), timeout=self._timeout) response_text = response.text response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except Exception: raise Exception(f"Error updating key: {response_text}") - def info(self, key: str, return_request: bool = False) -> dict[str, object] | requests.Request: + def info(self, key: str, return_request: bool = False) -> object | requests.Request: """ Get information about API keys. @@ -314,7 +317,7 @@ class KeysManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: redacted_message: Final = redact_string(str(e)) if e.response.status_code == 401: diff --git a/litellm/proxy/client/models.py b/litellm/proxy/client/models.py index 859322f061d..1a67ae1da15 100644 --- a/litellm/proxy/client/models.py +++ b/litellm/proxy/client/models.py @@ -5,9 +5,12 @@ from typing import Any, Final import requests from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.core_helpers import decoded_json_object + from .exceptions import NotFoundError, UnauthorizedError _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) +_DECODED_JSON: Final = TypeAdapter(object) class ModelsManagementClient: @@ -195,9 +198,11 @@ class ModelsManagementClient: # Find the matching model for model in models: - if (model_id and model.get("model_info", {}).get("id") == model_id) or ( - model_name and model.get("model_name") == model_name - ): + if ( + model_id + and decoded_json_object(_DECODED_JSON.validate_python(model.get("model_info", {}))).get("id") + == model_id + ) or (model_name and model.get("model_name") == model_name): return model # If we get here, no model was found diff --git a/litellm/proxy/client/teams.py b/litellm/proxy/client/teams.py index dc6fa1e228a..ad23892e998 100644 --- a/litellm/proxy/client/teams.py +++ b/litellm/proxy/client/teams.py @@ -4,9 +4,12 @@ import builtins from typing import Any, Final import requests +from pydantic import TypeAdapter from .exceptions import UnauthorizedError +_DECODED_JSON: Final = TypeAdapter(object) + class TeamsManagementClient: """Client for managing teams in LiteLLM proxy.""" @@ -80,7 +83,7 @@ class TeamsManagementClient: page_size: int = 10, sort_by: str | None = None, sort_order: str = "asc", - ) -> dict[str, object]: + ) -> object: """ Get a paginated list of teams with filtering and sorting options. @@ -125,7 +128,7 @@ class TeamsManagementClient: raise UnauthorizedError("Authentication failed. Check your API key.") response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) def get_available(self) -> builtins.list[dict[str, Any]]: """ diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 0ff094f6ad9..89f2ce87814 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -63,11 +63,14 @@ from litellm.litellm_core_utils.bug_report import ( strip_bug_report_notice, ) from litellm.litellm_core_utils.core_helpers import ( + attribute_or_none, finite_cost, get_or_create_metadata_bucket, + has_attributes, independent_snapshot, is_expected_client_error, redact_nested_match_and_regex_keys, + require_attributes, ) from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer from litellm.litellm_core_utils.get_supported_openai_params import ( @@ -674,6 +677,10 @@ class _ReadsHeaderValues(Protocol): def get(self, key: str, default: str = "") -> str: ... +class _NativeLoggingRelease(Protocol): + def release(self, success: bool, /) -> object: ... + + def _as_header_reader(headers: _ReadsHeaderValues) -> _ReadsHeaderValues: return headers @@ -693,6 +700,10 @@ def _as_success_dispatcher(logging_obj: _DispatchesSuccessHandlers) -> _Dispatch return logging_obj +class _ClosesAsync(Protocol): + def aclose(self) -> Awaitable[object]: ... + + def serialize_http_exception_detail( detail: object, ) -> tuple[str, dict | None]: @@ -4169,10 +4180,10 @@ class ProxyBaseLLMRequestProcessing: gating logic directly rather than reimplementing the finally block. """ if getattr(logging_obj, "call_type", None) in ("ocr", "aocr"): - pending: Final = getattr(logging_obj, "_native_pending_logging", None) + pending: Final = attribute_or_none(logging_obj, "_native_pending_logging") if pending is not None: logging_obj._native_pending_logging = None # rebind-ok: consume the native OCR release signal once - pending.release(not exception_raised) + require_attributes(pending, _NativeLoggingRelease, "release").release(not exception_raised) _enqueue_fn: Final = getattr(logging_obj, "enqueue_deferred_logging", None) if _enqueue_fn is None: return @@ -4604,7 +4615,7 @@ class ProxyBaseLLMRequestProcessing: async def finalize_streaming_generator_cleanup( request: Request | None, request_data: dict, - response: Any, + response: object, stream_completed: bool = False, client_disconnected: bool = False, user_api_key_dict: UserAPIKeyAuth | None = None, @@ -4641,7 +4652,7 @@ class ProxyBaseLLMRequestProcessing: ): await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict) - if hasattr(response, "aclose"): + if has_attributes(response, _ClosesAsync, "aclose"): try: await response.aclose() except BaseException as e: # noqa: BLE001 diff --git a/litellm/proxy/common_utils/cache_pydantic_utils.py b/litellm/proxy/common_utils/cache_pydantic_utils.py index 34d46be56cd..0beeacd4d87 100644 --- a/litellm/proxy/common_utils/cache_pydantic_utils.py +++ b/litellm/proxy/common_utils/cache_pydantic_utils.py @@ -17,7 +17,7 @@ are encoded; pass a Pydantic model or convert with e.g. ``dataclasses.asdict`` f from __future__ import annotations from collections.abc import Mapping, Sequence -from typing import Any, TypeAlias, TypeVar +from typing import TypeAlias, TypeVar from pydantic import BaseModel, ValidationError @@ -67,7 +67,7 @@ class CacheCodec: return value @staticmethod - def deserialize(cached: Any, model_type: type[T]) -> T | None: + def deserialize(cached: object, model_type: type[T]) -> T | None: """ Decode a cache entry to ``model_type``. diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index d1787f3ff8b..c195779aa0a 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -719,7 +719,7 @@ def strip_callback_config(metadata: dict[str, object] | None) -> dict[str, objec return {k: v for k, v in metadata.items() if k not in _CALLBACK_CONFIG_SLOTS} -def encrypt_callback_vars(metadata: object) -> Any: +def encrypt_callback_vars(metadata: object) -> object: """Return a deep copy of metadata with callback_vars values encrypted at rest. Idempotent: a value that already decrypts cleanly is left unchanged so diff --git a/litellm/proxy/common_utils/codex_catalog_source.py b/litellm/proxy/common_utils/codex_catalog_source.py new file mode 100644 index 00000000000..c36b7b0da22 --- /dev/null +++ b/litellm/proxy/common_utils/codex_catalog_source.py @@ -0,0 +1,181 @@ +"""Codex's own catalog for the Codex release asking for `/v1/models`, fetched from that release's tag on GitHub. + +Codex fetches its catalog from chatgpt.com, which needs a ChatGPT login the proxy never sees, so the proxy reads +the same two files at the release's public tag instead, and falls back to the copy bundled with LiteLLM when the +fetch fails, the release has no tag (dev builds report 0.0.0) or ``LITELLM_OFFLINE_MODE`` is on. +""" + +from __future__ import annotations + +import asyncio +import re +import time +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from functools import cache +from types import MappingProxyType +from typing import Final, Protocol, TypeVar + +import httpx + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.offline_mode import is_offline_mode +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared factory has untyped optional parameters +) +from litellm.proxy.common_utils.codex_model_catalog import ( + CODEX_BUNDLED_MODELS_PATH, + CodexStockCatalog, + CodexStockModel, + bundled_codex_models, + codex_base_instructions, +) +from litellm.types.llms.custom_http import httpxSpecialProvider + +CODEX_RELEASE_FILE_URL: Final = ( + "https://raw.githubusercontent.com/openai/codex/rust-v{version}/codex-rs/models-manager/{file}" +) +CODEX_FETCH_TIMEOUT_SECONDS: Final = 2.0 +CODEX_FETCH_FAILURE_TTL_SECONDS: Final = 3600.0 +CODEX_CACHED_VERSIONS: Final = 8 +CODEX_REMEMBERED_FAILURES: Final = 256 +CODEX_FETCHES_PER_MINUTE: Final = 6 + +_RELEASE_VERSION: Final = re.compile(r"[0-9]+\.[0-9]+\.[0-9]+(?:-[0-9A-Za-z.]+)?") + + +@dataclass(frozen=True, slots=True) +class CodexStock: + models: Mapping[str, CodexStockModel] + instructions: str + + +_Value = TypeVar("_Value") + + +class CodexReleaseClient(Protocol): + def get(self, url: str, *, timeout: float | None = None) -> Awaitable[httpx.Response]: ... + + +@cache +def bundled_codex_stock() -> CodexStock: + return CodexStock(models=bundled_codex_models(), instructions=codex_base_instructions()) + + +async def _release_file(client: CodexReleaseClient, version: str, file: str, timeout: float) -> str: + response: Final = await client.get(CODEX_RELEASE_FILE_URL.format(version=version, file=file), timeout=timeout) + return response.raise_for_status().content.decode("utf-8") + + +async def _release_stock(client: CodexReleaseClient, version: str, timeout: float) -> CodexStock: + models_json, instructions = await asyncio.wait_for( + asyncio.gather( + _release_file(client, version, "models.json", timeout), + _release_file(client, version, "prompt.md", timeout), + ), + timeout, + ) + catalog: Final = CodexStockCatalog.model_validate_json(models_json) + return CodexStock( + models=MappingProxyType({model.slug: model for model in catalog.models}), instructions=instructions + ) + + +class CodexCatalogSource: + """Codex's catalog per release: the one fetched from GitHub, else the bundled one. + + A fetched catalog is kept for good, a failed fetch is remembered for `failure_ttl` seconds so a blocked egress + costs one timeout per worker per version per hour. Only the `max_versions` most recently fetched catalogs and + the `max_failures` most recent failures are kept, apart, so versions that fail never push out ones that were + fetched. + Requests for a version already being fetched wait on that fetch, and a worker starts at most + `fetches_per_minute` fetches, serving the bundled catalog past that, so a caller cycling made-up versions + cannot turn every request into GitHub requests.""" + + def __init__( + self, + client: Callable[[], CodexReleaseClient], + *, + offline: Callable[[], bool] = is_offline_mode, + clock: Callable[[], float] = time.monotonic, + timeout: float = CODEX_FETCH_TIMEOUT_SECONDS, + failure_ttl: float = CODEX_FETCH_FAILURE_TTL_SECONDS, + max_versions: int = CODEX_CACHED_VERSIONS, + max_failures: int = CODEX_REMEMBERED_FAILURES, + fetches_per_minute: int = CODEX_FETCHES_PER_MINUTE, + ) -> None: + self._client: Final = client + self._offline: Final = offline + self._clock: Final = clock + self._timeout: Final = timeout + self._failure_ttl: Final = failure_ttl + self._max_versions: Final = max_versions + self._max_failures: Final = max_failures + self._fetches_per_minute: Final = fetches_per_minute + self._fetched: Mapping[str, CodexStock] = MappingProxyType({}) + self._failed_until: Mapping[str, float] = MappingProxyType({}) + self._in_flight: Mapping[str, asyncio.Future[CodexStock]] = MappingProxyType({}) + self._fetch_starts: tuple[float, ...] = () + + async def stock(self, client_version: str) -> CodexStock: + if self._offline() or _RELEASE_VERSION.fullmatch(client_version) is None: + return bundled_codex_stock() + known: Final = self._fetched.get(client_version) + if known is not None: + return known + if self._clock() < self._failed_until.get(client_version, float("-inf")): + return bundled_codex_stock() + pending: Final = self._in_flight.get(client_version) + if pending is not None: + return await asyncio.shield(pending) + if not self._take_fetch_slot(): + return bundled_codex_stock() + started: Final = asyncio.ensure_future(self._fetch_and_store(client_version)) + self._in_flight = MappingProxyType({**self._in_flight, client_version: started}) + return await asyncio.shield(started) + + def _take_fetch_slot(self) -> bool: + now: Final = self._clock() + recent: Final = tuple(start for start in self._fetch_starts if now - start < 60) + allowed: Final = len(recent) < self._fetches_per_minute + self._fetch_starts = (*recent, now) if allowed else recent + return allowed + + async def _fetch_and_store(self, version: str) -> CodexStock: + try: + fetched: Final = await self._fetch(version) + finally: + self._in_flight = MappingProxyType({key: task for key, task in self._in_flight.items() if key != version}) + if fetched is None: + self._failed_until = _latest( + self._failed_until, version, self._clock() + self._failure_ttl, self._max_failures + ) + return bundled_codex_stock() + self._fetched = _latest(self._fetched, version, fetched, self._max_versions) + return fetched + + async def _fetch(self, version: str) -> CodexStock | None: + try: + return await _release_stock(self._client(), version, self._timeout) + except Exception as e: # noqa: BLE001 # a failed best-effort fetch must never fail the model listing + verbose_proxy_logger.warning( + "Codex %s's model catalog could not be fetched from GitHub (%r), serving the catalog bundled " + "with LiteLLM (%s) for that version for the next %d seconds", + version, + e, + CODEX_BUNDLED_MODELS_PATH.name, + self._failure_ttl, + ) + return None + + +def _latest(stored: Mapping[str, _Value], version: str, value: _Value, limit: int) -> Mapping[str, _Value]: + others: Final = tuple((key, kept) for key, kept in stored.items() if key != version) + return MappingProxyType(dict((*others[max(0, len(others) - limit + 1) :], (version, value)))) + + +def _default_client() -> CodexReleaseClient: + return get_async_httpx_client(llm_provider=httpxSpecialProvider.CodexModelCatalog) + + +codex_catalog_source: Final = CodexCatalogSource(_default_client) diff --git a/litellm/proxy/common_utils/codex_model_catalog.py b/litellm/proxy/common_utils/codex_model_catalog.py index 462e372f3d8..6deee04e273 100644 --- a/litellm/proxy/common_utils/codex_model_catalog.py +++ b/litellm/proxy/common_utils/codex_model_catalog.py @@ -24,8 +24,10 @@ a newer tag. from __future__ import annotations +import hashlib +import json from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import asdict, dataclass from functools import cache from itertools import accumulate, pairwise from pathlib import Path @@ -33,7 +35,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, TypeAlias from packaging.version import Version -from pydantic import ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.proxy.common_utils.codex_model_translation import ( @@ -61,6 +63,8 @@ CODEX_CHAT_MODES: Final = frozenset({"chat", "responses"}) _BODY_PREFIX: Final = '{"models":[' _BODY_SUFFIX: Final = "]}" +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) +_JSON_OBJECTS: Final = TypeAdapter(tuple[dict[str, object], ...]) _NO_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) @@ -119,6 +123,7 @@ class CodexStockCatalog(LiteLLMBaseModel): CodexCatalogEntry = CodexFallbackModel | CodexStockModel CodexAliasDefaults: TypeAlias = Literal["codex", "deployment"] +_NATIVE_ALIAS_DEFAULTS: Final[CodexAliasDefaults] = "deployment" @dataclass(frozen=True, slots=True) @@ -408,6 +413,7 @@ def _entry( stock_model.upgrade if stock_model.upgrade is not None and stock_model.upgrade.model in served else None ) is_alias: Final = row.id != stock_model.slug + deployment_decides: Final = is_alias and alias_defaults == "deployment" stock_name: Final = row.id if is_alias else stock_model.display_name listing: Final = { "slug": row.id, @@ -417,12 +423,18 @@ def _entry( "supported_in_api": True, "upgrade": upgrade, **({} if search_tool is None else {"supports_search_tool": search_tool}), - **(_deployment_defaults(row, stock_model) if is_alias and alias_defaults == "deployment" else {}), + **(_deployment_defaults(row, stock_model) if deployment_decides else {}), **_configured_window(row, stock_model), } if configured is None: return stock_model.model_copy(update={**_priced_tiers(row, stock_model), **listing}) - return stock_model.model_copy(update={**listing, **_tier_fields(stock_model, configured)}) + return stock_model.model_copy( + update={ + **listing, + **_tier_fields(stock_model, configured), + **({"default_service_tier": None} if deployment_decides else {}), + } + ) def _serialized(entry: CodexCatalogEntry) -> str: @@ -435,13 +447,12 @@ def _bytes_left_after(left: int, size: int) -> int: return left - size if size <= left else left -def _fits(serialized: Sequence[str], byte_limit: int | None) -> tuple[bool, ...]: - """Whether each entry, taken in order, fits in what is left of a body of at most `byte_limit` - bytes: an entry larger than the bytes left is passed over and a smaller one after it is still - taken, each entry counted with its separating comma and the last one's comma given back.""" - if byte_limit is None: +def _fits(serialized: Sequence[str], budget: int | None) -> tuple[bool, ...]: + """Whether each entry, taken in order, fits in what is left of `budget` bytes: an entry larger + than the bytes left is passed over and a smaller one after it is still taken, each entry + counted with its separating comma.""" + if budget is None: return (True,) * len(serialized) - budget: Final = byte_limit - len(_BODY_PREFIX) - len(_BODY_SUFFIX) + 1 left: Final = tuple( accumulate((len(entry.encode()) + 1 for entry in serialized), _bytes_left_after, initial=budget) ) @@ -449,17 +460,34 @@ def _fits(serialized: Sequence[str], byte_limit: int | None) -> tuple[bool, ...] def _kept_positions( - entries: Sequence[CodexCatalogEntry], serialized: Sequence[str], byte_limit: int | None + entries: Sequence[CodexCatalogEntry], serialized: Sequence[str], budget: int | None ) -> frozenset[int]: """The listing positions that survive the byte limit: an entry offering a service tier is taken ahead of one offering none, each group in listing order, since a model without a tier is driven by name just as well off Codex's bundled catalog; an entry too large for the bytes left never costs the smaller ones after it their place.""" survival: Final = sorted(range(len(entries)), key=lambda position: (not entries[position].service_tiers, position)) - fits: Final = _fits(tuple(serialized[position] for position in survival), byte_limit) + fits: Final = _fits(tuple(serialized[position] for position in survival), budget) return frozenset(position for position, fit in zip(survival, fits, strict=True) if fit) +def _entries( + drivable: Sequence[CodexCatalogRow], + stock: Mapping[str, CodexStockModel] | None, + instructions: str | None, + first_priority: int, + alias_defaults: CodexAliasDefaults, + client_version: str | None, +) -> tuple[CodexCatalogEntry, ...]: + known: Final = bundled_codex_models() if stock is None else stock + prompt: Final = codex_base_instructions() if instructions is None else instructions + served: Final = frozenset(row.id for row in drivable) + return tuple( + _entry(first_priority + index, row, known, served, prompt, alias_defaults, client_version) + for index, row in enumerate(drivable) + ) + + @dataclass(frozen=True, slots=True) class CodexCatalogBody: """A `ModelsResponse` body, the ids it lists, and the chat-capable ids the byte limit left out.""" @@ -495,15 +523,12 @@ def codex_models_response_json( requesting Codex's version, which decides the reasoning levels a fallback entry may offer; None counts as a version too old for custom levels. """ - known: Final = bundled_codex_models() if stock is None else stock - prompt: Final = codex_base_instructions() if instructions is None else instructions drivable: Final = tuple(row for row in rows if _codex_can_drive(row.id, row.mode)) - served: Final = frozenset(row.id for row in drivable) - entries: Final = tuple( - _entry(index, row, known, served, prompt, alias_defaults, client_version) for index, row in enumerate(drivable) - ) + entries: Final = _entries(drivable, stock, instructions, 0, alias_defaults, client_version) serialized: Final = tuple(_serialized(entry) for entry in entries) - kept: Final = _kept_positions(entries, serialized, byte_limit) + kept: Final = _kept_positions( + entries, serialized, None if byte_limit is None else byte_limit - len(_BODY_PREFIX) - len(_BODY_SUFFIX) + 1 + ) return CodexCatalogBody( json=f"{_BODY_PREFIX}{','.join(serialized[position] for position in sorted(kept))}{_BODY_SUFFIX}", listed=tuple(row.id for position, row in enumerate(drivable) if position in kept), @@ -511,6 +536,74 @@ def codex_models_response_json( ) +def _priority(model: Mapping[str, object]) -> int: + priority: Final = model.get("priority") + return priority if isinstance(priority, int) else -1 + + +def _compact_json(value: object) -> str: + return json.dumps(value, separators=(",", ":")) + + +def _entry_json(entry: CodexCatalogEntry) -> Mapping[str, object]: + return _JSON_OBJECT.validate_python(entry.model_dump(mode="json", exclude_unset=isinstance(entry, CodexStockModel))) + + +def merged_codex_catalog( + upstream_body: bytes, + rows: Sequence[CodexCatalogRow], + *, + stock: Mapping[str, CodexStockModel] | None = None, + instructions: str | None = None, + byte_limit: int = CODEX_CATALOG_BYTE_LIMIT, + client_version: str | None = None, +) -> str | None: + try: + top: Final = _JSON_OBJECT.validate_json(upstream_body) + upstream: Final = _JSON_OBJECTS.validate_python(top.get("models")) + except ValidationError: + return None + slugs: Final = frozenset(model.get("slug") for model in upstream) + added: Final = tuple(row for row in rows if _codex_can_drive(row.id, row.mode) and row.id not in slugs) + first_priority: Final = max((_priority(model) for model in upstream), default=-1) + 1 + entries: Final = _entries(added, stock, instructions, first_priority, _NATIVE_ALIAS_DEFAULTS, client_version) + added_models: Final = tuple(_entry_json(entry) for entry in entries) + serialized: Final = tuple(_compact_json(model) for model in added_models) + budget: Final = byte_limit - len(_compact_json({**top, "models": upstream})) + kept: Final = tuple(sorted(_kept_positions(entries, serialized, budget))) + if len(kept) < len(entries): + verbose_proxy_logger.warning( + "Codex native model catalog left out %d of %d LiteLLM models to stay within Codex's %d byte limit", + len(entries) - len(kept), + len(entries), + byte_limit, + ) + if not kept: + return None + return _compact_json({**top, "models": [*upstream, *(added_models[position] for position in kept)]}) + + +def _canonical_value(value: object) -> object: + return value.model_dump(mode="json") if isinstance(value, BaseModel) else repr(value) + + +def codex_catalog_etag(upstream_etag: str, rows: Sequence[CodexCatalogRow]) -> str: + canonical: Final = json.dumps( + [asdict(row) for row in rows], sort_keys=True, separators=(",", ":"), default=_canonical_value + ) + digest: Final = hashlib.sha256(json.dumps([upstream_etag, canonical]).encode()).hexdigest() + return f'"{digest}"' + + +def upstream_codex_slugs(upstream_body: bytes) -> frozenset[str] | None: + try: + upstream: Final = _JSON_OBJECTS.validate_python(_JSON_OBJECT.validate_json(upstream_body).get("models")) + except ValidationError: + return None + slugs: Final = (model.get("slug") for model in upstream) + return frozenset(slug for slug in slugs if isinstance(slug, str)) + + def _group_reasoning_efforts( row: ModelInfoResponse, lookup_id: str, registry: RouterRegistry | None, team_id: str | None ) -> tuple[str, ...] | None: @@ -520,6 +613,11 @@ def _group_reasoning_efforts( return group.supported_reasoning_efforts if group is not None else None +def _shared_pinned_effort(pinned: tuple[object, ...]) -> str | None: + first: Final = next(iter(pinned), None) + return first if isinstance(first, str) and all(effort == first for effort in pinned) else None + + def deployment_capabilities(model_info: Mapping[str, object] | None, model_id: str) -> CodexDeploymentCapabilities: """One deployment's model info, merged with the cost map, as the capabilities a Codex entry describes. Model info that does not validate declares nothing, logged, so one bad value never fails the listing.""" @@ -562,6 +660,11 @@ def _catalog_row( registry.resolve(lookup_id, kind="upstream_provider", team_id=team_id) if registry is not None else None ), display_name=registry.group(label_id, field="display_name") if registry is not None else None, + reasoning_effort=( + _shared_pinned_effort(registry.group(lookup_id, field="pinned_reasoning_efforts", team_id=team_id)) + if registry is not None + else None + ), service_tiers=registry.group(lookup_id, field="service_tiers", team_id=team_id) if registry is not None else (), supported_reasoning_efforts=_group_reasoning_efforts(row, lookup_id, registry, team_id), search_tool_support=( @@ -606,10 +709,16 @@ def codex_model_list_body( label_entries: Sequence[tuple[str, str]] | None = None, *, client_version: str | None = None, + stock: Mapping[str, CodexStockModel] | None = None, + instructions: str | None = None, ) -> str: - """The `/v1/models?client_version=...` body, logging the models Codex's byte limit left out.""" + """The `/v1/models?client_version=...` body, logging the models Codex's byte limit left out. `stock` and + `instructions` are the requesting Codex release's catalog and prompt, the bundled ones when None.""" body: Final = codex_models_response_json( - codex_catalog_rows(rows, entries, llm_router, team_id, label_entries), client_version=client_version + codex_catalog_rows(rows, entries, llm_router, team_id, label_entries), + stock=stock, + instructions=instructions, + client_version=client_version, ) if body.left_out: verbose_proxy_logger.warning( @@ -619,3 +728,21 @@ def codex_model_list_body( ", ".join(body.left_out), ) return body.json + + +def codex_entry_names_default_reasoning_level( + model_id: str, llm_router: Router, team_id: str | None, client_version: str | None +) -> bool: + """Whether the native catalog entry for `model_id` names a `default_reasoning_level`. Codex's model picker + sends reasoning effort `none` for a model whose entry names none, while showing that choice as + "default".""" + listed: Final[ModelInfoResponse] = {"id": model_id, "object": "model", "created": 0, "owned_by": ""} + (entry,) = _entries( + (_catalog_row(listed, model_id, model_id, llm_router, team_id),), + None, + "", + 0, + _NATIVE_ALIAS_DEFAULTS, + client_version, + ) + return _entry_json(entry).get("default_reasoning_level") is not None diff --git a/litellm/proxy/common_utils/codex_model_translation.py b/litellm/proxy/common_utils/codex_model_translation.py index 2c4ea8896e3..694dc422994 100644 --- a/litellm/proxy/common_utils/codex_model_translation.py +++ b/litellm/proxy/common_utils/codex_model_translation.py @@ -16,7 +16,7 @@ from types import NoneType from typing import Annotated, Final, Literal, TypeAlias from packaging.version import Version -from pydantic import ConfigDict, Field, StringConstraints +from pydantic import ConfigDict, Field, StringConstraints, field_serializer from litellm.litellm_core_utils.reasoning_effort_strength import nearest_declared_reasoning_effort from litellm.types.llms.base_model import LiteLLMBaseModel @@ -117,6 +117,10 @@ class CodexDeploymentCapabilities(LiteLLMBaseModel): prices_tokens: bool = False priced_service_tier_keys: frozenset[str] = frozenset() + @field_serializer("priced_service_tier_keys") + def _sorted_priced_service_tier_keys(self, keys: frozenset[str]) -> tuple[str, ...]: + return tuple(sorted(keys)) + @cache def tier_price_keys() -> tuple[str, ...]: diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index e210c337922..08461aa3a44 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -2,7 +2,7 @@ import json import re from collections.abc import Collection, Mapping from types import MappingProxyType, UnionType -from typing import Annotated, Any, Final, Literal, Union, get_args, get_origin +from typing import Annotated, Any, Final, Literal, Protocol, Union, get_args, get_origin import orjson from fastapi import HTTPException, Request, UploadFile, status @@ -21,6 +21,7 @@ from litellm.constants import ( MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB, ) from litellm.integrations.otel.runtime import phase_event +from litellm.litellm_core_utils.core_helpers import as_object_list, has_attributes, require_attributes from litellm.proxy._types import ProxyException from litellm.proxy.common_utils.callback_utils import ( get_metadata_variable_name_from_kwargs, @@ -37,6 +38,18 @@ _DECODED_JSON: Final = TypeAdapter(object) _JSON_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object]) +class _ReadableUpload(Protocol): + async def read(self) -> bytes: ... + + +class _UploadedFileInfo(Protocol): + @property + def filename(self) -> str | None: ... + + @property + def content_type(self) -> str | None: ... + + def resolve_realtime_route_model(model: str | None, intent: str | None, is_translation: bool) -> str | None: if model is not None: return model @@ -208,7 +221,7 @@ def _parse_form_metadata(raw_metadata: str) -> object: 400 because no provider or database can encode it. ``json.loads`` accepts the escape, so the multipart ``metadata`` field gets the same check.""" try: - parsed: Final[object] = json.loads(raw_metadata) + parsed: Final = _DECODED_JSON.validate_python(json.loads(raw_metadata)) json.dumps(parsed, ensure_ascii=False).encode("utf-8") except (json.JSONDecodeError, UnicodeEncodeError) as e: raise ProxyException( @@ -602,7 +615,7 @@ async def get_form_data(request: Request) -> dict[str, Any]: async def convert_upload_files_to_file_data( - form_data: dict[str, Any], + form_data: Mapping[str, object], ) -> dict[str, Any]: """ Convert FastAPI UploadFile objects to file data tuples for litellm. @@ -625,21 +638,24 @@ async def convert_upload_files_to_file_data( """ data: Final = {} for key, value in form_data.items(): - if isinstance(value, list): + values = as_object_list(value) + if values is not None: # Check if it's a list of UploadFile objects - if value and hasattr(value[0], "read"): + if values and has_attributes(values[0], _ReadableUpload, "read"): files = [] - for f in value: - file_content = await f.read() + for f in values: + file_content = await require_attributes(f, _ReadableUpload, "read").read() + file_info = require_attributes(f, _UploadedFileInfo, "filename", "content_type") # Create tuple: (filename, content, content_type) - files.append((f.filename, file_content, f.content_type)) + files.append((file_info.filename, file_content, file_info.content_type)) data[key] = files else: data[key] = value - elif hasattr(value, "read"): + elif has_attributes(value, _ReadableUpload, "read"): # Single UploadFile object - read and convert to list for consistency file_content = await value.read() - data[key] = [(value.filename, file_content, value.content_type)] + upload_info = require_attributes(value, _UploadedFileInfo, "filename", "content_type") + data[key] = [(upload_info.filename, file_content, upload_info.content_type)] else: # Regular form field data[key] = value diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index c7c73cfc2a8..91529ee36fc 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -24,6 +24,7 @@ T = TypeVar("T", bound=BaseModel) from litellm.constants import AUTH_OBJECTS_TARGET as AUTH_OBJECTS_TARGET # noqa: E402, PLC0414 # public re-export _HASHED_TOKEN_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}") +_CACHED_VALUE: Final = TypeAdapter(object) _BUDGET_WINDOWS_ABSENT_PREFIX: Final = "budget_windows_absent:" AuthCacheFactory: TypeAlias = Callable[[RedisCache, AuthCacheSettings], AuthCacheStore] @@ -217,7 +218,7 @@ class UserApiKeyCache(DualCache): if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) shared: Final = self._shared - cached: Final = ( + cached: Final = _CACHED_VALUE.validate_python( self._get_shared_sync(shared, key, local_only) if shared is not None else self.key_object_cache.get_cache( @@ -260,7 +261,7 @@ class UserApiKeyCache(DualCache): if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) shared: Final = self._shared - cached: Final = ( + cached: Final = _CACHED_VALUE.validate_python( await self._get_shared(shared, key, local_only) if shared is not None else await self.key_object_cache.async_get_cache( diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index a2ffdd4d59d..a15ffd57333 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -15,6 +15,7 @@ from fastapi import APIRouter, Depends, Request, Response from fastapi.responses import ORJSONResponse from pydantic import TypeAdapter +from litellm.llms.custom_httpx.container_handler import EndpointsConfig from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -28,13 +29,14 @@ from litellm.proxy.container_endpoints.ownership import ( ) _UPLOADED_FILE_FIELD: Final = TypeAdapter(object) +_ENDPOINTS_CONFIG: Final = TypeAdapter(EndpointsConfig) -def _load_endpoints_config() -> dict: +def _load_endpoints_config() -> EndpointsConfig: """Load the endpoints configuration from JSON file.""" config_path: Final = Path(__file__).parent.parent.parent / "containers" / "endpoints.json" with open(config_path) as f: - return json.load(f) + return _ENDPOINTS_CONFIG.validate_python(json.load(f)) def get_all_route_types() -> list[str]: @@ -57,7 +59,7 @@ def _get_container_provider_config(custom_llm_provider: str): def _create_handler_for_path_params( - path_params: list[str], + path_params: Sequence[str], route_type: str, returns_binary: bool = False, is_multipart: bool = False, diff --git a/litellm/proxy/db/health_check_latest.py b/litellm/proxy/db/health_check_latest.py index 024be7603d6..bc9120e3b4e 100644 --- a/litellm/proxy/db/health_check_latest.py +++ b/litellm/proxy/db/health_check_latest.py @@ -13,10 +13,9 @@ import json from collections.abc import Iterable, Mapping, Sequence from datetime import datetime, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final from pydantic import ConfigDict, JsonValue, TypeAdapter, field_validator -from typing_extensions import LiteralString from litellm._logging import verbose_proxy_logger from litellm.proxy.db.db_span import db_span @@ -78,17 +77,11 @@ class LatestHealthCheckRow(LiteLLMBaseModel): _ROWS_ADAPTER: Final = TypeAdapter(tuple[LatestHealthCheckRow, ...]) -class _HealthCheckRowsQuery(Protocol): - async def __call__(self, query: LiteralString, *args: object) -> Sequence[Mapping[str, object]]: ... - - -def _health_check_rows_query(prisma_client: PrismaClient) -> _HealthCheckRowsQuery: - return prisma_client.db.query_raw - - async def query_latest_health_checks(prisma_client: PrismaClient) -> tuple[LatestHealthCheckRow, ...]: + from litellm.proxy.utils import query_raw_rows + async with db_span("latest_health_checks", "LiteLLM_HealthCheckTable"): - rows: Final = await _health_check_rows_query(prisma_client)(LATEST_HEALTH_CHECKS_SQL) + rows: Final = await query_raw_rows(prisma_client.db, LATEST_HEALTH_CHECKS_SQL) return _ROWS_ADAPTER.validate_python(rows) @@ -106,10 +99,10 @@ async def fetch_latest_health_checks_for_models( if not model_names: return () try: + from litellm.proxy.utils import query_raw_rows + async with db_span("latest_health_checks", "LiteLLM_HealthCheckTable"): - rows: Final = await _health_check_rows_query(prisma_client)( - LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, list(model_names) - ) + rows: Final = await query_raw_rows(prisma_client.db, LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, list(model_names)) return _ROWS_ADAPTER.validate_python(rows) except Exception as query_err: # noqa: BLE001 # a paged model list must not fail on its health decoration verbose_proxy_logger.error("Error getting latest health checks for models: %s", query_err) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 9d6214dc06d..bae64c883d3 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -12,7 +12,7 @@ import subprocess import time import urllib import urllib.parse -from collections.abc import Awaitable, Callable, Coroutine +from collections.abc import Awaitable, Callable, Coroutine, Mapping from datetime import datetime, timedelta from typing import TYPE_CHECKING, Final, Protocol, TypeVar @@ -36,8 +36,11 @@ from litellm.secret_managers.main import str_to_bool if TYPE_CHECKING: from prisma import Prisma + from prisma import models as prisma_models from prisma.types import HttpConfig + from litellm.repositories.prisma_protocols import TableActions + _ENVIRONMENT_VARIABLES_REVISION_MIGRATION: Final = "20261015120000_add_environment_variables_revision_trigger" __all__ = ( @@ -904,6 +907,21 @@ class PrismaWrapper: async def disconnect(self, timeout: float | timedelta | None = None) -> None: ... + async def execute_raw(self, query: str, *args: object) -> int: ... + + async def query_first(self, query: str, *args: object) -> Mapping[str, object] | None: ... + + @property + def litellm_deprecatedverificationtoken( + self, + ) -> TableActions[prisma_models.LiteLLM_DeprecatedVerificationToken]: ... + + @property + def litellm_mcpservertable(self) -> TableActions[prisma_models.LiteLLM_MCPServerTable]: ... + + @property + def litellm_shadowevaljob(self) -> TableActions[prisma_models.LiteLLM_ShadowEvalJob]: ... + def __getattr__(self, name: str): """ Proxy attribute access to the underlying Prisma client. diff --git a/litellm/proxy/db/proxy_worker_heartbeat.py b/litellm/proxy/db/proxy_worker_heartbeat.py index 3214c0974b4..d0f1b11ada7 100644 --- a/litellm/proxy/db/proxy_worker_heartbeat.py +++ b/litellm/proxy/db/proxy_worker_heartbeat.py @@ -97,11 +97,13 @@ async def count_live_proxy_workers(prisma_client: PrismaClient) -> int | None: counts on the primary: a lagging read replica must never undercount. """ try: + from litellm.proxy.utils import query_raw_rows + db: Final = prisma_client.db primary_db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) else db async with db_span("count_live_proxy_workers", "LiteLLM_ProxyWorkerHeartbeat"): rows: Final = _RAW_QUERY_RESULT.validate_python( - await primary_db.query_raw(COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) + await query_raw_rows(primary_db, COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) ) return _COUNT_ROWS_ADAPTER.validate_python(rows)[0]["live_workers"] except Exception as count_err: # noqa: BLE001 # an unknown count must degrade to "warn", never to a 503 diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 50e6aa54da4..c71e2206f45 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -6,7 +6,7 @@ otherwise PrismaClient uses the writer-only PrismaWrapper directly. import asyncio import os -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import timedelta from typing import TYPE_CHECKING, Any, Final @@ -16,8 +16,11 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy.db.prisma_client import PrismaWrapper if TYPE_CHECKING: + from prisma import models as prisma_models from prisma.types import HttpConfig + from litellm.repositories.prisma_protocols import TableActions + # Per-model action methods that read from the database. These are routed to # the read replica when one is configured. _MODEL_READ_METHODS: Final = frozenset( @@ -302,6 +305,23 @@ class RoutingPrismaWrapper: raise RuntimeError("DATABASE_URL_READ_REPLICA not set; cannot recreate read replica client") await self._reader.recreate_prisma_client(reader_url, http_client=http_client) + if TYPE_CHECKING: + + async def execute_raw(self, query: str, *args: object) -> int: ... + + async def query_first(self, query: str, *args: object) -> Mapping[str, object] | None: ... + + @property + def litellm_deprecatedverificationtoken( + self, + ) -> TableActions[prisma_models.LiteLLM_DeprecatedVerificationToken]: ... + + @property + def litellm_mcpservertable(self) -> TableActions[prisma_models.LiteLLM_MCPServerTable]: ... + + @property + def litellm_shadowevaljob(self) -> TableActions[prisma_models.LiteLLM_ShadowEvalJob]: ... + def __getattr__(self, name: str) -> Any: if name in _TOP_LEVEL_READ_METHODS: return _DELEGATED_ATTRIBUTE.validate_python(getattr(self.read_target, name)) diff --git a/litellm/proxy/guardrails/exception_utils.py b/litellm/proxy/guardrails/exception_utils.py index 518c61d1fd3..68a29e5c2ef 100644 --- a/litellm/proxy/guardrails/exception_utils.py +++ b/litellm/proxy/guardrails/exception_utils.py @@ -13,16 +13,22 @@ def is_fastapi_http_exception(e: Exception, block_status_codes: Collection[int]) def enrich_http_exception_with_guardrail_context(exc: BaseException, callback: object) -> None: + if not hasattr(exc, "mcp_block_behavior"): + setattr( # noqa: B010 # exception types do not declare per-deployment MCP context + exc, "mcp_block_behavior", getattr(callback, "mcp_block_behavior", None) + ) try: from fastapi.exceptions import HTTPException except ImportError: return if not isinstance(exc, HTTPException): return + guardrail_name: Final[object] = getattr(callback, "guardrail_name", None) detail: Final = getattr(exc, "detail", None) if not isinstance(detail, dict): + if isinstance(guardrail_name, str): + setattr(exc, "guardrail_name", guardrail_name) # noqa: B010 # HTTPException does not declare guardrail_name return - guardrail_name: Final[object] = getattr(callback, "guardrail_name", None) if guardrail_name: detail.setdefault("guardrail_name", guardrail_name) event_hook: Final[object] = getattr(callback, "event_hook", None) diff --git a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py index c4615a3cebf..cefcd4993eb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py @@ -14,12 +14,14 @@ import sys from typing import TYPE_CHECKING, Final, Literal from fastapi import HTTPException +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.logging_utils import ( convert_litellm_response_object_to_str, ) @@ -31,6 +33,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks GUARDRAIL_NAME: Final = "aporia" +_DECODED_JSON: Final = TypeAdapter(object) if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -128,8 +131,8 @@ class AporiaGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Aporia AI response: %s", response.text) if response.status_code == 200: # check if the response was flagged - _json_response: Final = response.json() - action: str = _json_response.get("action") # possible values are modify, passthrough, block, rephrase + _json_response: Final = decoded_json_object(_DECODED_JSON.validate_python(response.json())) + action: Final = _json_response.get("action") # possible values are modify, passthrough, block, rephrase if action == "block": raise HTTPException( status_code=400, diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py index 5f57969c310..dc1b9bb84a3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py @@ -3,6 +3,7 @@ from collections.abc import Mapping from typing import Any, Final, cast import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -29,6 +30,7 @@ JAVELIN_API_VERSION_STORED_BY_OLDER_RELEASES: Final = "v1" RESPONSES_API_CALL_TYPES: Final = frozenset({CallTypes.responses, CallTypes.aresponses}) _RESPONSES_API_CALL_TYPES: Final = RESPONSES_API_CALL_TYPES +_DECODED_JSON: Final = TypeAdapter(object) def resolve_content_safety_api_version(configured: str | None) -> str: @@ -62,9 +64,7 @@ class AzureGuardrailBase: self.api_base = api_base self.api_version: str | None = kwargs.get("api_version") - async def _post_to_content_safety( - self, endpoint_path: str, request_body: dict[str, object] - ) -> Mapping[str, object]: + async def _post_to_content_safety(self, endpoint_path: str, request_body: dict[str, object]) -> object: """POST to an Azure Content Safety endpoint with standard auth headers. Args: @@ -89,7 +89,7 @@ class AzureGuardrailBase: json=request_body, timeout=self.timeout, ) - response_json: Final[dict[str, object]] = response.json() + response_json: Final = _DECODED_JSON.validate_python(response.json()) verbose_proxy_logger.debug("Azure Content Safety response [%s]: %s", endpoint_path, response_json) return response_json diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 6a773f41f01..c3a979ad1bc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -16,7 +16,7 @@ import json import re import sys import time -from collections.abc import AsyncGenerator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence from datetime import datetime, timezone from itertools import accumulate, groupby from types import MappingProxyType @@ -33,7 +33,12 @@ from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.exceptions import ModifyResponseException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route -from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys +from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, + decoded_json_object, + iter_items, + redact_nested_match_and_regex_keys, +) from litellm.litellm_core_utils.litellm_logging import ( get_masked_values, ) @@ -1314,12 +1319,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) if httpx_response.status_code == 200: - _json_response: Final = httpx_response.json() + raw_json_response: Final = httpx_response.json() # check if the response was flagged verbose_proxy_logger.debug( "Bedrock AI response : %s", - redact_nested_match_and_regex_keys(_json_response), + redact_nested_match_and_regex_keys(raw_json_response), ) + _json_response: Final[BedrockGuardrailResponse] = raw_json_response bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response) if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response): self._log_apply_guardrail_attempt( @@ -1350,7 +1356,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): def _log_apply_guardrail_attempt( self, httpx_response: httpx.Response, - json_response: dict, # mutable-ok: raw AWS JSON payload + json_response: BedrockGuardrailResponse, request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper event_type: GuardrailEventHooks, start_time: "datetime", @@ -1365,7 +1371,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): billed_usages: Final[tuple[BedrockGuardrailUsage, ...]] = tuple(completed_chunk_usages) + ( (blocking_usage,) if isinstance(blocking_usage, dict) else () ) - logged_json_response: Final = ( + logged_json_response: Final[BedrockGuardrailResponse] = ( { **json_response, "usage": self._sum_usage_counters(billed_usages), @@ -2106,12 +2112,14 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): response: BedrockGuardrailChecksResponse, ) -> dict[str, object]: """Strip PII location offsets from a checks response before it is logged.""" - sanitized: Final[dict[str, Any]] = copy.deepcopy(dict(response)) - sensitive: Final = (sanitized.get("results") or {}).get("sensitiveInformation") or {} - for entry in sensitive.get("results") or []: - if isinstance(entry, dict): + sanitized: Final = copy.deepcopy(dict(response)) + sensitive: Final = decoded_json_object( + decoded_json_object(sanitized.get("results") or {}).get("sensitiveInformation") or {} + ) + for entry in iter_items(sensitive.get("results") or []): + if isinstance(pii_entry := as_str_mapping(entry), dict): for key in _BEDROCK_CHECKS_PII_LOCATION_KEYS: - entry.pop(key, None) + pii_entry.pop(key, None) return sanitized @staticmethod @@ -2205,7 +2213,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check if the guardrail would block content try: - _json_response: Final = response.json() + _json_response: Final[BedrockGuardrailResponse] = response.json() bedrock_guardrail_response: Final = BedrockGuardrailResponse(**_json_response) if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response): return "guardrail_intervened" @@ -2918,7 +2926,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: Any, + response: AsyncIterable[object], request_data: dict, ) -> AsyncGenerator[ModelResponseStream, None]: """ @@ -2970,7 +2978,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Collect all chunks to process them together started_at: Final = time.monotonic() - all_chunks: Final[list[ModelResponseStream]] = [] + all_chunks: Final[list[object]] = [] async for chunk in response: all_chunks.append(chunk) @@ -3079,18 +3087,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Return the reconstructed stream async for chunk in mock_response: yield chunk - elif raw_sse: - # Forwarding an unscannable stream would silently disable the guardrail, so fail closed. - # A raise cannot reach the client once a keepalive ping has flushed the headers, so the - # refusal travels as a frame, matching how a block is delivered above - for error_frame in anthropic_sse_error_frames( - f"{self.guardrail_name}: streamed response could not be assembled for scanning, blocking it" - ): - yield error_frame - return else: - for chunk in all_chunks: - yield chunk + for unassembled_chunk in ( + # Forwarding an unscannable stream would silently disable the guardrail, so fail closed. + # A raise cannot reach the client once a keepalive ping has flushed the headers, so the + # refusal travels as a frame, matching how a block is delivered above + anthropic_sse_error_frames( + f"{self.guardrail_name}: streamed response could not be assembled for scanning, blocking it" + ) + if raw_sse + else all_chunks + ): + yield unassembled_chunk def _extract_masked_texts_from_response(self, bedrock_guardrail_response: BedrockGuardrailResponse) -> list[str]: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py index 9b9e0949a1a..60905346fa1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -37,6 +37,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.prompt_templates.factory import ( get_attribute_or_key, get_tool_calls_from_response, @@ -73,7 +74,7 @@ if TYPE_CHECKING: ) BYPASS_HEADER: Final = "x-compresr-bypass" -_AGENTIC_LOOP_MODEL: Final = TypeAdapter(object) +_DECODED_JSON: Final = TypeAdapter(object) COMPRESR_RETRIEVE_TOOL_NAME: Final = "compresr_retrieve" DEFAULT_API_BASE: Final = "https://api.compresr.ai" DEFAULT_COMPRESSION_MODEL: Final = "latte_v2" @@ -150,6 +151,10 @@ class _JsonResponse(Protocol): def json(self) -> object: ... +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + def _json_body(response: _JsonResponse) -> object: return response.json() @@ -1121,8 +1126,13 @@ class CompresrGuardrail(CustomGuardrail): full_model_name = model if logging_obj is not None: - agentic_params: Final = getattr(logging_obj, "model_call_details", {}).get("agentic_loop_params", {}) - candidate: Final = _AGENTIC_LOOP_MODEL.validate_python(agentic_params.get("model", model)) + model_call_details: Final = require_attributes( + _DECODED_JSON.validate_python(getattr(logging_obj, "model_call_details", {})), _SupportsGet, "get" + ) + agentic_params: Final = require_attributes( + model_call_details.get("agentic_loop_params", {}), _SupportsGet, "get" + ) + candidate: Final = agentic_params.get("model", model) if isinstance(candidate, str) and candidate: full_model_name = candidate diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py index d507df24faa..0420bb2bb2a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py @@ -133,7 +133,7 @@ class DynamoAIGuardrails(CustomGuardrail): timeout=self.timeout, ) response.raise_for_status() - response_json: Final = response.json() + response_json: Final[DynamoAIResponse] = response.json() end_time = datetime.now() duration = (end_time - start_time).total_seconds() diff --git a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py index 98db3822092..1b9bfc725db 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py @@ -126,7 +126,7 @@ class EnkryptAIGuardrails(CustomGuardrail): timeout=self.timeout, ) response.raise_for_status() - response_json: Final = response.json() + response_json: Final[EnkryptAIResponse] = response.json() end_time = datetime.now() duration = (end_time - start_time).total_seconds() diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index e57d6e78378..6d4fb30dede 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -15,6 +15,7 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.llms.base_llm.guardrail_translation.utils import ( @@ -51,6 +52,10 @@ class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): pass +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + class _GraySwanMonitorHTTPResponse(Protocol): def raise_for_status(self) -> object: ... @@ -614,9 +619,12 @@ class GraySwanGuardrail(CustomGuardrail): if not headers: headers = (request_data.get("metadata") or {}).get("headers") if not headers and logging_obj and getattr(logging_obj, "model_call_details", None): - headers = ( - (logging_obj.model_call_details or {}).get("litellm_params", {}).get("metadata", {}).get("headers") + litellm_params: Final = require_attributes( + _DECODED_JSON.validate_python((logging_obj.model_call_details or {}).get("litellm_params", {})), + _SupportsGet, + "get", ) + headers = require_attributes(litellm_params.get("metadata", {}), _SupportsGet, "get").get("headers") if not isinstance(headers, dict): return None diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py index cf6592a3e58..c2455123521 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py @@ -83,7 +83,8 @@ class GuardrailsAI(CustomGuardrail): timeout=self.timeout, ) verbose_proxy_logger.debug("guardrails_ai response: %s", response) - _json_response: Final = GuardrailsAIResponse(**response.json()) + decoded_response: Final[GuardrailsAIResponse] = response.json() + _json_response: Final = GuardrailsAIResponse(**decoded_response) if _json_response.get("validationPassed") is False: raise HTTPException( status_code=400, @@ -130,7 +131,8 @@ class GuardrailsAI(CustomGuardrail): }, ) - _json_response: Final = GuardrailsAIResponsePreCall(**response.json()) + decoded_response: Final[GuardrailsAIResponsePreCall] = response.json() + _json_response: Final = GuardrailsAIResponsePreCall(**decoded_response) response = _json_response.get("outputs", [])[0].get("data", [])[0] return response diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 0b762b2d5fe..f08d9112c33 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -22,6 +22,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.prompt_templates.factory import ( get_attribute_or_key, get_tool_calls_from_response, @@ -76,12 +77,17 @@ _HASH_CACHE_TTL_SECONDS: Final = 15 * 60 # untranslated messages can be read with concrete types (values pass through by # reference, so this is a shallow top-level reconstruction). _REQUEST_DATA_ADAPTER: Final = TypeAdapter(dict[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) class _JsonResponse(Protocol): def json(self) -> object: ... +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + def _json_body(response: _JsonResponse) -> object: return response.json() @@ -985,7 +991,12 @@ class HeadroomGuardrail(CustomGuardrail): full_model_name = model if logging_obj is not None: - agentic_params: Final = getattr(logging_obj, "model_call_details", {}).get("agentic_loop_params", {}) + model_call_details: Final = require_attributes( + _DECODED_JSON.validate_python(getattr(logging_obj, "model_call_details", {})), _SupportsGet, "get" + ) + agentic_params: Final = require_attributes( + model_call_details.get("agentic_loop_params", {}), _SupportsGet, "get" + ) candidate: Final = agentic_params.get("model", model) if isinstance(candidate, str) and candidate: full_model_name = candidate diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 3e5c54b4501..d59c4c316b3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -414,7 +414,8 @@ class LakeraAIGuardrail(CustomGuardrail): verbose_proxy_logger.debug( "Lakera AI v2 guard response: %s", _DECODED_JSON.validate_python(response.json()) ) - lakera_response = LakeraAIResponse(**response.json()) + decoded_response: Final[LakeraAIResponse] = response.json() + lakera_response = LakeraAIResponse(**decoded_response) return lakera_response, masked_entity_count except Exception as e: status = "guardrail_failed_to_respond" diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 1877b0c97a9..76b5716a312 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -30,6 +30,7 @@ from litellm.constants import ( DEFAULT_MAX_RECURSE_DEPTH, ) from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.path_utils import is_within, try_safe_join from litellm.proxy.guardrails.content_filter_data import ( @@ -147,10 +148,10 @@ class _CategoryConfigView(TypedDict): def _streamed_choice_content(choice: StreamingChoices | TextChoices) -> str | None: if isinstance(choice, TextChoices): return choice.text if isinstance(choice.text, str) else "" - delta: Final = getattr(choice, "delta", None) + delta: Final = attribute_or_none(choice, "delta") if not delta: return None - content: Final = getattr(delta, "content", None) + content: Final = attribute_or_none(delta, "content") return content if isinstance(content, str) else "" diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py index da937bd0325..0faddb47cc3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py @@ -11,7 +11,7 @@ import re from collections.abc import Iterator, Sequence from enum import Enum from re import Pattern -from typing import Any, Final +from typing import Final from pydantic import ConfigDict, TypeAdapter, with_config from typing_extensions import ReadOnly, TypedDict @@ -33,7 +33,21 @@ class _PatternsFile(TypedDict): patterns: ReadOnly[Sequence[_PatternEntry]] +class _CategoryFileSummary(TypedDict, total=False): + category_name: ReadOnly[str] + display_name: ReadOnly[str] + description: ReadOnly[str] + default_action: ReadOnly[str] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _PatternExtraConfig(TypedDict, total=False): + keyword_pattern: ReadOnly[str] + allow_word_numbers: ReadOnly[bool] + + _PATTERNS_FILE: Final = TypeAdapter(_PatternsFile) +_PATTERN_EXTRAS: Final = TypeAdapter(_PatternExtraConfig) def _load_patterns_from_json() -> _PatternsFile: @@ -72,9 +86,11 @@ KNOWN_PATTERN_KEYS: Final = { "description", } -PATTERN_EXTRA_CONFIG: Final[dict[str, dict[str, Any]]] = {} +PATTERN_EXTRA_CONFIG: Final[dict[str, _PatternExtraConfig]] = {} for pattern_data in _PATTERNS_DATA["patterns"]: - extra_config = {key: value for key, value in pattern_data.items() if key not in KNOWN_PATTERN_KEYS} + extra_config = _PATTERN_EXTRAS.validate_python( + {key: value for key, value in pattern_data.items() if key not in KNOWN_PATTERN_KEYS} + ) PATTERN_EXTRA_CONFIG[pattern_data["name"]] = extra_config @@ -154,7 +170,7 @@ def _category_entry(categories_dir: str, filename: str) -> dict[str, str] | None if filename.endswith((".yaml", ".yml")): try: with open(category_file_path, "r") as f: - category_data = yaml.safe_load(f) + category_data: Final[_CategoryFileSummary] = yaml.safe_load(f) except Exception as e: from litellm._logging import verbose_proxy_logger diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py index 7c2d0dbc2fc..209a82f756f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py @@ -24,7 +24,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" "mode='pre_mcp_call'. JWT injection only fires for MCP tool calls." ) - optional_params: Final = getattr(litellm_params, "optional_params", None) + optional_params: Final[object] = getattr(litellm_params, "optional_params", None) def _get(key): if optional_params is not None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index 2e3f7120726..9cb929dec92 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -81,6 +81,7 @@ import jwt from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey +from pydantic import TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack from litellm._logging import verbose_proxy_logger @@ -89,6 +90,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypesLiteral @@ -96,6 +98,8 @@ from litellm.types.utils import CallTypesLiteral if TYPE_CHECKING: from jwt.types import Options +_DECODED_JSON: Final = TypeAdapter(object) + class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): pass @@ -496,7 +500,7 @@ class MCPJWTSigner(CustomGuardrail): payload: Final[dict[str, object]] = jwt.decode(raw_token, signing_jwk.key, **decode_kwargs) return payload - async def _introspect_opaque_token(self, token: str) -> dict[str, object]: + async def _introspect_opaque_token(self, token: str) -> Mapping[str, object]: """ Perform RFC 7662 token introspection for opaque (non-JWT) tokens. @@ -522,7 +526,7 @@ class MCPJWTSigner(CustomGuardrail): timeout=self.timeout, ) resp.raise_for_status() - result: Final[dict[str, object]] = resp.json() + result: Final = decoded_json_object(_DECODED_JSON.validate_python(resp.json())) if not result.get("active", False): raise jwt.exceptions.ExpiredSignatureError( "MCPJWTSigner: incoming token is inactive (introspection returned active=false)" @@ -821,7 +825,7 @@ class MCPJWTSigner(CustomGuardrail): # ------------------------------------------------------------------ # FR-5: Verify incoming token before re-signing # ------------------------------------------------------------------ - jwt_claims: dict[str, object] | None = None + jwt_claims: Mapping[str, object] | None = None raw_token: Final[str | None] = hook_data.get("incoming_bearer_token") if self.access_token_discovery_uri and raw_token: diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py index 1956170ba67..1790f335548 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py @@ -3,12 +3,13 @@ import time import uuid from collections import OrderedDict from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import httpx from typing_extensions import NotRequired, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items, json_contains from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) @@ -36,6 +37,17 @@ class GraphTokenResponse(TypedDict): expires_in: NotRequired[int] +def _policy_action_blocks(action: object) -> bool: + policy_action: Final = decoded_json_object(action) + odata_type: Final = policy_action.get("@odata.type", "") + action_field: Final = policy_action.get("action", "") + + if json_contains(odata_type, "restrictAccessAction") or action_field == "restrictAccess": + restriction: Final = policy_action.get("restrictionAction", "") + return restriction == "block" + return False + + class PurviewGuardrailBase: """ Base class for Microsoft Purview guardrails. @@ -364,16 +376,11 @@ class PurviewGuardrailBase: # ------------------------------------------------------------------ @staticmethod - def _should_block(response: dict[str, Any]) -> bool: + def _should_block(response: Mapping[str, object]) -> bool: """Return True if any policyAction requires blocking.""" - for action in response.get("policyActions", []): - odata_type = action.get("@odata.type", "") - action_field = action.get("action", "") - - if "restrictAccessAction" in odata_type or action_field == "restrictAccess": - restriction = action.get("restrictionAction", "") - if restriction == "block": - return True + for action in iter_items(response.get("policyActions", [])): + if _policy_action_blocks(action): + return True return False # ------------------------------------------------------------------ diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index b9871b1e276..721064d1169 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1,7 +1,8 @@ +import contextlib import time -from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Generator, Mapping, Sequence from enum import Enum, auto -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Protocol, runtime_checkable +from typing import TYPE_CHECKING, ClassVar, Final, Literal, Protocol, runtime_checkable import httpx from fastapi import HTTPException @@ -302,7 +303,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): def _log_request_debug( self, url: str, - body: dict, + body: Mapping[str, object], file_bytes: "bytes | None", file_type: "str | None", ) -> None: @@ -1162,11 +1163,17 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: Any, - request_data: dict, + response: AsyncIterable[object], + request_data: dict[str, object], ) -> AsyncGenerator[ModelResponseStream, None]: """Process streaming response chunks.""" + async with contextlib.aclosing(self._moderated_stream(response, request_data)) as moderated: + async for moderated_item in moderated: + yield moderated_item + async def _moderated_stream( + self, response: AsyncIterable[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, diff --git a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py index 7c26c4c4c56..021af0f7586 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py +++ b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py @@ -6,17 +6,20 @@ # +-------------------------------------------------------------+ import os import uuid +from collections.abc import Sequence from typing import TYPE_CHECKING, Final, Literal, Optional import httpx from fastapi import HTTPException from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -31,6 +34,11 @@ if TYPE_CHECKING: _DECODED_JSON: Final = TypeAdapter(object) +class _OnyxGuardVerdict(TypedDict, total=False): + allowed: ReadOnly[bool] + violated_rules: ReadOnly[Sequence[str]] + + class OnyxGuardrail(CustomGuardrail): @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: @@ -69,7 +77,7 @@ class OnyxGuardrail(CustomGuardrail): payload: object, input_type: Literal["request", "response"], conversation_id: str, - ) -> dict: + ) -> _OnyxGuardVerdict: """ Call external Onyx Guard server for validation """ @@ -85,8 +93,8 @@ class OnyxGuardrail(CustomGuardrail): }, ) response.raise_for_status() - result: Final = response.json() - if not result.get("allowed", True): + result: Final[_OnyxGuardVerdict] = response.json() + if not decoded_json_object(_DECODED_JSON.validate_python(result)).get("allowed", True): detection_message = "Unknown violation" if "violated_rules" in result: detection_message = ", ".join(result["violated_rules"]) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py index ddc6485b7d5..50a81506bb1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py @@ -21,6 +21,7 @@ _NO_TOOL_INPUT: Final[Mapping[str, object]] = MappingProxyType({}) _DATA_URL_RE: Final = re.compile(r"^data:(?P[^;,]+)?(?P(?:;[^;,]+)*?)(?P;base64)?,", re.IGNORECASE) _URLSAFE_TO_STANDARD_B64: Final = str.maketrans("-_", "+/") _OBJECT_LIST: Final = TypeAdapter(list[object]) +_DECODED_JSON: Final = TypeAdapter(object) def str_or_none(value: object) -> str | None: @@ -240,7 +241,7 @@ def _parsed_tool_input(raw_arguments: str) -> Mapping[str, object]: if not raw_arguments: return _NO_TOOL_INPUT try: - parsed: Final = json.loads(raw_arguments) + parsed: Final = _DECODED_JSON.validate_python(json.loads(raw_arguments)) except (ValueError, TypeError): return _NO_TOOL_INPUT return as_str_mapping(parsed) or _NO_TOOL_INPUT @@ -325,13 +326,17 @@ def _declarations_in(message_index: int, message: object) -> Iterator[_Declarati yield _Declaration(call_id, message_index, name) +def _declaration_call_id(declaration: _Declaration) -> str: + return declaration.call_id + + def _declarations_by_call_id(messages: Sequence[object]) -> Mapping[str, tuple[_Declaration, ...]]: """Every named declaration of each tool call id, in message order, built in one pass.""" declarations: Final = chain.from_iterable( _declarations_in(index, message) for index, message in enumerate(messages) ) ordered: Final = sorted(declarations, key=attrgetter("call_id", "message_index")) - return MappingProxyType({call_id: tuple(group) for call_id, group in groupby(ordered, key=attrgetter("call_id"))}) + return MappingProxyType({call_id: tuple(group) for call_id, group in groupby(ordered, key=_declaration_call_id)}) def _resolve_tool_name( diff --git a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py index 4ec3ed21359..152515c50ed 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py @@ -4,6 +4,7 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Final from fastapi import HTTPException +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache @@ -11,6 +12,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.llms.anthropic.common_utils import ( message_dict_with_server_tool_calls, message_dict_with_server_tool_calls_split_out, @@ -35,6 +37,8 @@ from litellm.types.utils import ( if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +_DECODED_JSON: Final = TypeAdapter(object) + def _choice_with_message(choice: object, rewrite: Callable[[object], object]) -> object: mapping: Final = validated_object_mapping(choice) @@ -150,9 +154,13 @@ class PangeaHandler(CustomGuardrail): response.raise_for_status() result: Final = response.json() + decoded_result: Final = decoded_json_object(_DECODED_JSON.validate_python(result)) + verdict: Final = decoded_json_object(decoded_result.get("result", {})) - if result.get("result", {}).get("blocked"): - verbose_proxy_logger.warning("Pangea Guardrail (%s): Request blocked. Response: %s", hook_name, result) + if verdict.get("blocked"): + verbose_proxy_logger.warning( + "Pangea Guardrail (%s): Request blocked. Response: %s", hook_name, decoded_result + ) raise HTTPException( status_code=400, # Bad Request, indicating violation detail={ @@ -161,7 +169,7 @@ class PangeaHandler(CustomGuardrail): }, ) verbose_proxy_logger.debug( - "Pangea Guardrail (%s): Request passed. Response: %s", hook_name, result.get("result", {}).get("detectors") + "Pangea Guardrail (%s): Request passed. Response: %s", hook_name, verdict.get("detectors") ) return result diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 379c2c4b075..73ccd82d37f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -12,6 +12,7 @@ import os import re from collections.abc import AsyncIterable, Mapping, Sequence from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias from urllib.parse import urlparse @@ -27,6 +28,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.llms.anthropic.common_utils import ( anthropic_server_tool_call_params, tool_calls_with_server_tool_calls, @@ -143,6 +145,7 @@ class _ResponsesInputItem(LiteLLMBaseModel): _ResponsesInput: TypeAlias = str | tuple[_ResponsesInputItem, ...] | None _RESPONSES_INPUT: Final[TypeAdapter[_ResponsesInput]] = TypeAdapter(_ResponsesInput) _DECODED_ARGUMENTS: Final = TypeAdapter(object) +_NO_METADATA: Final[Mapping[str, object]] = MappingProxyType({}) class _AirsToolEvent(TypedDict): @@ -901,7 +904,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): }, ) - def _prepare_metadata_from_request(self, data: dict[str, Any]) -> dict[str, object]: + def _prepare_metadata_from_request(self, data: Mapping[str, object]) -> dict[str, object]: """ Extract and prepare metadata from request data for PANW API call. @@ -915,8 +918,8 @@ class PanwPrismaAirsHandler(CustomGuardrail): the profile linked to the API key (configured in Strata Cloud Manager). If both are provided, PANW API uses profile_id (profile_id takes precedence). """ - user_metadata: Final = data.get("metadata", {}) or {} - requester_meta: Final = user_metadata.get("requester_metadata", {}) or {} + user_metadata: Final = decoded_json_object(data.get("metadata") or _NO_METADATA) + requester_meta: Final = user_metadata.get("requester_metadata") or _NO_METADATA metadata: Final[dict[str, object]] = { "user": data.get("user") or "litellm_user", "model": data.get("model") or "unknown", @@ -924,7 +927,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): # Pass through PANW API fields (check requester_metadata fallback for /v1/messages routes) for key in ("profile_name", "profile_id", "user_ip", "app_name", "app_user"): - val = user_metadata.get(key) or requester_meta.get(key) + val = user_metadata.get(key) or decoded_json_object(requester_meta).get(key) if val: metadata[key] = val @@ -939,7 +942,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): data.get("litellm_trace_id") or user_metadata.get("litellm_trace_id") or user_metadata.get("trace_id") - or requester_meta.get("litellm_trace_id") + or decoded_json_object(requester_meta).get("litellm_trace_id") ) if trace_id: metadata["litellm_trace_id"] = trace_id diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index e136b9baaa8..91d196f184f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -8,6 +8,7 @@ # Standard library imports import json import os +from collections.abc import Sized from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import quote @@ -73,6 +74,12 @@ class _PillarProtectHTTPClient(Protocol): ) -> _PillarProtectHTTPResponse: ... +def _json_len(value: object) -> int: + if isinstance(value, Sized): + return len(value) + raise TypeError(f"object of type '{type(value).__name__}' has no len()") + + def _encode_json_for_header(data: object) -> str: """ JSON-serialize and URL-encode data for safe header transmission. @@ -725,7 +732,7 @@ class PillarGuardrail(CustomGuardrail): ) return payload - async def _call_pillar_api(self, headers: dict[str, str], payload: dict[str, Any]) -> _PillarProtectResponse: + async def _call_pillar_api(self, headers: dict[str, str], payload: dict[str, object]) -> _PillarProtectResponse: """ Call the Pillar API and return the response. @@ -737,7 +744,7 @@ class PillarGuardrail(CustomGuardrail): Pillar API response as dictionary """ verbose_proxy_logger.debug( - "Pillar Guardrail: Scanning %s messages for security threats", len(payload.get("messages", [])) + "Pillar Guardrail: Scanning %s messages for security threats", _json_len(payload.get("messages", [])) ) response: Final = await self.async_handler.post( url=f"{self.api_base}/api/v1/protect", diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index c76e8e9c084..7bd72b7323d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -7,6 +7,7 @@ import json import os +from collections.abc import Mapping from typing import Any, Final, Literal from fastapi import HTTPException @@ -17,6 +18,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params from litellm.llms.custom_httpx.http_handler import ( @@ -247,7 +249,7 @@ class QualifireGuardrail(CustomGuardrail): return api_tools if api_tools else None - def _check_if_flagged(self, result: dict[str, Any]) -> bool: + def _check_if_flagged(self, result: Mapping[str, object]) -> bool: """ Check if the Qualifire evaluation result indicates flagged content. @@ -255,12 +257,12 @@ class QualifireGuardrail(CustomGuardrail): A high score (close to 100) indicates GOOD content, low score indicates problems. """ # Check evaluation results for any flagged items - evaluation_results: Final = result.get("evaluationResults", []) or [] + evaluation_results: Final = result.get("evaluationResults") or () - for eval_result in evaluation_results: - results = eval_result.get("results", []) or [] - for r in results: - if r.get("flagged"): + for eval_result in iter_items(evaluation_results): + results = decoded_json_object(eval_result).get("results") or () + for r in iter_items(results): + if decoded_json_object(r).get("flagged"): return True return False @@ -384,7 +386,7 @@ class QualifireGuardrail(CustomGuardrail): timeout=self.timeout, ) response.raise_for_status() - result: Final = response.json() + result: Final = decoded_json_object(_DECODED_JSON.validate_python(response.json())) # Extract response info for logging qualifire_response: Final[dict[str, object]] = { diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py index fa5251ba558..2e712f027cb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -110,7 +110,7 @@ class SingulrGuardrail(CustomGuardrail): return SingulrGuardrailConfigModel @staticmethod - def _metadata_containers(request_data: Mapping[str, Any]) -> tuple[Mapping[str, Any], ...]: + def _metadata_containers(request_data: Mapping[str, Any]) -> tuple[Mapping[str, object], ...]: litellm_params: Final = request_data.get("litellm_params") or _EMPTY_MAPPING return tuple( container @@ -124,7 +124,7 @@ class SingulrGuardrail(CustomGuardrail): ) @classmethod - def _resolve_metadata_value(cls, request_data: Mapping[str, object], key: str) -> str | None: + def _resolve_metadata_value(cls, request_data: Mapping[str, object], key: str) -> object: for container in cls._metadata_containers(request_data=request_data): value = container.get(key) if value: @@ -140,7 +140,7 @@ class SingulrGuardrail(CustomGuardrail): return None @classmethod - def _build_metadata(cls, request_data: Mapping[str, object]) -> Mapping[str, str] | None: + def _build_metadata(cls, request_data: Mapping[str, object]) -> Mapping[str, object] | None: fields: Final = ( "user_api_key_alias", "user_api_key_user_id", @@ -225,7 +225,7 @@ class SingulrGuardrail(CustomGuardrail): inputs: GenericGuardrailAPIInputs, texts: Sequence[str], structured_messages: Sequence[AllMessageValues], - request_data: Mapping[str, Any], + request_data: Mapping[str, object], ) -> GenericGuardrailAPIInputs: messages: Final = ( tuple(structured_messages) @@ -242,14 +242,16 @@ class SingulrGuardrail(CustomGuardrail): metadata: Final = self._build_metadata(request_data=request_data) - singulr_req_obj = SingulrGuardrailPayload( - correlation_id=request_data.get("litellm_call_id"), - model_name=inputs.get("model"), - guardrail_scope="request", - messages=messages, - images=images, - tools=tools, - metadata=metadata, + singulr_req_obj = SingulrGuardrailPayload.model_validate( + { + "correlation_id": request_data.get("litellm_call_id"), + "model_name": inputs.get("model"), + "guardrail_scope": "request", + "messages": messages, + "images": images, + "tools": tools, + "metadata": metadata, + } ) payload = singulr_req_obj.model_dump(mode="json") guardrail_resp = await self._call_api(payload) @@ -267,7 +269,7 @@ class SingulrGuardrail(CustomGuardrail): return inputs @staticmethod - def _mcp_tool_name(request_data: Mapping[str, Any]) -> str | None: + def _mcp_tool_name(request_data: Mapping[str, object]) -> object: return request_data.get("mcp_tool_name") or request_data.get("name") @staticmethod @@ -283,15 +285,17 @@ class SingulrGuardrail(CustomGuardrail): model: Final = request_data.get("model") return "mcp_tool_name" in request_data or (isinstance(model, str) and model.startswith(_MCP_MODEL_PREFIX)) - async def _apply_guardrail_on_mcp_request(self, request_data: Mapping[str, Any]) -> None: + async def _apply_guardrail_on_mcp_request(self, request_data: Mapping[str, object]) -> None: metadata: Final = self._build_metadata(request_data=request_data) - singulr_mcp_obj = SingulrMcpGuardrailPayload( - guardrail_scope="mcp_request", - tool_name=self._mcp_tool_name(request_data), - tool_arguments=self._mcp_arguments(request_data), - mcp_server_name=request_data.get("mcp_server_name"), - metadata=metadata, + singulr_mcp_obj = SingulrMcpGuardrailPayload.model_validate( + { + "guardrail_scope": "mcp_request", + "tool_name": self._mcp_tool_name(request_data), + "tool_arguments": self._mcp_arguments(request_data), + "mcp_server_name": request_data.get("mcp_server_name"), + "metadata": metadata, + } ) payload = singulr_mcp_obj.model_dump(mode="json") guardrail_resp = await self._call_api(payload) @@ -308,18 +312,20 @@ class SingulrGuardrail(CustomGuardrail): ) async def _apply_guardrail_on_mcp_response( - self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, Any] + self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, object] ) -> GenericGuardrailAPIInputs: if not texts: return inputs metadata: Final = self._build_metadata(request_data=request_data) - singulr_mcp_obj = SingulrMcpGuardrailPayload( - model_name=request_data.get("model"), - guardrail_scope="mcp_response", - tool_result=texts, - metadata=metadata, + singulr_mcp_obj = SingulrMcpGuardrailPayload.model_validate( + { + "model_name": request_data.get("model"), + "guardrail_scope": "mcp_response", + "tool_result": texts, + "metadata": metadata, + } ) payload = singulr_mcp_obj.model_dump(mode="json") guardrail_resp = await self._call_api(payload) @@ -358,7 +364,7 @@ class SingulrGuardrail(CustomGuardrail): ) async def _apply_guardrail_on_response( - self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, Any] + self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, object] ) -> GenericGuardrailAPIInputs: combined_texts: Final = "\n".join(texts) if texts else None @@ -377,14 +383,16 @@ class SingulrGuardrail(CustomGuardrail): metadata: Final = self._build_metadata(request_data=request_data) - singulr_resp_obj = SingulrGuardrailPayload( - correlation_id=request_data.get("litellm_call_id"), - guardrail_scope="response", - model_name=request_data.get("model"), - messages=request_data.get("messages"), - images=inputs.get("images"), - response=assistant_message, - metadata=metadata, + singulr_resp_obj = SingulrGuardrailPayload.model_validate( + { + "correlation_id": request_data.get("litellm_call_id"), + "guardrail_scope": "response", + "model_name": request_data.get("model"), + "messages": request_data.get("messages"), + "images": inputs.get("images"), + "response": assistant_message, + "metadata": metadata, + } ) payload = singulr_resp_obj.model_dump(mode="json") diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py index fbf1f2576ff..1feb6d533f4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -27,7 +27,10 @@ from pydantic import TypeAdapter from typing_extensions import TypedDict, Unpack from litellm._logging import verbose_proxy_logger -from litellm.completion_extras.litellm_responses_transformation.transformation import replayed_tool_search_tool_names +from litellm.completion_extras.litellm_responses_transformation.transformation import ( + nested_tool_member_names, + replayed_tool_search_tool_names, +) from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, @@ -47,6 +50,14 @@ class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" +def _function_tool_name(tool: object) -> str | None: + match tool: + case {"function": {"name": str(name)}}: + return name + case _: + return None + + def _attribute(source: object, name: str) -> object: return getattr(source, name, None) @@ -140,17 +151,18 @@ class ToolPolicyGuardrail(CustomGuardrail): """ if input_type == "request": tools: Final = inputs.get("tools") or [] - function_tool_names: Final = [ - t["function"]["name"] - for t in tools - if isinstance(t, dict) and isinstance(t.get("function"), dict) and t["function"].get("name") - ] + function_tool_names: Final = [name for name in map(_function_tool_name, tools) if name] route: Final = _get_request_route_from_data(request_data) - replayed_tool_names: Final = replayed_tool_search_tool_names(_REQUEST_DATA.validate_python(request_data)) - tool_names = ( - extract_request_tool_names(route, request_data) - if route and not function_tool_names - else [*function_tool_names, *replayed_tool_names] + request: Final = _REQUEST_DATA.validate_python(request_data) + tool_names = list( + dict.fromkeys( + ( + *function_tool_names, + *(extract_request_tool_names(route, request_data) if route else ()), + *replayed_tool_search_tool_names(request), + *nested_tool_member_names(request), + ) + ) ) else: tool_calls: Final = inputs.get("tool_calls") or [] diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 24b30e65e2b..2f2294db522 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -644,6 +644,9 @@ def _configure_callback_scoping( ) logging_only_scope: Final = None if logging_only_scope_error is not None else litellm_params.logging_only_scope custom_guardrail_callback.logging_only_scope = logging_only_scope + custom_guardrail_callback.mcp_block_behavior = ( # rebind-ok: initialize provider-independent callback configuration + litellm_params.mcp_block_behavior + ) custom_guardrail_callback.logging_only_continue_on_input_failure = ( logging_only_continue_error is None and logging_only_scope is None diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index c4fe6d6506b..b4d343b3932 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -39,6 +39,7 @@ from litellm.batches.batch_utils import ( from litellm.constants import BATCH_TPD_DESCRIPTOR_SUFFIX, BATCH_TPD_WINDOW_SECONDS from litellm.exceptions import RateLimitErrorCategory from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.proxy._types import ( ProxyErrorTypes, ProxyException, @@ -100,6 +101,7 @@ else: _BATCH_BODY_ADAPTER: Final = TypeAdapter(dict[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) _WINDOW_START_ADAPTER: Final[TypeAdapter[int | float | str | None]] = TypeAdapter(int | float | str | None) @@ -985,7 +987,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): for raw_line in iter_batch_input_lines(file_content_bytes): request_count += 1 try: - entry = json.loads(raw_line) + entry = _DECODED_JSON.validate_python(json.loads(raw_line)) except Exception: entry_total_tokens = estimate_batch_entry_tokens(raw_line) entry_output_tokens = self.parallel_request_limiter.no_max_tokens_output_floor( @@ -1008,7 +1010,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): output_tokens += entry_output_tokens try: - entry_total_tokens = count_entry_tokens(entry) + entry_total_tokens = count_entry_tokens(decoded_json_object(entry)) except Exception: entry_total_tokens = estimate_batch_entry_tokens(raw_line) total_tokens += entry_total_tokens diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 8f9aebecf30..8784c02ac41 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -2,13 +2,14 @@ import asyncio import json from collections.abc import Sequence from datetime import datetime, timezone -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import attribute_or_none, has_attributes from litellm.proxy._types import ( GenerateKeyRequest, GenerateKeyResponse, @@ -34,6 +35,10 @@ if TYPE_CHECKING: LITELLM_PREFIX_STORED_VIRTUAL_KEYS: Final = "litellm/" +class _ModelDumpable(Protocol): + def model_dump(self) -> object: ... + + class KeyManagementEventHooks: @staticmethod async def async_key_generated_hook( @@ -476,12 +481,13 @@ class KeyManagementEventHooks: verbose_proxy_logger.debug("Unable to load team metadata for team_id=%s: %s", team_id, exc) return None - metadata = getattr(team_obj, "metadata", None) - if metadata is None: + raw_metadata: Final = attribute_or_none(team_obj, "metadata") + if raw_metadata is None: return None - if hasattr(metadata, "model_dump"): - metadata = metadata.model_dump() + metadata: Final = ( + raw_metadata.model_dump() if has_attributes(raw_metadata, _ModelDumpable, "model_dump") else raw_metadata + ) if not isinstance(metadata, dict): return None diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index 6bd8417eb5e..576e3997db0 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -36,6 +36,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import has_attributes, iter_items, require_attributes from litellm.llms.litellm_proxy.skills.constants import LITELLM_SKILL_ID_PREFIX from litellm.llms.litellm_proxy.skills.prompt_injection import ( SkillPromptInjectionHandler, @@ -84,6 +85,11 @@ class _ChatCompletion(Protocol): def choices(self) -> Sequence[_ChatChoice]: ... +class _HasContent(Protocol): + @property + def content(self) -> object: ... + + @with_config(ConfigDict(extra="allow", strict=True)) class _CodeExecutionToolArgs(TypedDict, total=False): code: ReadOnly[str] @@ -105,6 +111,10 @@ def _first_choice(response: _ChatCompletion) -> _ChatChoice: return response.choices[0] +def _dict_or_none(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, dict) else None + + class SkillsInjectionHook(CustomLogger): """ Pre/Post-call hook that processes skills from container.skills parameter. @@ -456,26 +466,28 @@ class SkillsInjectionHook(CustomLogger): skill_files=all_skill_files, ) - def _extract_tool_calls(self, response: Any) -> list[dict[str, Any]]: + def _extract_tool_calls(self, response: object) -> list[dict[str, Any]]: """Extract tool calls from response, handling both formats.""" tool_calls: Final = [] # Get content - handle both dict and object responses - content = None - if isinstance(response, dict): - content = response.get("content", []) - elif hasattr(response, "content"): - content = response.content + response_mapping: Final = _dict_or_none(response) + content: Final = ( + response_mapping.get("content", []) + if response_mapping is not None + else (response.content if has_attributes(response, _HasContent, "content") else None) + ) # Anthropic/messages API format: response has "content" list with tool_use blocks if content: - for block in content: - if isinstance(block, dict) and block.get("type") == "tool_use": + for block in iter_items(content): + block_mapping = _dict_or_none(block) + if block_mapping is not None and block_mapping.get("type") == "tool_use": tool_calls.append( { - "id": block.get("id"), - "name": block.get("name"), - "input": block.get("input", {}), + "id": block_mapping.get("id"), + "name": block_mapping.get("name"), + "input": block_mapping.get("input", {}), } ) elif hasattr(block, "type") and getattr(block, "type", None) == "tool_use": @@ -488,7 +500,7 @@ class SkillsInjectionHook(CustomLogger): ) # OpenAI format: response has choices[0].message.tool_calls - if not tool_calls and hasattr(response, "choices") and response.choices: + if not tool_calls and has_attributes(response, _ChatCompletion, "choices") and response.choices: msg: Final = response.choices[0].message if hasattr(msg, "tool_calls") and msg.tool_calls: for tc in msg.tool_calls: @@ -770,11 +782,11 @@ print('No executable skill module found') executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) generated_files: Final[list[dict[str, object]]] = [] - current_response: Any = response + current_response: object = response for iteration in range(self.max_iterations): # OpenAI format response has choices[0].message - choice: _ChatChoice = _first_choice(current_response) + choice: _ChatChoice = _first_choice(require_attributes(current_response, _ChatCompletion, "choices")) assistant_message: _ChatMessage = choice.message stop_reason: str | None = choice.finish_reason diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py index 5b7d7d04e24..f12756bc1e6 100644 --- a/litellm/proxy/hooks/max_iterations_limiter.py +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -191,7 +191,7 @@ class _PROXY_MaxIterationsHandler(CustomLogger): """ if self.increment_script is not None: try: - result: Final = await self.increment_script( + result: Final[int] = await self.increment_script( keys=[cache_key], args=[self.ttl], ) @@ -207,7 +207,7 @@ class _PROXY_MaxIterationsHandler(CustomLogger): async def _in_memory_increment(self, cache_key: str) -> int: """Increment counter in in-memory cache with TTL.""" - current: Final = await self.internal_usage_cache.async_get_cache( + current: Final[int | None] = await self.internal_usage_cache.async_get_cache( key=cache_key, litellm_parent_otel_span=None, local_only=True, diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index c6e26999766..1830961dd06 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -6,7 +6,7 @@ Reduces context window size and improves tool selection accuracy. """ from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence -from typing import TYPE_CHECKING, Final, Optional +from typing import TYPE_CHECKING, Final, Optional, Protocol from fastapi import HTTPException from typing_extensions import ReadOnly, TypedDict @@ -18,6 +18,7 @@ from litellm.constants import ( DEFAULT_MCP_SEMANTIC_FILTER_TOP_K, ) from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import as_str_mapping, has_attributes from litellm.proxy.mcp_server.semantic_tool_filter import ( SemanticToolFilterContextWindowError, ) @@ -50,6 +51,35 @@ def _truncate_csv_at_tool_name_boundary(tool_names_csv: str, max_length: int) -> return head.rsplit(",", 1)[0] +class _PydanticTool(Protocol): + def model_dump(self, *, exclude_none: bool) -> Mapping[str, object]: ... + + +class _PydanticV1Tool(Protocol): + def dict(self, *, exclude_none: bool) -> Mapping[str, object]: ... + + +def _tool_as_dict(tool: object) -> object: + if has_attributes(tool, _PydanticTool, "model_dump"): + tool_dict: Final = tool.model_dump(exclude_none=True) + verbose_proxy_logger.debug( + "Converted Pydantic tool to dict: %s -> dict with keys: %s", + type(tool).__name__, + list(tool_dict.keys()), + ) + return tool_dict + if has_attributes(tool, _PydanticV1Tool, "dict"): + legacy_tool_dict: Final = tool.dict(exclude_none=True) + verbose_proxy_logger.debug("Converted Pydantic tool (v1) to dict: %s -> dict", type(tool).__name__) + return legacy_tool_dict + tool_mapping: Final = as_str_mapping(tool) + if tool_mapping is not None and isinstance(tool, dict): + verbose_proxy_logger.debug("Tool is already a dict with keys: %s", list(tool_mapping.keys())) + return tool_mapping + verbose_proxy_logger.warning("Tool is unknown type: %s, passing as-is", type(tool)) + return tool + + class SemanticToolFilterHook(CustomLogger): """ Pre-call hook that filters MCP tools semantically. @@ -92,7 +122,7 @@ class SemanticToolFilterHook(CustomLogger): self, tools: Iterable[Mapping[str, object]], user_api_key_dict: "UserAPIKeyAuth", - ) -> list[dict[str, object]]: + ) -> list[object]: """ Expand MCP references to actual tool definitions. @@ -119,26 +149,7 @@ class SemanticToolFilterHook(CustomLogger): ) # Convert Pydantic models to dicts for compatibility - openai_tools_as_dicts: Final[list[dict[str, object]]] = [] - for tool in openai_tools: - if hasattr(tool, "model_dump"): - tool_dict = tool.model_dump(exclude_none=True) - verbose_proxy_logger.debug( - "Converted Pydantic tool to dict: %s -> dict with keys: %s", - type(tool).__name__, - list(tool_dict.keys()), - ) - openai_tools_as_dicts.append(tool_dict) - elif hasattr(tool, "dict"): - tool_dict = tool.dict(exclude_none=True) - verbose_proxy_logger.debug("Converted Pydantic tool (v1) to dict: %s -> dict", type(tool).__name__) - openai_tools_as_dicts.append(tool_dict) - elif isinstance(tool, dict): - verbose_proxy_logger.debug("Tool is already a dict with keys: %s", list(tool.keys())) - openai_tools_as_dicts.append(tool) - else: - verbose_proxy_logger.warning("Tool is unknown type: %s, passing as-is", type(tool)) - openai_tools_as_dicts.append(tool) + openai_tools_as_dicts: Final = list(map(_tool_as_dict, openai_tools)) verbose_proxy_logger.debug( "Expanded %s MCP reference(s) to %s tools (all as dicts)", len(mcp_tools), len(openai_tools_as_dicts) @@ -149,8 +160,8 @@ class SemanticToolFilterHook(CustomLogger): async def _filter_expanded_tools( self, data: dict, - expanded_tools: list[dict[str, object]], - ) -> list[dict[str, object]]: + expanded_tools: list[object], + ) -> list[object]: """ Apply the semantic filter to expanded MCP tool definitions. diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 33a14bfbeda..6b36a6ac8bc 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -4743,7 +4743,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reserved_tokens=reserved_tokens, ) ) - ptu_skips_socket_usage: Final = kwargs.get("call_type") == CallTypes.aresponses_websocket.value + call_type: Final[object] = kwargs.get("call_type") + ptu_skips_socket_usage: Final = call_type == CallTypes.aresponses_websocket.value pipeline_operations.extend( self._build_team_ptu_tpm_ops( standard_logging_metadata, # pyright: ignore[reportUnknownArgumentType] # untyped logging metadata diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 0269a784445..0ab3d1ca7af 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -5,9 +5,9 @@ This hook uses the DBSpendUpdateWriter to batch-write response IDs to the databa instead of writing immediately on each request. """ -from collections.abc import AsyncGenerator, Callable, Mapping +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypeVar +from typing import TYPE_CHECKING, Final, TypeVar from fastapi import HTTPException from pydantic import TypeAdapter, ValidationError @@ -460,8 +460,8 @@ class ResponsesIDSecurity(CustomLogger): return response async def async_post_call_streaming_iterator_hook( - self, user_api_key_dict: "UserAPIKeyAuth", response: Any, request_data: dict[str, object] - ) -> AsyncGenerator[BaseLiteLLMOpenAIResponseObject, None]: + self, user_api_key_dict: "UserAPIKeyAuth", response: AsyncIterable[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: advertise: Final = self._streamed_response_id_advertiser(user_api_key_dict, request_data) async for chunk in response: yield ( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..546521f9048 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -14,6 +14,7 @@ from typing import ( Any, Final, Literal, + Protocol, TypeAlias, cast, # noqa: TID251 # narrows untyped deployment and credential payloads ) @@ -49,7 +50,13 @@ from litellm.constants import ( X_LITELLM_DISABLE_CALLBACKS, ) from litellm.integrations.custom_guardrail import merged_guardrail_entries, without_server_streaming_classification -from litellm.litellm_core_utils.core_helpers import is_codex_user_agent, normalize_drop_params +from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, + attribute_or_none, + is_codex_user_agent, + normalize_drop_params, + require_attributes, +) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( TRUSTED_CALLBACK_VARS_FIELD, @@ -2664,11 +2671,11 @@ async def add_litellm_data_to_request( _raw_headers: Final[dict[str, str]] = RedactedDict(safe_get_request_headers(request)) - forward_llm_auth = False - if general_settings: - 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) + general_settings_view: Final = as_str_mapping(general_settings) + forward_llm_auth: Final = bool( + (general_settings_view.get("forward_llm_provider_auth_headers", False) if general_settings_view else False) + or attribute_or_none(litellm, "forward_llm_provider_auth_headers") + ) _headers: Final[dict[str, str]] = clean_headers( request.headers, litellm_key_header_name=( @@ -3562,6 +3569,10 @@ def _warn_stale_team_alias_once(warning_key: str, message: str, *args: str) -> N verbose_proxy_logger.warning(message, *args) +class _SupportsStartswith(Protocol): + def startswith(self, prefix: str, /) -> bool: ... + + def _update_model_if_team_alias_exists( data: dict, user_api_key_dict: UserAPIKeyAuth, @@ -3589,14 +3600,15 @@ def _update_model_if_team_alias_exists( exist (e.g. a gateway-level model group shared with the team). """ _model: Final = data.get("model") - if not _model or not user_api_key_dict.team_model_aliases or _model not in user_api_key_dict.team_model_aliases: + team_model_aliases: Final = as_str_mapping(user_api_key_dict.team_model_aliases) + if not _model or not team_model_aliases or _model not in team_model_aliases: return from litellm.proxy.proxy_server import llm_router # Skip alias rewrite if this model resolves to team-specific deployments # (team models use team_public_model_name, not model_aliases) - aliased_target: Final = user_api_key_dict.team_model_aliases[_model] + aliased_target: Final = team_model_aliases[_model] # Optional bypass for stale aliases from pre-PR deployments: # only enabled via feature flag to preserve backwards compatibility. @@ -3607,7 +3619,9 @@ def _update_model_if_team_alias_exists( enable_stale_alias_bypass: Final = _ENABLE_TEAM_STALE_ALIAS_BYPASS # Check if the alias points to a team-scoped UUID name # (format: "model_name_{team_id}_{uuid}") - is_stale_team_alias: Final = aliased_target.startswith(f"model_name_{user_api_key_dict.team_id}_") + is_stale_team_alias: Final = require_attributes(aliased_target, _SupportsStartswith, "startswith").startswith( + f"model_name_{user_api_key_dict.team_id}_" + ) if is_stale_team_alias and llm_router: if aliased_target not in llm_router.registry().index("model_name"): _warn_stale_team_alias_once( diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 42321fe1b1d..1e43a8375a6 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -236,8 +236,10 @@ _RAW_ROWS: Final = TypeAdapter(Sequence[dict[str, object]], config=ConfigDict(st async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) -> Sequence[Mapping[str, object]]: + from litellm.proxy.utils import query_raw_rows + async with db_span("auto_router_report_query", sql_relation(query)): - return _RAW_ROWS.validate_python(await prisma_client.db.query_raw(query, *args)) + return _RAW_ROWS.validate_python(await query_raw_rows(prisma_client.db, query, *args)) async def _execute_raw(prisma_client: "PrismaClient", query: str, *args: object) -> int: @@ -939,7 +941,7 @@ def _idle_router_groups( zero: Final = _summed_agg_row(()) idle: Final = frozenset( key - for key in (_strategy_router_key(deployment) for deployment in llm_router.registry().all() or ()) + for key in map(_strategy_router_key, llm_router.registry().all() or ()) if key is not None and key not in covered ) return tuple( @@ -1112,6 +1114,7 @@ def _is_configured_pre_routing_strategy(llm_router: "Router", router_name: str) router_registry.plugins("complexity"), router_registry.plugins("adaptive"), router_registry.plugins("quality"), + router_registry.plugins("switchyard"), ) ) diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index aaf5d30d2e1..a09045829b8 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -12,10 +12,11 @@ import asyncio import json from collections.abc import Mapping from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Final, Protocol from fastapi import APIRouter, Depends, Header, HTTPException from pydantic import Field, TypeAdapter +from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger from litellm._redis import _redis_kwargs_from_environment @@ -84,7 +85,11 @@ _REDACTED_VALUE: Final = "***REDACTED***" _URL_OVERRIDDEN_CONNECTION_FIELDS: Final[frozenset[str]] = frozenset({"host", "port", "db", "password", "username"}) -def _resolve_cache_url_precedence(settings: Mapping[str, object]) -> dict[str, Any]: +class _CacheSettings(TypedDict, total=False, extra_items=object): + pass + + +def _resolve_cache_url_precedence(settings: Mapping[str, object]) -> _CacheSettings: """Return cache settings with the url-vs-discrete-fields ambiguity resolved. When a full ``url`` is supplied it wins: the discrete @@ -97,8 +102,8 @@ def _resolve_cache_url_precedence(settings: Mapping[str, object]) -> dict[str, A url: Final = settings.get("url") has_url: Final = isinstance(url, str) and url.strip() != "" if not has_url or settings.get("redis_startup_nodes"): - return dict(settings) - return {k: v for k, v in settings.items() if k not in _URL_OVERRIDDEN_CONNECTION_FIELDS} + return _CacheSettings(**settings) + return _CacheSettings(**{k: v for k, v in settings.items() if k not in _URL_OVERRIDDEN_CONNECTION_FIELDS}) def _parse_stored_settings(cache_settings_value: object) -> dict[str, object]: diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index bbd6a1c9021..475608e15b6 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -5,7 +5,7 @@ from dataclasses import dataclass, replace from datetime import date, datetime, timedelta from itertools import groupby from types import MappingProxyType -from typing import Final, Literal, NoReturn, Protocol +from typing import Final, Literal, NoReturn, Protocol, TypeVar from fastapi import HTTPException, status from typing_extensions import ReadOnly, TypedDict, assert_never @@ -14,6 +14,7 @@ from litellm import constants from litellm._logging import verbose_proxy_logger from litellm.constants import PTU_SENTINEL_API_KEY, SUBSCRIPTION_BACKED_PROVIDERS from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_details, recover_cli_session_key_metadata, @@ -56,6 +57,8 @@ from litellm.types.repositories.daily_activity import ( SpendLogsWindow, ) +_T = TypeVar("_T") + @dataclass(frozen=True, slots=True) class ScopeDenied: @@ -188,9 +191,11 @@ class _AggregatedSpendData(TypedDict): totals: ReadOnly[SpendMetrics] -def _key_metadata(api_key_metadata: Mapping[str, KeyMetadataRow], api_key: str | None) -> KeyMetadata: +def _key_metadata(api_key_metadata: Mapping[str, KeyMetadataRow] | None, api_key: str | None) -> KeyMetadata: if not api_key: return KeyMetadata(key_alias="Unattributed", key_exists=False) + if api_key_metadata is None: + return KeyMetadata() meta: Final = api_key_metadata.get(api_key) return KeyMetadata( key_alias=meta.key_alias if meta is not None else None, @@ -306,7 +311,7 @@ def update_breakdown_metrics( record: DailySpendRecord, model_metadata: Mapping[str, dict[str, object]], provider_metadata: Mapping[str, dict[str, object]], - api_key_metadata: Mapping[str, KeyMetadataRow], + api_key_metadata: Mapping[str, KeyMetadataRow] | None, entity_id_field: str | None = None, entity_metadata_field: Mapping[str, dict[str, object]] | None = None, ) -> BreakdownMetrics: @@ -626,7 +631,7 @@ async def get_api_key_metadata( def _aggregate_spend_records_sync( *, records: Sequence[DailySpendRecord], - api_key_metadata: Mapping[str, KeyMetadataRow], + api_key_metadata: Mapping[str, KeyMetadataRow] | None, entity_id_field: str | None, entity_metadata_field: Mapping[str, dict[str, object]] | None, ) -> _AggregatedSpendData: @@ -685,22 +690,10 @@ async def _aggregate_spend_records( The per-row loop is offloaded to a worker thread via asyncio.to_thread so a large result set doesn't peg the event loop. """ - api_keys: Final[set[str]] = { - record.api_key for record in records if record.api_key and record.api_key != PTU_SENTINEL_API_KEY - } - - api_key_metadata: Final[Mapping[str, KeyMetadataRow]] = ( - await repository.key_metadata( - frozenset(api_keys), spend_logs_window(frozenset(record.date for record in records)) - ) - if api_keys - else MappingProxyType({}) - ) - return await asyncio.to_thread( _aggregate_spend_records_sync, records=records, - api_key_metadata=api_key_metadata, + api_key_metadata=await _key_metadata_unless_database_fails(repository, records), entity_id_field=entity_id_field, entity_metadata_field=entity_metadata_field, ) @@ -783,7 +776,7 @@ def _record_to_spend_metrics(record: RollupMetricsRow) -> SpendMetrics: def _aggregate_grouping_sets_records_sync( *, records: Sequence[GroupingSetsRow], - api_key_metadata: Mapping[str, KeyMetadataRow], + api_key_metadata: Mapping[str, KeyMetadataRow] | None, ) -> _AggregatedSpendData: """Build the response from rollup rows produced by the GROUPING SETS query. @@ -919,24 +912,40 @@ def _aggregate_grouping_sets_records_sync( return {"results": results, "totals": total_metrics} +async def unless_database_fails(lookup: Awaitable[_T], fallback: _T, left_out: str) -> _T: + try: + return await lookup + except Exception as exc: + if not PrismaDBExceptionHandler.is_database_service_unavailable_error(exc): + raise + verbose_proxy_logger.warning("Returning daily activity without %s: %s", left_out, exc) + return fallback + + +async def _key_metadata_unless_database_fails( + repository: DailyActivityRepository, + rows: Sequence[GroupingSetsRow] | Sequence[EntityRollupRow] | Sequence[DailySpendRecord], +) -> Mapping[str, KeyMetadataRow] | None: + api_keys: Final = frozenset(row.api_key for row in rows if row.api_key and row.api_key != PTU_SENTINEL_API_KEY) + if not api_keys: + return MappingProxyType({}) + return await unless_database_fails( + repository.key_metadata(api_keys, spend_logs_window(frozenset(row.date for row in rows))), + None, + "its key metadata", + ) + + async def _aggregate_grouping_sets_records( *, repository: DailyActivityRepository, records: Sequence[GroupingSetsRow], ) -> _AggregatedSpendData: """Async wrapper: fetch api_key_metadata, then dispatch on a worker thread.""" - api_keys: Final[set[str]] = {r.api_key for r in records if r.api_key and r.api_key != PTU_SENTINEL_API_KEY} - - api_key_metadata: Final[Mapping[str, KeyMetadataRow]] = ( - await repository.key_metadata(frozenset(api_keys), spend_logs_window(frozenset(r.date for r in records))) - if api_keys - else MappingProxyType({}) - ) - return await asyncio.to_thread( _aggregate_grouping_sets_records_sync, records=records, - api_key_metadata=api_key_metadata, + api_key_metadata=await _key_metadata_unless_database_fails(repository, records), ) @@ -991,7 +1000,13 @@ async def get_daily_activity( if resolve_entity_metadata is not None: resolved_entity_metadata = { **(entity_metadata_field or {}), - **(await resolve_entity_metadata(daily_spend_data)), + **( + await unless_database_fails( + resolve_entity_metadata(daily_spend_data), + MappingProxyType[str, dict[str, object]]({}), + "its entity metadata", + ) + ), } aggregated: Final = await _aggregate_spend_records( repository=repository, @@ -1038,7 +1053,7 @@ def _fold_entity_rollups_sync( *, results: Sequence[DailySpendData], entity_rows: Sequence[EntityRollupRow], - api_key_metadata: Mapping[str, KeyMetadataRow], + api_key_metadata: Mapping[str, KeyMetadataRow] | None, entity_metadata_field: Mapping[str, Mapping[str, object]] | None, ) -> None: """Write breakdown.entities onto the already-built per-day results.""" @@ -1107,7 +1122,7 @@ async def _guardrail_spend_or_empty( from prisma.errors import PrismaError try: - return await repository.guardrail_spend(scope) + return await unless_database_fails(repository.guardrail_spend(scope), (), "its guardrail breakdown") except PrismaError as exc: verbose_proxy_logger.warning("Returning daily activity without its guardrail breakdown: %s", exc) return () @@ -1117,7 +1132,9 @@ async def _user_entity_metadata( repository: DailyActivityRepository, entity_rows: Sequence[EntityRollupRow] ) -> Mapping[str, Mapping[str, object]]: user_ids: Final = frozenset(row.entity_id for row in entity_rows if row.entity_id) - users: Final = await repository.user_metadata(user_ids) + users: Final = await unless_database_fails( + repository.user_metadata(user_ids), MappingProxyType({}), "its entity user metadata" + ) return MappingProxyType( { user_id: MappingProxyType({"user_email": user.user_email, "user_alias": user.user_alias}) @@ -1160,16 +1177,7 @@ async def get_daily_activity_aggregated( ) if aggregated_rows.entity_rows: entity_records: Final = aggregated_rows.entity_rows - entity_api_keys: Final = frozenset( - row.api_key for row in entity_records if row.api_key and row.api_key != PTU_SENTINEL_API_KEY - ) - entity_key_metadata: Final[Mapping[str, KeyMetadataRow]] = ( - await repository.key_metadata( - entity_api_keys, spend_logs_window(frozenset(row.date for row in entity_records)) - ) - if entity_api_keys - else MappingProxyType({}) - ) + entity_key_metadata: Final = await _key_metadata_unless_database_fails(repository, entity_records) user_entity_metadata: Final[Mapping[str, Mapping[str, object]]] = ( await _user_entity_metadata(repository, entity_records) if scope.table is DailyActivityTable.USER diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 130b92f5e73..0fbfd76e5ba 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -35,7 +35,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( end_user_cache_key, end_user_restricted_registry_cache_key, ) -from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity +from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity, unless_database_fails from litellm.proxy.management_endpoints.common_utils import validate_budget_duration from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports _set_object_permission, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -522,10 +522,12 @@ async def resolve_customer_daily_activity_scope( prisma_client: "PrismaClient", ) -> _CustomerDailyActivityScope: end_user_table: Final = _typed_table(EndUserRepository(prisma_client)) - end_user_aliases: Final = ( - await find_many_in(end_user_table, "user_id", end_user_ids) + end_user_aliases: Final = await unless_database_fails( + find_many_in(end_user_table, "user_id", end_user_ids) if end_user_ids is not None - else await end_user_table.find_many(where={}) + else end_user_table.find_many(where={}), + (), + "its customer aliases", ) metadata: Final = MappingProxyType({end_user.user_id: {"alias": end_user.alias} for end_user in end_user_aliases}) return _CustomerDailyActivityScope(end_user_ids, metadata) diff --git a/litellm/proxy/management_endpoints/daily_activity_routes.py b/litellm/proxy/management_endpoints/daily_activity_routes.py index 1a3445aae52..b5af673585c 100644 --- a/litellm/proxy/management_endpoints/daily_activity_routes.py +++ b/litellm/proxy/management_endpoints/daily_activity_routes.py @@ -4,6 +4,7 @@ import json from collections.abc import AsyncIterator, Mapping, Sequence from dataclasses import asdict, fields, replace from datetime import date, datetime +from types import MappingProxyType from typing import Annotated, Final, Literal from fastapi import APIRouter, Depends, HTTPException, Query @@ -23,6 +24,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( raise_public, spend_logs_window, sum_spend_metrics, + unless_database_fails, ) from litellm.proxy.management_endpoints.daily_activity_scopes import ( AGENT_RESOLVER, @@ -149,9 +151,10 @@ async def _key_activity_rows( spend_window: Final[tuple[datetime, datetime] | None] = spend_logs_window( frozenset((resolved_scope.scope.start_date, resolved_scope.scope.end_date)) ) - metadata: Final[Mapping[str, KeyMetadataRow]] = await repository.key_metadata( - frozenset(row.api_key for row in rows if row.api_key), - spend_window, + metadata: Final[Mapping[str, KeyMetadataRow]] = await unless_database_fails( + repository.key_metadata(frozenset(row.api_key for row in rows if row.api_key), spend_window), + MappingProxyType({}), + "its key metadata", ) return [_key_activity_row(row, metadata) for row in rows] @@ -176,8 +179,10 @@ async def _user_activity_rows( repository: DailyActivityRepository, rows: Sequence[UserSpendRow], ) -> Sequence[UserActivityRow]: - metadata: Final[Mapping[str, UserMetadataRow]] = await repository.user_metadata( - frozenset(user_id for row in rows if (user_id := row.user_id) is not None) + metadata: Final[Mapping[str, UserMetadataRow]] = await unless_database_fails( + repository.user_metadata(frozenset(user_id for row in rows if (user_id := row.user_id) is not None)), + MappingProxyType({}), + "its user metadata", ) return [_user_activity_row(row, metadata) for row in rows] diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index a76dad8d02a..6c6118b2505 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -85,7 +85,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: ) from litellm.proxy.management_helpers.team_roster_sync import user_tx_db from litellm.proxy.management_helpers.utils import management_endpoint_wrapper -from litellm.proxy.utils import handle_exception_on_proxy, hash_password, run_off_loop_if_fips +from litellm.proxy.utils import handle_exception_on_proxy, hash_password, query_raw_rows, run_off_loop_if_fips from litellm.redis.stores.spend_counters import ( user_counter, ) @@ -1240,7 +1240,7 @@ async def user_info_v2( async def _fetch_admin_teams_and_keys_rows( prisma_client: "PrismaClient", sql_query: str ) -> Sequence[_AdminTeamAndKeyRows]: - return _ADMIN_TEAM_AND_KEY_ROWS_ADAPTER.validate_python(await prisma_client.db.query_raw(sql_query)) + return _ADMIN_TEAM_AND_KEY_ROWS_ADAPTER.validate_python(await query_raw_rows(prisma_client.db, sql_query)) async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 4cdb318d370..e7597fe70ac 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -40,6 +40,7 @@ from litellm.constants import ( MINIMUM_CUSTOM_KEY_LENGTH, UI_SESSION_TOKEN_TEAM_ID, ) +from litellm.litellm_core_utils.core_helpers import as_object_list from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.models.credentials import CredentialItem @@ -65,7 +66,6 @@ from litellm.proxy.auth.auth_utils import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( evict_and_broadcast, - publish_auth_cache_invalidation, ) from litellm.proxy.common_utils.callback_config_validation import logging_metadata_config_error from litellm.proxy.common_utils.callback_utils import ( @@ -171,7 +171,6 @@ from litellm.proxy.utils import ( require_enterprise_license, ) from litellm.redis.stores.spend_counters import ( - SPEND_COUNTERS_TARGET, key_counter, key_window_counter, ) @@ -2060,7 +2059,7 @@ async def _check_project_key_limits( # Validate key models are a subset of project models if data.models and len(project_obj.models) > 0: - for m in data.models: + for m in as_object_list(data.models) or (): if m not in project_obj.models and m not in _INHERITED_MODEL_SENTINELS: raise HTTPException( status_code=400, @@ -6809,38 +6808,6 @@ def _validate_reset_spend_value(reset_to: object, key_in_db: LiteLLM_Verificatio return reset_to -@with_service_target(SPEND_COUNTERS_TARGET) -async def _set_spend_counter_with_floor_and_broadcast(counter_key: str, value: float) -> None: - """ - Set a Redis-backed spend counter to `value`, mirror it into the short-lived - spend_db_floor marker `_authoritative_floor_spend` reads, and broadcast both - to every worker (LIT-3803 pattern: setting, not deleting, means a worker's - own self-delivered broadcast still carries the reset value forward). - - Without the floor marker, `_authoritative_floor_spend` can re-derive a - stale, pre-reset value from a marker another worker cached moments earlier - and raise the just-reset counter right back up via `_repair_stale_spend_counter`. - Without the broadcast, a worker that already cached the pre-reset key object - or floor marker keeps enforcing against it until its own TTL expires. - """ - from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache, spend_counter_store - - try: - await spend_counter_store.set(counter_key, value, ttl=60) - except Exception as redis_err: - verbose_proxy_logger.warning( - "Failed to update spend counter %s in Redis: %s. Budget checks may use stale value until counter expires.", - counter_key, - redis_err, - ) - - floor_key: Final = f"spend_db_floor:{counter_key}" - spend_counter_cache.in_memory_cache.set_cache(key=floor_key, value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) - - await publish_auth_cache_invalidation(cache_key=counter_key, new_value=value, ttl=60) - await publish_auth_cache_invalidation(cache_key=floor_key, new_value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) - - def _budget_limit_windows(budget_limits: Sequence[object] | str | None) -> tuple[Mapping[str, object], ...]: """Coerce a key's stored `budget_limits` into a tuple of plain window dicts. @@ -6925,6 +6892,8 @@ async def _reset_key_budget_windows( pre-reset boundary, re-sum the unchanged historical spend, and put the counter right back where it was before the write ever landed. """ + from litellm.proxy.proxy_server import set_spend_counter_with_floor_and_broadcast + windows: Final = _budget_limit_windows(budget_limits) if not windows: return @@ -6943,7 +6912,7 @@ async def _reset_key_budget_windows( duration = window.get("budget_duration") if isinstance(duration, str) and duration: counter_key = key_window_counter(hashed_api_key, duration) - await _set_spend_counter_with_floor_and_broadcast(counter_key=counter_key, value=0.0) + await set_spend_counter_with_floor_and_broadcast(counter_key=counter_key, value=0.0) @router.post( @@ -6966,6 +6935,7 @@ async def reset_key_spend_fn( hash_token, prisma_client, proxy_logging_obj, + set_spend_counter_with_floor_and_broadcast, user_api_key_cache, ) @@ -7017,7 +6987,7 @@ async def reset_key_spend_fn( # amount for every enforcement check immediately instead of the stale # pre-reset value. _counter_key: Final = key_counter(hashed_api_key) - await _set_spend_counter_with_floor_and_broadcast(counter_key=_counter_key, value=reset_to) + await set_spend_counter_with_floor_and_broadcast(counter_key=_counter_key, value=reset_to) await _reset_key_budget_windows( prisma_client=prisma_client, hashed_api_key=hashed_api_key, diff --git a/litellm/proxy/management_endpoints/management_v1/budgets.py b/litellm/proxy/management_endpoints/management_v1/budgets.py index 19a5408496a..4c8d8fdf302 100644 --- a/litellm/proxy/management_endpoints/management_v1/budgets.py +++ b/litellm/proxy/management_endpoints/management_v1/budgets.py @@ -35,7 +35,7 @@ from litellm.proxy.list_api.list_framework import ( where_sql, ) from litellm.proxy.management_endpoints.management_v1.common import MANAGEMENT_V1_PREFIX -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import PrismaClient, query_raw_rows from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.proxy.management_endpoints.management_v1 import ( ListResponse, @@ -113,7 +113,7 @@ class PrismaBudgetListExecutor: async def count(self, where: tuple[Predicate, ...]) -> int: clauses, params = where_sql(where) sql: Final = f"SELECT COUNT(*) AS count FROM {BUDGET_TABLE}" + (f" WHERE {clauses}" if clauses else "") - counted: Final = _ROW_COUNTS.validate_python(await self.prisma_client.db.query_raw(sql, *params)) + counted: Final = _ROW_COUNTS.validate_python(await query_raw_rows(self.prisma_client.db, sql, *params)) return counted[0].count if counted else 0 async def find_many(self, plan: QueryPlan) -> Sequence[BudgetListItem]: @@ -130,7 +130,9 @@ class PrismaBudgetListExecutor: + f" ORDER BY {ordering}" + f" LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}" ) - return _BUDGET_ROWS.validate_python(await self.prisma_client.db.query_raw(sql, *params, plan.take, plan.skip)) + return _BUDGET_ROWS.validate_python( + await query_raw_rows(self.prisma_client.db, sql, *params, plan.take, plan.skip) + ) def _serialize(row: BudgetListItem) -> BudgetListItem: diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 2c7453ebb1c..b3c1f48bf99 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -30,6 +30,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import LITELLM_PROXY_ADMIN_NAME, TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS +from litellm.litellm_core_utils.core_helpers import as_object_list, decoded_json_object from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.ptu_pricing import ( @@ -142,6 +143,7 @@ from litellm.router_strategy.auto_router.auto_router_model_naming import ( count_capability_routers, gated_capabilities_of, is_complexity_router_model, + strategy_router_dependencies, validate_complexity_router_config_placement, validate_complexity_router_config_write, validate_strategy_router_model_write, @@ -209,6 +211,7 @@ class KubernetesPodsResponse(BaseModel): CLEARABLE_LITELLM_PARAMS: Final = frozenset({"cache_control_injection_points", "litellm_credential_name"}) NULL_CLEARABLE_LITELLM_PARAMS: Final = frozenset((*SPECIAL_MODEL_INFO_PARAMS, *CLEARABLE_LITELLM_PARAMS)) +_STR_OBJECT_DICT: Final = TypeAdapter(dict[str, object]) async def update_team(*args, **kwargs): @@ -245,12 +248,12 @@ def _model_write_response( ) -> _ProxyModelRow | Mapping[str, object]: if member_write is None: return row - payload: Final = TypeAdapter(dict[str, object]).validate_json(row.model_dump_json()) + payload: Final = _STR_OBJECT_DICT.validate_json(row.model_dump_json()) stored_params: Final = payload.get("litellm_params") params: Final = ( - TypeAdapter(dict[str, object]).validate_json(stored_params) + _STR_OBJECT_DICT.validate_json(stored_params) if isinstance(stored_params, str) - else TypeAdapter(dict[str, object]).validate_python(stored_params) + else _STR_OBJECT_DICT.validate_python(stored_params) ) redacted: Final = redact_credentials_in_payload(params) return { @@ -355,6 +358,22 @@ async def get_db_model(model_id: str, prisma_client: PrismaClient) -> Deployment return deployment_pydantic_obj +def _self_routing_violation( + model_name: str | None, + incoming_params: GenericLiteLLMParams | None, + existing_params: GenericLiteLLMParams | None, +) -> str | None: + if not model_name: + return None + effective_params: Final = { + **(_decrypted_litellm_params(existing_params) if existing_params is not None else {}), + **(_dump_without_none(incoming_params) if incoming_params is not None else {}), + } + if any(dependency.model_name == model_name for dependency in strategy_router_dependencies(effective_params)): + return f"A strategy router cannot route to itself: {model_name} is listed as one of its own models" + return None + + def _strategy_router_write_violation( incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None, @@ -374,6 +393,11 @@ def _strategy_router_write_violation( """ if incoming_params is None: return None + from litellm.router_strategy.switchyard_router.config import validate_switchyard_config_write + + switchyard_violation: Final = validate_switchyard_config_write(incoming_params.switchyard_router_config) + if switchyard_violation is not None: + return switchyard_violation config_violation: Final = validate_complexity_router_config_write( complexity_router_config=( _effective_complexity_router_config(incoming_params, existing_params) @@ -406,8 +430,11 @@ def _strategy_router_write_violation( def _raise_on_strategy_router_write_violation( incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None, + model_name: str | None = None, ) -> None: - violation = _strategy_router_write_violation(incoming_params=incoming_params, existing_params=existing_params) + violation: Final = _self_routing_violation(model_name, incoming_params, existing_params) or ( + _strategy_router_write_violation(incoming_params=incoming_params, existing_params=existing_params) + ) if violation is None: return raise ProxyException( @@ -1346,7 +1373,9 @@ def legacy_update_model_info_changes( value that differs from the stored override resets that override (``None``). """ stored: Final = db_model.model_info.model_dump(exclude_none=True) - requested: Final = incoming_model_info.model_dump(exclude_unset=True, exclude_none=True) + requested: Final = _STR_OBJECT_DICT.validate_python( + incoming_model_info.model_dump(exclude_unset=True, exclude_none=True) + ) sent: Final = {k: v for k, v in requested.items() if k not in _LEGACY_UPDATE_IGNORED_MODEL_INFO_FIELDS and v != ""} echoed: Final = frozenset( echoed_cost_map_fields( @@ -1356,7 +1385,7 @@ def legacy_update_model_info_changes( discovered_model_info, ) ) - kept: Final = without_server_derived_pricing(sent) + kept: Final[Mapping[str, object]] = without_server_derived_pricing(sent) return MappingProxyType( { **{k: getattr(incoming_model_info, k) for k in kept if k not in echoed}, @@ -1396,20 +1425,16 @@ async def _locked_legacy_update_model_info_json( return None stored_model_info: Final[object] = payload.get("model_info") stored: Final = ( - TypeAdapter(dict[str, object]).validate_json(stored_model_info) + _STR_OBJECT_DICT.validate_json(stored_model_info) if isinstance(stored_model_info, str) - else TypeAdapter(dict[str, object]).validate_python(stored_model_info or {}) - ) - return ( - TypeAdapter(dict[str, object]) - .dump_json( - { - **{k: v for k, v in stored.items() if k not in writes}, - **{k: v for k, v in writes.items() if v is not None}, - } - ) - .decode() + else _STR_OBJECT_DICT.validate_python(stored_model_info or {}) ) + return _STR_OBJECT_DICT.dump_json( + { + **{k: v for k, v in stored.items() if k not in writes}, + **{k: v for k, v in writes.items() if v is not None}, + } + ).decode() def _param_keys_to_remove(patch: updateDeployment) -> tuple[str, ...]: @@ -1518,7 +1543,7 @@ def update_db_model( prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel( model_name=merged_model_name, litellm_params=json.dumps(merged_litellm_params), - model_info=TypeAdapter(dict[str, object]).dump_json(merged_model_info).decode(), + model_info=_STR_OBJECT_DICT.dump_json(merged_model_info).decode(), ) if updated_patch.blocked is not None: @@ -1751,6 +1776,7 @@ async def patch_model( _raise_on_strategy_router_write_violation( incoming_params=incoming_params, existing_params=remaining_params, + model_name=patch_data.model_name or db_model.model_name, ) written_team_id: Final = ( @@ -2054,7 +2080,7 @@ async def add_model_to_db( slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None, ) -> "_ProxyModelRow | LiteLLM_ProxyModelTable": # encrypt litellm params # - _litellm_params_dict: Final = TypeAdapter(dict[str, object]).validate_python( + _litellm_params_dict: Final = _STR_OBJECT_DICT.validate_python( model_params.litellm_params.model_dump(exclude_none=True) ) if "complexity_router_config" in _litellm_params_dict: @@ -3120,6 +3146,7 @@ async def add_new_model( _raise_on_strategy_router_write_violation( incoming_params=model_params.litellm_params, existing_params=None, + model_name=model_params.model_name, ) _raise_if_rate_limits_required_but_missing( @@ -3325,6 +3352,7 @@ async def update_model( _raise_on_strategy_router_write_violation( incoming_params=model_params.litellm_params, existing_params=deployment.litellm_params, + model_name=model_params.model_name or deployment.model_name, ) renamed_to: Final = ( model_params.model_name @@ -3972,19 +4000,20 @@ async def clear_cache() -> ReconcileOutcome: db_router_names: Final = set() - for model in current_models: - model_info = model.get("model_info", {}) + for model in as_object_list(current_models) or (): + deployment = decoded_json_object(model) + model_info = decoded_json_object(deployment.get("model_info", {})) if model_info.get("db_model", False): db_model_ids.append(model_info.get("id")) # Auto-router deployments (and only those) are wiped here, in the # same pass, so the reload rebuilds them -- see the comment below. - model_name = model.get("model_name") - if model_name is not None and str(model.get("litellm_params", {}).get("model", "")).startswith( - "auto_router/" - ): - db_router_names.add(model_name) + deployment_name = deployment.get("model_name") + if isinstance(deployment_name, str) and str( + decoded_json_object(deployment.get("litellm_params", {})).get("model", "") + ).startswith("auto_router/"): + db_router_names.add(deployment_name) router_model_id = model_info.get("id") - if router_model_id is not None: + if isinstance(router_model_id, str): llm_router.apply(ChangeSet(delete=[router_model_id])) else: # This is a config model, preserved by the reconcile below @@ -4028,6 +4057,7 @@ async def clear_cache() -> ReconcileOutcome: llm_router.registry().plugins("complexity").pop(model_name, None) llm_router.registry().plugins("adaptive").pop(model_name, None) llm_router.registry().plugins("quality").pop(model_name, None) + llm_router.registry().plugins("switchyard").pop(model_name, None) # Reload only DB models. _add_deployment_locked, not add_deployment: this # coroutine already holds MODEL_RECONCILE_LOCK and asyncio.Lock is not diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 4e511aee24a..15d05633679 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -48,7 +48,7 @@ from litellm.proxy.management_endpoints.budget_management_endpoints import ( new_budget, update_budget, ) -from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity +from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity, unless_database_fails from litellm.proxy.management_endpoints.common_utils import ( # noqa: F401 # legacy module exports _set_object_metadata_field, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _user_has_admin_view, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -655,10 +655,12 @@ async def resolve_organization_daily_activity_scope( ) organization_table: Final = _table(OrganizationRepository(prisma_client)) - organization_rows: Final = ( - await find_many_in(organization_table, "organization_id", resolved_organization_ids) + organization_rows: Final = await unless_database_fails( + find_many_in(organization_table, "organization_id", resolved_organization_ids) if resolved_organization_ids is not None - else await organization_table.find_many(where={}) + else organization_table.find_many(where={}), + (), + "its organization aliases", ) metadata: Final = MappingProxyType( { diff --git a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py index 1096954536a..d3854aa3ef3 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py @@ -6,11 +6,16 @@ based on user-provided attack examples and descriptions. import json from typing import Final +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items from litellm.proxy._types import ProxyErrorTypes, ProxyException +_DECODED_JSON: Final = TypeAdapter(object) + SUGGEST_TOOL: Final = { "type": "function", "function": { @@ -97,14 +102,19 @@ class AiPolicySuggester: "explanation": "No templates could be matched to your requirements.", } - result: Final = json.loads(tool_calls[0].function.arguments) + decoded: Final = _DECODED_JSON.validate_python(json.loads(tool_calls[0].function.arguments)) valid_ids: Final = {t["id"] for t in templates} - result["selected_templates"] = [ - s for s in result.get("selected_templates", []) if s.get("template_id") in valid_ids - ] - - return result + result: Final = decoded_json_object(decoded) + suggestion: Final = { + **result, + "selected_templates": [ + s + for s in iter_items(result.get("selected_templates", [])) + if decoded_json_object(s).get("template_id") in valid_ids + ], + } + return suggestion except Exception as e: verbose_proxy_logger.error("AI policy suggestion failed: %s", e) raise diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 3ec92165a84..0b7a7ca43cd 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -30,6 +30,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, get_daily_activity, ) +from litellm.proxy.management_endpoints.common_utils import validate_finite_spend from litellm.proxy.management_helpers.utils import handle_budget_for_entity from litellm.redis.stores.spend_counters import ( tag_counter, @@ -415,6 +416,7 @@ async def update_tag( Parameters: - name: str - The name of the tag to update + - spend: Optional[float] - Updated spend for the tag - description: Optional[str] - Updated description - models: List[str] - Updated list of allowed LLM models - budget_id: Optional[str] - The id for a budget to associate with the tag @@ -428,11 +430,13 @@ async def update_tag( - model_max_budget: Optional[dict] - Max budget for a specific model - budget_duration: Optional[str] - Frequency of resetting tag budget """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import prisma_client, set_spend_counter_with_floor_and_broadcast if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") + validate_finite_spend(tag.spend) + try: # Check if tag exists existing_tag: Final = await _table(TagRepository(prisma_client)).find_unique(where={"tag_name": tag.name}) @@ -459,12 +463,10 @@ async def update_tag( "description": tag.description, "models": tag.models or [], "model_info": json.dumps(model_info), + **({"spend": tag.spend} if tag.spend is not None else {}), + **({"budget_id": budget_id} if budget_id != existing_tag.budget_id else {}), } - # Add budget_id if it changed - if budget_id != existing_tag.budget_id: - update_data["budget_id"] = budget_id - # Update tag in database updated_tag_record: Final = await _table(TagRepository(prisma_client)).update( where={"tag_name": tag.name}, @@ -473,6 +475,9 @@ async def update_tag( await _evict_tag_cache_keys((tag_cache_key(tag.name),)) + if tag.spend is not None: + await set_spend_counter_with_floor_and_broadcast(counter_key=f"spend:tag:{tag.name}", value=tag.spend) + # Build response tag_config: Final = TagConfig( name=updated_tag_record.tag_name, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index cf2ae392677..59fa8a18ac7 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -155,6 +155,7 @@ from litellm.proxy.management.teams.dependencies import get_team_access from litellm.proxy.management_endpoints.common_daily_activity import ( InvalidDateRange, parse_canonical_date_range, + unless_database_fails, ) from litellm.proxy.management_endpoints.common_utils import ( # noqa: F401 # legacy module exports _check_disable_global_guardrails_caller_permission, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -7303,7 +7304,12 @@ async def resolve_team_daily_activity_scope( where_condition: Final[_TeamIdInFilter] = {} if team_ids_list: where_condition["team_id"] = {"in": list(team_ids_list)} - team_aliases: Final = await _team_db(prisma_client).find_many(where=where_condition) + team_alias_lookup: Final = _team_db(prisma_client).find_many(where=where_condition) + team_aliases: Final = ( + await unless_database_fails(team_alias_lookup, [], "its team aliases") + if user_api_key_has_admin_view(user_api_key_dict) + else await team_alias_lookup + ) team_alias_metadata: Final = {t.team_id: {"team_alias": _as_object(t.team_alias)} for t in team_aliases} # Check if user is team admin or has /team/daily/activity permission diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index aaa261eb19b..80df6a95cc2 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -65,7 +65,7 @@ from litellm.constants import ( MICROSOFT_USER_ID_ATTRIBUTE, MICROSOFT_USER_LAST_NAME_ATTRIBUTE, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -4856,7 +4856,7 @@ class MicrosoftSSOHandler: # Decode the JWT without signature verification # (signature is already verified by fastapi_sso) - decoded_token: Final = jwt.decode(id_token, options={"verify_signature": False}) + decoded_token: Final = decoded_json_object(jwt.decode(id_token, options={"verify_signature": False})) # Extract app_roles claim from the token ## check for both 'roles' and 'app_roles' claims diff --git a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py index 2938348cf9b..b4d863672c0 100644 --- a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py +++ b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py @@ -158,7 +158,9 @@ def _typed_table(repo: DailyTagSpendRepository | VerificationTokenRepository | U async def _query_raw(prisma_client: "PrismaClient", sql_query: str, *params: object) -> object: - return await prisma_client.db.query_raw(sql_query, *params) + from litellm.proxy.utils import query_raw_rows + + return await query_raw_rows(prisma_client.db, sql_query, *params) @router.get( diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index a719a5c67d4..27d7cff0572 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -139,6 +139,14 @@ def _object_permission_mapping(raw: object) -> Mapping[str, object]: return mapping if mapping is not None and isinstance(decoded, dict) else {} +def _decoded_tool_permissions(raw: str) -> Mapping[str, object]: + decoded: Final = _DECODED_JSON.validate_python(json.loads(raw)) + mapping: Final = as_str_mapping(decoded) + if mapping is not None: + return mapping + raise AttributeError(f"{type(decoded).__name__!r} object has no attribute 'keys'", name="keys", obj=decoded) + + async def handle_update_object_permission_common( data_json: dict, existing_object_permission_id: str | None, @@ -460,10 +468,11 @@ async def _resolve_team_allowed_mcp_servers( access_group_servers: Final[list[str]] = await MCPRequestHandler.get_mcp_servers_from_access_groups( team_object_permission.mcp_access_groups or [] ) - raw_tool_perms = team_object_permission.mcp_tool_permissions or {} - if isinstance(raw_tool_perms, str): - raw_tool_perms = json.loads(raw_tool_perms) - tool_perm_servers: Final[list[str]] = list(raw_tool_perms.keys()) + stored_tool_perms: Final = team_object_permission.mcp_tool_permissions or {} + tool_perms: Final = ( + _decoded_tool_permissions(stored_tool_perms) if isinstance(stored_tool_perms, str) else stored_tool_perms + ) + tool_perm_servers: Final[list[str]] = list(tool_perms.keys()) raw_servers: Final = set(direct_servers + access_group_servers + tool_perm_servers) resolved_servers: Final = await _resolve_mcp_server_identifiers_to_ids( identifiers=raw_servers, @@ -555,7 +564,7 @@ async def _get_grandfathered_key_mcp_server_ids( return frozenset() raw_tool_perms: Final = existing_object_permission.mcp_tool_permissions or {} tool_perm_keys: Final[frozenset[str]] = frozenset( - json.loads(raw_tool_perms).keys() if isinstance(raw_tool_perms, str) else raw_tool_perms.keys() + _decoded_tool_permissions(raw_tool_perms).keys() if isinstance(raw_tool_perms, str) else raw_tool_perms.keys() ) identifiers: Final = (frozenset(existing_object_permission.mcp_servers or []) | tool_perm_keys) - { SpecialMCPServerNames.no_mcp_servers.value, diff --git a/litellm/proxy/management_helpers/team_list_sorting.py b/litellm/proxy/management_helpers/team_list_sorting.py index caddd8afc08..a3a35857328 100644 --- a/litellm/proxy/management_helpers/team_list_sorting.py +++ b/litellm/proxy/management_helpers/team_list_sorting.py @@ -6,7 +6,7 @@ from typing import Final from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import PrismaClient, query_raw_rows COMPUTED_TEAM_SORTS: Final = MappingProxyType( { @@ -105,5 +105,5 @@ async def sorted_team_ids( f"ORDER BY {expression} {direction} NULLS LAST, t.team_id ASC " f"LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}" ) - rows: Final = _TEAM_IDS.validate_python(await prisma_client.db.query_raw(sql, *params, take, skip)) + rows: Final = _TEAM_IDS.validate_python(await query_raw_rows(prisma_client.db, sql, *params, take, skip)) return tuple(row.team_id for row in rows) diff --git a/litellm/proxy/management_helpers/user_list_sorting.py b/litellm/proxy/management_helpers/user_list_sorting.py index 407a188ff1a..aa77417311f 100644 --- a/litellm/proxy/management_helpers/user_list_sorting.py +++ b/litellm/proxy/management_helpers/user_list_sorting.py @@ -4,7 +4,7 @@ from typing import Final, Literal from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm.constants import UI_SESSION_TOKEN_TEAM_ID -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import PrismaClient, query_raw_rows class _UserId(BaseModel): @@ -63,7 +63,8 @@ async def sorted_user_ids( LIMIT $9 OFFSET $10 """ rows: Final = _USER_IDS.validate_python( - await prisma_client.db.query_raw( + await query_raw_rows( + prisma_client.db, sql, role or None, ids if user_ids else None, diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index c9fdc75c4ec..3b37b9d6e02 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -217,7 +217,7 @@ async def handle_budget_for_entity( ) # Get all budget field names - budget_params: Final = LiteLLM_BudgetTable.model_fields.keys() + budget_params: Final = LiteLLM_BudgetTable.model_fields.keys() - {"budget_id"} # Extract budget fields from data _json_data: Final = _BUDGET_DATA_MAPPING.validate_python( diff --git a/litellm/proxy/mcp_server/guardrail_translation/handler.py b/litellm/proxy/mcp_server/guardrail_translation/handler.py index 28c51148da5..b268ba893f5 100644 --- a/litellm/proxy/mcp_server/guardrail_translation/handler.py +++ b/litellm/proxy/mcp_server/guardrail_translation/handler.py @@ -19,7 +19,7 @@ payload (name + arguments) so we just build the tool definition. """ from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from fastapi import HTTPException from mcp.types import Tool as MCPTool @@ -139,7 +139,7 @@ class MCPGuardrailTranslationHandler(BaseTranslation): async def process_input_messages( self, - data: dict[str, Any], + data: dict[str, object], guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> dict[str, object]: @@ -154,10 +154,12 @@ class MCPGuardrailTranslationHandler(BaseTranslation): # Convert MCP input via transform_mcp_tool_to_openai_tool, then map to litellm # ChatCompletionToolParam (openai SDK type has incompatible strict/cache_control). - mcp_tool: Final = MCPTool( - name=mcp_tool_name, - description=mcp_tool_description or "", - input_schema=dict(mcp_input_schema) if isinstance(mcp_input_schema, Mapping) else {}, + mcp_tool: Final = MCPTool.model_validate( + { + "name": mcp_tool_name, + "description": mcp_tool_description or "", + "input_schema": dict(mcp_input_schema) if isinstance(mcp_input_schema, Mapping) else {}, + } ) openai_tool: Final = transform_mcp_tool_to_openai_tool(mcp_tool) fn: Final = openai_tool["function"] diff --git a/litellm/proxy/mcp_server/oauth2_flow_backfill.py b/litellm/proxy/mcp_server/oauth2_flow_backfill.py index 1ef67df498f..7b72ba2a26f 100644 --- a/litellm/proxy/mcp_server/oauth2_flow_backfill.py +++ b/litellm/proxy/mcp_server/oauth2_flow_backfill.py @@ -39,15 +39,17 @@ from collections import Counter from collections.abc import Mapping, Sequence from typing import Final, Literal, Protocol -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy.mcp_server.db import ( # noqa: F401 # legacy module exports _decode_oauth_payload, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export decode_oauth_payload, decrypt_credentials, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import MCPUserCredentialsRepository from litellm.types.mcp import MCPCredentials OAuth2Flow = Literal["client_credentials", "authorization_code"] @@ -60,6 +62,7 @@ BackfillRule = Literal[ ] _BACKFILL_AUDIT_ACTOR: Final = "oauth2_flow_backfill" +_DECODED_JSON: Final = TypeAdapter(object) class _MCPServerRow(Protocol): @@ -108,21 +111,21 @@ def _mcp_server_table(prisma_client: PrismaClient) -> _MCPServerTable: def _mcp_user_credentials_table(prisma_client: PrismaClient) -> _MCPUserCredentialsTable: """The per-user MCP credential table, typed so the untyped prisma client surface stops here.""" - return prisma_client.db.litellm_mcpusercredentials + return MCPUserCredentialsRepository(prisma_client).table def _decrypted_credentials(raw_credentials: str | Mapping[str, JsonValue] | None) -> MCPCredentials | None: if raw_credentials is None: return None - parsed: JsonValue | Mapping[str, JsonValue] + parsed: Mapping[str, object] | None if isinstance(raw_credentials, str): try: - parsed = json.loads(raw_credentials) + parsed = as_str_mapping(_DECODED_JSON.validate_python(json.loads(raw_credentials))) except (ValueError, TypeError): return None else: - parsed = raw_credentials - if not isinstance(parsed, dict): + parsed = raw_credentials if isinstance(raw_credentials, dict) else None + if parsed is None: return None return decrypt_credentials(credentials=dict(parsed)) diff --git a/litellm/proxy/mcp_server/oauth2_token_cache.py b/litellm/proxy/mcp_server/oauth2_token_cache.py index 1e6e83c0f4b..4bb61f6a5bb 100644 --- a/litellm/proxy/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/mcp_server/oauth2_token_cache.py @@ -43,6 +43,7 @@ if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer _DECODED_JSON: Final = TypeAdapter(object) +_CACHED_TOKEN: Final = TypeAdapter(object) class MCPOAuth2TokenCache(InMemoryCache): @@ -259,8 +260,8 @@ class MCPPerUserTokenCache: from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 key: Final = self._cache_key(user_id, server_id) - encrypted: Final = await user_api_key_cache.async_get_cache(key) - if encrypted is None: + encrypted: Final = _CACHED_TOKEN.validate_python(await user_api_key_cache.async_get_cache(key)) + if not isinstance(encrypted, str): return None return self._codec().decode(encrypted) except Exception as exc: diff --git a/litellm/proxy/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/mcp_server/openapi_to_mcp_generator.py index 0fa904cb582..7282e59d746 100644 --- a/litellm/proxy/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/mcp_server/openapi_to_mcp_generator.py @@ -45,6 +45,7 @@ OPENAPI_TOOL_NAME_MAX_LEN: Final = 128 _OPENAPI_TOOL_NAME_MAX_LEN: Final = OPENAPI_TOOL_NAME_MAX_LEN _OPENAPI_DOCUMENT: Final = TypeAdapter(dict[str, object]) +_DECODED_SPEC: Final = TypeAdapter(object) _MAX_YAML_INT_LENGTH: Final = 1024 # YAML 1.2 core scalars; isolate these from PyYAML's global YAML 1.1 resolver. _YAML_SCALARS: Final = { @@ -225,6 +226,22 @@ def load_openapi_spec(filepath: str) -> _OpenAPISpec: return asyncio.run(load_openapi_spec_async(filepath)) +def _has_nonspecific_tag(event: ScalarEvent) -> bool: + match event: + case ScalarEvent(tag="!"): + return True + case _: + return False + + +def _scalar_text(node: ScalarNode) -> str: + match node: + case ScalarNode(value=str() as text): + return text + case _: + raise ValueError("Invalid YAML OpenAPI spec") + + class _OpenAPISafeLoader(yaml.SafeLoader): yaml_implicit_resolvers: ClassVar[Mapping[str | None, Sequence[tuple[str, re.Pattern[str]]]]] = { # pyright: ignore[reportIncompatibleVariableOverride] # PyYAML resolver tables are class state None: list(_YAML_SCALARS.items()) @@ -236,7 +253,7 @@ class _OpenAPISafeLoader(yaml.SafeLoader): raise ValueError("YAML aliases are not supported in OpenAPI specs") # noqa: TRY004 # callers handle malformed documents as ValueError node: Final = super().compose_node(parent, index) # PyYAML otherwise implicitly resolves even an explicit nonspecific ! tag. - if isinstance(event, ScalarEvent) and event.tag == "!" and isinstance(node, ScalarNode): + if isinstance(event, ScalarEvent) and _has_nonspecific_tag(event) and isinstance(node, ScalarNode): node.tag = "tag:yaml.org,2002:str" return node @@ -245,7 +262,7 @@ class _OpenAPISafeLoader(yaml.SafeLoader): if pattern is not None: if not isinstance(node, ScalarNode): raise ValueError("Invalid YAML OpenAPI spec") - scalar: Final[str] = node.value + scalar: Final = _scalar_text(node) if node.tag == "tag:yaml.org,2002:int" and len(scalar) > _MAX_YAML_INT_LENGTH: raise ValueError("YAML integer is too long in OpenAPI spec") if pattern.fullmatch(scalar) is None: @@ -281,10 +298,10 @@ def _validate_openapi_spec(value: object) -> _OpenAPISpec: def _parse_openapi_spec(content: str | bytes) -> _OpenAPISpec: try: - return _validate_openapi_spec(json.loads(content)) + return _validate_openapi_spec(_DECODED_SPEC.validate_python(json.loads(content))) except json.JSONDecodeError: try: - return _validate_openapi_spec(yaml.load(content, Loader=_OpenAPISafeLoader)) + return _validate_openapi_spec(_DECODED_SPEC.validate_python(yaml.load(content, Loader=_OpenAPISafeLoader))) except (yaml.YAMLError, RecursionError) as exc: raise ValueError("Invalid YAML OpenAPI spec") from exc diff --git a/litellm/proxy/mcp_server/operations.py b/litellm/proxy/mcp_server/operations.py index 66e46c73f98..dfc45417814 100644 --- a/litellm/proxy/mcp_server/operations.py +++ b/litellm/proxy/mcp_server/operations.py @@ -49,7 +49,7 @@ from mcp.types import ( ) from mcp.types import Tool as MCPTool from pydantic import AnyUrl, ConfigDict, Field, JsonValue, SecretStr, TypeAdapter -from typing_extensions import ReadOnly, TypedDict, assert_never +from typing_extensions import ReadOnly, TypedDict, Unpack, assert_never from litellm._logging import verbose_logger from litellm.constants import ( @@ -562,6 +562,7 @@ async def _dispatch_virtual_mcp_tool( is_error=True, ) + assert user_api_key_auth is not None # guaranteed by the flag check above args: Final = arguments or {} if name == MCP_TOOL_SEARCH_TOOL_NAME: return await handle_mcp_tool_search( @@ -576,7 +577,6 @@ async def _dispatch_virtual_mcp_tool( raw_headers=raw_headers, ) - assert user_api_key_auth is not None # guaranteed by the flag check above if name == AGENT_SEARCH_TOOL_NAME: return await handle_agent_search( query=str(args.get("query", "")), @@ -2926,6 +2926,10 @@ async def resolve_scoped_mcp_server( return _single_scoped_server(mcp_servers, allowed_mcp_servers) +class _CallMcpToolOptions(TypedDict, total=False, extra_items=object): + litellm_logging_obj: ReadOnly[LiteLLMLoggingObj | None] + + @client @catalog_operation(lambda: global_mcp_server_manager) async def call_mcp_tool( @@ -2943,7 +2947,7 @@ async def call_mcp_tool( request_state: str | None = None, incoming_bearer_token: SecretStr | None | EllipsisType = ..., guardrail_auth: UserAPIKeyAuth | None = None, - **kwargs: Any, + **kwargs: Unpack[_CallMcpToolOptions], ) -> CallToolResult | InputRequiredResult: """ Call a specific tool with the provided arguments (handles prefixed tool names). diff --git a/litellm/proxy/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/mcp_server/outbound_credentials/sso_assertion_store.py index fb4d838aff9..65ab6f6c609 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/mcp_server/outbound_credentials/sso_assertion_store.py @@ -30,6 +30,7 @@ from pydantic import ConfigDict, SecretStr, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, MCP_SSO_ASSERTION_CACHE_TTL_SECONDS +from litellm.repositories.table_repositories import MCPServerRepository, PrismaTableRepository from litellm.types.llms.base import LiteLLMBaseModel if TYPE_CHECKING: @@ -68,14 +69,18 @@ class _MCPServerTable(Protocol): async def find_first(self, *, where: Mapping[str, str]) -> object | None: ... +class _SSOAssertionRepository(PrismaTableRepository["LiteLLM_SSOIdentityAssertion"]): + table_name = "litellm_ssoidentityassertion" + + def _assertion_table(prisma_client: PrismaClient) -> _SSOAssertionTable: """The SSO assertion table, typed so the untyped prisma client surface stops here.""" - return prisma_client.db.litellm_ssoidentityassertion + return _SSOAssertionRepository(prisma_client).table def _mcp_server_table(prisma_client: PrismaClient) -> _MCPServerTable: """The MCP server table, typed so the untyped prisma client surface stops here.""" - return prisma_client.db.litellm_mcpservertable + return MCPServerRepository(prisma_client).table class SSOIdentityAssertion(LiteLLMBaseModel): diff --git a/litellm/proxy/mcp_server/utils.py b/litellm/proxy/mcp_server/utils.py index 3f72804ff35..92bea6053b0 100644 --- a/litellm/proxy/mcp_server/utils.py +++ b/litellm/proxy/mcp_server/utils.py @@ -16,6 +16,7 @@ from urllib.parse import quote from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.types.mcp_server.mcp_server_manager import MCPServer if typing.TYPE_CHECKING: @@ -702,8 +703,12 @@ class MCPMissingUserEnvVarsError(Exception): _ENV_VAR_PATTERN: Final = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}") +class _ModelDumpable(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + def parse_admin_env_vars( - env_vars: Iterable[Any] | None, + env_vars: Iterable[object] | None, ) -> tuple[dict[str, str], list[dict[str, Any]]]: """Split admin-configured env var entries into globals and per-user specs. @@ -723,7 +728,7 @@ def parse_admin_env_vars( for raw in env_vars: if raw is None: continue - if hasattr(raw, "model_dump"): + if has_attributes(raw, _ModelDumpable, "model_dump"): entry: Mapping[str, object] = raw.model_dump() elif isinstance(raw, dict): entry = raw diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index a6eee889fea..1ae0d475d8c 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -23,6 +23,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.constants import MAX_FILE_LIST_LIMIT from litellm.litellm_core_utils.batch_utils import batch_cost_is_final +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import ProxyException from litellm.repositories.managed_file_repository import ManagedFileRepository from litellm.repositories.table_repositories import ( @@ -597,7 +598,7 @@ def get_team_provider_credentials( return isinstance(public_model_name, str) and _key_allowlist_covers(public_model_name, key_model_allowlist) def _provider_credentials(model_id: str) -> dict[str, object] | None: - credentials: Final = registry.credentials(model_id, form="with_provider", team_id=team_id) + credentials: Final = as_str_mapping(registry.credentials(model_id, form="with_provider", team_id=team_id)) if credentials is not None and credentials.get("custom_llm_provider") == custom_llm_provider: return {key: value for key, value in credentials.items() if key != "model"} return None @@ -1545,6 +1546,10 @@ async def get_batch_from_database( return None, None +class _BatchCostJobScheduler(Protocol): + def get_job(self, job_id: str, /) -> object: ... + + def batch_cost_poller_is_active() -> bool: """ Whether the CheckBatchCost poller will account for a managed batch's cost itself. @@ -1565,10 +1570,10 @@ def batch_cost_poller_is_active() -> bool: try: import litellm.proxy.proxy_server as proxy_server_module - scheduler = getattr(proxy_server_module, "scheduler", None) + scheduler: Final[_BatchCostJobScheduler | None] = getattr(proxy_server_module, "scheduler", None) if scheduler is None: return False - job = scheduler.get_job("check_batch_cost_job") + job: Final = scheduler.get_job("check_batch_cost_job") if job is None: return False poller = getattr(getattr(job, "func", None), "__self__", None) diff --git a/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py index 3087faa9f67..d321f5b5861 100644 --- a/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py @@ -1,15 +1,94 @@ +import asyncio +import hashlib import hmac +import json +import re +from collections.abc import Iterator, Mapping +from dataclasses import dataclass from types import MappingProxyType -from typing import Annotated, Final, Literal +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol +import httpx from fastapi import APIRouter, Depends, HTTPException, Request, Response +from fastapi.responses import StreamingResponse +from pydantic import TypeAdapter, ValidationError +from starlette.types import Message +import litellm +from litellm._internal_context import with_service_target +from litellm._logging import verbose_proxy_logger +from litellm.exceptions import BadRequestError +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_llm_provider from litellm.llms.chatgpt.common_utils import CHATGPT_API_BASE -from litellm.proxy._types import UserAPIKeyAuth +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared client factory has untyped params +) +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import check_model_access_helper, check_tools_allowlist +from litellm.proxy.auth.handle_jwt import JWTAuthManager +from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.codex_model_catalog import ( + CodexCatalogRow, + bundled_codex_models, + codex_catalog_etag, + codex_entry_names_default_reasoning_level, + merged_codex_catalog, + upstream_codex_slugs, +) +from litellm.proxy.common_utils.http_parsing_utils import ( + get_client_requested_model, + read_request_body, # pyright: ignore[reportUnknownVariableType] # returns a bare dict, validated at the call +) +from litellm.proxy.litellm_pre_call_utils import should_auto_drop_params_for_agentic_cli from litellm.proxy.pass_through_endpoints.pass_through_endpoints import pass_through_request +from litellm.proxy.response_api_endpoints.endpoints import responses_api +from litellm.responses.additional_tools import with_additional_tools_hoisted +from litellm.responses.litellm_completion_transformation.reasoning_items import is_other_provider_reasoning_item +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.custom_http import httpxSpecialProvider +from litellm.types.router import Deployment +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager, supports_none_reasoning_effort + +if TYPE_CHECKING: + from litellm.router import Router router: Final = APIRouter() +_CHATGPT_PROVIDER: Final = "chatgpt" +_CHATGPT_CREDENTIAL_HEADERS: Final = frozenset({b"authorization", b"chatgpt-account-id"}) +_MERGED_CATALOG_DROPPED_HEADERS: Final = frozenset({"etag", "content-length", "content-type"}) +_ETAG_HEADER: Final = "etag" +_MODELS_ETAG_HEADER: Final = "x-models-etag" +_ENCRYPTED_CONTENT_AFFINITY_FLAG: Final = "encrypted_content_affinity_enabled" +_LITELLM_RESPONSES_PATH: Final = "/v1/responses" +_CODEX_NO_DEFAULT_EFFORT: Final = "none" +_NONE_EFFORT_SUPPORT_FLAG: Final = "supports_none_reasoning_effort" +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) +_INPUT_ITEMS: Final = TypeAdapter(tuple[object, ...]) +_CACHED_SLUGS: Final = TypeAdapter(frozenset[str]) +_CACHED_MODELS_ETAG: Final = TypeAdapter(dict[Literal["etag"], str]) +_MODEL_NAMES: Final = TypeAdapter(list[str]) +_HEADER_DERIVED_STATE: Final = frozenset({"_cached_headers"}) +_ACCOUNT_CATALOG_HEADERS: Final = frozenset({"authorization", "chatgpt-account-id", "originator", "user-agent"}) +_CODEX_USER_AGENT_PRODUCT: Final = re.compile(r"([^/]+)/(\d+\.\d+\.\d+(?:[-+][0-9A-Za-z.+-]+)?)") +_ACCOUNT_CATALOG_TIMEOUT_SECONDS: Final = 5.0 +_SHARED_CACHE_TTL_SECONDS: Final = 3600 +_ACCOUNT_SLUGS_UNAVAILABLE: Final = "unavailable" +_ACCOUNT_SLUGS_UNAVAILABLE_TTL_SECONDS: Final = 60 +_ACCOUNT_SLUGS_CACHE_PREFIX: Final = "codex_account_slugs:" +_ACCOUNT_SLUGS_TARGET: Final = "codex_account_slugs" +_ACCOUNT_SLUGS_FETCHES_CACHE_PREFIX: Final = "codex_account_slugs_fetches:" +_ACCOUNT_SLUGS_FETCHES_PER_CALLER: Final = 5 +_ACCOUNT_SLUGS_FETCHES_WINDOW_SECONDS: Final = 60 +_MODELS_ETAG_CACHE_PREFIX: Final = "codex_models_etag:" +_MODELS_ETAG_TTL_SECONDS: Final = 300 +_OPENAI_API_PROVIDERS: Final = frozenset({LlmProviders.OPENAI.value, LlmProviders.AZURE.value}) +_CODEX_TOOL_PROVIDERS: Final = frozenset( + {*_OPENAI_API_PROVIDERS, LlmProviders.FIREWORKS_AI.value, LlmProviders.BEDROCK_MANTLE.value} +) +_CHAT_COMPLETIONS_BRIDGE: Final = "use_chat_completions_api" _HEADERS: Final = frozenset( { "authorization", @@ -29,10 +108,637 @@ _HEADERS: Final = frozenset( ) -@router.get("/backend-api/codex/{endpoint}", tags=["ChatGPT Pass-through", "pass-through"]) +def _is_chatgpt(deployment: Deployment) -> bool: + return ( + deployment.litellm_params.custom_llm_provider == _CHATGPT_PROVIDER + or deployment.litellm_params.model.partition("/")[0] == _CHATGPT_PROVIDER + ) + + +def _serves_by_name(deployment: Deployment, names: frozenset[str]) -> bool: + return not _is_chatgpt(deployment) and ( + deployment.model_name in names or deployment.model_info.team_public_model_name in names + ) + + +def _candidate_deployments(model: str, llm_router: "Router", team_id: str | None) -> tuple[Deployment, ...]: + registry: Final = llm_router.registry() + candidates: Final = tuple( + registry.deployment(model_id) for model_id in registry.names(model, kind="route_candidates", team_id=team_id) + ) + return tuple(deployment for deployment in candidates if deployment is not None) + + +def routes_to_litellm(model: object, llm_router: "Router | None", team_id: str | None) -> bool: + if ( + not isinstance(model, str) + or llm_router is None + or model in bundled_codex_models() + or f"{_CHATGPT_PROVIDER}/{model}" in litellm.model_cost # pyright: ignore[reportUnknownMemberType] # registry is dynamically typed + ): + return False + deployments: Final = _candidate_deployments(model, llm_router, team_id) + names: Final = frozenset({model, llm_router.registry().resolve(model, kind="alias") or model}) + return bool(deployments) and all(_serves_by_name(deployment, names) for deployment in deployments) + + +def _key_and_team_allow(model: str, llm_router: "Router", user_api_key_dict: UserAPIKeyAuth) -> bool: + allowed_lists: Final = ( + _MODEL_NAMES.validate_python(user_api_key_dict.models), + _MODEL_NAMES.validate_python(user_api_key_dict.team_models), + ) + return all( + check_model_access_helper(model, llm_router, allowed, team_id=user_api_key_dict.team_id) + for allowed in allowed_lists + ) + + +def native_chatgpt_model(model: str, llm_router: "Router", user_api_key_dict: UserAPIKeyAuth) -> str | None: + if model.startswith(f"{_CHATGPT_PROVIDER}/") or _key_and_team_allow(model, llm_router, user_api_key_dict): + return None + chatgpt_model: Final = f"{_CHATGPT_PROVIDER}/{model}" + deployments: Final = _candidate_deployments(chatgpt_model, llm_router, user_api_key_dict.team_id) + if not deployments or not all(_is_chatgpt(deployment) for deployment in deployments): + return None + return chatgpt_model if _key_and_team_allow(chatgpt_model, llm_router, user_api_key_dict) else None + + +class _SharedCache(Protocol): + async def async_get_cache(self, key: str) -> object: ... + + async def async_set_cache(self, key: str, value: object, *, ttl: int) -> None: ... + + async def async_increment_cache(self, key: str, value: float, *, ttl: int) -> float | None: ... + + +def codex_shared_cache() -> _SharedCache: + from litellm.proxy.proxy_server import proxy_logging_obj + + return proxy_logging_obj.internal_usage_cache.dual_cache + + +def _key_may_call_litellm_responses(user_api_key_dict: UserAPIKeyAuth) -> bool: + if not RouteChecks.is_allowed_on_proxy(_LITELLM_RESPONSES_PATH): + return False + try: + _ = RouteChecks.should_call_route(route=_LITELLM_RESPONSES_PATH, valid_token=user_api_key_dict) + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + RouteChecks.custom_admin_only_route_check(route=_LITELLM_RESPONSES_PATH) + except HTTPException: + return False + return True + + +async def _jwt_policy_admits_litellm_responses( + gateway_token: str, model: str, request_headers: Mapping[str, str] +) -> bool: + from litellm.proxy.proxy_server import ( + general_settings_view, + jwt_handler, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + try: + _ = await JWTAuthManager.auth_builder( + api_key=gateway_token, + jwt_handler=jwt_handler, + request_data={"model": model}, + general_settings=dict(general_settings_view()), + route=_LITELLM_RESPONSES_PATH, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + request_headers=request_headers, + request_method="POST", + ) + except Exception: # noqa: BLE001 # any JWT policy refusal keeps the caller on ChatGPT's models + return False + return True + + +@dataclass(frozen=True, slots=True) +class _LitellmResponsesAccess: + key_admits: bool + jwt_token: str | None + request_headers: Mapping[str, str] + + async def admits(self, model: str) -> bool: + if not self.key_admits: + return False + if self.jwt_token is None: + return True + return await _jwt_policy_admits_litellm_responses(self.jwt_token, model, self.request_headers) + + +def _litellm_responses_access(user_api_key_dict: UserAPIKeyAuth, request: Request) -> _LitellmResponsesAccess: + standard_jwt: Final = bool(user_api_key_dict.jwt_claims) and not user_api_key_dict.token + return _LitellmResponsesAccess( + key_admits=_key_may_call_litellm_responses(user_api_key_dict), + jwt_token=_gateway_token(request) if standard_jwt else None, + request_headers=request.headers, + ) + + +def _account_slugs_cache_key(headers: Mapping[str, str]) -> str | None: + account_id: Final = headers.get("chatgpt-account-id") + if not account_id: + return None + caller: Final = json.dumps([account_id, headers.get("authorization", "")]) + return f"{_ACCOUNT_SLUGS_CACHE_PREFIX}{hashlib.sha256(caller.encode()).hexdigest()}" + + +def _cached_slugs(cached: object) -> frozenset[str] | None: + if cached is None: + return None + try: + return _CACHED_SLUGS.validate_python(cached) + except ValidationError: + return None + + +async def _remember_account_slugs(cache_key: str, slugs: frozenset[str] | None, cache: _SharedCache) -> None: + if slugs is None: + await cache.async_set_cache(cache_key, _ACCOUNT_SLUGS_UNAVAILABLE, ttl=_ACCOUNT_SLUGS_UNAVAILABLE_TTL_SECONDS) + return + await cache.async_set_cache(cache_key, tuple(sorted(slugs)), ttl=_SHARED_CACHE_TTL_SECONDS) + + +def _codex_client_version(headers: Mapping[str, str]) -> str | None: + product: Final = _CODEX_USER_AGENT_PRODUCT.fullmatch(headers.get("user-agent", "").partition(" ")[0]) + if product is None or product.group(1) != headers.get("originator"): + return None + return product.group(2) + + +async def _fetched_account_slugs(headers: Mapping[str, str], client_version: str) -> frozenset[str] | None: + client: Final = get_async_httpx_client( + llm_provider=httpxSpecialProvider.PassThroughEndpoint, + params={"timeout": _ACCOUNT_CATALOG_TIMEOUT_SECONDS}, + ) + try: + response: Final = await client.client.get( + f"{CHATGPT_API_BASE}/models", + params={"client_version": client_version}, + headers={name: value for name, value in headers.items() if name in _ACCOUNT_CATALOG_HEADERS}, + ) + except httpx.HTTPError: + return None + return upstream_codex_slugs(response.content) if response.status_code == 200 else None + + +async def _admits_account_slugs_fetch(caller: str, cache: _SharedCache) -> bool: + fetches: Final = await cache.async_increment_cache( + f"{_ACCOUNT_SLUGS_FETCHES_CACHE_PREFIX}{caller}", 1, ttl=_ACCOUNT_SLUGS_FETCHES_WINDOW_SECONDS + ) + return fetches is not None and fetches <= _ACCOUNT_SLUGS_FETCHES_PER_CALLER + + +@with_service_target(_ACCOUNT_SLUGS_TARGET) +async def _account_slugs(headers: Mapping[str, str], caller: str, cache: _SharedCache) -> frozenset[str] | None: + cache_key: Final = _account_slugs_cache_key(headers) + if cache_key is None: + return None + cached_value: Final = await cache.async_get_cache(cache_key) + if cached_value == _ACCOUNT_SLUGS_UNAVAILABLE: + return None + cached: Final = _cached_slugs(cached_value) + if cached is not None: + return cached + client_version: Final = _codex_client_version(headers) + if client_version is None or not await _admits_account_slugs_fetch(caller, cache): + return None + fetched: Final = await _fetched_account_slugs(headers, client_version) + await _remember_account_slugs(cache_key, fetched, cache) + return fetched + + +async def _serves_through_litellm( + model: object, + llm_router: "Router | None", + user_api_key_dict: UserAPIKeyAuth, + access: _LitellmResponsesAccess, + caller: str, + headers: Mapping[str, str], + cache: _SharedCache, +) -> bool: + if not isinstance(model, str) or not routes_to_litellm(model, llm_router, user_api_key_dict.team_id): + return False + if not await access.admits(model): + return False + account_slugs: Final = await _account_slugs(headers, caller, cache) + return account_slugs is not None and model not in account_slugs + + +def _gateway_token(request: Request) -> str: + gateway_parts: Final = request.headers.get("x-litellm-api-key", "").split() + return ( + gateway_parts[-1] + if len(gateway_parts) == 1 or (len(gateway_parts) == 2 and gateway_parts[0].lower() == "bearer") + else "" + ) + + +def _chatgpt_bearer(request: Request) -> str: + scheme, _, token = request.headers.get("authorization", "").partition(" ") + return token.strip() if scheme.lower() == "bearer" else "" + + +def _forwarded_headers(request: Request) -> Mapping[str, str]: + return MappingProxyType({name: value for name, value in request.headers.items() if name in _HEADERS}) + + +async def native_turn_serves_through_litellm( + model: str, llm_router: "Router", request: Request, user_api_key_dict: UserAPIKeyAuth +) -> bool: + chatgpt_bearer: Final = _chatgpt_bearer(request) + gateway_token: Final = _gateway_token(request) + if not chatgpt_bearer or not gateway_token or hmac.compare_digest(chatgpt_bearer.encode(), gateway_token.encode()): + return False + return await _serves_through_litellm( + model, + llm_router, + user_api_key_dict, + _litellm_responses_access(user_api_key_dict, request), + _models_etag_caller(user_api_key_dict, gateway_token), + _forwarded_headers(request), + codex_shared_cache(), + ) + + +def _json_object(value: object) -> Mapping[str, object]: + try: + return _JSON_OBJECT.validate_python(value) + except ValidationError: + return MappingProxyType({}) + + +def _scope_state(request: Request) -> Mapping[str, object]: + return _json_object(request.scope.get("state")) + + +def _as_litellm_responses_request(request: Request, body: bytes, parsed_body: Mapping[str, object]) -> Request: + replayed: Final[Iterator[Message]] = iter(({"type": "http.request", "body": body, "more_body": False},)) + + async def receive() -> Message: + return next(replayed, None) or await request.receive() + + return Request( + { + **request.scope, + "path": _LITELLM_RESPONSES_PATH, + "raw_path": _LITELLM_RESPONSES_PATH.encode(), + "headers": [ + (name, value) for name, value in request.headers.raw if name not in _CHATGPT_CREDENTIAL_HEADERS + ], + "parsed_body": (tuple(parsed_body), parsed_body), + "state": { + name: value for name, value in _scope_state(request).items() if name not in _HEADER_DERIVED_STATE + }, + }, + receive, + ) + + +def _drops_params(body: Mapping[str, object], model: str, llm_router: "Router", team_id: str | None) -> bool: + deployments: Final = _candidate_deployments(model, llm_router, team_id) + return drops_unsupported_params(body.get("drop_params")) or ( + bool(deployments) + and all(normalize_drop_params(deployment.litellm_params.drop_params) is True for deployment in deployments) + ) + + +def _deployment_supports_no_effort(deployment: Deployment) -> bool: + configured: Final = _json_object(deployment.model_info.model_extra).get(_NONE_EFFORT_SUPPORT_FLAG) + if isinstance(configured, bool): + return configured + return supports_none_reasoning_effort( + model=deployment.litellm_params.model, custom_llm_provider=deployment.litellm_params.custom_llm_provider + ) + + +def _supports_no_effort(model: str, llm_router: "Router", team_id: str | None) -> bool: + deployments: Final = _candidate_deployments(model, llm_router, team_id) + return bool(deployments) and all(_deployment_supports_no_effort(deployment) for deployment in deployments) + + +def _with_provider_default_reasoning( + body: Mapping[str, object], llm_router: "Router | None", team_id: str | None, client_version: str | None +) -> Mapping[str, object]: + model: Final = body.get("model") + try: + reasoning: Final = _JSON_OBJECT.validate_python(body.get("reasoning")) + except ValidationError: + return body + if ( + reasoning.get("effort") != _CODEX_NO_DEFAULT_EFFORT + or not isinstance(model, str) + or llm_router is None + or not _drops_params(body, model, llm_router, team_id) + or codex_entry_names_default_reasoning_level(model, llm_router, team_id, client_version) + or _supports_no_effort(model, llm_router, team_id) + ): + return body + return {**body, "reasoning": {key: value for key, value in reasoning.items() if key != "effort"}} + + +def _with_wrapped_reasoning(body: Mapping[str, object]) -> Mapping[str, object]: + return { + **body, + "litellm_metadata": {**_json_object(body.get("litellm_metadata")), _ENCRYPTED_CONTENT_AFFINITY_FLAG: True}, + } + + +def _with_reasoning_the_targets_can_decrypt( + body: Mapping[str, object], model: str, llm_router: "Router", deployments: tuple[Deployment, ...] +) -> Mapping[str, object]: + try: + items: Final = _INPUT_ITEMS.validate_python(body.get("input")) + except ValidationError: + return body + model_group: Final = llm_router.registry().resolve(model, kind="alias") or model + kept: Final = llm_router.registry().decryptable_reasoning(model_group, items, deployments) + if body.get("store") is True: + return {**body, "input": list(kept)} + resolvable: Final = (_without_unpersisted_reasoning_reference(item) for item in kept) + return {**body, "input": [item for item in resolvable if item is not None]} + + +def _without_unpersisted_reasoning_reference(item: object) -> object | None: + if not _is_unpersisted_reasoning_reference(item): + return item + return ResponsesAPIRequestUtils.readable_reasoning(_JSON_OBJECT.validate_python(item)) + + +def _model_and_provider(deployment: Deployment) -> tuple[str, str] | None: + try: + model, provider, _, _ = get_llm_provider( + model=deployment.litellm_params.model, custom_llm_provider=deployment.litellm_params.custom_llm_provider + ) + except BadRequestError: + return None + return model, provider + + +def _is_openai_api(deployment: Deployment) -> bool: + resolved: Final = _model_and_provider(deployment) + return resolved is not None and resolved[1] in _OPENAI_API_PROVIDERS + + +def _takes_codex_tools(deployment: Deployment) -> bool: + resolved: Final = _model_and_provider(deployment) + if resolved is None: + return False + model, provider = resolved + return ( + provider in _CODEX_TOOL_PROVIDERS + or ProviderConfigManager.get_provider_responses_api_config( + provider=provider, model=model, api_base=deployment.litellm_params.api_base + ) + is None + ) + + +def _with_tools_the_targets_accept( + body: Mapping[str, object], deployments: tuple[Deployment, ...], drop_params: bool +) -> Mapping[str, object]: + if all(_is_openai_api(deployment) for deployment in deployments): + return body + if all(_takes_codex_tools(deployment) for deployment in deployments): + return with_additional_tools_hoisted(body, drop_params=drop_params) + return {**body, _CHAT_COMPLETIONS_BRIDGE: True} + + +def _drops_params_for_client( + body: Mapping[str, object], model: str, llm_router: "Router", team_id: str | None, user_agent: str +) -> bool: + from litellm.proxy.proxy_server import proxy_config + + return _drops_params(body, model, llm_router, team_id) or should_auto_drop_params_for_agentic_cli( + user_agent, "drop_params" in body, proxy_config + ) + + +def _litellm_body( + body: Mapping[str, object], llm_router: "Router | None", team_id: str | None, headers: Mapping[str, str] +) -> Mapping[str, object]: + wrapped: Final = _with_wrapped_reasoning( + _with_provider_default_reasoning(body, llm_router, team_id, _codex_client_version(headers)) + ) + model: Final = body.get("model") + if not isinstance(model, str) or llm_router is None: + return wrapped + deployments: Final = _candidate_deployments(model, llm_router, team_id) + return _with_tools_the_targets_accept( + _with_reasoning_the_targets_can_decrypt(wrapped, model, llm_router, deployments), + deployments, + _drops_params_for_client(body, model, llm_router, team_id, headers.get("user-agent", "")), + ) + + +def _is_unpersisted_reasoning_reference(item: object) -> bool: + try: + fields: Final = _JSON_OBJECT.validate_python(item) + except ValidationError: + return False + item_id: Final = fields.get("id") + encrypted_content: Final = fields.get("encrypted_content") + return ( + fields.get("type") == "reasoning" + and isinstance(item_id, str) + and bool(item_id) + and not (isinstance(encrypted_content, str) and bool(encrypted_content)) + ) + + +def _rejected_by_chatgpt(item: object, stored: bool) -> bool: + return is_other_provider_reasoning_item(item) or (not stored and _is_unpersisted_reasoning_reference(item)) + + +def _input_chatgpt_accepts(body: Mapping[str, object]) -> tuple[object, ...] | None: + try: + items: Final = _INPUT_ITEMS.validate_python(body.get("input")) + except ValidationError: + return None + stored: Final = body.get("store") is True + kept: Final = tuple(item for item in items if not _rejected_by_chatgpt(item, stored)) + return kept if len(kept) < len(items) else None + + +def _chatgpt_body(body: Mapping[str, object], client_model: str | None) -> Mapping[str, object] | None: + kept_input: Final = _input_chatgpt_accepts(body) + restored_model: Final = client_model is not None and body.get("model") == f"{_CHATGPT_PROVIDER}/{client_model}" + if kept_input is None and not restored_model: + return None + return { + **body, + **({"input": list(kept_input)} if kept_input is not None else {}), + **({"model": client_model} if restored_model else {}), + } + + +async def _litellm_catalog_rows( + user_api_key_dict: UserAPIKeyAuth, access: _LitellmResponsesAccess, llm_router: "Router | None" +) -> tuple[CodexCatalogRow, ...]: + from litellm.proxy.proxy_server import codex_catalog_rows_for_key + + if llm_router is None or not access.key_admits: + return () + try: + rows: Final = await codex_catalog_rows_for_key(user_api_key_dict) + except Exception as error: # noqa: BLE001 # a failed listing must never fail the ChatGPT answer it extends + verbose_proxy_logger.warning( + "Codex catalog keeps only ChatGPT's models: listing the key's LiteLLM models raised %s", + type(error).__name__, + ) + return () + routed: Final = tuple(row for row in rows if routes_to_litellm(row.id, llm_router, user_api_key_dict.team_id)) + admitted: Final = await asyncio.gather(*(access.admits(row.id) for row in routed)) + return tuple(row for row, row_admitted in zip(routed, admitted, strict=True) if row_admitted) + + +def _models_etag_caller(user_api_key_dict: UserAPIKeyAuth, gateway_token: str) -> str: + return user_api_key_dict.token or hashlib.sha256(gateway_token.encode()).hexdigest() + + +def _models_etag_cache_key(caller: str, upstream_etag: str) -> str: + return f"{_MODELS_ETAG_CACHE_PREFIX}{hashlib.sha256(json.dumps([caller, upstream_etag]).encode()).hexdigest()}" + + +async def _remembered_models_etag( + caller: str, upstream_etag: str, rows: tuple[CodexCatalogRow, ...], cache: _SharedCache +) -> str: + etag: Final = codex_catalog_etag(upstream_etag, rows) if rows else upstream_etag + await cache.async_set_cache( + _models_etag_cache_key(caller, upstream_etag), {"etag": etag}, ttl=_MODELS_ETAG_TTL_SECONDS + ) + return etag + + +async def _cached_models_etag(caller: str, upstream_etag: str, cache: _SharedCache) -> str | None: + try: + return _CACHED_MODELS_ETAG.validate_python( + await cache.async_get_cache(_models_etag_cache_key(caller, upstream_etag)) + )["etag"] + except ValidationError: + return None + + +async def _key_models_etag( + upstream_etag: str, + user_api_key_dict: UserAPIKeyAuth, + access: _LitellmResponsesAccess, + caller: str, + llm_router: "Router | None", + cache: _SharedCache, +) -> str: + cached: Final = await _cached_models_etag(caller, upstream_etag, cache) + if cached is not None: + return cached + rows: Final = await _litellm_catalog_rows(user_api_key_dict, access, llm_router) + return await _remembered_models_etag(caller, upstream_etag, rows, cache) + + +def _with_etag(response: Response, content: bytes | str, etag: str | None) -> Response: + kept_headers: Final = { + name: value for name, value in response.headers.items() if name not in _MERGED_CATALOG_DROPPED_HEADERS + } + return Response( + content=content, + status_code=response.status_code, + headers={**kept_headers, **({_ETAG_HEADER: etag} if etag is not None else {})}, + media_type="application/json", + ) + + +async def _catalog( + request: Request, + user_api_key_dict: UserAPIKeyAuth, + access: _LitellmResponsesAccess, + caller: str, + headers: Mapping[str, str], + cache: _SharedCache, +) -> Response: + from litellm.proxy.proxy_server import llm_router + + upstream: Final = await pass_through_request( + request=request, + target=f"{CHATGPT_API_BASE}/models", + custom_headers=dict(headers), + user_api_key_dict=user_api_key_dict, + custom_llm_provider=_CHATGPT_PROVIDER, + upstream_headers_filter=lambda _: headers, + ) + if upstream.status_code != 200 or isinstance(upstream, StreamingResponse): + return upstream + upstream_body: Final = bytes(upstream.body) + slugs: Final = upstream_codex_slugs(upstream_body) + if slugs is None: + return upstream + account_slugs_key: Final = _account_slugs_cache_key(headers) + if account_slugs_key is not None: + await _remember_account_slugs(account_slugs_key, slugs, cache) + rows: Final = await _litellm_catalog_rows(user_api_key_dict, access, llm_router) + upstream_etag: Final = upstream.headers.get(_ETAG_HEADER) + etag: Final = None if upstream_etag is None else await _remembered_models_etag(caller, upstream_etag, rows, cache) + if not rows: + return upstream + merged: Final = await asyncio.to_thread( + merged_codex_catalog, upstream_body, rows, client_version=request.query_params.get("client_version") + ) + return _with_etag(upstream, merged if merged is not None else upstream_body, etag) + + +async def _responses( + request: Request, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth, + access: _LitellmResponsesAccess, + caller: str, + headers: Mapping[str, str], + cache: _SharedCache, +) -> object: + from litellm.proxy.proxy_server import llm_router + + body: Final = _JSON_OBJECT.validate_python(await read_request_body(request)) + if await _serves_through_litellm(body.get("model"), llm_router, user_api_key_dict, access, caller, headers, cache): + litellm_body: Final = _litellm_body(body, llm_router, user_api_key_dict.team_id, headers) + await check_tools_allowlist( + request_body=dict(litellm_body), + valid_token=user_api_key_dict, + team_object=None, + route=_LITELLM_RESPONSES_PATH, + ) + return await responses_api( + _as_litellm_responses_request(request, json.dumps(litellm_body).encode(), litellm_body), + fastapi_response, + user_api_key_dict.model_copy(update={"request_route": _LITELLM_RESPONSES_PATH}), + ) + chatgpt_body: Final = _chatgpt_body(body, get_client_requested_model(request)) + upstream: Final = await pass_through_request( + request=request, + target=f"{CHATGPT_API_BASE}/responses", + custom_headers=dict(headers), + user_api_key_dict=user_api_key_dict, + custom_body=dict(chatgpt_body) if chatgpt_body is not None else None, + custom_llm_provider=_CHATGPT_PROVIDER, + upstream_headers_filter=lambda _: headers, + ) + upstream_models_etag: Final = upstream.headers.get(_MODELS_ETAG_HEADER) + if upstream_models_etag is not None: + upstream.headers[_MODELS_ETAG_HEADER] = await _key_models_etag( + upstream_models_etag, user_api_key_dict, access, caller, llm_router, cache + ) + return upstream + + +@router.get("/backend-api/codex/{endpoint}", tags=["ChatGPT Pass-through", "pass-through"], response_model=None) @router.post( "/backend-api/codex/{endpoint}", tags=["ChatGPT Pass-through", "pass-through"], + response_model=None, openapi_extra={ "requestBody": { "required": True, @@ -41,7 +747,7 @@ _HEADERS: Final = frozenset( }, responses={ 200: { - "description": "Unmodified Codex response or server-sent events", + "description": "Codex response, the model catalog with the key's LiteLLM models, or server-sent events", "content": { "application/json": {"schema": {"type": "object", "additionalProperties": True}}, "text/event-stream": {"schema": {"type": "string"}}, @@ -52,36 +758,29 @@ _HEADERS: Final = frozenset( async def codex_passthrough_route( endpoint: Literal["models", "responses"], request: Request, + fastapi_response: Response, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], -) -> Response: - """Forward a client's ChatGPT subscription with separate LiteLLM authentication.""" + cache: Annotated[_SharedCache, Depends(codex_shared_cache)], +) -> object: + """Forward a client's ChatGPT subscription with separate LiteLLM authentication, listing and serving the + key's other chat models through LiteLLM.""" expected_method: Final = "GET" if endpoint == "models" else "POST" if request.method != expected_method: raise HTTPException(status_code=405, detail="Method not allowed", headers={"Allow": expected_method}) - authorization: Final = request.headers.get("authorization", "") - scheme, _, token = authorization.partition(" ") - gateway_key: Final = request.headers.get("x-litellm-api-key", "") - gateway_parts: Final = gateway_key.split() - gateway_token: Final = ( - gateway_parts[-1] - if len(gateway_parts) == 1 or (len(gateway_parts) == 2 and gateway_parts[0].lower() == "bearer") - else "" - ) - if scheme.lower() != "bearer" or not token.strip() or not gateway_token: + chatgpt_bearer: Final = _chatgpt_bearer(request) + gateway_token: Final = _gateway_token(request) + if not chatgpt_bearer or not gateway_token: raise HTTPException( status_code=401, detail="Send the ChatGPT bearer in Authorization and the LiteLLM key in x-litellm-api-key", ) - if hmac.compare_digest(token.strip().encode(), gateway_token.encode()): + if hmac.compare_digest(chatgpt_bearer.encode(), gateway_token.encode()): raise HTTPException(status_code=401, detail="ChatGPT and LiteLLM credentials must be separate") - headers: Final = MappingProxyType({name: value for name, value in request.headers.items() if name in _HEADERS}) - return await pass_through_request( - request=request, - target=f"{CHATGPT_API_BASE}/{endpoint}", - custom_headers=dict(headers), - user_api_key_dict=user_api_key_dict, - custom_llm_provider="chatgpt", - upstream_headers_filter=lambda _: headers, - ) + headers: Final = _forwarded_headers(request) + caller: Final = _models_etag_caller(user_api_key_dict, gateway_token) + access: Final = _litellm_responses_access(user_api_key_dict, request) + if endpoint == "models": + return await _catalog(request, user_api_key_dict, access, caller, headers, cache) + return await _responses(request, fastapi_response, user_api_key_dict, access, caller, headers, cache) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index aabb83b3a5e..e4e3fec855d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -131,6 +131,7 @@ def _responses_sub_route(path: str) -> tuple[str, ...] | None: _STR_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) class _AgentConfig(BaseModel): @@ -550,8 +551,11 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): ) = None handler_instance: Final = OpenAIPassthroughLoggingHandler() + logged_custom_llm_provider: Final = _DECODED_JSON.validate_python( + logging_obj.model_call_details.get("custom_llm_provider", "openai") + ) custom_llm_provider: Final = TypeAdapter(str).validate_python( - kwargs.get("custom_llm_provider", logging_obj.model_call_details.get("custom_llm_provider", "openai")) + kwargs.get("custom_llm_provider", logged_custom_llm_provider) ) if is_chat_completions: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index fde616c34b4..d9f0e56ffce 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -68,6 +68,7 @@ from litellm.litellm_core_utils.core_helpers import ( bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, + require_attributes, ) from litellm.litellm_core_utils.initialize_dynamic_callback_params import validate_no_callback_env_reference from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY @@ -194,6 +195,13 @@ def _is_mergeable_client_metadata(key: str, value: object) -> bool: return not key.startswith("user_api_key_") and (key != "model_info" or isinstance(value, Mapping)) +def _request_body_fields(parsed: object) -> Mapping[str, object]: + fields: Final = as_str_mapping(parsed) + if fields is None: + raise TypeError(f"unsupported operand type(s) for |: '{type(parsed).__name__}' and 'dict'") + return fields + + class _RegisteredPassThroughRoute(TypedDict): endpoint_id: ReadOnly[str] path: ReadOnly[str | None] @@ -287,9 +295,9 @@ async def chat_completion_pass_through_endpoint( body: Final = await request.body() body_str: Final = body.decode() try: - data = ast.literal_eval(body_str) | data + data = {**_request_body_fields(_DECODED_JSON.validate_python(ast.literal_eval(body_str))), **data} except Exception: - data = json.loads(body_str) | data + data = {**_request_body_fields(_DECODED_JSON.validate_python(json.loads(body_str))), **data} data["adapter_id"] = adapter_id @@ -2237,15 +2245,13 @@ def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> di return restored_metadata -class _PassThroughRequestEnvelope(TypedDict, total=False): - query_params: Mapping[str, object] | None - custom_body: Mapping[str, object] | None - stream: bool | None +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... async def _parse_request_data_by_content_type( request: Request, -) -> tuple[object, object, None, bool | None]: +) -> tuple[object, object, None, object]: """ Parse request data based on content type. @@ -2264,10 +2270,12 @@ async def _parse_request_data_by_content_type( if "application/json" in content_type: # ✅ Handle JSON try: - body: _PassThroughRequestEnvelope = await request.json() - query_params_data = body.get("query_params") - custom_body_data = body.get("custom_body") - stream = body.get("stream") + json_body: Final = require_attributes( + _DECODED_JSON.validate_python(await request.json()), _SupportsGet, "get" + ) + query_params_data = json_body.get("query_params") + custom_body_data = json_body.get("custom_body") + stream = json_body.get("stream") except json.JSONDecodeError: # Handle requests with no body (e.g., DELETE requests) pass @@ -2275,14 +2283,16 @@ async def _parse_request_data_by_content_type( # ✅ Try to parse as JSON first (handles misconfigured clients sending JSON with multipart content-type) # If that fails, skip parsing - pass_through_request will handle actual multipart try: - body = await request.json() + multipart_json_body: Final = require_attributes( + _DECODED_JSON.validate_python(await request.json()), _SupportsGet, "get" + ) # Successfully parsed as JSON - treat as JSON body - query_params_data = body.get("query_params") - custom_body_data = body.get("custom_body") - stream = body.get("stream") + query_params_data = multipart_json_body.get("query_params") + custom_body_data = multipart_json_body.get("custom_body") + stream = multipart_json_body.get("stream") # If custom_body is not set, use the entire body - if custom_body_data is None and body: - custom_body_data = body + if custom_body_data is None and multipart_json_body: + custom_body_data = multipart_json_body except (json.JSONDecodeError, Exception): # Not JSON - this is actual multipart data # Skip parsing here to avoid consuming the request body stream diff --git a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py index bd8390d704e..07d4f7dae93 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py @@ -7,7 +7,7 @@ Handles guardrail execution for passthrough endpoints with: - Automatic inheritance from org/team/key levels when enabled """ -from typing import Any, Final +from typing import Final from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import merged_guardrail_entries @@ -243,14 +243,15 @@ class PassthroughGuardrailHandler: verbose_proxy_logger.debug("Passthrough guardrails config is empty, skipping") return None - temp_data: Final[dict[str, Any]] = {"metadata": {}} + inherited_metadata: Final = dict[str, list[object]]() + temp_data: Final = {"metadata": inherited_metadata} add_guardrails_from_key_or_team_metadata( key_metadata=user_api_key_dict.metadata, team_metadata=user_api_key_dict.team_metadata, data=temp_data, metadata_variable_name="metadata", ) - inherited_guardrails: Final[list[object]] = temp_data["metadata"].get("guardrails", []) + inherited_guardrails: Final = inherited_metadata.get("guardrails", []) guardrails_to_run: Final = merged_guardrail_entries(normalized_config, inherited_guardrails) verbose_proxy_logger.debug("Collected guardrails for passthrough endpoint: %s", guardrails_to_run) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 0c89627e74d..41aedf4b5cd 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -264,6 +264,7 @@ WORKER_APP: Final = "litellm_cli.main:proxy_worker_app" WORKER_LOG_CONFIG_ENV: Final = "LITELLM_WORKER_LOG_CONFIG" _LOG_CONFIG: Final = TypeAdapter[str | Mapping[str, object]](str | Mapping[str, object]) _LOG_CONFIG_FILE: Final = TypeAdapter[Mapping[str, object]](Mapping[str, object]) +_DECODED_JSON: Final = TypeAdapter[object](object) class ProxyInitializationHelpers: @@ -1332,8 +1333,7 @@ def run_server( ProxyInitializationHelpers.run_test_chat_completion(host, port, model, test) return else: - if headers: - headers = json.loads(headers) + worker_headers: Final = _DECODED_JSON.validate_python(json.loads(headers) if headers else headers) save_worker_config( model=model, alias=alias, @@ -1347,7 +1347,7 @@ def run_server( max_budget=max_budget, drop_params=drop_params, add_function_to_prompt=add_function_to_prompt, - headers=headers, + headers=worker_headers, save=save, config=config, use_queue=use_queue, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f5369bc34b0..3297a298af4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -162,6 +162,7 @@ from litellm.proxy._types import ( ProxyException, ProxyLifespanState, SpecialModelNames, + SSOUserDefinedValues, SupportedDBObjectType, TeamDefaultSettings, TokenCountRequest, @@ -221,7 +222,12 @@ from litellm.types.utils import ( is_server_derived_pricing_key, pricing_override_fields, ) -from litellm.utils import acredential_updates, cost_map_omits_token_price, with_named_credential +from litellm.utils import ( + acredential_updates, + cost_map_omits_token_price, + native_structured_output_support, + with_named_credential, +) if TYPE_CHECKING: from aiohttp import ClientSession @@ -231,6 +237,7 @@ if TYPE_CHECKING: from litellm.integrations.opentelemetry import OpenTelemetry from litellm.proxy.health_check_utils.shared_health_check_manager import SharedHealthCheckManager + from litellm.types.proxy.model_listing import ModelInfoResponse Span = _Span | Any else: @@ -512,9 +519,11 @@ from litellm.proxy.common_utils.anthropic_error_payload import ( ) from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AuthCacheInvalidationSubscriber, + publish_auth_cache_invalidation, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy -from litellm.proxy.common_utils.codex_model_catalog import codex_model_list_body +from litellm.proxy.common_utils.codex_catalog_source import codex_catalog_source +from litellm.proxy.common_utils.codex_model_catalog import CodexCatalogRow, codex_catalog_rows, codex_model_list_body from litellm.proxy.common_utils.config_includes import resolve_include_file_path, resolve_includes from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.credential_hydration import decrypted_or_stored @@ -553,6 +562,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( from litellm.proxy.common_utils.load_config_utils import get_config_from_bucket from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations from litellm.proxy.common_utils.model_listing_utils import ( + CallerAliases, ClaudeCodeRoutingNames, DeployedModelGroups, TeamModelNameTranslator, @@ -722,6 +732,7 @@ from litellm.proxy.litellm_pre_call_utils import ( SCHEDULER_PRIORITY_HEADER, LiteLLMProxyRequestSetup, add_litellm_data_to_request, + sanitize_for_log, ) from litellm.proxy.logging_endpoints.callback_logs_endpoints import ( rust_control_plane_router, @@ -1127,7 +1138,6 @@ from litellm.types.utils import ( SelectTokenizerResponse, StandardLoggingPayload, ) -from litellm.types.utils import ModelInfo as ModelMapInfo from litellm.utils import add_custom_logger_callback_to_specific_event try: @@ -3191,16 +3201,23 @@ def _apply_background_response_settings(value: SettingsJsonValue) -> None: ) -user_custom_auth = None -user_custom_key_generate = None +_CustomKeyHook: TypeAlias = Callable[..., Awaitable[Mapping[str, object]]] + + +class _CustomSsoHandler(Protocol): + async def __call__(self, sso_response: object) -> SSOUserDefinedValues | None: ... + + +user_custom_auth: Callable[..., Awaitable[object]] | None = None +user_custom_key_generate: _CustomKeyHook | None = None # Sentinel: prevents PKCE-no-Redis advisory from re-logging on config hot-reload. # Tests that need to reset it can patch 'litellm.proxy.proxy_server._pkce_no_redis_warning_emitted'. _pkce_no_redis_warning_emitted: bool = False _cp_no_redis_warning_emitted: bool = False -user_custom_key_update = None -user_custom_key_policy = None -user_custom_sso = None -user_custom_ui_sso_sign_in_handler = None +user_custom_key_update: _CustomKeyHook | None = None +user_custom_key_policy: _CustomKeyHook | None = None +user_custom_sso: _CustomSsoHandler | None = None +user_custom_ui_sso_sign_in_handler: object | None = None use_background_health_checks = None use_shared_health_check = None use_queue = False @@ -4420,6 +4437,36 @@ async def invalidate_spend_counter(counter_key: str) -> None: _invalidate_spend_counter: Final = invalidate_spend_counter +@with_service_target(SPEND_COUNTERS_TARGET) +async def set_spend_counter_with_floor_and_broadcast(counter_key: str, value: float) -> None: + """ + Set a Redis-backed spend counter to `value`, mirror it into the short-lived + spend_db_floor marker `_authoritative_floor_spend` reads, and broadcast both + to every worker (LIT-3803 pattern: setting, not deleting, means a worker's + own self-delivered broadcast still carries the reset value forward). + + Without the floor marker, `_authoritative_floor_spend` can re-derive a + stale, pre-reset value from a marker another worker cached moments earlier + and raise the just-reset counter right back up via `_repair_stale_spend_counter`. + Without the broadcast, a worker that already cached the pre-reset key object + or floor marker keeps enforcing against it until its own TTL expires. + """ + try: + await spend_counter_store.set(counter_key, value, ttl=60) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to update spend counter %s in Redis: %s. Budget checks may use stale value until counter expires.", + sanitize_for_log(counter_key), + redis_err, + ) + + floor_key: Final = f"spend_db_floor:{counter_key}" + spend_counter_cache.in_memory_cache.set_cache(key=floor_key, value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) + + await publish_auth_cache_invalidation(cache_key=counter_key, new_value=value, ttl=60) + await publish_auth_cache_invalidation(cache_key=floor_key, new_value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) + + async def _apply_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> None: if _defer_spend_counter_increments(pending): return @@ -4507,9 +4554,13 @@ async def _add_spend_to_shared_global_proxy_spend(cache: UserApiKeyCache, respon async def _add_spend_to_global_proxy_spend(cache: UserApiKeyCache, response_cost: float) -> None: - local_before: Final = cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True) + local_before: Final = _UNTYPED_VALUE.validate_python( + cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True) + ) shared_total: Final = await _add_spend_to_shared_global_proxy_spend(cache, response_cost) - local_total: Final = cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True) + local_total: Final = _UNTYPED_VALUE.validate_python( + cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True) + ) if not isinstance(local_total, (int, float)): return if shared_total is not None and local_total != local_before and shared_total < local_total: @@ -5840,8 +5891,16 @@ class _DbCallbackRegistration: displaced: tuple[tuple[str, object], ...] = () +class _CallbackList(Protocol): + def __contains__(self, value: object, /) -> bool: ... + + def append(self, value: object, /) -> None: ... + + def _restore_displaced_callback(list_name: str, entry: object) -> None: - callback_list: Final = getattr(litellm, list_name) + callback_list: Final = require_attributes( + _UNTYPED_VALUE.validate_python(getattr(litellm, list_name)), _CallbackList, "append" + ) if entry not in callback_list: callback_list.append(entry) @@ -7336,9 +7395,11 @@ class ProxyConfig: custom_ui_sso_sign_in_handler: Final = general_settings.get("custom_ui_sso_sign_in_handler", None) if custom_ui_sso_sign_in_handler is not None: - user_custom_ui_sso_sign_in_handler = get_instance_fn( - value=custom_ui_sso_sign_in_handler, - config_file_path=config_file_path, + user_custom_ui_sso_sign_in_handler = _UNTYPED_VALUE.validate_python( + get_instance_fn( + value=custom_ui_sso_sign_in_handler, + config_file_path=config_file_path, + ) ) if enterprise_proxy_config is not None: @@ -8640,7 +8701,7 @@ class ProxyConfig: environment_read: Final = next(self._environment_read_numbers) sections: Final = tuple(self._settings_stores) responses: Final = await asyncio.gather(*(get_config_param(prisma_client, section) for section in sections)) - general_settings_response: Final = responses[sections.index("general_settings")] + general_settings_response: Final[_ConfigParamValueRow | None] = responses[sections.index("general_settings")] general_settings_value: Final[Mapping[str, SettingsJsonValue]] = ( self._prepared_db_settings_values("general_settings", general_settings_response.param_value) if general_settings_response is not None and general_settings_response.param_value is not None @@ -10997,7 +11058,7 @@ async def async_data_generator( await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( request=request, request_data=request_data, - response=response, + response=_UNTYPED_VALUE.validate_python(response), stream_completed=stream_completed, client_disconnected=client_disconnected, user_api_key_dict=user_api_key_dict, @@ -12579,6 +12640,116 @@ async def _deployment_hidden_by_listing_callbacks(deployment: Deployment, user_a return listed_name in await _names_hidden_by_listing_callbacks(user_api_key_dict, (listed_name,)) +async def _caller_aliases(user_api_key_dict: UserAPIKeyAuth, team_id: str | None) -> CallerAliases: + caller_router_settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict, prisma_client, proxy_logging_obj + ) + return caller_alias_maps( + user_api_key_dict.aliases, + user_api_key_dict.team_model_aliases, + user_api_key_dict.team_id, + team_id, + caller_router_settings.get("model_group_alias") if isinstance(caller_router_settings, Mapping) else None, + ) + + +async def _listing_hidden_names(healthy_only: bool | None) -> frozenset[str]: + blocked_names: Final = llm_router.health().fully_blocked_model_names() if llm_router is not None else set[str]() + unhealthy_names: Final = await get_hidden_unhealthy_model_names( + healthy_only=healthy_only, + general_settings=general_settings_view(), + llm_router=llm_router, + ) + return frozenset(blocked_names | unhealthy_names) + + +@dataclass(frozen=True, slots=True) +class _KeyModelListing: + rows: "tuple[ModelInfoResponse, ...]" + entries: tuple[tuple[str, str], ...] + labels: tuple[tuple[str, str], ...] + + +def _listed_model_info( + response_id: str, lookup_id: str, include_metadata: bool, fallback_type: str | None +) -> "ModelInfoResponse": + from litellm.proxy.utils import create_model_info_response + + model_info: Final = create_model_info_response( + model_id=lookup_id, + provider="openai", + include_metadata=include_metadata, + fallback_type=fallback_type, + llm_router=llm_router, + ) + model_info["id"] = response_id + return model_info + + +async def _key_model_listing( + user_api_key_dict: UserAPIKeyAuth, + caller_aliases: CallerAliases, + hidden_names: frozenset[str], + *, + team_id: str | None = None, + include_model_access_groups: bool = False, + only_model_access_groups: bool = False, + include_wildcard_routes: bool = False, + include_metadata: bool = False, + fallback_type: str | None = None, + offered_model_names: Sequence[str] = (), +) -> _KeyModelListing: + from litellm.proxy.utils import get_available_models_for_user + + available: Final = await get_available_models_for_user( + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + general_settings=general_settings, + user_model=user_model, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + team_id=team_id, + include_model_access_groups=include_model_access_groups, + only_model_access_groups=only_model_access_groups, + return_wildcard_routes=include_wildcard_routes, + user_api_key_cache=user_api_key_cache, + ) + names: Final = (*available, *(name for name in offered_model_names if name not in available)) + excluded: Final = hidden_names | undiscoverable_model_names( + names, llm_router, user_api_key_dict, team_id or user_api_key_dict.team_id + ) + listed_entries: Final = await _entries_kept_by_listing_callbacks( + TeamModelNameTranslator.listing_entries( + [name for name in names if name not in excluded], llm_router, general_settings_view() + ), + user_api_key_dict, + ) + entries: Final = alias_listing_entries( + listed_entries, + caller_aliases, + _describable_model_groups(listed_entries, hidden_names, user_api_key_dict, team_id), + ) + return _KeyModelListing( + rows=tuple( + _listed_model_info(response_id, lookup_id, include_metadata, fallback_type) + for response_id, lookup_id in entries + ), + entries=entries, + labels=label_entries(listed_entries, entries), + ) + + +async def codex_catalog_rows_for_key(user_api_key_dict: UserAPIKeyAuth) -> tuple[CodexCatalogRow, ...]: + listing: Final = await _key_model_listing( + user_api_key_dict, + await _caller_aliases(user_api_key_dict, None), + await _listing_hidden_names(None), + ) + return await asyncio.to_thread( + codex_catalog_rows, listing.rows, listing.entries, llm_router, user_api_key_dict.team_id, listing.labels + ) + + @router.get("/v1/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"]) @router.get( "/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"] @@ -12646,10 +12817,7 @@ async def model_list( from litellm.proxy.management_endpoints.common_utils import ( user_has_admin_privileges, ) - from litellm.proxy.utils import ( - create_model_info_response, - get_available_models_for_user, - ) + from litellm.proxy.utils import create_model_info_response from litellm.types.proxy.model_listing import ModelInfoResponse http_request: Final = cast(Request | None, request) # cast-ok: in-process callers pass no request @@ -12657,16 +12825,7 @@ async def model_list( http_request is not None and http_request.headers.get("anthropic-version") is not None ) client_headers: Final[Mapping[str, str]] = http_request.headers if http_request is not None else _EMPTY_HEADERS - caller_router_settings: Final = await proxy_config.get_hierarchical_router_settings( - user_api_key_dict, prisma_client, proxy_logging_obj - ) - caller_aliases: Final = caller_alias_maps( - user_api_key_dict.aliases, - user_api_key_dict.team_model_aliases, - user_api_key_dict.team_id, - team_id, - caller_router_settings.get("model_group_alias") if isinstance(caller_router_settings, Mapping) else None, - ) + caller_aliases: Final = await _caller_aliases(user_api_key_dict, team_id) routing_names: Final = ClaudeCodeRoutingNames( llm_router, team_id or user_api_key_dict.team_id, @@ -12692,18 +12851,7 @@ async def model_list( proxy_logging_obj=proxy_logging_obj, ) - # Compute once — used in both branches below to hide paused models from the listing. - blocked_names: Final = llm_router.health().fully_blocked_model_names() if llm_router is not None else set[str]() - - # Opt-in: also hide models whose deployments are all unhealthy per background - # health checks. Empty when health state is unavailable or stale (fail open). - unhealthy_names: Final = await get_hidden_unhealthy_model_names( - healthy_only=healthy_only, - general_settings=settings, - llm_router=llm_router, - ) - - hidden_names: Final = blocked_names | unhealthy_names + hidden_names: Final = await _listing_hidden_names(healthy_only) from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import typesafe_allowlist_model_names @@ -12773,6 +12921,7 @@ async def model_list( admin_listing: Final = cast(Sequence[ModelInfoResponse], model_data) # cast-ok: rows built above if client_version is not None: + admin_codex_stock: Final = await codex_catalog_source.stock(client_version) return Response( content=await asyncio.to_thread( codex_model_list_body, @@ -12781,6 +12930,8 @@ async def model_list( llm_router, team_id or user_api_key_dict.team_id, client_version=client_version, + stock=admin_codex_stock.models, + instructions=admin_codex_stock.instructions, ), media_type="application/json", ) @@ -12801,59 +12952,26 @@ async def model_list( object="list", ) - # Otherwise, use the normal behavior (current implementation) - # Get available models for the user - available_models: Final = await get_available_models_for_user( - user_api_key_dict=user_api_key_dict, - llm_router=llm_router, - general_settings=general_settings, - user_model=user_model, - prisma_client=prisma_client, - proxy_logging_obj=proxy_logging_obj, - team_id=team_id, - include_model_access_groups=include_model_access_groups or False, - only_model_access_groups=only_model_access_groups or False, - return_wildcard_routes=include_wildcard_routes, - user_api_key_cache=user_api_key_cache, - ) offers_typesafe_allowlist_names: Final = ( asks_for_allowlist_names and team_id is None and user_api_key_has_admin_view(user_api_key_dict) ) - typesafe_names_to_offer: Final = typesafe_allowlist_model_names() if offers_typesafe_allowlist_names else () - all_models = [*available_models, *(name for name in typesafe_names_to_offer if name not in available_models)] - - undiscoverable_names: Final = undiscoverable_model_names( - all_models, llm_router, user_api_key_dict, team_id or user_api_key_dict.team_id - ) - if hidden_names or undiscoverable_names: - all_models = [m for m in all_models if m not in hidden_names and m not in undiscoverable_names] - - # Surface the public team name by default; legacy internal keys via flag. - # The internal routing key drives the metadata/fallback lookup, while the - # public name is what the client sees as the model id. - model_data = [] - listed_entries: Final = await _entries_kept_by_listing_callbacks( - TeamModelNameTranslator.listing_entries(all_models, llm_router, settings), user_api_key_dict - ) - entries: Final = alias_listing_entries( - listed_entries, + key_listing: Final = await _key_model_listing( + user_api_key_dict, caller_aliases, - _describable_model_groups(listed_entries, hidden_names, user_api_key_dict, team_id), + hidden_names, + team_id=team_id, + include_model_access_groups=include_model_access_groups or False, + only_model_access_groups=only_model_access_groups or False, + include_wildcard_routes=include_wildcard_routes, + include_metadata=include_metadata or False, + fallback_type=fallback_type, + offered_model_names=typesafe_allowlist_model_names() if offers_typesafe_allowlist_names else (), ) - labels: Final = label_entries(listed_entries, entries) - for response_id, lookup_id in entries: - model_info = create_model_info_response( - model_id=lookup_id, - provider="openai", - include_metadata=include_metadata or False, - fallback_type=fallback_type, - llm_router=llm_router, - ) - model_info["id"] = response_id - model_data.append(model_info) - - listing: Final = cast(Sequence[ModelInfoResponse], model_data) # cast-ok: rows built above + listing: Final = key_listing.rows + entries: Final = key_listing.entries + labels: Final = key_listing.labels if client_version is not None: + codex_stock: Final = await codex_catalog_source.stock(client_version) return Response( content=await asyncio.to_thread( codex_model_list_body, @@ -12863,6 +12981,8 @@ async def model_list( team_id or user_api_key_dict.team_id, labels, client_version=client_version, + stock=codex_stock.models, + instructions=codex_stock.instructions, ), media_type="application/json", ) @@ -12879,7 +12999,7 @@ async def model_list( ) return dict( - data=model_data, + data=list(listing), object="list", ) @@ -15928,10 +16048,9 @@ async def _count_request_tokens( deployment: dict[str, Any] | None = None if routed is None else dict(routed) litellm_model_name = None - model_info: ModelMapInfo | None = None + model_info: Final = None if routed is None else routed.get("model_info", {}) if deployment is not None: litellm_model_name = deployment.get("litellm_params", {}).get("model") - model_info = deployment.get("model_info", {}) deployment = { # rebind-ok: merge the named credential into a copy **deployment, "litellm_params": with_named_credential( @@ -16109,7 +16228,7 @@ async def _count_request_tokens( if model_info is not None: custom_tokenizer = cast( CustomHuggingfaceTokenizer | None, - model_info.get("custom_tokenizer", None), + _with_get(model_info).get("custom_tokenizer", None), ) system_message: Final = _system_message(local_system) prompt_turn: Final = prompt_only and not (system_message is None and not tools) @@ -16733,7 +16852,7 @@ def _enrich_model_info_with_litellm_data( stamped_model_info: Final = MappingProxyType( {**model_info, **_pricing_override_stamps(model_info, model.get("litellm_params") or MappingProxyType({}))} ) - model["model_info"] = { + merged_model_info: Final = { **stamped_model_info, **{ k: None if unpriced and k in ("input_cost_per_token", "output_cost_per_token") else v @@ -16741,6 +16860,18 @@ def _enrich_model_info_with_litellm_data( if k not in stamped_model_info or (stamped_model_info[k] is None and k in discovered_model_info) }, } + native_output_info: Final = { + **merged_model_info, + "supports_bedrock_native_structured_output_by_default": litellm_model_info.get( + "supports_bedrock_native_structured_output_by_default" + ), + } + model["model_info"] = ( + {**merged_model_info, "supports_native_structured_output": native_structured_output_support(native_output_info)} + if "supports_native_structured_output" in merged_model_info + and "supports_native_structured_output" not in stamped_model_info + else merged_model_info + ) # don't return the api key / vertex credentials # don't return the llm credentials model = remove_sensitive_info_from_deployment(model, excluded_keys={"litellm_credential_name"}) @@ -20618,7 +20749,7 @@ async def update_config_general_settings( detail={"error": CommonProxyErrors.not_allowed_access.value}, ) - requested_value: Final[object] = data.field_value + requested_value: Final = _UNTYPED_VALUE.validate_python(data.field_value) if data.field_name in _GENERAL_SETTINGS_UI_LITELLM_FIELDS: return await _persist_general_settings_ui_litellm_field(data.field_name, requested_value, user_api_key_dict) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 8fd4fbfcb8b..df156ccedd6 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -21,7 +21,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.proxy.db.db_span import db_span, db_spanned -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import PrismaClient, query_raw_rows from litellm.repositories.chunked_in import find_many_in from litellm.repositories.user_repository import UserRepository from litellm.types.llms.base import LiteLLMBaseModel @@ -207,7 +207,9 @@ async def _reverse_hash_key_metadata( warning: str, ) -> Mapping[str, KeyMetadataDict]: rows: Final = await _db_or_empty( - lambda: db_spanned("recover_key_metadata", table, lambda: prisma_client.db.query_raw(sql, sorted(wanted))), + lambda: db_spanned( + "recover_key_metadata", table, lambda: query_raw_rows(prisma_client.db, sql, sorted(wanted)) + ), warning, len(wanted), ) diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index 64b67ec55dd..e4a98619bfd 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -33,6 +33,7 @@ from litellm.constants import ( PTU_ROLLUP_MAX_BACKFILL_DAYS, PTU_SENTINEL_API_KEY, ) +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.ptu_pricing import ptu_terms from litellm.proxy.db.db_transaction_queue.pod_lock_manager import POD_LOCK_TARGET from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled @@ -79,6 +80,10 @@ class _SentinelBatchManager(Protocol): async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + def _sentinel_batch(prisma_client: "PrismaClient") -> _SentinelBatchManager: batch: Final[_SentinelBatchManager] = prisma_client.db.batch_() return batch @@ -171,7 +176,7 @@ class _PTUDeployment: model_info: Mapping[str, object] -def _router_deployment(deployment: Mapping[str, object]) -> _PTUDeployment | None: +def _router_deployment(deployment: _SupportsGet) -> _PTUDeployment | None: """A router ``model_list`` entry in the shape the parser reads, else None. An id is required rather than defaulted because it keys the sentinel row: every @@ -421,8 +426,8 @@ def _config_deployments(router: "Router | None", *, owned_by_db: frozenset[str]) credential clone carries ``original_model_id`` and reuses its source's PTU config under a fresh id, so pricing it would bill one reservation once per distinct client key. """ - entries: Final = tuple((router.registry().all() if router is not None else None) or ()) - records: Final = tuple(_router_deployment(entry) for entry in entries) + entries: Final[tuple[object, ...]] = tuple((router.registry().all() if router is not None else None) or ()) + records: Final = tuple(_router_deployment(require_attributes(entry, _SupportsGet, "get")) for entry in entries) return tuple( record for record in records diff --git a/litellm/proxy/spend_tracking/spend_capture_rate.py b/litellm/proxy/spend_tracking/spend_capture_rate.py index 443010b65da..0e6bed9017e 100644 --- a/litellm/proxy/spend_tracking/spend_capture_rate.py +++ b/litellm/proxy/spend_tracking/spend_capture_rate.py @@ -89,8 +89,14 @@ async def captured_spend_by_day( end_date: date, ) -> Mapping[str, float]: """LiteLLM's tracked spend per UTC day (ISO date) for the given ``custom_llm_provider`` values.""" - rows: Final = await prisma_client.db.query_raw( - _CAPTURED_SPEND_BY_DAY_SQL, start_date.isoformat(), end_date.isoformat(), tuple(litellm_providers) + from litellm.proxy.utils import query_raw_rows + + rows: Final = await query_raw_rows( + prisma_client.db, + _CAPTURED_SPEND_BY_DAY_SQL, + start_date.isoformat(), + end_date.isoformat(), + tuple(litellm_providers), ) return MappingProxyType({row.date: row.spend for row in _CAPTURED_SPEND_ROWS.validate_python(rows)}) diff --git a/litellm/proxy/types_utils/utils.py b/litellm/proxy/types_utils/utils.py index 62b87dbbf36..4f9b7787682 100644 --- a/litellm/proxy/types_utils/utils.py +++ b/litellm/proxy/types_utils/utils.py @@ -206,8 +206,8 @@ async def _download_gcs_file_wrapper(bucket_name: str, object_key: str, local_fi def validate_custom_validate_return_type( - fn: Callable[..., Any] | None, -) -> Callable[..., Literal[True]] | None: + fn: Callable[..., object] | None, +) -> Callable[..., object] | None: if fn is None: return None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b6538eecfe0..cf25816ada7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -75,6 +75,7 @@ from litellm.constants import ( VLLM_CHOICE_FIELDS, ) from litellm.cost_calculator import BaseTokenUsageProcessor +from litellm.integrations.otel.runtime import seed_request_identity from litellm.litellm_core_utils.bug_report import ( bug_report_notice, should_report_bug, @@ -162,10 +163,12 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import add_langfuse_trace_id_to_alert from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.litellm_core_utils.core_helpers import ( + attribute_or_none, coerce_token_limit, get_or_create_metadata_bucket, independent_snapshot, is_expected_client_error, + require_attributes, ) from litellm.litellm_core_utils.deployment_pricing import pricing_models from litellm.litellm_core_utils.duration_parser import get_budget_window_start @@ -315,7 +318,6 @@ if TYPE_CHECKING: from mcp.types import CallToolResult from opentelemetry.trace import Span as _Span from prisma import models as prisma_models - from prisma.actions import LiteLLM_DeprecatedVerificationTokenActions from prisma.client import TransactionManager from prisma.errors import UniqueViolationError from prisma.models import LiteLLM_DeprecatedVerificationToken @@ -364,6 +366,7 @@ _VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0 _VIEW_SETUP_GATE_TABLE: Final = '"LiteLLM_SpendLogs"' _VIEW_SETUP_GATE_PROBE_ROWS: Final = TypeAdapter(tuple[Mapping[str, bool], ...]) _DELEGATED_ATTRIBUTE: Final = TypeAdapter(object) +_DECODED_JSON: Final = TypeAdapter(object) _ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"] _ViewSetupAttempt: TypeAlias = Literal["ready", "table_missing"] | Exception @@ -405,7 +408,7 @@ def _update_many_batch(db: PrismaWrapper | RoutingPrismaWrapper) -> _UpdateManyB return db.batch_() -def _query_raw_rows( +def query_raw_rows( db: PrismaWrapper | RoutingPrismaWrapper, query: str, *args: object ) -> Awaitable[Sequence[Mapping[str, object]]]: return db.query_raw(query, *args) @@ -645,6 +648,24 @@ class _StreamIteratorHook(Protocol[_T]): def __call__(self, *, response: AsyncIterator[_T]) -> AsyncIterator[_T]: ... +class _BudgetAlertEmailer(Protocol): + async def budget_alerts( + self, + type: Literal[ + "token_budget", + "user_budget", + "soft_budget", + "max_budget_alert", + "team_budget", + "organization_budget", + "proxy_budget", + "projected_limit_exceeded", + "project_budget", + ], + user_info: CallInfo, + ) -> object: ... + + def _is_client_error_exception(exc: Exception) -> bool: if isinstance(exc, HTTPException): return exc.status_code < 500 @@ -738,6 +759,7 @@ def _drop_token_ids_ahead_of_chunk_check( _PLAIN_STREAM_FIELDS: Final = TypeAdapter(dict[str, object]) +_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object]) _EMPTY_STREAM_VALUES: Final[tuple[object, ...]] = (None, "", [], {}) @@ -1941,7 +1963,7 @@ class ProxyLogging: alerting=self.alerting, internal_usage_cache=self.internal_usage_cache.dual_cache, ) - self.email_logging_instance: Any | None = None + self.email_logging_instance: _BudgetAlertEmailer | None = None if BaseEmailLogger is not None: email_logger_class: Final = _get_email_logger_class() if email_logger_class is not None: @@ -2660,14 +2682,18 @@ class ProxyLogging: if callback is None: raise ValueError(f"No callback found for guardrail: {guardrail_name}") - return await self._execute_guardrail_hook( - callback=callback, - hook_type=hook_type, - data=data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - response=response, - ) + try: + return await self._execute_guardrail_hook( + callback=callback, + hook_type=hook_type, + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + response=response, + ) + except Exception as e: + enrich_http_exception_with_guardrail_context(e, callback) + raise async def _process_guardrail_callback( self, @@ -3201,6 +3227,12 @@ class ProxyLogging: if data is None: return None + if not guardrails_only: + seed_request_identity( + user_api_key_dict, + request_data=cast("Mapping[str, object]", data), # cast-ok: untyped request body, read with .get + ) + litellm_logging_obj: Final = cast(Optional["LiteLLMLoggingObj"], data.get("litellm_logging_obj", None)) prompt_id: Final[str | None] = data.get("prompt_id", None) @@ -4394,6 +4426,7 @@ class ProxyLogging: response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call") + request_view: Final = _REQUEST_FIELDS.validate_python(data) guardrail_callbacks, other_callbacks = _partition_post_call_callbacks() try: # Merge model-level guardrails before checking which guardrails to run @@ -4424,6 +4457,9 @@ class ProxyLogging: ): continue + if callback.post_call_already_checked(request_view, response) is True: + continue + guardrail_response: Any | None = None if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks: @@ -4497,10 +4533,13 @@ class ProxyLogging: immediately before awaiting, and the unified hook pops it before its first suspension point, so concurrent guardrails never race on that key. """ + request_view: Final = _REQUEST_FIELDS.validate_python(data) async def _run_one(callback: CustomGuardrail) -> None: if callback.should_run_guardrail(data=guardrail_data, event_type=GuardrailEventHooks.post_call) is not True: return + if callback.post_call_already_checked(request_view, response) is True: + return if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks: data["guardrail_to_apply"] = callback await self._run_guardrail_with_metrics( @@ -5358,9 +5397,9 @@ async def _lookup_deprecated_key( _deprecated_key_cache.pop(hashed_token, None) try: - deprecated_keys_table: Final[ - LiteLLM_DeprecatedVerificationTokenActions[LiteLLM_DeprecatedVerificationToken] - ] = db.litellm_deprecatedverificationtoken + deprecated_keys_table: Final[TableActions[LiteLLM_DeprecatedVerificationToken]] = ( + db.litellm_deprecatedverificationtoken + ) deprecated_row: Final = await deprecated_keys_table.find_first( where={ "token": hashed_token, @@ -5411,11 +5450,20 @@ def _config_cache_key(param_name: str) -> str: return f"litellm_config:param:{param_name}" -def _pack_config_row(row: Any) -> dict[str, object]: +class _PackableConfigRow(Protocol): + @property + def param_name(self) -> object: ... + + @property + def param_value(self) -> object: ... + + +def _pack_config_row(row: object) -> dict[str, object]: + config_row: Final = require_attributes(row, _PackableConfigRow, "param_name", "param_value") return { - "param_name": row.param_name, - "param_value": row.param_value, - "reload_revision": getattr(row, "reload_revision", None), + "param_name": config_row.param_name, + "param_value": config_row.param_value, + "reload_revision": attribute_or_none(row, "reload_revision"), } @@ -6243,7 +6291,7 @@ class PrismaClient: LIMIT $1 OFFSET $2 """ - response = await _query_raw_rows(self.db, sql_query, limit, offset) + response = await query_raw_rows(self.db, sql_query, limit, offset) return response elif table_name == "spend": verbose_proxy_logger.debug("PrismaClient: get_data: table_name == 'spend'") @@ -7553,7 +7601,7 @@ class PrismaClient: # caller's own retried query, a stronger check than SELECT 1, # and a reader that fails to come back sets `_reader_unavailable` # so reads fall through to the writer just recreated here. - await _query_raw_rows(self.writer_db, "SELECT 1") + await query_raw_rows(self.writer_db, "SELECT 1") # A recreate can decline: the optimistic-lock guard no-ops when # the writer generation moved since cycle entry, and the routing # wrapper then leaves the reader untouched as well. Callers that @@ -7903,7 +7951,7 @@ class PrismaClient: async def _view_setup_gate_table_present(self) -> bool: rows: Final = _VIEW_SETUP_GATE_PROBE_ROWS.validate_python( - await _query_raw_rows(self.db, "SELECT to_regclass($1) IS NOT NULL AS present", _VIEW_SETUP_GATE_TABLE) + await query_raw_rows(self.db, "SELECT to_regclass($1) IS NOT NULL AS present", _VIEW_SETUP_GATE_TABLE) ) return rows[0]["present"] @@ -7913,7 +7961,7 @@ class PrismaClient: await asyncio.sleep(self._db_health_watchdog_interval_seconds) async with db_span("db_health_watchdog", None): await asyncio.wait_for( - _query_raw_rows(self.db, "SELECT 1"), + query_raw_rows(self.db, "SELECT 1"), timeout=self._db_health_watchdog_probe_timeout_seconds, ) if isinstance(self.db, RoutingPrismaWrapper) and self.db.writer_unavailable: @@ -7973,7 +8021,7 @@ class PrismaClient: async def _writer_is_read_only(self, writer: PrismaWrapper) -> bool: """True iff the pooled writer session answers reads but rejects writes (SQLSTATE 25006).""" rows: Final = _WRITER_WRITABILITY_PROBE_ROWS.validate_python( - await _query_raw_rows(writer, _WRITER_WRITABILITY_PROBE_SQL) + await query_raw_rows(writer, _WRITER_WRITABILITY_PROBE_SQL) ) return any(row.get("transaction_read_only") == "on" for row in rows) @@ -8000,7 +8048,7 @@ class PrismaClient: """ sql_query: Final = "SELECT 1" async with db_span("health_check", None): - response: Final[object] = await _query_raw_rows(wrapper, sql_query) + response: Final[object] = await query_raw_rows(wrapper, sql_query) return response async def _probe_answers_now(self, wrapper: PrismaWrapper) -> bool: @@ -8168,12 +8216,12 @@ class PrismaClient: verbose_proxy_logger.warning("Invalid response_time_ms value: %s", response_time_ms) return None - def _clean_details(self, details: dict | None) -> dict | None: + def _clean_details(self, details: dict | None) -> object: """Clean and validate details JSON""" if not isinstance(details, dict): return None try: - return safe_json_loads(safe_dumps(details)) + return _DECODED_JSON.validate_python(safe_json_loads(safe_dumps(details))) except Exception as e: verbose_proxy_logger.warning("Failed to clean details JSON: %s", e) return None diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index b6a0ebe3c58..c0d7ab8d6e1 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -1,18 +1,21 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import ( + TYPE_CHECKING, Annotated, - Any, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict Final, + Protocol, cast, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict ) from fastapi import APIRouter, Depends, HTTPException, Request, Response +from pydantic import TypeAdapter from litellm.constants import VECTOR_STORE_CALLER_METADATA_KWARG from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) +from litellm.litellm_core_utils.core_helpers import as_object_list, decoded_json_object, require_attributes from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -30,8 +33,25 @@ from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse from litellm.vector_stores.utils import VectorStoreWriteBody, callers_metadata_snapshot from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry +if TYPE_CHECKING: + from litellm import Router + router: Final = APIRouter() +_DECODED_JSON: Final = TypeAdapter(object) + + +class _CreatesManagedVectorStores(Protocol): + async def acreate_vector_store( + self, + create_request: Mapping[str, object], + llm_router: "Router", + target_model_names_list: Sequence[object], + litellm_parent_otel_span: object, + user_api_key_dict: UserAPIKeyAuth, + ) -> object: ... + + BLOCKED_QUERY_EMBEDDING_SELECTION_PARAMS: Final = frozenset( { "embedding_model", @@ -247,22 +267,23 @@ async def vector_store_create( ) # Check for target_model_names parameter - target_model_names: Final = data.pop("target_model_names", None) + target_model_names: Final = _DECODED_JSON.validate_python(data.pop("target_model_names", None)) if target_model_names: # Use managed vector stores for multi-model support - if isinstance(target_model_names, str): - target_model_names_list = [m.strip() for m in target_model_names.split(",")] - elif isinstance(target_model_names, list): - target_model_names_list = target_model_names - else: + target_model_names_list: Final = ( + [m.strip() for m in target_model_names.split(",")] + if isinstance(target_model_names, str) + else as_object_list(target_model_names) + ) + if target_model_names_list is None: raise HTTPException( status_code=400, detail="target_model_names must be a comma-separated string or list of model names", ) # Get managed vector stores hook - managed_vector_stores: Final[Any] = proxy_logging_obj.get_proxy_hook("managed_vector_stores") + managed_vector_stores: Final = proxy_logging_obj.get_proxy_hook("managed_vector_stores") if managed_vector_stores is None: raise HTTPException( status_code=500, @@ -276,8 +297,10 @@ async def vector_store_create( ) # Create vector store across multiple models - response: Final[object] = await managed_vector_stores.acreate_vector_store( - create_request=data, + response: Final = await require_attributes( + managed_vector_stores, _CreatesManagedVectorStores, "acreate_vector_store" + ).acreate_vector_store( + create_request=decoded_json_object(_DECODED_JSON.validate_python(data)), llm_router=llm_router, target_model_names_list=target_model_names_list, litellm_parent_otel_span=user_api_key_dict.parent_otel_span, diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 6bfccd2426b..4b9547300db 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -308,9 +308,7 @@ async def new_vector_store( # Extract and validate metadata metadata: Final = vector_store.get("vector_store_metadata") - validated_metadata: dict | None = None - if metadata is not None and isinstance(metadata, dict): - validated_metadata = metadata + validated_metadata: Final = metadata if isinstance(metadata, dict) else None new_vector_store: Final = await create_vector_store_in_db( vector_store_id=vector_store_id, @@ -382,7 +380,6 @@ async def list_vector_stores( from litellm.proxy.proxy_server import prisma_client vector_store_map: Final[dict[str, LiteLLM_ManagedVectorStore]] = {} - db_vector_store_ids: Final[set] = set() try: # Get vector stores from database first (source of truth) @@ -393,7 +390,7 @@ async def list_vector_stores( vector_store_id = vector_store.get("vector_store_id", None) if vector_store_id: vector_store_map[vector_store_id] = vector_store - db_vector_store_ids.add(vector_store_id) + db_vector_store_ids: Final = frozenset(vector_store_map) # Process in-memory vector stores if litellm.vector_store_registry is not None: @@ -647,8 +644,9 @@ async def update_vector_store( ) # Handle metadata serialization - if update_data.get("vector_store_metadata") is not None: - update_data["vector_store_metadata"] = safe_dumps(update_data["vector_store_metadata"]) + submitted_metadata: Final[object] = update_data.get("vector_store_metadata") + if submitted_metadata is not None: + update_data["vector_store_metadata"] = safe_dumps(submitted_metadata) # Handle litellm_params if provided. As with the create path, the # embedding-config auto-resolve previously persisted cleartext diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index 9d8318d5a83..696ad4ab170 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -1,7 +1,7 @@ from typing import TYPE_CHECKING, Any, Final import orjson -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm.types.videos.main import DecodedVideoId, VideoObject from litellm.types.videos.utils import ( @@ -14,6 +14,8 @@ from litellm.types.videos.utils import ( if TYPE_CHECKING: from litellm.router import Router +_DECODED_JSON: Final = TypeAdapter(object) + def extract_model_from_target_model_names(target_model_names: object) -> str | None: if isinstance(target_model_names, str): @@ -30,7 +32,7 @@ def video_reference_to_id(video_ref: object) -> str: if not isinstance(video_ref, str): return "" try: - parsed_ref: Final = orjson.loads(video_ref) + parsed_ref: Final = _DECODED_JSON.validate_python(orjson.loads(video_ref)) except orjson.JSONDecodeError: return video_ref return parsed_ref.get("id", "") if isinstance(parsed_ref, dict) else video_ref @@ -44,7 +46,7 @@ def get_custom_provider_from_data(data: dict[str, Any]) -> str | None: extra_body = data.get("extra_body") if isinstance(extra_body, str): try: - parsed_extra_body: Final = orjson.loads(extra_body) + parsed_extra_body: Final = _DECODED_JSON.validate_python(orjson.loads(extra_body)) if isinstance(parsed_extra_body, dict): extra_body = parsed_extra_body except Exception: diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 77bf6868bf1..2d86f4344db 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -36,6 +36,7 @@ from litellm.types.rag import ( RAGIngestResponse, ) from litellm.types.utils import ModelResponse +from litellm.types.vector_stores import VectorStoreSearchResponse from litellm.utils import client if TYPE_CHECKING: @@ -257,7 +258,7 @@ async def _execute_query_pipeline( {**provider_search_params, **kwargs, **filter_search_params, **store_search_params} ) with internal_sub_call(parent_request_id): - search_response: Final = await litellm.vector_stores.asearch( + search_response: Final[VectorStoreSearchResponse] = await litellm.vector_stores.asearch( vector_store_id=retrieval_config["vector_store_id"], query=query_text, max_num_results=retrieval_config.get("top_k", 10), diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py index 9f608dbc99a..d35a8ddc137 100644 --- a/litellm/rag/rag_query.py +++ b/litellm/rag/rag_query.py @@ -1,5 +1,5 @@ from collections.abc import Sequence -from typing import Any, Final +from typing import Any, Final, TypeVar from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -7,6 +7,8 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessag from litellm.types.utils import ModelResponse from litellm.types.vector_stores import VectorStoreSearchResponse +_ItemT = TypeVar("_ItemT") + class _ResultContentView(TypedDict): """Content entry carried by a vector store search result.""" @@ -22,12 +24,13 @@ class _SearchResultView(TypedDict): text: ReadOnly[NotRequired[str]] -class _SearchDataView(TypedDict): - results: ReadOnly[Sequence[_SearchResultView]] +class _StoreContentView(TypedDict): + type: ReadOnly[NotRequired[str | None]] + text: ReadOnly[NotRequired[str | None]] -class _ContextChunksView(TypedDict): - chunks: ReadOnly[Sequence[_SearchResultView | str | None]] +class _StoreResultView(TypedDict): + content: ReadOnly[NotRequired[Sequence[_StoreContentView] | None]] class _RerankResultView(TypedDict): @@ -42,6 +45,18 @@ class _MessageView(TypedDict): message: ReadOnly[object] +def _iterable_items(items: Sequence[_ItemT] | None) -> Sequence[_ItemT]: + if items is None: + raise TypeError("'NoneType' object is not iterable") + return items + + +def _sized_items(items: Sequence[_ItemT] | None) -> Sequence[_ItemT]: + if items is None: + raise TypeError("object of type 'NoneType' has no len()") + return items + + class RAGQuery: CONTENT_PREFIX_STRING = "Context:\n\n" @@ -70,16 +85,17 @@ class RAGQuery: return None @staticmethod - def build_context_message(context_chunks: Sequence[_SearchResultView | str | None]) -> ChatCompletionUserMessage: + def build_context_message( + context_chunks: Sequence[_StoreResultView | _SearchResultView | str | None] | None, + ) -> ChatCompletionUserMessage: """ Process search results and build a context message. """ context_content = RAGQuery.CONTENT_PREFIX_STRING - chunks: Final[_ContextChunksView] = {"chunks": context_chunks} - for chunk in chunks["chunks"]: + for chunk in _iterable_items(context_chunks): if isinstance(chunk, dict): - result_content: Sequence[_ResultContentView] | None = chunk.get("content") + result_content: Sequence[_StoreContentView | _ResultContentView] | None = chunk.get("content") if result_content: for content_item in result_content: content_text: str | None = content_item.get("text") @@ -123,27 +139,27 @@ class RAGQuery: @staticmethod def extract_documents_from_search( - search_response: Any, + search_response: VectorStoreSearchResponse, ) -> list[str | dict[str, object]]: """Extract text documents from vector store search response.""" documents: Final[list[str | dict[str, object]]] = [] - search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])} - for result in search_data["results"]: - content_list = result.get("content", []) - for content in content_list: - if content.get("type") == "text" and content.get("text"): - documents.append(content["text"]) + results: Final[Sequence[_StoreResultView]] = _iterable_items(search_response.get("data", [])) + for result in results: + for content in _iterable_items(result.get("content", [])): + if content.get("type") == "text" and (content_text := content.get("text")): + documents.append(content_text) return documents @staticmethod - def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> list[_SearchResultView]: + def get_top_chunks_from_rerank( + search_response: VectorStoreSearchResponse, rerank_response: Any + ) -> list[_StoreResultView]: """Get the original search results corresponding to the top reranked results.""" - top_chunks: Final[list[_SearchResultView]] = [] - search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])} - original_results: Final = search_data["results"] + top_chunks: Final[list[_StoreResultView]] = [] + original_results: Final = search_response.get("data", []) reranked: Final[_RerankResultsView] = {"results": rerank_response.get("results", [])} for result in reranked["results"]: index = result.get("index") - if index is not None and index < len(original_results): - top_chunks.append(original_results[index]) + if index is not None and index < len(sized_results := _sized_items(original_results)): + top_chunks.append(sized_results[index]) return top_chunks diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index be1df4aa0f6..85d87554023 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -149,7 +149,7 @@ def _build_litellm_metadata(kwargs: Mapping[str, object]) -> dict[str, object]: _GUARDRAIL_OPT_OUT_METADATA_KEYS: Final = ("user_api_key_metadata", "user_api_key_team_metadata") -def _translation_litellm_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: +def _session_litellm_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: request_metadata: Final = _OBJECT_MAPPING.validate_python(kwargs.get("metadata") or {}) guardrail_opt_outs: Final = { key: request_metadata[key] for key in _GUARDRAIL_OPT_OUT_METADATA_KEYS if key in request_metadata @@ -218,7 +218,7 @@ def _translation_rejection_error( def _raise_if_http_translation_rejected(custom_llm_provider: str, model: str, kwargs: Mapping[str, object]) -> None: rejection: Final = _translation_rejection( - custom_llm_provider, _HTTP_TRANSLATION_PROVIDERS, _translation_litellm_metadata(kwargs) + custom_llm_provider, _HTTP_TRANSLATION_PROVIDERS, _session_litellm_metadata(kwargs) ) if rejection is None: return @@ -858,9 +858,7 @@ async def _arealtime( ) session_litellm_metadata: Final = ( - dict(_translation_litellm_metadata(kwargs)) - if realtime_mode == "translation" - else _build_litellm_metadata(kwargs) + _build_litellm_metadata(kwargs) if live_session_start is not None else dict(_session_litellm_metadata(kwargs)) ) translation_rejection: Final = ( _translation_rejection(_custom_llm_provider, _WEBSOCKET_TRANSLATION_PROVIDERS, session_litellm_metadata) @@ -1014,6 +1012,8 @@ async def _arealtime( aws_external_id=aws_external_id, aws_session_tags=aws_session_tags, drop_params=drop_params, + user_api_key_dict=kwargs.get("user_api_key_dict"), + litellm_metadata=session_litellm_metadata, ) elif _custom_llm_provider == "xai": api_base = ( diff --git a/litellm/repositories/autorouter_session_repository.py b/litellm/repositories/autorouter_session_repository.py index 640fd26ffe3..e39087167e4 100644 --- a/litellm/repositories/autorouter_session_repository.py +++ b/litellm/repositories/autorouter_session_repository.py @@ -2,31 +2,25 @@ Repository for the auto-router per-session rollup (LiteLLM_AutoRouterSession). """ -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final from litellm.models.autorouter_session import LiteLLM_AutoRouterSession from litellm.repositories.base_repository import BaseRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models -class _AutoRouterSessionDb(Protocol): - @property - def litellm_autoroutersession(self) -> TableActions["prisma_models.LiteLLM_AutoRouterSession"]: ... - - -class _PrismaClientView(Protocol): - @property - def db(self) -> _AutoRouterSessionDb: ... +class _AutoRouterSessionTable(PrismaTableRepository["prisma_models.LiteLLM_AutoRouterSession"]): + table_name = "litellm_autoroutersession" class AutoRouterSessionRepository(BaseRepository[LiteLLM_AutoRouterSession]): @property def table(self) -> TableActions["prisma_models.LiteLLM_AutoRouterSession"]: - client: Final[_PrismaClientView] = self.prisma_client - return client.db.litellm_autoroutersession + return _AutoRouterSessionTable(self._prisma_client).table @property def model_class(self) -> type[LiteLLM_AutoRouterSession]: diff --git a/litellm/repositories/base_repository.py b/litellm/repositories/base_repository.py index 92f8ba984f4..a0662a5fc33 100644 --- a/litellm/repositories/base_repository.py +++ b/litellm/repositories/base_repository.py @@ -45,12 +45,15 @@ class BaseRepository(ABC, Generic[T]): def __init__(self, prisma_client: object): self._prisma_client = prisma_client - @property - def prisma_client(self) -> Any: # any-ok: PrismaClient is an untyped runtime wrapper + def _connected_client(self) -> object: if self._prisma_client is None: raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") return self._prisma_client + @property + def prisma_client(self) -> Any: # any-ok: PrismaClient is an untyped runtime wrapper + return self._connected_client() + @property @abstractmethod def table(self) -> TableActions[DbRecord]: diff --git a/litellm/repositories/budget_repository.py b/litellm/repositories/budget_repository.py index 205646c8393..af6c8cfe6f5 100644 --- a/litellm/repositories/budget_repository.py +++ b/litellm/repositories/budget_repository.py @@ -3,28 +3,19 @@ Budget repository for database operations on LiteLLM_BudgetTable. """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final from litellm.models.budget import LiteLLM_BudgetTable from litellm.repositories.base_repository import BaseRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models -class _BudgetDb(Protocol): - """The single Prisma table this repository reaches for on ``prisma_client.db``.""" - - @property - def litellm_budgettable(self) -> TableActions["prisma_models.LiteLLM_BudgetTable"]: ... - - -class _PrismaClientView(Protocol): - """The one attribute this repository reads off the untyped Prisma client wrapper.""" - - @property - def db(self) -> _BudgetDb: ... +class _BudgetTable(PrismaTableRepository["prisma_models.LiteLLM_BudgetTable"]): + table_name = "litellm_budgettable" class BudgetRepository(BaseRepository[LiteLLM_BudgetTable]): @@ -32,8 +23,7 @@ class BudgetRepository(BaseRepository[LiteLLM_BudgetTable]): @property def table(self) -> TableActions["prisma_models.LiteLLM_BudgetTable"]: - client: Final[_PrismaClientView] = self.prisma_client - return client.db.litellm_budgettable + return _BudgetTable(self._prisma_client).table @property def model_class(self) -> type[LiteLLM_BudgetTable]: diff --git a/litellm/repositories/credentials_repository.py b/litellm/repositories/credentials_repository.py index cb1bc1d101b..d618b71ce52 100644 --- a/litellm/repositories/credentials_repository.py +++ b/litellm/repositories/credentials_repository.py @@ -7,11 +7,12 @@ so reads return the stored values verbatim. """ from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias from litellm.models.credentials import CredentialItem from litellm.repositories.base_repository import DbRecord, record_to_dict from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models @@ -29,27 +30,24 @@ class _PrismaClientView(Protocol): def db(self) -> _PrismaCredentialsDb: ... +class _CredentialsTableRepository(PrismaTableRepository["prisma_models.LiteLLM_CredentialsTable"]): + table_name = "litellm_credentialstable" + + class CredentialsRepository: """Repository for credentials database operations, keyed by credential name.""" - def __init__(self, prisma_client: Any): # any-ok: PrismaClient is an untyped runtime wrapper + def __init__(self, prisma_client: object): self._prisma_client = prisma_client @property def prisma_client(self) -> _PrismaClientView: - if self._prisma_client is None: - raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") - client: Final[_PrismaClientView] = self._prisma_client + client: Final[_PrismaClientView] = _CredentialsTableRepository(self._prisma_client).prisma_client return client @property def table(self) -> "_CredentialsTable": - from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync - - return wrap_table_actions_for_config_sync( - actions=self.prisma_client.db.litellm_credentialstable, - table_name="litellm_credentialstable", - ) + return _CredentialsTableRepository(self._prisma_client).table @staticmethod def _to_model(record: DbRecord | None) -> CredentialItem | None: diff --git a/litellm/repositories/object_permission_repository.py b/litellm/repositories/object_permission_repository.py index 99d6aeed88b..e90da64afcd 100644 --- a/litellm/repositories/object_permission_repository.py +++ b/litellm/repositories/object_permission_repository.py @@ -2,27 +2,19 @@ ObjectPermission repository for database operations on LiteLLM_ObjectPermissionTable. """ -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.repositories.base_repository import BaseRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models -class _ObjectPermissionDb(Protocol): - @property - def litellm_objectpermissiontable(self) -> TableActions["prisma_models.LiteLLM_ObjectPermissionTable"]: ... - - -class _PrismaClientView(Protocol): - @property - def db(self) -> _ObjectPermissionDb: ... - - @property - def writer_db(self) -> _ObjectPermissionDb: ... +class _ObjectPermissionTable(PrismaTableRepository["prisma_models.LiteLLM_ObjectPermissionTable"]): + table_name = "litellm_objectpermissiontable" class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): @@ -34,9 +26,7 @@ class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): @property def table(self) -> TableActions["prisma_models.LiteLLM_ObjectPermissionTable"]: - client: Final[_PrismaClientView] = self.prisma_client - database: Final = client.writer_db if self._use_writer else client.db - return database.litellm_objectpermissiontable + return _ObjectPermissionTable(self._prisma_client, use_writer=self._use_writer).table @property def model_class(self) -> type[LiteLLM_ObjectPermissionTable]: diff --git a/litellm/repositories/organization_repository.py b/litellm/repositories/organization_repository.py index 47eb8f4a609..1a773015414 100644 --- a/litellm/repositories/organization_repository.py +++ b/litellm/repositories/organization_repository.py @@ -3,28 +3,19 @@ Organization repository for database operations on LiteLLM_OrganizationTable. """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final from litellm.models.organization import LiteLLM_OrganizationTable from litellm.repositories.base_repository import BaseRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models -class _OrganizationDb(Protocol): - """The single Prisma table this repository reaches for on ``prisma_client.db``.""" - - @property - def litellm_organizationtable(self) -> TableActions["prisma_models.LiteLLM_OrganizationTable"]: ... - - -class _PrismaClientView(Protocol): - """The one attribute this repository reads off the untyped Prisma client wrapper.""" - - @property - def db(self) -> _OrganizationDb: ... +class _OrganizationTable(PrismaTableRepository["prisma_models.LiteLLM_OrganizationTable"]): + table_name = "litellm_organizationtable" class OrganizationRepository(BaseRepository[LiteLLM_OrganizationTable]): @@ -32,8 +23,7 @@ class OrganizationRepository(BaseRepository[LiteLLM_OrganizationTable]): @property def table(self) -> TableActions["prisma_models.LiteLLM_OrganizationTable"]: - client: Final[_PrismaClientView] = self.prisma_client - return client.db.litellm_organizationtable + return _OrganizationTable(self._prisma_client).table @property def model_class(self) -> type[LiteLLM_OrganizationTable]: diff --git a/litellm/repositories/project_repository.py b/litellm/repositories/project_repository.py index 0ea5009a5b6..ab19027a3d0 100644 --- a/litellm/repositories/project_repository.py +++ b/litellm/repositories/project_repository.py @@ -3,24 +3,19 @@ Project repository for database operations on LiteLLM_ProjectTable. """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final from litellm.models.project import LiteLLM_ProjectTable from litellm.repositories.base_repository import BaseRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models -class _ProjectDb(Protocol): - @property - def litellm_projecttable(self) -> TableActions["prisma_models.LiteLLM_ProjectTable"]: ... - - -class _PrismaClientView(Protocol): - @property - def db(self) -> _ProjectDb: ... +class _ProjectTable(PrismaTableRepository["prisma_models.LiteLLM_ProjectTable"]): + table_name = "litellm_projecttable" class ProjectRepository(BaseRepository[LiteLLM_ProjectTable]): @@ -28,8 +23,7 @@ class ProjectRepository(BaseRepository[LiteLLM_ProjectTable]): @property def table(self) -> TableActions["prisma_models.LiteLLM_ProjectTable"]: - client: Final[_PrismaClientView] = self.prisma_client - return client.db.litellm_projecttable + return _ProjectTable(self._prisma_client).table @property def model_class(self) -> type[LiteLLM_ProjectTable]: diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index d010455faa3..aa33da54572 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -10,10 +10,11 @@ methods; richer repositories live in their own modules. from collections.abc import Mapping, Sequence from datetime import timedelta from itertools import groupby -from typing import TYPE_CHECKING, Any, Final, Generic +from typing import TYPE_CHECKING, Any, Final, Generic, Protocol from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.repositories.accounting_delivery_repository import LedgerDatabase, LedgerTransaction from litellm.repositories.prisma_protocols import RowT_co, TableActions @@ -22,6 +23,16 @@ if TYPE_CHECKING: from prisma import models as prisma_models # noqa: F401 # used by quoted base-class subscripts +class _PrismaReaderClientView(Protocol): + @property + def db(self) -> object: ... + + +class _PrismaWriterClientView(Protocol): + @property + def writer_db(self) -> object: ... + + class PrismaTableRepository(Generic[RowT_co]): """Base for repositories that expose a single Prisma table.""" @@ -31,19 +42,26 @@ class PrismaTableRepository(Generic[RowT_co]): self._prisma_client = prisma_client self._use_writer = use_writer - @property - def prisma_client(self) -> Any: + def _connected_client(self) -> object: if self._prisma_client is None: raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") return self._prisma_client + @property + def prisma_client(self) -> Any: + return self._connected_client() + + def _selected_database(self) -> object: + client: Final = self._connected_client() + if self._use_writer: + return require_attributes(client, _PrismaWriterClientView, "writer_db").writer_db + return require_attributes(client, _PrismaReaderClientView, "db").db + @property def table(self) -> TableActions[RowT_co]: from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync - actions: Final[TableActions[RowT_co]] = getattr( - self.prisma_client.writer_db if self._use_writer else self.prisma_client.db, self.table_name - ) + actions: Final[TableActions[RowT_co]] = getattr(self._selected_database(), self.table_name) return wrap_table_actions_for_config_sync(actions=actions, table_name=self.table_name) diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index b89d87aa354..d076d148783 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.models.team import LiteLLM_TeamTable, Member from litellm.repositories.base_repository import ( BaseRepository, @@ -66,10 +67,12 @@ class _PrismaTeamDb(_TeamTables, Protocol): def tx(self) -> _TeamTransactionManager: ... -class _PrismaClientView(Protocol): +class _PrismaReaderClientView(Protocol): @property def db(self) -> _PrismaTeamDb: ... + +class _PrismaWriterClientView(Protocol): @property def writer_db(self) -> _PrismaTeamDb: ... @@ -95,8 +98,10 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): @property def _db(self) -> _PrismaTeamDb: - client: Final[_PrismaClientView] = self.prisma_client - return client.writer_db if self._use_writer else client.db + client: Final = self._connected_client() + if self._use_writer: + return require_attributes(client, _PrismaWriterClientView, "writer_db").writer_db + return require_attributes(client, _PrismaReaderClientView, "db").db @property def table(self) -> TableActions["prisma_models.LiteLLM_TeamTable"]: diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py index 28fba2de564..f39d28a3039 100644 --- a/litellm/repositories/user_repository.py +++ b/litellm/repositories/user_repository.py @@ -12,10 +12,12 @@ from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.models.user import LiteLLM_UserTable, SCIMPlaceholder from litellm.repositories.base_repository import BaseRepository, DbRecord, record_to_dict from litellm.repositories.chunked_in import IN_LIST_CHUNK_SIZE from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import Prisma @@ -62,15 +64,13 @@ class _UserTransaction(_UserDb, Protocol): class _PrismaClientView(Protocol): - @property - def db(self) -> _UserDb: ... - - @property - def writer_db(self) -> _UserDb: ... - def tx(self) -> AbstractAsyncContextManager[_UserTransaction]: ... +class _UserTable(PrismaTableRepository["prisma_models.LiteLLM_UserTable"]): + table_name = "litellm_usertable" + + async def _count_billable_users(table: TableActions["prisma_models.LiteLLM_UserTable"]) -> int: from prisma import Json @@ -88,9 +88,7 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): @property def table(self) -> TableActions["prisma_models.LiteLLM_UserTable"]: - client: Final[_PrismaClientView] = self.prisma_client - database: Final = client.writer_db if self._use_writer else client.db - return database.litellm_usertable + return _UserTable(self._prisma_client, use_writer=self._use_writer).table @property def model_class(self) -> type[LiteLLM_UserTable]: @@ -184,7 +182,7 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): where={"user_id": data["user_id"]}, data={"create": dict(data), "update": {}} ) - client: Final[_PrismaClientView] = self.prisma_client + client: Final = require_attributes(self._connected_client(), _PrismaClientView, "tx") async with client.tx() as transaction: await transaction.execute_raw("SET TRANSACTION ISOLATION LEVEL READ COMMITTED") await transaction.query_raw("SELECT 1 AS locked FROM pg_advisory_xact_lock(468530001)") diff --git a/litellm/repositories/verification_token_repository.py b/litellm/repositories/verification_token_repository.py index fb1f1c89e69..4b7ceca3212 100644 --- a/litellm/repositories/verification_token_repository.py +++ b/litellm/repositories/verification_token_repository.py @@ -9,6 +9,7 @@ from types import TracebackType from typing import TYPE_CHECKING, Final, Protocol from litellm.constants import UI_SESSION_TOKEN_TEAM_ID +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.models.verification_token import ( LiteLLM_VerificationToken, ) @@ -76,8 +77,7 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): @property def _db(self) -> _PrismaVerificationTokenDb: - client: Final[_PrismaClientView] = self.prisma_client - return client.db + return require_attributes(self._connected_client(), _PrismaClientView, "db").db @property def table(self) -> TableActions["PrismaVerificationToken"]: diff --git a/litellm/responses/additional_tools.py b/litellm/responses/additional_tools.py index 9c7a3a5f72e..879bc8ba41c 100644 --- a/litellm/responses/additional_tools.py +++ b/litellm/responses/additional_tools.py @@ -1,21 +1,20 @@ import json from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass -from itertools import groupby +from itertools import chain, groupby from types import MappingProxyType from typing import Final, cast # noqa: TID251 # validating the openai tool union strips vendor keys from raw tools -from pydantic import ValidationError +from pydantic import TypeAdapter, ValidationError +from typing_extensions import TypeIs # noqa: TID251 # narrows a client input item before its type is read from litellm._logging import verbose_logger from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS, ResponseInputParam ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools" - - -class _InputItemType(LiteLLMBaseModel): - type: str = "" +_NAMESPACE_TOOL_TYPE: Final = "namespace" +_ITEMS: Final = TypeAdapter(tuple[object, ...]) class _ToolList(LiteLLMBaseModel): @@ -29,11 +28,14 @@ class HoistedAdditionalTools: hoisted: tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...] +def _is_an_input_item( + item: object, +) -> TypeIs[Mapping[str, object]]: # guard-ok: auth scans every list input, and a pydantic error per item is slow + return isinstance(item, Mapping) + + def _is_additional_tools_item(item: object) -> bool: - try: - return _InputItemType.model_validate(item).type == ADDITIONAL_TOOLS_INPUT_ITEM_TYPE - except ValidationError: - return False + return _is_an_input_item(item) and item.get("type") == ADDITIONAL_TOOLS_INPUT_ITEM_TYPE def _tool_list_of(value: object) -> tuple[Mapping[str, object], ...]: @@ -61,7 +63,7 @@ def _members_of_copies(copies: Sequence[Mapping[str, object]]) -> Iterator[Mappi def _merged(copies: Sequence[Mapping[str, object]]) -> Mapping[str, object]: newest: Final = copies[-1] - if newest.get("type") != "namespace" or len(copies) == 1: + if newest.get("type") != _NAMESPACE_TOOL_TYPE or len(copies) == 1: return newest members: Final = tuple(_members_of_copies(copies)) return {**newest, "tools": list(_unique_tools(members))} if members else newest @@ -86,6 +88,11 @@ def _unique_tools(tools: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, o return tuple(_merged(copies_by_identity[identity]) for identity in first_seen_order) +def additional_tools_in_input(input: object) -> tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]: + items: Final = _ITEMS.validate_python(input) if isinstance(input, list) else () + return tuple(chain.from_iterable(_tools_of_item(item) for item in items if _is_additional_tools_item(item))) + + def hoist_additional_tools( input: str | ResponseInputParam, tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None, @@ -112,3 +119,26 @@ def hoist_additional_tools( for tool in _unique_tools((*existing, *hoisted)) ) return HoistedAdditionalTools(input=remaining_input, tools=unique, hoisted=hoisted) + + +def _top_level_tools(request: Mapping[str, object]) -> tuple[object, ...] | None: + try: + return _ITEMS.validate_python(request.get("tools") or ()) + except ValidationError: + return None + + +def with_additional_tools_hoisted(request: Mapping[str, object], *, drop_params: bool) -> Mapping[str, object]: + request_input: Final = request.get("input") + top_level_tools: Final = _top_level_tools(request) + if not isinstance(request_input, list) or top_level_tools is None: + return request + input_items: Final = cast(ResponseInputParam, request_input) # cast-ok: forwarded as sent, as inference does + hoisted: Final = hoist_additional_tools( + input_items, + cast("tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]", top_level_tools), # cast-ok: forwarded as sent + drop_params=drop_params, + ) + if not hoisted.hoisted: + return request + return {**request, "input": hoisted.input, "tools": list(hoisted.tools)} diff --git a/litellm/responses/litellm_completion_transformation/count_request.py b/litellm/responses/litellm_completion_transformation/count_request.py index 3deceac07aa..2141b55850b 100644 --- a/litellm/responses/litellm_completion_transformation/count_request.py +++ b/litellm/responses/litellm_completion_transformation/count_request.py @@ -12,7 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.file_id_utils import ( ) from litellm.litellm_core_utils.tool_fields import DropPolicy from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body -from litellm.responses.additional_tools import hoist_additional_tools +from litellm.responses.additional_tools import with_additional_tools_hoisted from litellm.responses.file_search.emulated_handler import replace_file_search_tools from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary from litellm.responses.litellm_completion_transformation.transformation import ( @@ -20,7 +20,7 @@ from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) from litellm.responses.main import previous_response_client_ids -from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS, ResponseInputParam +from litellm.types.llms.openai import ResponseInputParam from litellm.utils import ProviderConfigManager _COUNTED_CHAT_PARAMS: Final = ( @@ -114,20 +114,6 @@ def with_provider_file_ids( } -def _with_hoisted_additional_tools(responses_request: Mapping[str, object], drop_params: bool) -> Mapping[str, object]: - responses_input: Final = responses_request.get("input") - if not isinstance(responses_input, list): - return responses_request - input_items: Final = cast(ResponseInputParam, responses_input) # cast-ok: forwarded as sent, as inference does - top_level_tools: Final = cast( # cast-ok: tools are forwarded as sent, as inference does - "tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]", _ITEMS.validate_python(responses_request.get("tools") or ()) - ) - hoisted: Final = hoist_additional_tools(input_items, top_level_tools, drop_params=drop_params) - if not hoisted.hoisted: - return responses_request - return {**responses_request, "input": hoisted.input, "tools": list(hoisted.tools)} - - def _caller_input(responses_request: Mapping[str, object]) -> ResponseInputParam | None: responses_input: Final = responses_request.get("input") return ( @@ -205,7 +191,9 @@ async def chat_count_request( chat_body: Final = await _with_previous_turns( responses_batch_body_to_chat_body( _with_compaction_as_summary_text( - _with_emulated_file_search(_with_hoisted_additional_tools(responses_request, policy.drop_params)) + _with_emulated_file_search( + with_additional_tools_hoisted(responses_request, drop_params=policy.drop_params) + ) ), custom_llm_provider, policy.drop_params, diff --git a/litellm/responses/litellm_completion_transformation/reasoning_items.py b/litellm/responses/litellm_completion_transformation/reasoning_items.py index ab982222902..dda313bfc1b 100644 --- a/litellm/responses/litellm_completion_transformation/reasoning_items.py +++ b/litellm/responses/litellm_completion_transformation/reasoning_items.py @@ -6,6 +6,7 @@ from typing import Final from pydantic import BaseModel, TypeAdapter, ValidationError REASONING_ITEM_ID_PREFIX: Final = "rs_" +WRAPPED_ENCRYPTED_CONTENT_PREFIX: Final = "litellm_enc:" _JSON_LIST: Final = TypeAdapter(list[object]) _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) @@ -58,16 +59,35 @@ def is_minted_reasoning_item_id(item_id: object) -> bool: return parsed.version == 4 and str(parsed) == suffix -def is_litellm_minted_reasoning_item(item: object) -> bool: +def _reasoning_fields(item: object) -> Mapping[str, object] | None: try: fields: Final = _JSON_OBJECT.validate_python( item.model_dump(exclude_none=True) if isinstance(item, BaseModel) else item ) except ValidationError: - return False - if fields.get("type") != "reasoning": - return False + return None + return fields if fields.get("type") == "reasoning" else None + + +def _is_minted(fields: Mapping[str, object]) -> bool: return ( is_minted_reasoning_item_id(fields.get("id")) or decode_thinking_blocks(fields.get("encrypted_content")) is not None ) + + +def is_litellm_minted_reasoning_item(item: object) -> bool: + fields: Final = _reasoning_fields(item) + return fields is not None and _is_minted(fields) + + +def is_other_provider_reasoning_item(item: object) -> bool: + fields: Final = _reasoning_fields(item) + if fields is None: + return False + encrypted_content: Final = fields.get("encrypted_content") + return ( + _is_minted(fields) + or bool(fields.get("content")) + or (isinstance(encrypted_content, str) and encrypted_content.startswith(WRAPPED_ENCRYPTED_CONTENT_PREFIX)) + ) diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 7e6ee84e2b2..9da89d96e73 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -4,13 +4,13 @@ from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, require_attributes from litellm.responses.litellm_completion_transformation.client_tool_search import ( bridge_client_tool_search, client_tool_search_loaded_tools, @@ -46,12 +46,24 @@ COLD_STORAGE_HANDLER: Final = ColdStorageHandler() ######################################################## +class _CarriesArguments(Protocol): + arguments: object + + +def _normalize_redacted_function_arguments(tool_call: object) -> None: + function: Final = attribute_or_none(tool_call, "function") + if function is None: + return + carrier: Final = require_attributes(function, _CarriesArguments, "arguments") + if carrier.arguments == REDACTED_BY_LITELLM: + carrier.arguments = REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER + + def _normalize_redacted_tool_call_arguments(message: Message) -> None: """Redaction stores the bare sentinel (invalid JSON) in tool-call arguments; normalize replayed history to "{}" so provider converters can parse it.""" for tool_call in message.tool_calls or []: - if (function := getattr(tool_call, "function", None)) is not None and function.arguments == REDACTED_BY_LITELLM: - function.arguments = REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER + _normalize_redacted_function_arguments(tool_call) function_call: Final = message.function_call if function_call is not None and function_call.arguments == REDACTED_BY_LITELLM: function_call.arguments = REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 0f8ba6ef98a..8b6c53a272d 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -3,10 +3,10 @@ import uuid from collections.abc import Mapping, Sequence from itertools import filterfalse from types import MappingProxyType -from typing import Any, Final, cast +from typing import Any, Final, Protocol, cast import litellm -from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_items +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_items, require_attributes from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params, served_from_cache from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend from litellm.llms.anthropic.common_utils import anthropic_server_call_fields, anthropic_server_tool_calls @@ -84,7 +84,11 @@ def _index_of_output_item_type(items: Sequence[object], item_type: str) -> int | ) -def _output_items_with_id(items: tuple[Any, ...], item_type: str, item_id: str | None) -> tuple[object, ...]: +class _CopiesWithUpdate(Protocol): + def model_copy(self, *, update: Mapping[str, object]) -> object: ... + + +def _output_items_with_id(items: tuple[object, ...], item_type: str, item_id: str | None) -> tuple[object, ...]: if item_id is None: return items @@ -93,7 +97,10 @@ def _output_items_with_id(items: tuple[Any, ...], item_type: str, item_id: str | return items return tuple( - item.model_copy(update={"id": item_id}) if index == target_index else item for index, item in enumerate(items) + require_attributes(item, _CopiesWithUpdate, "model_copy").model_copy(update={"id": item_id}) + if index == target_index + else item + for index, item in enumerate(items) ) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 79e060cdd6b..d8651754d8d 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -12,7 +12,6 @@ from itertools import accumulate, chain, groupby from types import MappingProxyType from typing import ( TYPE_CHECKING, - Any, Final, Literal, Protocol, @@ -22,12 +21,6 @@ from typing import ( runtime_checkable, ) -from openai.types.chat.chat_completion_named_tool_choice_param import ( - ChatCompletionNamedToolChoiceParam, -) -from openai.types.chat.chat_completion_named_tool_choice_param import ( - Function as NamedToolChoiceFunction, -) from openai.types.responses import ResponseFunctionToolCall, ResponseFunctionWebSearch, ResponseOutputRefusal from openai.types.responses.response_compaction_item import ResponseCompactionItem from openai.types.responses.response_create_params import ResponseInputParam @@ -56,7 +49,12 @@ from litellm.integrations.anthropic_cache_control_hook import ( whole_messages_within_cache_control_limit, with_cache_control, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params, normalize_drop_params +from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, + drops_unsupported_params, + normalize_drop_params, + require_attributes, +) from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, ) @@ -66,6 +64,7 @@ from litellm.litellm_core_utils.hidden_params import ( set_hidden_params, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import is_str_keyed_dict from litellm.litellm_core_utils.model_param_helper import embeds_thought_signature_in_tool_call_id from litellm.litellm_core_utils.prompt_templates.common_utils import ( invalid_type_message, @@ -89,6 +88,7 @@ from litellm.responses.litellm_completion_transformation.session_handler import ResponsesSessionHandler, StoredTurns, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.anthropic import AnthropicToolSearchToolRegex, claude_hosted_web_tool_call from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.openai import ( @@ -198,6 +198,8 @@ ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool | AnthropicToolSearchToolRegex NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n" NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: Final = frozenset({"function", "custom"}) +RESPONSES_WEB_SEARCH_TOOL_TYPES: Final = frozenset({"web_search", "web_search_preview"}) +RESPONSES_TOOL_TYPES_WITHOUT_CHAT_FORM: Final = frozenset({"computer_use", "image_generation", "local_shell", "shell"}) _INCOMPLETE_REASON_BY_FINISH_REASON: Final[Mapping[str, Literal["max_output_tokens", "content_filter"]]] = ( MappingProxyType({"length": "max_output_tokens", "content_filter": "content_filter", "refusal": "content_filter"}) ) @@ -383,6 +385,10 @@ class _HasId(Protocol): id: object +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + class _ResponsesToolCallItem(Protocol): name: object arguments: object @@ -981,8 +987,8 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_tool_choice( - tool_choice: Any, - ) -> str | dict[str, object] | None: + tool_choice: object, + ) -> object: """ Transform tool_choice from various formats to OpenAI Chat Completion format. @@ -1006,11 +1012,13 @@ class LiteLLMCompletionResponsesConfig: if isinstance(tool_choice, str): return tool_choice - if isinstance(tool_choice, dict): + if is_str_keyed_dict(tool_choice): tool_choice_type: Final = tool_choice.get("type") # If it has a function with name, it's standard OpenAI format - pass through - if tool_choice.get("function") and tool_choice.get("function", {}).get("name"): + if tool_choice.get("function") and require_attributes( + tool_choice.get("function", {}), _SupportsGet, "get" + ).get("name"): return tool_choice # Handle Cursor IDE dict formats without function name @@ -1025,17 +1033,13 @@ class LiteLLMCompletionResponsesConfig: elif tool_choice_type == "function": function_name: Final = tool_choice.get("name") if function_name: - return ChatCompletionNamedToolChoiceParam( - type="function", function=NamedToolChoiceFunction(name=function_name) - ) + return {"type": "function", "function": {"name": function_name}} return "required" elif tool_choice_type == "custom": custom: Final = tool_choice.get("custom") - custom_name = tool_choice.get("name") or (custom.get("name") if isinstance(custom, dict) else None) + custom_name = tool_choice.get("name") or (custom.get("name") if is_str_keyed_dict(custom) else None) if custom_name: - return ChatCompletionNamedToolChoiceParam( - type="function", function=NamedToolChoiceFunction(name=custom_name) - ) + return {"type": "function", "function": {"name": custom_name}} return "required" # Return as-is for unknown formats @@ -1059,7 +1063,7 @@ class LiteLLMCompletionResponsesConfig: return ToolChoiceCustomParam(type="custom", name=custom_name) case _, {"type": "function", "function": {"name": str(function_name)}}: return ToolChoiceFunctionParam(type="function", name=function_name) - case _, "none" | "auto" | "required" as normalized: + case _, normalized if normalized in ("none", "auto", "required"): return normalized case _, _: return "auto" @@ -2556,7 +2560,12 @@ class LiteLLMCompletionResponsesConfig: Returns None for anything this deployment did not write, so a genuinely opaque blob is still skipped rather than forwarded as garbage. """ - decoded: Final = decode_thinking_blocks(input_item.get("encrypted_content")) + encrypted_content: Final = input_item.get("encrypted_content") + decoded: Final = decode_thinking_blocks( + ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id(encrypted_content)[1] + if isinstance(encrypted_content, str) + else encrypted_content + ) if decoded is None: return None return tuple( @@ -2942,7 +2951,7 @@ class LiteLLMCompletionResponsesConfig: tool_type: Final = tool.get("type") if tool_type == "mcp": return ResponsesToolChatForm(chat_tools=(cast(OpenAIMcpServerTool, tool),), web_search_options=None) - if tool_type == "web_search_preview" or tool_type == "web_search": + if tool_type in RESPONSES_WEB_SEARCH_TOOL_TYPES: _search_context_size: Final[Literal["low", "medium", "high"]] = cast( Literal["low", "medium", "high"], tool.get("search_context_size") ) @@ -2995,7 +3004,7 @@ class LiteLLMCompletionResponsesConfig: return ResponsesToolChatForm(chat_tools=(client_tool_search_chat_tool(tool),), web_search_options=None) if tool_type == "tool_search": return ResponsesToolChatForm(chat_tools=(anthropic_tool_search_chat_tool(),), web_search_options=None) - if tool_type in ("computer_use", "image_generation", "local_shell", "shell"): + if tool_type in RESPONSES_TOOL_TYPES_WITHOUT_CHAT_FORM: verbose_logger.warning( "Dropping Responses API tool of type '%s': it has no Chat Completions " "equivalent and the target provider would reject the request.", diff --git a/litellm/responses/main.py b/litellm/responses/main.py index ae7c3fccffa..e085c4d7a2e 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -31,7 +31,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 drops_unsupported_params, normalize_drop_params +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params, require_attributes from litellm.litellm_core_utils.get_llm_provider_logic import ( declared_authenticating_provider, get_provider_affinity_header, @@ -1147,8 +1147,12 @@ def _apply_prompt_management_to_responses_call( ) -> tuple[str | ResponseInputParam, str, str | None, object]: """Returns the prompt-managed input, model and provider, plus the deployment metadata that still describes the upstream (``None`` once the prompt manager moved the request to another provider).""" - async_outcome: Final[_AsyncPromptManagementOutcome | None] = kwargs.pop("_async_prompt_merged_params", None) - if async_outcome is not None: + carried_outcome: Final[object] = kwargs.get("_async_prompt_merged_params") + kwargs.pop("_async_prompt_merged_params", None) + if carried_outcome is not None: + async_outcome: Final = require_attributes( + carried_outcome, _AsyncPromptManagementOutcome, "merged_optional_params" + ) for key, value in async_outcome.merged_optional_params.items(): local_vars[key] = value return input, model, custom_llm_provider, async_outcome.deployment_model_info diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 1104aee353e..d045689cb56 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -455,7 +455,7 @@ class LiteLLM_Proxy_MCP_Handler: mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], litellm_trace_id: str | None = None, request_tags: list[str] | None = None, - ) -> tuple[list[Any], dict[str, str]]: + ) -> tuple[list[FunctionToolParam], dict[str, str]]: """ Centralized method to process MCP tools through the complete pipeline. @@ -964,6 +964,8 @@ class LiteLLM_Proxy_MCP_Handler: error=e, ) verbose_logger.error("BlockedPiiEntityError in MCP tool call: %s", e) + if LiteLLM_Proxy_MCP_Handler.should_raise_mcp_guardrail_block(e): + raise error_message = f"Tool call blocked: PII entity '{getattr(e, 'entity_type', 'unknown')}' detected by guardrail '{getattr(e, 'guardrail_name', 'unknown')}'. {e}" tool_results.append( { @@ -980,6 +982,8 @@ class LiteLLM_Proxy_MCP_Handler: error=e, ) verbose_logger.error("GuardrailRaisedException in MCP tool call: %s", e) + if LiteLLM_Proxy_MCP_Handler.should_raise_mcp_guardrail_block(e): + raise error_message = ( f"Tool call blocked: Guardrail '{getattr(e, 'guardrail_name', 'unknown')}' violation. {e}" ) @@ -998,6 +1002,8 @@ class LiteLLM_Proxy_MCP_Handler: error=e, ) verbose_logger.error("HTTPException in MCP tool call: %s", e) + if LiteLLM_Proxy_MCP_Handler.should_raise_mcp_guardrail_block(e): + raise error_message = f"Tool call failed: {str(e.detail) if hasattr(e, 'detail') else str(e)}" tool_results.append( { @@ -1026,6 +1032,10 @@ class LiteLLM_Proxy_MCP_Handler: _execute_tool_calls = execute_tool_calls + @staticmethod + def should_raise_mcp_guardrail_block(error: Exception) -> bool: + return getattr(error, "mcp_block_behavior", None) == "raise" + @staticmethod def create_follow_up_messages_for_chat( original_messages: list[object], diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 189366feab4..ea6064b554d 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -1,6 +1,9 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast +from typing_extensions import ReadOnly +from typing_extensions import TypedDict as ExtraItemsTypedDict + from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.agentic_loop_settings import validated_max_agentic_loops @@ -39,6 +42,11 @@ else: MAX_MCP_TOOL_CALL_ROUNDS: Final = 5 +class _FollowUpRequest(ExtraItemsTypedDict, total=False, extra_items=object): + input: ReadOnly[str | ResponseInputParam] + model: ReadOnly[str] + + def _output_items(response: ResponsesAPIResponse) -> Sequence[object]: """Read a response's output items as plain objects; the field is a wide union of item models.""" return tuple(cast("Sequence[object]", response.output)) # cast-ok: items are only carried, never inspected @@ -550,6 +558,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Phase 4: Tool Execution Events if self.phase == "tool_execution": + if self._stream_error is not None: + self.phase = "finished" + return await self._anext_impl() # Emit any queued tool execution events if self.tool_execution_events: return self.tool_execution_events.pop(0) @@ -713,6 +724,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.phase = "tool_execution" await self._generate_tool_execution_events() + if self._stream_error is not None: + return None if not isinstance(response_obj, ResponsesAPIResponse): return chunk if self._follow_up_pending(): @@ -974,6 +987,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): except Exception as e: verbose_logger.error("Error in tool execution: %s", e) + if LiteLLM_Proxy_MCP_Handler.should_raise_mcp_guardrail_block(e): + self._stream_error = e + return import traceback traceback.print_exc() @@ -1035,7 +1051,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.max_tool_call_rounds, ) - follow_up_response: Final = await aresponses(**follow_up_params) + follow_up_response: Final = await aresponses(**_FollowUpRequest(**follow_up_params)) # Route the follow-up through the same base_iterator machinery as # the initial call so its completion is checked for further tool diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 288e4a7119d..0c002a770e8 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -250,6 +250,8 @@ _ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType( "server_error": 500, "rate_limit_exceeded": 429, "insufficient_quota": 429, + "resource_unavailable": 429, + "flex_unavailable": 429, "vector_store_timeout": 504, "invalid_prompt": 400, "invalid_image": 400, diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 0e7f40e11af..e21c7f883c7 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -13,7 +13,7 @@ from typing_extensions import TypeIs # noqa: TID251 # narrows untyped wire pay import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, iter_items +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, decoded_json_object, iter_items from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value, is_nested_path from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils from litellm.llms.base_llm.responses.transformation import ( @@ -50,6 +50,7 @@ from litellm.types.utils import ( _PROVIDER_SPECIFIC_HEADER_ADAPTER: Final[ TypeAdapter[ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None] ] = TypeAdapter(ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None) +_REASONING_FIELDS_ONLY_ITS_ORIGIN_RESOLVES: Final = frozenset({"encrypted_content", "id", "status"}) class _UsageExtraFields(ExtraItemsTypedDict, total=False, extra_items=object): @@ -90,8 +91,9 @@ def _is_object_dict( def _prompt_tokens_details_wrapper(details: object) -> PromptTokensDetailsWrapper: if _is_object_dict(details): input_details: Final = InputTokensDetails.model_validate(details) + input_fields: Final = decoded_json_object(input_details.model_dump(exclude={"cache_write_tokens_details"})) return PromptTokensDetailsWrapper( - **input_details.model_dump(exclude={"cache_write_tokens_details"}), + **input_fields, cache_creation_token_details=_cache_write_lifetimes(input_details.cache_write_tokens_details), ) return PromptTokensDetailsWrapper( @@ -678,7 +680,7 @@ class ResponsesAPIRequestUtils: return items: Final = cast(list[object], request_input) # cast-ok: untyped client json stripped: Final = tuple( - ResponsesAPIRequestUtils._without_encrypted_reasoning(reasoning) + ResponsesAPIRequestUtils.readable_reasoning(reasoning) if (reasoning := ResponsesAPIRequestUtils._encrypted_reasoning_item(item)) is not None and (should_strip is None or should_strip(reasoning)) else item @@ -696,14 +698,14 @@ class ResponsesAPIRequestUtils: return reasoning @staticmethod - def _without_encrypted_reasoning(reasoning: Mapping[str, object]) -> object | None: + def readable_reasoning(reasoning: Mapping[str, object]) -> Mapping[str, object] | None: readable: Final = any( ResponsesAPIRequestUtils._has_readable_text(reasoning.get(key)) for key in ("summary", "content") ) if not readable: return None kept: Final[dict[str, object]] = { # mutable-ok: request item rebuilt without the undecryptable keys - key: value for key, value in reasoning.items() if key not in ("encrypted_content", "id") + key: value for key, value in reasoning.items() if key not in _REASONING_FIELDS_ONLY_ITS_ORIGIN_RESOLVES } return kept diff --git a/litellm/router.py b/litellm/router.py index 8f53ae3056c..4bece1892db 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -15,6 +15,8 @@ from collections.abc import Callable, Mapping, Sequence from types import FunctionType, MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal +from pydantic import TypeAdapter + from litellm.router_config import RouterConfig, build_router_config from litellm.router_core import RouterCore, make_core from litellm.router_deprecated import DeprecatedRouterMembers @@ -202,6 +204,7 @@ def _router_init( _CORE_SIGNATURE: Final = inspect.signature(_router_init) +_PASSED_VALUE: Final = TypeAdapter(object) def router_arguments( @@ -213,7 +216,13 @@ def router_arguments( bound: Final = _CORE_SIGNATURE.bind_partial(None, *args, **kwargs).arguments except TypeError as error: raise TypeError(f"Router.__init__() {error}") from None - return MappingProxyType({name: value for name, value in bound.items() if name != "self"}) + return MappingProxyType( + { + name: value + for name, value in zip(bound, map(_PASSED_VALUE.validate_python, bound.values())) + if name != "self" + } + ) def _class_descriptor(cls: type, name: str) -> bool: diff --git a/litellm/router_api.py b/litellm/router_api.py index fa6ad69fab4..ca370488f5d 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -65,6 +65,7 @@ if TYPE_CHECKING: from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.quality_router.quality_router import QualityRouter + from litellm.router_strategy.switchyard_router.switchyard_router import SwitchyardRouter from litellm.router_utils.cooldown_cache import CooldownCache from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets as _AttemptedFallbackTargets from litellm.router_utils.health_state_cache import DeploymentHealthCache @@ -106,6 +107,7 @@ if TYPE_CHECKING: | dict[str, list[TaggedPreRoutingStrategy[ComplexityRouter]]] | dict[str, list[TaggedPreRoutingStrategy[QualityRouter]]] | dict[str, list[TaggedPreRoutingStrategy[AdaptiveRouter]]] + | dict[str, list[TaggedPreRoutingStrategy[SwitchyardRouter]]] | list[RoutingPlugin] | RouterBudgetLimiting | DualCache @@ -158,6 +160,7 @@ GroupField: TypeAlias = Literal[ "display_name", "mode", "service_tiers", + "pinned_reasoning_efforts", "max_input_tokens", "deployment_model_info", "token_limits", @@ -170,6 +173,7 @@ PluginKind: TypeAlias = Literal[ "complexity", "quality", "adaptive", + "switchyard", "routing", "budget_limiter", "prompt_caching_pins", @@ -939,7 +943,7 @@ class RouterRegistry: self, name: str, *, - field: Literal["service_tiers"], + field: Literal["service_tiers", "pinned_reasoning_efforts"], include_hidden: bool = False, team_id: str | None = None, ) -> tuple[object, ...]: ... @@ -1001,6 +1005,7 @@ class RouterRegistry: (were ``get_configured_display_name``, ``get_configured_mode``, ``get_configured_token_limits``, ``get_model_listing_info``) - ``service_tiers``: (was ``get_configured_service_tiers``) + - ``pinned_reasoning_efforts``: the ``reasoning_effort`` each routable deployment pins - ``max_input_tokens``: each routable deployment's own ``model_info.max_input_tokens`` - ``deployment_model_info``: each routable deployment's model info, merged with the pricing map """ @@ -1015,6 +1020,8 @@ class RouterRegistry: return r._core_get_configured_mode(name) if field == "service_tiers": return r._core_get_configured_service_tiers(name, team_id=team_id) + if field == "pinned_reasoning_efforts": + return r._core_get_pinned_reasoning_efforts(name, team_id=team_id) if field == "max_input_tokens": return r._core_get_configured_max_input_tokens(name, team_id=team_id) if field == "deployment_model_info": @@ -1102,6 +1109,11 @@ class RouterRegistry: self, kind: Literal["adaptive"] ) -> "dict[str, list[TaggedPreRoutingStrategy[AdaptiveRouter]]]": ... # mutable-ok: the live dict + @overload + def plugins( + self, kind: Literal["switchyard"] + ) -> "dict[str, list[TaggedPreRoutingStrategy[SwitchyardRouter]]]": ... # mutable-ok: the live dict + @overload def plugins(self, kind: Literal["routing"]) -> "list[RoutingPlugin]": ... # mutable-ok: the live list @@ -1128,6 +1140,8 @@ class RouterRegistry: return r._core_quality_routers if kind == "adaptive": return r._core_adaptive_routers + if kind == "switchyard": + return r._core_switchyard_routers if kind == "routing": return r._core_routing_plugins if kind == "budget_limiter": @@ -1136,6 +1150,18 @@ class RouterRegistry: return r._core_prompt_caching_pin_cache raise ValueError(f"unknown plugin kind {kind!r}") + def decryptable_reasoning( + self, model_group: str, request_input: Sequence[object], targets: "Sequence[Deployment]" + ) -> tuple[object, ...]: + """ + ``request_input`` without the encrypted reasoning that ``targets``, the deployments a + ``model_group`` request may go to, cannot decrypt, keeping each item's readable summary. + When encrypted content affinity is on for the group, the router sends the request to the + reasoning's origin, so reasoning one target can decrypt stays; otherwise every target must + decrypt it (was reading ``optional_callbacks`` for ``EncryptedContentAffinityCheck``). + """ + return self._router._core_decryptable_reasoning(model_group, request_input, targets) + # Settings a caller may read. Assigning one from outside the router goes through apply(ChangeSet(assign=...)). SETTING_FIELDS: Final[tuple[str, ...]] = ( diff --git a/litellm/router_config.py b/litellm/router_config.py index 36154af296c..874fb314320 100644 --- a/litellm/router_config.py +++ b/litellm/router_config.py @@ -13,6 +13,8 @@ from dataclasses import dataclass, field from types import MappingProxyType from typing import TYPE_CHECKING, Final +from pydantic import TypeAdapter + import litellm if TYPE_CHECKING: @@ -45,6 +47,7 @@ class RouterDynamicConfig: context_window_fallbacks: Sequence[object] | None content_policy_fallbacks: Sequence[object] | None drop_params: bool | None + disable_token_counter: bool provider_list: Sequence[str] expose_router_debug_in_errors: bool redact_user_api_key_info: bool | None @@ -102,6 +105,7 @@ def current_dynamic_config() -> RouterDynamicConfig: context_window_fallbacks=litellm.context_window_fallbacks, content_policy_fallbacks=litellm.content_policy_fallbacks, drop_params=litellm.drop_params, + disable_token_counter=litellm.disable_token_counter, provider_list=litellm.provider_list, expose_router_debug_in_errors=litellm.expose_router_debug_in_errors, redact_user_api_key_info=litellm.redact_user_api_key_info, @@ -118,6 +122,7 @@ def current_dynamic_config() -> RouterDynamicConfig: ) +_LIVE_VALUE: Final = TypeAdapter(object) _LIVE_SOURCES: Final = MappingProxyType( { "num_retries": "num_retries", @@ -130,6 +135,7 @@ _LIVE_SOURCES: Final = MappingProxyType( "context_window_fallbacks": "context_window_fallbacks", "content_policy_fallbacks": "content_policy_fallbacks", "drop_params": "drop_params", + "disable_token_counter": "disable_token_counter", "provider_list": "provider_list", "expose_router_debug_in_errors": "expose_router_debug_in_errors", "redact_user_api_key_info": "redact_user_api_key_info", @@ -156,7 +162,7 @@ class LiveDynamicConfig: source: Final = _LIVE_SOURCES.get(name) if source is None: raise AttributeError(name) - return getattr(litellm, source) + return _LIVE_VALUE.validate_python(getattr(litellm, source)) live_dynamic_config: Final = LiveDynamicConfig() diff --git a/litellm/router_deprecated.py b/litellm/router_deprecated.py index 82a53f001c5..0baaec70b1c 100644 --- a/litellm/router_deprecated.py +++ b/litellm/router_deprecated.py @@ -32,6 +32,8 @@ from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, S from types import MappingProxyType from typing import TYPE_CHECKING, Final, Generic, ParamSpec, TypeVar, overload +from pydantic import TypeAdapter + from litellm.router_api import ASSIGNABLE_ALIASES, SETTING_FIELDS, ChangeSet, RoutingContext from litellm.router_strategy.least_busy import picks_reserve @@ -67,6 +69,8 @@ DEPRECATED_ROUTER_MEMBER_USES: Final[Counter[str]] = Counter() # mutable-ok: th _INTERNAL_PREFIX: Final = "litellm.router" +_MEMBER_VALUE: Final = TypeAdapter(object) + # Router's own attribute hooks: a write or delete through an old name passes through them first. _ATTRIBUTE_HOOKS: Final = frozenset({"__getattribute__", "__setattr__", "__delattr__"}) @@ -935,7 +939,7 @@ def _name_members_after_router() -> None: # The old members were Router methods, so errors that name the method (a wrong-arity TypeError, # or "litellm.router.Router.() argument after ** must be a mapping") keep naming it the # way they did before the members moved here. - for member in vars(DeprecatedRouterMembers).values(): + for member in map(_MEMBER_VALUE.validate_python, vars(DeprecatedRouterMembers).values()): function = member.__func__ if isinstance(member, staticmethod) else member if inspect.isfunction(function): function.__qualname__ = f"Router.{function.__name__}" diff --git a/litellm/router_python_host.py b/litellm/router_python_host.py index 84e20270322..425d38d6233 100644 --- a/litellm/router_python_host.py +++ b/litellm/router_python_host.py @@ -19,6 +19,7 @@ from typing import ( cast, # noqa: TID251 # untyped litellm globals: the live config view and the callback lists ) +from pydantic import TypeAdapter from typing_extensions import overload import litellm @@ -53,6 +54,8 @@ CallbackList: TypeAlias = Literal[ Endpoint: TypeAlias = Callable[..., object] +_LITELLM_GLOBAL: Final = TypeAdapter(object) + class PythonHost: __slots__ = () @@ -126,7 +129,7 @@ class PythonHost: return litellm.logging_callback_manager def callback_list(self, name: CallbackList) -> list[object]: # mutable-ok: the process callback list itself - value: Final[object] = getattr(litellm, name) # pyright: ignore[reportAny] # an untyped litellm global + value: Final = _LITELLM_GLOBAL.validate_python(getattr(litellm, name)) return cast(list[object], value) # cast-ok: the process callback list is an untyped litellm global def set_callback_list(self, name: CallbackList, value: list[object]) -> None: # mutable-ok: assigned as given diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 40ae7e0e554..3e222313f32 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -1177,6 +1177,7 @@ _HOSTED_STATE: Final = frozenset( "_core_auto_routers", "_core_complexity_routers", "_core_quality_routers", + "_core_switchyard_routers", "_core_adaptive_routers", "_core_routing_plugins", "_core_deployment_latency_map", @@ -1722,6 +1723,7 @@ class RustCore( run_plugins=partial(self._run_plugins, request_kwargs, messages) if self._hosted._core_routing_plugins # pyright: ignore[reportPrivateUsage] # the live plugin list else None, + call_type=None, ) def _core_log_retry( @@ -2715,7 +2717,7 @@ class RustCore( web_search: Final = not plan.sync and self._picks_async() and is_web_search_request(request) # The generic handler `_route_capability_pin` names: a call only some providers serve # drops the others, and a Decisions call carrying media skips the deployments refusing it. - handler: Final = plan.host_op.rsplit(".", 1)[-1] + handler: Final = plan.handler decisions_media: Final = handler in DECISIONS_MEDIA_CALL_TYPES and bool( request.get("audio") or request.get("videos") ) @@ -3113,6 +3115,7 @@ class RustCore( specific_deployment: object = _NO_INPUT, *, pass_through: bool = False, + call_type: str | None = None, ) -> object: """The host's pre-routing of one pick (router_rust.steps.prerouting): the model group to resolve instead of the step's ``model``, a custom strategy's deployment, or ``None``. ``pass_through`` @@ -3166,6 +3169,7 @@ class RustCore( fallback_depth=depth if isinstance(depth, int) else 0, router_strategy=normalize_strategy(r.routing_strategy), pass_through=pass_through, + call_type=call_type, enable_tag_filtering=r.enable_tag_filtering, match_any=r.tag_filtering_match_any, public_router=public if public is not None else self, @@ -5266,7 +5270,9 @@ class _Driver: async def pre_route(self, step: Mapping[str, object]) -> object: request: Final = self._request try: - return await self._core._pre_route(step, request, *self._plan.pick_input(request)) # pyright: ignore[reportPrivateUsage] # see above + return await self._core._pre_route( # pyright: ignore[reportPrivateUsage] # see above + step, request, *self._plan.pick_input(request), call_type=self._plan.handler + ) except Exception as error: # noqa: BLE001 # see select return self._pick_raised(error, logged=True) diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index e2c621ccb7d..7a5b52761e2 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -244,6 +244,11 @@ class Plan: # the request (`{**data, name: value, **kwargs}`). named: tuple[str, ...] = () + @property + def handler(self) -> str: + """The handler name the host op ends in, which route capabilities and pre-routing strategies key on.""" + return self.host_op.rsplit(".", 1)[-1] + @property def on_its_own_loop(self) -> bool: """``function_with_fallbacks`` runs these sync endpoints' whole walk on a private event loop.""" diff --git a/litellm/router_rust/invoke/walk.py b/litellm/router_rust/invoke/walk.py index 17e20119f43..ff8bc7eef5a 100644 --- a/litellm/router_rust/invoke/walk.py +++ b/litellm/router_rust/invoke/walk.py @@ -39,6 +39,7 @@ _WALK_KEYS: Final = ( _SERVED_KEYS: Final = ("model_info", "deployment") # The walk state `run_async_fallback` hands every hop as kwargs. _HOP_KWARGS: Final = ("fallback_depth", "max_fallbacks", "attempted_targets", "include_fallback_errors") +_RETRY_RECORD_KEYS: Final = ("previous_models", "request_retry_count") def _reentrant_hop_stamps(request: Mapping[str, object], field: str) -> Mapping[str, object]: @@ -179,11 +180,13 @@ class WalkMetadata: return sibling: Final = request.get("metadata" if self._field == "litellm_metadata" else "litellm_metadata") if isinstance(sibling, dict): - sibling.pop("attempted_fallbacks", None) # pyright: ignore[reportUnknownMemberType] # the caller's metadata dict - sibling.pop("original_model_group", None) # pyright: ignore[reportUnknownMemberType] # see above + for key in ("attempted_fallbacks", "original_model_group", *_RETRY_RECORD_KEYS): + sibling.pop(key, None) # pyright: ignore[reportUnknownMemberType] # the caller's metadata dict bucket: Final = request.get(self._field) if not isinstance(bucket, dict): return + for key in _RETRY_RECORD_KEYS: + bucket.pop(key, None) # pyright: ignore[reportUnknownMemberType] # see above bucket["attempted_fallbacks"] = 0 model: Final = request.get("model") if model is not None: diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py index 75afe38548e..cd2d69d13a8 100644 --- a/litellm/router_rust/steps/callbacks.py +++ b/litellm/router_rust/steps/callbacks.py @@ -44,7 +44,10 @@ from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( ) from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck -from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck +from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( + PromptCachingDeploymentCheck, + new_responses_pins_in_flight, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging @@ -130,6 +133,7 @@ class OptionalChecks: redis_cache=cache.redis_cache, in_memory_cache=InMemoryCache(max_size_in_memory=PROMPT_CACHE_PIN_CACHE_MAX_SIZE), ) + self.prompt_caching_pins_in_flight: Final = new_responses_pins_in_flight() self._request_scoped_selectors: tuple[CustomLogger, ...] = () self._request_scoped_policies: tuple[CustomLogger, ...] = () # ``_slack_alerting_loggers``: the Slack alerting logger ``alerting_config`` builds, if any. @@ -167,6 +171,9 @@ class OptionalChecks: group: Final = hosted.get_routing_group(resolved) return group is not None and group.routing_strategy == "priority" + def _group_has_alternatives(self, model: str) -> bool: + return len(self._hosted._core_deployments_named(model) or ()) > 1 # pyright: ignore[reportPrivateUsage] # the host registry's name index + async def repin_session_after_success( self, litellm_params: Mapping[str, object], response: object, cooldowns: "CooldownCache" ) -> None: @@ -291,7 +298,10 @@ class OptionalChecks: callback: CustomLogger | None = None if name == "prompt_caching": callback = PromptCachingDeploymentCheck( - cache=self.prompt_caching_pin_cache, is_priority_group=self._is_priority_group + cache=self.prompt_caching_pin_cache, + is_priority_group=self._is_priority_group, + responses_pins_in_flight=self.prompt_caching_pins_in_flight, + group_has_alternatives=self._group_has_alternatives, ) elif name == "router_budget_limiting": if self.budget_limiter() is not None: @@ -407,6 +417,28 @@ class OptionalChecks: bound: Final = request.get("litellm_logging_obj") in self._bound_requests return {"success": bound and not sync, "failure": bound, "make_call": make_call} + def decryptable_reasoning( + self, model_group: str, request_input: Sequence[object], targets: "Sequence[Deployment]" + ) -> tuple[object, ...]: + """``registry().decryptable_reasoning``: one target decrypting is enough only when encrypted + content affinity pins ``model_group`` to the reasoning's origin.""" + target_dicts: Final = tuple(target.model_dump(exclude_none=True) for target in targets) + pinned: Final = any( + isinstance(callback, EncryptedContentAffinityCheck) and callback.is_enabled_for_model_group(model_group) + for callback in self.optional_callbacks or () + ) + target_groups: Final = (target_dicts,) if pinned else tuple((target,) for target in target_dicts) + kept: Final = list(request_input) + for group in target_groups: + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + self._ctx, # pyright: ignore[reportArgumentType] # the core's RoutingContext serves registry() + kept, + None, + group, + unmarked_origin=None, + ) + return tuple(kept) + def routing_callbacks(self) -> Iterable[object]: """``_routing_callbacks``: the host's callbacks, or after ``observe(request_scoped=True)`` the host's callbacks with the router's own selectors and policies in place of theirs.""" @@ -426,16 +458,21 @@ class OptionalChecks: deployments: list[dict[str, object]], # mutable-ok: CustomLogger deployment-list contract messages: object, request: Mapping[str, object], - ) -> list[object]: # mutable-ok: as the filters return it + ) -> tuple[list[object], str | None]: # mutable-ok: as the filters return it """``async_callback_filter_deployments``: each routing callback narrows the deployments in - turn; a filter's error is raised for the Rust core to decide on.""" + turn, then the prompt caching check names the deployment it prefers among the rest, which + the Rust core takes only if its later filters keep it; a filter's error is raised for the + Rust core to decide on.""" remaining: list[object] = list(deployments) # mutable-ok: filter contract # rebind-ok: narrowed base: Final[object] = CustomLogger.async_filter_deployments # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped deployment dicts request_kwargs: Final = dict(request) - for callback in tuple(self.routing_callbacks()): + callbacks: Final = tuple(self.routing_callbacks()) + for callback in callbacks: # The base filter returns its input unchanged: only overriding filters are run. if not isinstance(callback, CustomLogger) or type(callback).async_filter_deployments is base: # pyright: ignore[reportUnknownMemberType] # see above continue + if isinstance(callback, PromptCachingDeploymentCheck): + continue remaining = await callback.async_filter_deployments( # pyright: ignore[reportUnknownMemberType, reportAssignmentType, reportUnknownVariableType] # see above model=model, healthy_deployments=remaining, @@ -443,7 +480,15 @@ class OptionalChecks: request_kwargs=request_kwargs, parent_otel_span=None, ) - return remaining + prompt_caching: Final = next((cb for cb in callbacks if isinstance(cb, PromptCachingDeploymentCheck)), None) + if prompt_caching is None: + return remaining, None + return remaining, await prompt_caching.async_preferred_deployment_id( + model, + remaining, # pyright: ignore[reportArgumentType] # RustCore's deployment dicts + messages, # pyright: ignore[reportArgumentType] # the request's messages, as given + request_kwargs, + ) # -- release -------------------------------------------------------------------------------- @@ -648,15 +693,15 @@ class RustOptionalChecks: async def _answer_callback_filters( self, step: Mapping[str, object], request: Mapping[str, object], messages: object ) -> tuple[object, BaseException | None]: - """A ``callback_filters`` step: the remaining ids, or a filter's error as an attempt - failure with the error itself, which the caller keeps to raise.""" + """A ``callback_filters`` step: the remaining ids and the one a filter prefers, or a + filter's error as an attempt failure with the error itself, which the caller keeps to raise.""" from litellm.router_rust.core import attempt_failure candidates: Final = [str(id) for id in step.get("candidates") or ()] # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType, reportUnknownArgumentType] # the step's id list records: Final = self._records_of(candidates) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # RustCore's deployment dicts narrowed: Final = _request_with_failed_ids(request, step) try: - remaining: Final = await self._opt.filter_deployments( + remaining, preferred_id = await self._opt.filter_deployments( str(step.get("model_group")), records, # pyright: ignore[reportUnknownArgumentType] # see above messages, @@ -667,7 +712,8 @@ class RustOptionalChecks: # A filter's pin lands on the request kwargs the pick goes on to read. if narrowed is not request and narrowed.get(_ENCRYPTED_PIN) and isinstance(request, dict): request[_ENCRYPTED_PIN] = True # rebind-ok: the request dict the pick goes on to read - return filtered_ids(remaining, candidates), None # pyright: ignore[reportUnknownArgumentType] # see above + kept: Final = filtered_ids(remaining, candidates) # pyright: ignore[reportUnknownArgumentType] # see above + return (kept, preferred_id), None __all__ = [ diff --git a/litellm/router_rust/steps/precall.py b/litellm/router_rust/steps/precall.py index 227bb8de66f..c68c5fd5982 100644 --- a/litellm/router_rust/steps/precall.py +++ b/litellm/router_rust/steps/precall.py @@ -13,7 +13,9 @@ deployment records rather than a step, and those need not be in the Rust registr """ from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from functools import partial +from itertools import accumulate, groupby from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol @@ -21,6 +23,7 @@ from litellm._logging import verbose_router_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider +from litellm.litellm_core_utils.prompt_templates.tpm_media_utils import is_plain_embedding_text from litellm.litellm_core_utils.token_counter import ( offload_token_count, served_openai_models, @@ -38,6 +41,17 @@ if TYPE_CHECKING: # The chat params a deployment must take when the request sends them (`_PRE_CALL_CHECK_FILTERED_PARAMS`). _FILTERED_PARAMS: Final = frozenset({"response_format"}) +_CHARS_PER_OFFLOADED_COUNT: Final = 1 << 16 + + +@dataclass(frozen=True, slots=True) +class _TextBatch: + texts: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class _TokenIdBatch: + longest: int class PreCallRegistry(Protocol): @@ -59,35 +73,111 @@ def count_tokens( request_kwargs: Mapping[str, object] | None, tokenizer_model: str, deployments: Sequence[Mapping[str, object]] = (), + fits_within: int = 0, ) -> int: """``_count_pre_call_check_tokens``: messages, or Responses input normalized to messages, plus the tools and Anthropic ``system`` block the message list never carries. Embedding media an item of ``deployments`` reads as media is left out of the input: it never counts as text.""" from litellm.llms.anthropic.messages.native.utils import anthropic_system_to_openai_message - extras: Final[Mapping[str, object]] = request_kwargs if request_kwargs is not None else MappingProxyType({}) - raw_instructions: Final = extras.get("instructions") - instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None - raw_tools: Final = extras.get("tools") - tools: Final = raw_tools if _is_list(raw_tools) and raw_tools else None + if messages is None and input is not None: + return _input_tokens(host, _text_input(input, deployments), request_kwargs, tokenizer_model, fits_within) + if messages is None: + raise ValueError("Either messages or input must be provided to count tokens") + extras: Final = _extras(request_kwargs) system_message: Final = anthropic_system_to_openai_message(extras.get("system")) counter: Final = host.invoke("token_counter") # pyright: ignore[reportUnknownMemberType] # the host's overloads name litellm's untyped returns - if messages is not None: - counted_messages: Final = (system_message, *_items(messages)) if system_message is not None else messages - return counter(model=tokenizer_model, messages=counted_messages, tools=tools) - if input is not None: - from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, - ) + counted_messages: Final = (system_message, *_items(messages)) if system_message is not None else messages + return counter(model=tokenizer_model, messages=counted_messages, tools=_tools(extras)) - text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments)) - input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # its input type is untyped - input=text_input, # pyright: ignore[reportArgumentType] # str | list, as the request carries it - responses_api_request={"instructions": instructions} if instructions is not None else {}, - only_since_compaction=True, - ) - return counter(model=tokenizer_model, messages=input_messages, tools=tools) - raise ValueError("Either messages or input must be provided to count tokens") + +def _input_tokens( + host: "PythonHost", + text_input: object, + request_kwargs: Mapping[str, object] | None, + tokenizer_model: str, + fits_within: int, +) -> int: + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + counter: Final = host.invoke("token_counter") # pyright: ignore[reportUnknownMemberType] # see count_tokens + match _embedding_batch(text_input): + case _TextBatch(texts): + return _longest_text_tokens(counter, tokenizer_model, texts, fits_within) + case _TokenIdBatch(longest): + return _token_id_count(host, longest) + case None: + pass + extras: Final = _extras(request_kwargs) + raw_instructions: Final = extras.get("instructions") + instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None + input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # its input type is untyped + input=text_input, # pyright: ignore[reportArgumentType] # str | list, as the request carries it + responses_api_request={"instructions": instructions} if instructions is not None else {}, + only_since_compaction=True, + ) + return counter(model=tokenizer_model, messages=input_messages, tools=_tools(extras)) + + +def _extras(request_kwargs: Mapping[str, object] | None) -> Mapping[str, object]: + return request_kwargs if request_kwargs is not None else MappingProxyType({}) + + +def _tools(extras: Mapping[str, object]) -> object: + raw_tools: Final = extras.get("tools") + return raw_tools if _is_list(raw_tools) and raw_tools else None + + +def _text_input(input: object, deployments: Sequence[Mapping[str, object]]) -> object: # the request field's name + text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments)) + return text_input + + +def _embedding_batch(text_input: object) -> _TextBatch | _TokenIdBatch | None: + items: Final = _items(text_input) if _is_list(text_input) else () + if not items: + return None + texts: Final = tuple(item for item in items if isinstance(item, str)) + if len(texts) == len(items): + return _TextBatch(texts) + if all(_is_token_id(item) for item in items): + return _TokenIdBatch(len(items)) + if all(_is_token_id_list(item) for item in items): + return _TokenIdBatch(max(len(_items(item)) for item in items)) + return None + + +def _longest_text_tokens( + counter: Callable[..., int], tokenizer_model: str, texts: Sequence[str], fits_within: int +) -> int: + return max(_text_tokens(counter, tokenizer_model, text, fits_within) for text in texts) + + +def _offload_slices(texts: Sequence[str]) -> tuple[tuple[str, ...], ...]: + ends: Final = accumulate(len(text) for text in texts) + slices: Final = groupby( + zip(ends, texts, strict=True), key=lambda end_and_text: (end_and_text[0] - 1) // _CHARS_PER_OFFLOADED_COUNT + ) + return tuple(tuple(text for _, text in texts_in_slice) for _, texts_in_slice in slices) + + +def _text_tokens(counter: Callable[..., int], tokenizer_model: str, text: str, fits_within: int) -> int: + size: Final = len(text.encode(errors="surrogatepass")) + return size if size <= fits_within else counter(model=tokenizer_model, text=text) + + +def _is_token_id(value: object) -> bool: + return isinstance(value, int) and not isinstance(value, bool) + + +def _is_token_id_list(value: object) -> bool: + return _is_list(value) and all(_is_token_id(token) for token in _items(value)) + + +def _token_id_count(host: "PythonHost", count: int) -> int: + return 0 if host.dynamic().disable_token_counter else count def _max_input_tokens(registry: PreCallRegistry, model: str, deployment: Mapping[str, object]) -> int | None: @@ -246,16 +336,31 @@ class PreCallFacts: def _candidate_records(self) -> Sequence[Mapping[str, object]]: return tuple(self._record(id) for id in self._candidates) + def _fits_within(self) -> int: + return min((limit for limit in self._windows.values() if limit is not None), default=0) + + def _input_is_a_list(self) -> bool: + return self._messages is None and _is_list(self._input) + def _count(self) -> int | None: try: - return count_tokens( - self._host, - self._messages, - self._input, - self._request, - self._tokenizer_model(), - self._candidate_records(), - ) + tokenizer_model: Final = self._tokenizer_model() + if not self._input_is_a_list(): + return count_tokens( + self._host, + self._messages, + self._input, + self._request, + tokenizer_model, + self._candidate_records(), + self._fits_within(), + ) + text_input: Final = _text_input(self._input, self._candidate_records()) + match _embedding_batch(text_input): + case _TextBatch() if not tokenizer_model: + return 0 + case _: + return _input_tokens(self._host, text_input, self._request, tokenizer_model, self._fits_within()) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request verbose_router_logger.error( "litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e @@ -264,14 +369,39 @@ class PreCallFacts: async def _count_async(self) -> int | None: try: - return await offload_token_count(count_tokens)( - self._host, - self._messages, - self._input, - self._request, - self._tokenizer_model(), - self._candidate_records(), + tokenizer_model: Final = self._tokenizer_model() + if not self._input_is_a_list(): + return await offload_token_count(count_tokens)( + self._host, + self._messages, + self._input, + self._request, + tokenizer_model, + self._candidate_records(), + self._fits_within(), + ) + text_input: Final = ( + self._input + if all(map(is_plain_embedding_text, _items(self._input))) + else await offload_token_count(_text_input)(self._input, self._candidate_records()) ) + match _embedding_batch(text_input): + case _TextBatch() if not tokenizer_model: + return 0 + case _TextBatch(texts): + counter: Final = self._host.invoke("token_counter") # pyright: ignore[reportUnknownMemberType] # see count_tokens + fits_within: Final = self._fits_within() + count_slice: Final = offload_token_count(_longest_text_tokens) + return max( + [ + await count_slice(counter, tokenizer_model, texts_in_slice, fits_within) + for texts_in_slice in _offload_slices(texts) + ] + ) + case _: + return await offload_token_count(_input_tokens)( + self._host, text_input, self._request, tokenizer_model, self._fits_within() + ) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request verbose_router_logger.error( "litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e diff --git a/litellm/router_rust/steps/prerouting.py b/litellm/router_rust/steps/prerouting.py index 63d15a1a80a..47bd1e69046 100644 --- a/litellm/router_rust/steps/prerouting.py +++ b/litellm/router_rust/steps/prerouting.py @@ -197,9 +197,13 @@ class PreRoutingStrategies: def quality(self) -> StrategyRegistry: return self._hosted._core_quality_routers # pyright: ignore[reportPrivateUsage] # see above + @property + def switchyard(self) -> StrategyRegistry: + return self._hosted._core_switchyard_routers # pyright: ignore[reportPrivateUsage] # see above + def all(self) -> tuple[StrategyRegistry, ...]: """In the order ``_select_pre_routing_strategy`` searches them.""" - return (self.auto, self.complexity, self.adaptive, self.quality) + return (self.switchyard, self.auto, self.complexity, self.adaptive, self.quality) def any_registered(self) -> bool: return any(self.all()) @@ -211,6 +215,7 @@ class PreRoutingStrategies: self._hosted._core_quality_routers = {} # pyright: ignore[reportPrivateUsage] # see above self._hosted._core_complexity_routers = {} # pyright: ignore[reportPrivateUsage] # see above self._hosted._core_auto_routers = {} # pyright: ignore[reportPrivateUsage] # see above + self._hosted._core_switchyard_routers = {} # pyright: ignore[reportPrivateUsage] # see above def register(self, deployment: Deployment) -> None: """The end of ``_add_deployment``: an auto, complexity or quality marker registers its router @@ -222,6 +227,8 @@ class PreRoutingStrategies: self._init_complexity_router(deployment) case "quality": self._init_quality_router(deployment) + case "switchyard": + self._init_switchyard_router(deployment) case _: pass @@ -231,7 +238,7 @@ class PreRoutingStrategies: return model_name: Final = deployment.model_name tags: Final = _deployment_tags(deployment) - for registry in (self.auto, self.complexity, self.quality): + for registry in (self.auto, self.complexity, self.quality, self.switchyard): _ = _unregister(registry, model_name, tags) if _unregister(self.adaptive, model_name, tags): self._sync_adaptive_router_hooks() @@ -463,6 +470,32 @@ class PreRoutingStrategies: ) self._register(self.quality, deployment, router, "Quality-router") + def _init_switchyard_router(self, deployment: Deployment) -> None: + from litellm.router_strategy.switchyard_router.config import SwitchyardRouterConfig, build_algorithm + from litellm.router_strategy.switchyard_router.switchyard_router import SwitchyardRouter + from litellm.types.decisions import OpenAIDecisionResponse + from litellm.types.utils import ModelResponse + + config: Final = SwitchyardRouterConfig.model_validate(deployment.litellm_params.switchyard_router_config) + if deployment.model_name in (config.models.capable, config.models.efficient, config.models.judge): + raise ValueError("Switchyard cannot route to itself") + ctx: Final = self._ctx + + async def complete_judge(request: Mapping[str, object]) -> ModelResponse: + result: Final = await ctx.acompletion(**request) # pyright: ignore[reportAttributeAccessIssue] # the core's RoutingContext + if not isinstance(result, ModelResponse): + raise TypeError("Switchyard judge must return a non-streaming completion") + return result + + async def decide(request: Mapping[str, object]) -> OpenAIDecisionResponse: + result: Final = await ctx.adecisions(**request) # pyright: ignore[reportAttributeAccessIssue] # see above + if not isinstance(result, OpenAIDecisionResponse): + raise TypeError("Switchyard decision judge must return an OpenAI-shaped decision response") + return result + + router: Final = SwitchyardRouter(deployment.model_name, config, build_algorithm(config), complete_judge, decide) + self._register(self.switchyard, deployment, router, "Switchyard") + def _init_adaptive_router(self, deployment: Deployment) -> None: from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook @@ -741,6 +774,7 @@ class PreRouter: public_router: object, run_plugins: Callable[[str], Awaitable[None]] | None = None, pass_through: bool = False, + call_type: str | None = None, ) -> str | None: """The start of ``_async_get_available_deployment``: run the hook, record its pick and pin the tier's litellm_params. The model group to resolve, or ``None`` to resolve ``model`` itself. @@ -759,6 +793,7 @@ class PreRouter: match_any=match_any, public_router=public_router, run_plugins=run_plugins, + call_type=call_type, ) routed: Final = response.model if response is not None else model if response is not None: @@ -782,6 +817,7 @@ class PreRouter: match_any: bool, public_router: object, run_plugins: Callable[[str], Awaitable[None]] | None, + call_type: str | None, ) -> "PreRoutingHookResponse | None": _restore_client_output_ceiling(request_kwargs) _record_routing_decision(request_kwargs, None) @@ -861,6 +897,7 @@ class PreRouter: messages=routing_messages if routing_messages is not None else messages, # pyright: ignore[reportArgumentType] # the request's messages, as given input=input, # pyright: ignore[reportArgumentType] # the request's input, as given specific_deployment=False, + call_type=call_type, ) hooked: Final[PreRoutingHookResponse | None] = routed traced: Final = ( diff --git a/litellm/router_rust/surface/model_info.py b/litellm/router_rust/surface/model_info.py index b1649617557..508d2bb1145 100644 --- a/litellm/router_rust/surface/model_info.py +++ b/litellm/router_rust/surface/model_info.py @@ -46,6 +46,7 @@ from litellm.litellm_core_utils.add_retry_fallback_headers import ( from litellm.litellm_core_utils.batch_utils import deployment_supported_endpoints from litellm.litellm_core_utils.core_helpers import as_str_mapping, coerce_token_limit from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.litellm_core_utils.litellm_logging import resolve_vertex_location_for_cost from litellm.litellm_core_utils.reasoning_effort_capability import ( deployment_is_catalog_mapped, intersect_supported_reasoning_efforts, @@ -56,11 +57,13 @@ from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage, UsageCounter from litellm.router_python_host import python_host +from litellm.router_utils.model_group_pricing import deployment_pricing from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import build_io_token_rate_limit_headers from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata from litellm.types.router import ( Deployment, DeploymentModelListingInfo, + DeploymentPricing, DeploymentTypedDict, DiscoveredDeploymentModelInfo, LiteLLM_Params, @@ -78,6 +81,7 @@ if TYPE_CHECKING: ) _MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) # The key router-core's merge reads as an opaque value (``router-core/src/model_info.rs``). _OPAQUE_KEY: Final = "__litellm_python_value__" @@ -140,6 +144,11 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) +def _configured_deployment_id(deployment: Mapping[str, object]) -> str | None: + deployment_id: Final = _configured_model_info(deployment).get("id") + return deployment_id if isinstance(deployment_id, str) else None + + def _operator_max_input_tokens(deployment: Mapping[str, object]) -> int | None: """``_operator_max_input_tokens``: the deployment's own limit, skipping one a saved ``/model/info`` response copied from the cost map.""" @@ -187,6 +196,7 @@ class _GroupBuild: # Filled as the lookups arrive. params: dict[int, LiteLLM_Params] = dataclasses.field(default_factory=dict) # mutable-ok: see above info: ModelGroupInfo | None = None + pricing: tuple[DeploymentPricing, ...] = () def _usage_total(values: Sequence[object]) -> int | None: @@ -358,6 +368,7 @@ class HostModelInfo: return None for field, value in (as_str_mapping(answer["assignments"]) or {}).items(): setattr(group.info, field, _group_field(field, wire.from_wire(value))) + group.info.deployment_pricing = group.pricing return group.info def _answer_group_lookup(self, step: Mapping[str, object], group: "_GroupBuild", wire: "_Wire") -> object: @@ -384,6 +395,17 @@ class HostModelInfo: efforts: Final = resolve_supported_reasoning_efforts( model_info, deployment_is_mapped=step["mapped"] is True ) + pricing: Final = deployment_pricing( + llm_provider, + _configured_deployment_id(deployment), + group.params[int(str(step["index"]))].model, + model_info, + resolve_vertex_location_for_cost( + llm_provider, deployment.get("litellm_params"), None, litellm_model + ), + ) + if pricing is not None and pricing not in group.pricing: + group.pricing = (*group.pricing, pricing) # rebind-ok: the merge's state return { "supported_endpoints": wire.to_wire( list( @@ -740,6 +762,20 @@ class HostModelInfo: ) return tuple(wire.from_wire(tier) for tier in tiers) + def _core_get_pinned_reasoning_efforts(self, model_name: str, team_id: str | None = None) -> tuple[object, ...]: + """ + Return the litellm_params.reasoning_effort each deployment a request + from team_id to model_name can route to (see _routable_deployments) + pins, unvalidated and in model_list order, None for a deployment that + pins none; the caller decides how the deployments combine. + + Returns an empty tuple for wildcard-expanded or unknown names. + """ + return tuple( + _LITELLM_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]).get("reasoning_effort") + for deployment in self._routable_deployments(model_name, team_id) + ) + def _core_get_configured_max_input_tokens( self, model_name: str, team_id: str | None = None ) -> tuple[int | None, ...]: diff --git a/litellm/router_rust/surface/registry.py b/litellm/router_rust/surface/registry.py index eaeac10cb5f..820db52d5a1 100644 --- a/litellm/router_rust/surface/registry.py +++ b/litellm/router_rust/surface/registry.py @@ -658,6 +658,7 @@ class HostRegistry(HostModelInfo): self._core_complexity_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: same self._core_adaptive_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: same self._core_quality_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: same + self._core_switchyard_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: same self.strategies: PreRoutingStrategies | None = None # The router's optional pre-call checks (deployment budgets, prompt caching pins); # RustCore installs its ``OptionalChecks`` (litellm/router_rust/steps/callbacks.py). @@ -1510,6 +1511,14 @@ class HostRegistry(HostModelInfo): return self.__getattr__("_core_prompt_caching_pin_cache") # pyright: ignore[reportReturnType] # raises RustCoreUnsupported return self.checks.prompt_caching_pin_cache + def _core_decryptable_reasoning( + self, model_group: str, request_input: Sequence[object], targets: Sequence[Deployment] + ) -> tuple[object, ...]: + """The input ``targets`` can decrypt (``registry().decryptable_reasoning``).""" + if self.checks is None: + return self.__getattr__("_core_decryptable_reasoning") # pyright: ignore[reportReturnType] # raises RustCoreUnsupported + return self.checks.decryptable_reasoning(model_group, request_input, targets) + def replay_model_cost_registrations(self) -> None: """Re-register every served deployment's cost-map entries (``replay_model_cost_registrations``).""" provider_unresolved: Final = self._provider_unresolved_deployments diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index c3a65d6aa91..43869fdf832 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -184,6 +184,7 @@ class AdaptiveRouter: messages: list[dict[str, object]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, + call_type: str | None = None, ) -> PreRoutingHookResponse | None: """ Plugin entry point invoked by `Router.async_pre_routing_hook` when the diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index 3427b81be68..ed402fd5262 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -14,10 +14,13 @@ from __future__ import annotations import hashlib import json from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable + +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, iter_items, require_attributes from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter from litellm.router_strategy.adaptive_router.classifier import classify_prompt from litellm.router_strategy.adaptive_router.config import ( @@ -41,7 +44,12 @@ _IDENTITY_FIELDS: Final = ( ) -def _resolve_session_key(kwargs: dict[str, Any]) -> str | None: +class _LoggedCallKwargs(TypedDict, total=False): + litellm_params: ReadOnly[Mapping[str, object]] + messages: ReadOnly[Sequence[Mapping[str, object]] | None] + + +def _resolve_session_key(kwargs: _LoggedCallKwargs) -> str | None: """Pick a stable per-conversation key for owner-cache attribution. Order: @@ -135,38 +143,70 @@ def _recent_tool_results( return results -def _assistant_content_and_tool_calls(response_obj: Any) -> tuple[object, Sequence[Mapping[str, object]]]: +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, index: int, /) -> object: ... + + +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + +class _SupportsModelDump(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +def _first_item(value: object) -> object: + if isinstance(value, _SupportsGetItem): + return value[0] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + + +def _dict_fields(value: object) -> Mapping[str, object] | None: + fields: Final = as_str_mapping(value) + return fields if isinstance(value, dict) else None + + +def _attribute_or_dict_field(source: object, name: str) -> object: + attribute: Final = attribute_or_none(source, name) + fields: Final = _dict_fields(source) + return fields.get(name) if attribute is None and fields is not None else attribute + + +def _tool_call_fields(tool_call: object) -> Mapping[str, object]: + fields: Final = _dict_fields(tool_call) + if fields is not None: + return fields + try: + return require_attributes(tool_call, _SupportsModelDump, "model_dump").model_dump() + except Exception: + return {"name": getattr(tool_call, "name", ""), "arguments": ""} + + +def _assistant_content_and_tool_calls(response_obj: object) -> tuple[object, Sequence[Mapping[str, object]]]: """Return (assistant_text, tool_calls_list) extracted from a ModelResponse-ish object.""" if response_obj is None: return None, [] try: - choices: Final = getattr(response_obj, "choices", None) or response_obj.get("choices") + choices: Final = attribute_or_none(response_obj, "choices") or require_attributes( + response_obj, _SupportsGet, "get" + ).get("choices") except Exception: return None, [] if not choices: return None, [] - msg = choices[0] - msg = getattr(msg, "message", None) or (msg.get("message") if isinstance(msg, dict) else None) + first_choice: Final = _first_item(choices) + choice_fields: Final = _dict_fields(first_choice) + msg: Final = attribute_or_none(first_choice, "message") or ( + choice_fields.get("message") if choice_fields is not None else None + ) if msg is None: return None, [] - content = getattr(msg, "content", None) - if content is None and isinstance(msg, dict): - content = msg.get("content") - - raw_tool_calls = getattr(msg, "tool_calls", None) - if raw_tool_calls is None and isinstance(msg, dict): - raw_tool_calls = msg.get("tool_calls") - tool_calls: Final[list[dict[str, object]]] = [] - for tc in raw_tool_calls or []: - if isinstance(tc, dict): - tool_calls.append(tc) - else: - try: - tool_calls.append(tc.model_dump()) - except Exception: - tool_calls.append({"name": getattr(tc, "name", ""), "arguments": ""}) + content: Final = _attribute_or_dict_field(msg, "content") + raw_tool_calls: Final = _attribute_or_dict_field(msg, "tool_calls") + tool_calls: Final = [_tool_call_fields(tc) for tc in iter_items(raw_tool_calls or [])] return content, tool_calls @@ -214,7 +254,7 @@ class AdaptiveRouterPostCallHook(CustomLogger): async def _record( self, - kwargs: dict[str, Any], + kwargs: _LoggedCallKwargs, response_obj: object, response_status: int, ) -> None: diff --git a/litellm/router_strategy/adaptive_router/signals.py b/litellm/router_strategy/adaptive_router/signals.py index c7fffd05632..9159bc6c2c0 100644 --- a/litellm/router_strategy/adaptive_router/signals.py +++ b/litellm/router_strategy/adaptive_router/signals.py @@ -15,8 +15,9 @@ from __future__ import annotations import re from collections.abc import Mapping, Sequence from dataclasses import dataclass, field -from typing import Any, Final +from typing import Final, Protocol +from litellm.litellm_core_utils.core_helpers import as_str_mapping, require_attributes from litellm.router_strategy.adaptive_router.config import ( LOOP_REPEAT_THRESHOLD, MIN_TURNS_FOR_CLEAN_CREDIT, @@ -174,14 +175,23 @@ def _detect_failure(tool_results: Sequence[Mapping[str, object]]) -> bool: return False -def _signature(call: Mapping[str, Any]) -> str: +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + +def _function_field(call: Mapping[str, object], key: str) -> object: + return require_attributes(call.get("function", {}), _SupportsGet, "get").get(key, "") + + +def _signature(call: Mapping[str, object]) -> str: """Stable signature for loop detection: name + sorted JSON-ish args.""" - name: Final = call.get("name") or call.get("function", {}).get("name", "") - call_args = call.get("arguments") - if call_args is None: - call_args = call.get("function", {}).get("arguments", "") - if isinstance(call_args, dict): - call_args = ",".join(f"{k}={call_args[k]}" for k in sorted(call_args.keys())) + name: Final = call.get("name") or _function_field(call, "name") + raw_args: Final = call.get("arguments") + call_args: Final = _function_field(call, "arguments") if raw_args is None else raw_args + args_mapping: Final = as_str_mapping(call_args) + if args_mapping is not None and isinstance(call_args, dict): + joined: Final = ",".join(f"{k}={args_mapping[k]}" for k in sorted(args_mapping.keys())) + return f"{name}({joined})" return f"{name}({call_args})" diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index b92c53fa93c..4e119c414a0 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -6,12 +6,13 @@ import asyncio from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, NamedTuple, Optional -from pydantic import ConfigDict +from pydantic import ConfigDict, TypeAdapter from litellm._logging import verbose_router_logger from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS from litellm.exceptions import Timeout as LiteLLMTimeout from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items from litellm.litellm_core_utils.internal_call_metadata import ( effective_turn_off_message_logging, forwarded_internal_call_metadata, @@ -41,6 +42,9 @@ else: LiteLLMRouterEncoder = Any +_DECODED_JSON: Final = TypeAdapter(object) + + class _CallerMetadata(LiteLLMBaseModel): model_config = ConfigDict(frozen=True, extra="ignore") @@ -123,14 +127,15 @@ class AutoRouter(CustomLogger): if self.auto_router_config is None: raise ValueError("No auto router config provided") auto_router_routes: Final[list[Route]] = [] - loaded_config: Final = json.loads(self.auto_router_config) - for route in loaded_config.get("routes", []): + loaded_config: Final = decoded_json_object(_DECODED_JSON.validate_python(json.loads(self.auto_router_config))) + for route in iter_items(loaded_config.get("routes", [])): + route_config = decoded_json_object(route) auto_router_routes.append( Route( - name=route.get("name"), - description=route.get("description"), - utterances=route.get("utterances", []), - score_threshold=route.get("score_threshold"), + name=route_config.get("name"), + description=route_config.get("description"), + utterances=route_config.get("utterances", []), + score_threshold=route_config.get("score_threshold"), ) ) return auto_router_routes @@ -204,6 +209,7 @@ class AutoRouter(CustomLogger): messages: list[dict[str, object]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, + call_type: str | None = None, ) -> Optional["PreRoutingHookResponse"]: """ This hook is called before the routing decision is made. diff --git a/litellm/router_strategy/auto_router/auto_router_model_naming.py b/litellm/router_strategy/auto_router/auto_router_model_naming.py index f5aeccd9305..47daa175c99 100644 --- a/litellm/router_strategy/auto_router/auto_router_model_naming.py +++ b/litellm/router_strategy/auto_router/auto_router_model_naming.py @@ -26,7 +26,7 @@ from litellm.router_strategy.complexity_router.config import ( AUTO_ROUTER_MODEL_PREFIX: Final = "auto_router/" -StrategyRouterKind = Literal["semantic", "complexity", "adaptive", "quality"] +StrategyRouterKind: TypeAlias = Literal["semantic", "complexity", "adaptive", "quality", "switchyard"] StrategyRouterDependencyRole: TypeAlias = Literal[ "tier", "default", "classifier", "embedding", "evaluation", "compactor" @@ -43,6 +43,7 @@ class StrategyRouterDependency: STRATEGY_ROUTER_PARAM_FIELDS: Final[frozenset[str]] = frozenset( { + "switchyard_router_config", "auto_router_config", "auto_router_config_path", "auto_router_default_model", @@ -57,6 +58,7 @@ STRATEGY_ROUTER_PARAM_FIELDS: Final[frozenset[str]] = frozenset( ) _REQUIRED_FIELD_GROUPS: Final[Mapping[StrategyRouterKind, tuple[tuple[str, ...], ...]]] = { + "switchyard": (("switchyard_router_config",),), "semantic": ( ("auto_router_config", "auto_router_config_path"), ("auto_router_default_model",), @@ -78,6 +80,8 @@ def classify_strategy_router_model(model: str) -> StrategyRouterKind | None: if not model.startswith(AUTO_ROUTER_MODEL_PREFIX): return None remainder: Final = model[len(AUTO_ROUTER_MODEL_PREFIX) :] + if remainder.startswith("switchyard"): + return "switchyard" if remainder.startswith("complexity_router"): return "complexity" if remainder.startswith("adaptive_router"): @@ -137,6 +141,13 @@ def strategy_router_dependencies( kind: Final = classify_strategy_router_model(str(litellm_params.get("model", ""))) if kind is None: return () + if kind == "switchyard": + models: Final = _mapping(_mapping(litellm_params.get("switchyard_router_config")).get("models")) + return ( + _named(models.get("capable"), "tier") + + _named(models.get("efficient"), "tier") + + _named(models.get("judge"), "classifier") + ) if kind == "semantic": return _named(litellm_params.get("auto_router_default_model"), "default") + _named( litellm_params.get("auto_router_embedding_model"), "embedding" diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 98dd39b40e3..c65820b9892 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -22,11 +22,11 @@ import asyncio import builtins import logging import time -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Mapping, MutableMapping, Sequence from datetime import datetime, timedelta, timezone from itertools import groupby from types import MappingProxyType -from typing import Any, Final +from typing import Any, Final, Protocol from pydantic import TypeAdapter @@ -44,6 +44,7 @@ from litellm.litellm_core_utils.batch_utils import ( ) from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, + require_attributes, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds_from from litellm.router_strategy.tag_based_routing import get_tags_from_request_kwargs @@ -51,7 +52,7 @@ from litellm.router_utils.cooldown_callbacks import ( get_prometheus_logger_from_callbacks, ) from litellm.types.llms.openai import AllMessageValues -from litellm.types.router import DeploymentTypedDict, LiteLLM_Params, RouterErrors +from litellm.types.router import LiteLLM_Params, RouterErrors from litellm.types.utils import BudgetConfig, GenericBudgetConfigType, StandardLoggingPayload from litellm.types.utils import BudgetConfig as GenericBudgetInfo @@ -60,6 +61,32 @@ _CALLBACK_FIELDS: Final = TypeAdapter(dict[str, object]) _JOB_CHARGE_VALUES: Final = TypeAdapter(tuple[object, float]) +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + +def _model_info_id(deployment: Mapping[str, object]) -> object: + return require_attributes(deployment.get("model_info", {}), _SupportsGet, "get").get("id") + + +def _deployment_budget(deployment: Mapping[str, object]) -> tuple[object, GenericBudgetInfo] | None: + litellm_params: Final = deployment.get("litellm_params", {}) + model_id: Final = require_attributes(deployment.get("model_info") or {}, _SupportsGet, "get").get("id") + params: Final = require_attributes(litellm_params, _SupportsGet, "get") + max_budget: Final = params.get("max_budget") + budget_duration: Final = params.get("budget_duration") + + verbose_router_logger.debug( + "Init Deployment Budget: max_budget: %s, budget_duration: %s, model_id: %s", + max_budget, + budget_duration, + model_id, + ) + if max_budget is None or budget_duration is None or model_id is None: + return None + return model_id, GenericBudgetInfo(time_period=budget_duration, budget_limit=max_budget) + + class _LiteLLMParamsDictView: """ Lightweight attribute view over `litellm_params` dict. @@ -124,7 +151,7 @@ class RouterBudgetLimiting(CustomLogger): self, dual_cache: DualCache, provider_budget_config: dict | None, - model_list: list[DeploymentTypedDict | dict[str, Any]] | None = None, + model_list: Sequence[Mapping[str, object]] | None = None, clock: Callable[[], float] = time.time, ): self.dual_cache = dual_cache @@ -137,7 +164,7 @@ class RouterBudgetLimiting(CustomLogger): self._retired = False self._flush_after_retire: asyncio.Task[None] | None = None self.provider_budget_config: GenericBudgetConfigType | None = provider_budget_config - self.deployment_budget_config: GenericBudgetConfigType | None = None + self.deployment_budget_config: Mapping[object, GenericBudgetInfo] | None = None self.tag_budget_config: GenericBudgetConfigType | None = None self._init_provider_budgets() self._init_deployment_budgets(model_list=model_list) @@ -251,9 +278,9 @@ class RouterBudgetLimiting(CustomLogger): def _filter_out_deployments_above_budget( self, potential_deployments: list[dict[str, object]], - healthy_deployments: list[dict[str, Any]], + healthy_deployments: Sequence[dict[str, object]], provider_configs: dict[str, GenericBudgetInfo], - deployment_configs: dict[str, GenericBudgetInfo], + deployment_configs: Mapping[object, GenericBudgetInfo], deployment_providers: list[str | None], spend_map: dict[str, float], request_tags: list[str], @@ -300,9 +327,9 @@ class RouterBudgetLimiting(CustomLogger): # Check deployment budget if self.deployment_budget_config and is_within_budget: _model_name = deployment.get("model_name") - _litellm_params = deployment.get("litellm_params") or {} + _litellm_params = require_attributes(deployment.get("litellm_params") or {}, _SupportsGet, "get") _litellm_model_name = _litellm_params.get("model") - model_id = deployment.get("model_info", {}).get("id") + model_id = _model_info_id(deployment) if model_id in deployment_configs: config = deployment_configs[model_id] current_spend = spend_map.get(f"deployment_spend:{model_id}:{config.budget_duration}", 0.0) @@ -334,12 +361,12 @@ class RouterBudgetLimiting(CustomLogger): async def _async_get_cache_keys_for_router_budget_limiting( self, - healthy_deployments: list[dict[str, Any]], + healthy_deployments: Sequence[Mapping[str, object]], request_kwargs: dict | None = None, ) -> tuple[ list[str], dict[str, GenericBudgetInfo], - dict[str, GenericBudgetInfo], + dict[object, GenericBudgetInfo], list[str | None], ]: """ @@ -354,7 +381,7 @@ class RouterBudgetLimiting(CustomLogger): """ cache_keys: Final[list[str]] = [] provider_configs: Final[dict[str, GenericBudgetInfo]] = {} - deployment_configs: Final[dict[str, GenericBudgetInfo]] = {} + deployment_configs: Final[dict[object, GenericBudgetInfo]] = {} deployment_providers: Final[list[str | None]] = [] # Resolve tags once before the loop (loop-invariant) @@ -378,7 +405,7 @@ class RouterBudgetLimiting(CustomLogger): # Check deployment budgets if self.deployment_budget_config: - model_id = deployment.get("model_info", {}).get("id") + model_id = _model_info_id(deployment) if model_id is not None: budget_config = self._get_budget_config_for_deployment(model_id) if budget_config is not None: @@ -768,7 +795,7 @@ class RouterBudgetLimiting(CustomLogger): def _get_budget_config_for_deployment( self, - model_id: str, + model_id: object, ) -> GenericBudgetInfo | None: if self.deployment_budget_config is None: return None @@ -784,7 +811,7 @@ class RouterBudgetLimiting(CustomLogger): return None return self.tag_budget_config.get(tag, None) - def _get_llm_provider_for_deployment(self, deployment: dict) -> str | None: + def _get_llm_provider_for_deployment(self, deployment: Mapping[str, object]) -> str | None: try: deployment_litellm_params: Final = deployment.get("litellm_params") or {} @@ -893,7 +920,7 @@ class RouterBudgetLimiting(CustomLogger): @staticmethod def should_init_router_budget_limiter( provider_budget_config: dict | None, - model_list: list[DeploymentTypedDict] | list[dict[str, Any]] | None = None, + model_list: Sequence[Mapping[str, object]] | None = None, ): """ Returns `True` if the router budget routing settings are set and RouterBudgetLimiting should be initialized @@ -913,7 +940,7 @@ class RouterBudgetLimiting(CustomLogger): return False for _model in model_list: - _litellm_params = _model.get("litellm_params", {}) + _litellm_params = require_attributes(_model.get("litellm_params", {}), _SupportsGet, "get") if _litellm_params.get("max_budget") or _litellm_params.get("budget_duration") is not None: return True return False @@ -943,37 +970,20 @@ class RouterBudgetLimiting(CustomLogger): def _init_deployment_budgets( self, - model_list: list[DeploymentTypedDict | dict[str, Any]] | None = None, + model_list: Sequence[Mapping[str, object]] | None = None, ): if model_list is None: return - for _model in model_list: - _litellm_params = _model.get("litellm_params", {}) - _model_info: dict = _model.get("model_info") or {} - _model_id = _model_info.get("id") - _max_budget = _litellm_params.get("max_budget") - _budget_duration = _litellm_params.get("budget_duration") - - verbose_router_logger.debug( - "Init Deployment Budget: max_budget: %s, budget_duration: %s, model_id: %s", - _max_budget, - _budget_duration, - _model_id, - ) - if _max_budget is not None and _budget_duration is not None and _model_id is not None: - _budget_config = GenericBudgetInfo( - time_period=_budget_duration, - budget_limit=_max_budget, - ) - if self.deployment_budget_config is None: - self.deployment_budget_config = {} - self.deployment_budget_config[_model_id] = _budget_config + for model_id, budget in filter(None, map(_deployment_budget, model_list)): + if not isinstance(self.deployment_budget_config, MutableMapping): + self.deployment_budget_config = dict(self.deployment_budget_config or {}) + self.deployment_budget_config[model_id] = budget verbose_router_logger.debug("Initialized Deployment Budget Config: %s", self.deployment_budget_config) def register_deployment_budget( self, - deployment: dict[str, Any] | DeploymentTypedDict, + deployment: Mapping[str, object], ) -> None: """ Register or refresh deployment-level budget config for a runtime-added deployment. @@ -981,7 +991,7 @@ class RouterBudgetLimiting(CustomLogger): self._init_deployment_budgets(model_list=[deployment]) def unregister_deployment_budget(self, model_id: str) -> None: - if self.deployment_budget_config is None: + if not isinstance(self.deployment_budget_config, MutableMapping): return self.deployment_budget_config.pop(model_id, None) if len(self.deployment_budget_config) == 0: diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 05dd4706b62..600b9d7c455 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -28,7 +28,7 @@ from functools import partial from itertools import accumulate, chain, islice, takewhile from threading import Lock from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, cast from pydantic import BaseModel, Field, TypeAdapter, ValidationError, create_model from pydantic_core import ErrorDetails @@ -52,6 +52,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_parent_otel_span_from_kwargs, is_codex_user_agent, + require_attributes, ) from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata @@ -144,7 +145,7 @@ if TYPE_CHECKING: from litellm.router_api import RoutingContext from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter from litellm.router_strategy.savings_baseline import Baseline - from litellm.types.router import PreRoutingHookResponse + from litellm.types.router import DeploymentTypedDict, PreRoutingHookResponse else: RoutingContext = Any PreRoutingHookResponse = Any @@ -184,6 +185,15 @@ def _built_in_tier_or_none(tier_name: str) -> ComplexityTier | None: return ComplexityTier.__members__.get(tier_name) +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + +def _declared_vision_support(deployment: Mapping[str, object]) -> object: + model_info: Final = deployment.get("model_info") + return require_attributes(model_info, _SupportsGet, "get").get("supports_vision") if model_info else None + + _CLASSIFICATION_TIER_CRITERIA: Final[Mapping[ComplexityTier, str]] = MappingProxyType( { ComplexityTier.NON_REASONING: ( @@ -3800,11 +3810,12 @@ class ComplexityRouter(CustomLogger): return True return False if is_vision_explicitly_disabled(model) else None - def deployment_verdict(deployment: Mapping[str, Any]) -> bool | None: - declared: Final = (deployment.get("model_info") or EMPTY_MAPPING).get("supports_vision") + def deployment_verdict(deployment: DeploymentTypedDict) -> bool | None: + declared: Final = _declared_vision_support(deployment) if declared is not None: return declared is True - return model_verdict((deployment.get("litellm_params") or EMPTY_MAPPING).get("model") or model_name) + litellm_params: Final = deployment.get("litellm_params") + return model_verdict((litellm_params.get("model") if litellm_params else None) or model_name) deployments: Final = self.ctx.registry().deployments(model_name=model_name) if not deployments: @@ -4490,6 +4501,7 @@ class ComplexityRouter(CustomLogger): messages: list[dict[str, object]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, + call_type: str | None = None, ) -> PreRoutingHookResponse | None: """ Pre-routing hook called before the routing decision. diff --git a/litellm/router_strategy/quality_router/quality_router.py b/litellm/router_strategy/quality_router/quality_router.py index 2d9ad63e9de..c3b038b5f26 100644 --- a/litellm/router_strategy/quality_router/quality_router.py +++ b/litellm/router_strategy/quality_router/quality_router.py @@ -16,7 +16,7 @@ then cheapest `model_info.input_cost_per_token`). """ import math -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Optional from litellm._logging import verbose_router_logger @@ -162,7 +162,7 @@ class QualityRouter(CustomLogger): `_model_quality`, and `_model_cost`. Raises if any listed model is missing `litellm_routing_preferences`. """ - model_list: Final = self.ctx.registry().all() or [] + model_list: Final[Sequence[object]] = self.ctx.registry().all() or [] available: Final = set(self.config.available_models) # Track which available models we've matched so we can error on missing. @@ -330,6 +330,7 @@ class QualityRouter(CustomLogger): messages: list[dict[str, object]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, + call_type: str | None = None, ) -> Optional["PreRoutingHookResponse"]: """Try keyword override first; fall back to complexity-tier routing.""" from litellm.types.router import PreRoutingHookResponse diff --git a/litellm/router_strategy/switchyard_router/__init__.py b/litellm/router_strategy/switchyard_router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/router_strategy/switchyard_router/config.py b/litellm/router_strategy/switchyard_router/config.py new file mode 100644 index 00000000000..7e0d3ad3a94 --- /dev/null +++ b/litellm/router_strategy/switchyard_router/config.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Discriminator, Field, Tag, TypeAdapter, model_validator + +if TYPE_CHECKING: + from switchyard_rust.libsy import Algorithm, TaskClassifierConfig + + +class LlmClassifierConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + kind: Literal["llm"] = "llm" + base_threshold: float = Field(default=0.5, ge=0, le=1) + threshold_step: float = Field(default=0.0, ge=0, le=1) + recent_turn_window: int | None = Field(default=None, gt=0) + max_output_tokens: int = Field(default=4096, gt=0) + response_format_type: Literal["json_schema", "json_object"] = "json_schema" + prompt: str | None = None + + +class DecisionClassifierConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + kind: Literal["decision"] + cutoff: float = Field(default=0.6, ge=0, le=1) + instructions: str | Mapping[str, object] | None = None + evidence: Mapping[str, object] | None = None + + +_KIND: Final = TypeAdapter(Mapping[str, object]) + + +def _classifier_kind(value: object) -> str: + match value: + case LlmClassifierConfig() | DecisionClassifierConfig(): + return value.kind + case Mapping(): + return str(_KIND.validate_python(value).get("kind", "llm")) + case _: + return "llm" + + +ClassifierConfig: TypeAlias = Annotated[ + Annotated[LlmClassifierConfig, Tag("llm")] | Annotated[DecisionClassifierConfig, Tag("decision")], + Discriminator(_classifier_kind), +] + + +class StageOptions(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + picker: Literal["efficient_first", "capable_first"] = "efficient_first" + confidence_threshold: float | None = Field(default=None, ge=0, le=1) + recent_window: int | None = Field(default=None, gt=0) + tool_semantics: Mapping[str, tuple[str, ...]] | None = None + only_on_wrong_signal_escalation: bool = True + + +class StageModels(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", str_strip_whitespace=True) + + capable: str = Field(min_length=1) + efficient: str = Field(min_length=1) + judge: str | None = Field(default=None, min_length=1) + + @model_validator(mode="after") + def reject_recursive_targets(self) -> StageModels: + if any(name.startswith("auto_router/") for name in (self.capable, self.efficient, self.judge or "")): + raise ValueError("Switchyard targets must be concrete model groups, not auto-router markers") + if self.capable == self.efficient: + raise ValueError("The capable and efficient models must be different model groups") + return self + + +class SwitchyardRouterConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + algorithm: Literal["stage_router"] = "stage_router" + models: StageModels + options: StageOptions = Field(default_factory=StageOptions) + classifier: ClassifierConfig | None = None + + @model_validator(mode="after") + def validate_judge(self) -> SwitchyardRouterConfig: + if (self.classifier is None) != (self.models.judge is None): + raise ValueError("A judge model is required exactly when the classifier is enabled") + return self + + @property + def confidence_threshold(self) -> float: + if self.options.confidence_threshold is not None: + return self.options.confidence_threshold + return 0.7 if self.classifier is not None else 0.5 + + def runtime_models(self) -> Mapping[str, tuple[str, ...]]: + return { + "any": (self.models.capable, self.models.efficient), + "capable": (self.models.capable,), + "efficient": (self.models.efficient,), + **({"judge": (self.models.judge,)} if self.models.judge is not None else {}), + } + + +def _task_classifier(config: SwitchyardRouterConfig) -> TaskClassifierConfig | None: + from switchyard_rust.libsy import TaskClassifierConfig + + match config.classifier: + case None: + return None + case LlmClassifierConfig() as llm: + return TaskClassifierConfig( + base_threshold=llm.base_threshold, + threshold_step=llm.threshold_step, + recent_turn_window=llm.recent_turn_window, + max_output_tokens=llm.max_output_tokens, + response_format_type=llm.response_format_type, + prompt=llm.prompt, + ) + case DecisionClassifierConfig() as decision: + return TaskClassifierConfig.decision( + cutoff=decision.cutoff, + candidates={"capable": config.models.capable, "efficient": config.models.efficient}, + evidence=dict(decision.evidence or {}), + instructions=decision.instructions, + ) + + +def build_algorithm(config: SwitchyardRouterConfig) -> Algorithm: + from switchyard_rust.libsy import LlmFallback, stage_router + + classifier: Final = _task_classifier(config) + return stage_router( + picker=config.options.picker, + confidence_threshold=config.confidence_threshold, + recent_window=config.options.recent_window, + tool_semantics=dict(config.options.tool_semantics) if config.options.tool_semantics is not None else None, + only_on_wrong_signal_escalation=config.options.only_on_wrong_signal_escalation, + classifier=LlmFallback(config=classifier) if classifier is not None else None, + ) + + +def validate_switchyard_config_write(config: Mapping[str, object] | None) -> str | None: + if config is None: + return None + try: + build_algorithm(SwitchyardRouterConfig.model_validate(config)) + except ImportError: + return "Switchyard requires the optional dependency. Install litellm[switchyard] (nemo-switchyard==0.4.0a1)." + except (ValueError, TypeError) as error: + return f"Invalid switchyard_router_config: {error}" + return None diff --git a/litellm/router_strategy/switchyard_router/switchyard_router.py b/litellm/router_strategy/switchyard_router/switchyard_router.py new file mode 100644 index 00000000000..cb455bab337 --- /dev/null +++ b/litellm/router_strategy/switchyard_router/switchyard_router.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +import hashlib +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Final, Literal, TypeAlias, TypeVar + +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs +from litellm.litellm_core_utils.internal_call_metadata import ( + effective_turn_off_message_logging, + forwarded_internal_call_metadata, + parent_session_kwargs, +) +from litellm.router_strategy.switchyard_router.config import ( + DecisionClassifierConfig, + LlmClassifierConfig, + SwitchyardRouterConfig, +) +from litellm.router_strategy.switchyard_router.translation import ( + chat_messages, + decision_request, + decision_response, + judge_request, + judge_response, + mapping, + request_to_switchyard, +) +from litellm.types.decisions import OpenAIDecisionResponse +from litellm.types.router import PreRoutingHookResponse +from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN, ModelResponse, StandardLoggingRoutingDecision + +if TYPE_CHECKING: + from switchyard_rust.libsy import Algorithm, DecisionCall, ModelCall, RoutingOutcome + +JudgeCompletion: TypeAlias = Callable[[Mapping[str, object]], Awaitable[ModelResponse]] +DecisionCompletion: TypeAlias = Callable[[Mapping[str, object]], Awaitable[OpenAIDecisionResponse]] +FallbackReason: TypeAlias = Literal[ + "unsupported_request_type", + "unsupported_content", + "libsy_error", + "unexpected_judge_call", + "unexpected_decision_call", + "target_outside_config", + "no_decision", +] +_T = TypeVar("_T") +_SESSION_HEADER: Final = "x-switchyard-session-id" +_FORWARDED_HEADER_PREFIXES: Final = ("x-switchyard-", "x-codex-", "x-nemo-relay-", "x-dynamo-") +_FORWARDED_HEADERS: Final = frozenset( + {"x-session-id", "session-id", "thread-id", "x-task-id", "x-request-id", "x-client-request-id"} +) + + +async def _resolved(value: _T) -> _T: + return value + + +_SESSION_ALIASES: Final = frozenset( + {_SESSION_HEADER, "x-claude-code-session-id", "x-nemo-relay-session-id", "x-session-id", "session-id"} +) +_CODEX_TURN_METADATA_HEADER: Final = "x-codex-turn-metadata" + + +def _forwarded_header(name: str, session_scoped: bool) -> bool: + lowered: Final = name.lower() + if lowered in _SESSION_ALIASES or (lowered == _CODEX_TURN_METADATA_HEADER and not session_scoped): + return False + return lowered in _FORWARDED_HEADERS or lowered.startswith(_FORWARDED_HEADER_PREFIXES) + + +def _caller_key_hash(request_kwargs: Mapping[str, object]) -> str: + internal: Final = mapping(request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs)) or {}) + key: Final = internal.get("user_api_key_hash") + return str(key) if key is not None else "" + + +def _scoped_session(request_kwargs: Mapping[str, object]) -> str | None: + session: Final = parent_session_kwargs(request_kwargs).get("litellm_session_id") + if not session: + return None + return hashlib.sha256(f"{_caller_key_hash(request_kwargs)}\0{session}".encode()).hexdigest() + + +def switchyard_headers(request_kwargs: Mapping[str, object]) -> Mapping[str, str]: + proxy_request: Final = mapping(request_kwargs.get("proxy_server_request") or {}) + client: Final = mapping(proxy_request.get("headers") or {}) + session: Final = _scoped_session(request_kwargs) + forwarded: Final = { + name.lower(): value + for name, value in client.items() + if isinstance(value, str) and _forwarded_header(name, session is not None) + } + return {**forwarded, **({_SESSION_HEADER: session} if session else {})} + + +class SwitchyardRouter: + def __init__( + self, + model_name: str, + config: SwitchyardRouterConfig, + algorithm: Algorithm, + completion: JudgeCompletion, + decide: DecisionCompletion, + ) -> None: + self.model_name = model_name + self.config = config + self.algorithm = algorithm + self.completion = completion + self.decide = decide + + def _fallback( + self, reason: FallbackReason, messages: Sequence[Mapping[str, object]] | None + ) -> PreRoutingHookResponse: + return PreRoutingHookResponse( + model=self.config.models.capable, + messages=[dict(message) for message in messages] if messages is not None else None, + routing_decision={ + "router_type": "switchyard", + "router_model_name": self.model_name, + "routed_model": self.config.models.capable, + "switchyard_algorithm": self.config.algorithm, + "switchyard_fallback_reason": reason, + }, + ) + + def _internal_call( + self, + request: Mapping[str, object], + caller: Mapping[str, object], + bucket: Literal["metadata", "litellm_metadata"], + ) -> Mapping[str, object]: + caller_metadata: Final = mapping(caller.get(get_metadata_variable_name_from_kwargs(caller)) or {}) + return { + **request, + bucket: { + **forwarded_internal_call_metadata(caller_metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN), + INTERNAL_CALL_ORIGIN_METADATA_KEY: AUTOROUTER_CLASSIFIER_CALL_ORIGIN, + }, + "proxy_server_request": {"body": request}, + "num_retries": 0, + "turn_off_message_logging": effective_turn_off_message_logging(caller), + **parent_session_kwargs(caller), + } + + def _decision( + self, pick: str, evidence: object, messages: Sequence[Mapping[str, object]] | None + ) -> PreRoutingHookResponse: + decision: Final[StandardLoggingRoutingDecision] = { + "router_type": "switchyard", + "router_model_name": self.model_name, + "routed_model": pick, + "switchyard_algorithm": self.config.algorithm, + "switchyard_evidence": evidence, + } + return PreRoutingHookResponse( + model=pick, + messages=[dict(message) for message in messages] if messages is not None else None, + routing_decision=decision, + ) + + async def _answer_judge(self, call: ModelCall, caller: Mapping[str, object]) -> None: + try: + converted: Final = judge_response( + await self.completion( + self._internal_call(judge_request(call.request, call.models[0]), caller, "metadata") + ) + ) + except Exception as error: # noqa: BLE001 # judge/provider failures are inputs to Switchyard's fall-open decision + await call.fail(error) + else: + await call.respond(_resolved(converted)) + + async def _answer_decision(self, call: DecisionCall, caller: Mapping[str, object]) -> None: + try: + answer: Final = decision_response( + await self.decide( + self._internal_call(decision_request(call.request, call.model), caller, "litellm_metadata") + ), + call.request, + ) + except Exception as error: # noqa: BLE001 # decision/provider failures are inputs to Switchyard's fall-open decision + await call.fail(error) + else: + await call.respond(_resolved(answer)) + + def _outcome( + self, outcome: RoutingOutcome, messages: Sequence[Mapping[str, object]] | None + ) -> PreRoutingHookResponse: + pick: Final = outcome.selected_model_ids[0] + if pick not in (self.config.models.capable, self.config.models.efficient): + return self._fallback("target_outside_config", messages) + metadata: Final = outcome.metadata + evidence: Final[object] = metadata.evidence if metadata is not None else None + return self._decision(pick, evidence, messages) + + async def async_pre_routing_hook( + self, + model: str, + request_kwargs: Mapping[str, object], + messages: Sequence[Mapping[str, object]] | None = None, + input: object = None, + specific_deployment: bool | None = False, + call_type: str | None = None, + ) -> PreRoutingHookResponse | None: + from switchyard_rust.libsy import LibsyError, Step + + if model != self.model_name: + return None + try: + chat: Final = chat_messages(call_type, model, messages, input, request_kwargs) + request: Final = request_to_switchyard(chat) if chat is not None else None + except (ValueError, KeyError, TypeError): + return self._fallback("unsupported_content", messages) + if request is None: + return self._fallback("unsupported_request_type", messages) + judge: Final = self.config.models.judge + try: + async for step in self.algorithm.run_stream( + request, self.config.runtime_models(), headers=switchyard_headers(request_kwargs) + ): + match step: + case Step.CallModel(call) if isinstance( + self.config.classifier, LlmClassifierConfig + ) and call.models[:1] == [judge]: + await self._answer_judge(call, request_kwargs) + case Step.CallDecision(call) if ( + isinstance(self.config.classifier, DecisionClassifierConfig) and call.model == judge + ): + await self._answer_decision(call, request_kwargs) + case Step.CallModel(): + return self._fallback("unexpected_judge_call", messages) + case Step.CallDecision(): + return self._fallback("unexpected_decision_call", messages) + case Step.Done(outcome): + return self._outcome(outcome, messages) + except LibsyError: + return self._fallback("libsy_error", messages) + return self._fallback("no_decision", messages) diff --git a/litellm/router_strategy/switchyard_router/translation.py b/litellm/router_strategy/switchyard_router/translation.py new file mode 100644 index 00000000000..246f2ba3e41 --- /dev/null +++ b/litellm/router_strategy/switchyard_router/translation.py @@ -0,0 +1,232 @@ +from __future__ import annotations + +import json +from collections.abc import Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Final, Protocol + +from pydantic import TypeAdapter + +from litellm.types.decisions import ( + OpenAIChoiceAnswer, + OpenAIDecisionResponse, + systemone_choice_key, +) +from litellm.types.utils import CallTypes, ModelResponse + +if TYPE_CHECKING: + from switchyard_rust.libsy import LlmResponse + +_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_SEQUENCE: Final = TypeAdapter(tuple[object, ...]) +_CALL_TYPES: Final = frozenset(call_type.value for call_type in CallTypes) + + +def mapping(value: object) -> Mapping[str, object]: + return _MAPPING.validate_python(value) + + +def sequence(value: object) -> tuple[object, ...]: + return _SEQUENCE.validate_python(value) + + +def text_content(value: object) -> tuple[Mapping[str, object], ...]: + if value is None: + return () + if isinstance(value, str): + return ({"type": "text", "text": value},) + blocks: Final = tuple(mapping(block) for block in sequence(value)) + if any(block.get("type") != "text" or not isinstance(block.get("text"), str) for block in blocks): + raise ValueError("Switchyard routing supports text content only; media and audio are unsupported") + return tuple({"type": "text", "text": block["text"]} for block in blocks) + + +def tool_call(value: object) -> Mapping[str, object]: + call: Final = mapping(value) + function: Final = mapping(call.get("function")) + arguments: Final = function.get("arguments") + if call.get("type") != "function" or not isinstance(arguments, str): + raise ValueError("Switchyard requires function calls with JSON arguments") + return { + "type": "tool_call", + "id": call["id"], + "name": function["name"], + "arguments": _MAPPING.validate_json(arguments), + } + + +def message_to_switchyard(value: object) -> Mapping[str, object]: + message: Final = mapping(value) + if message.get("audio") is not None: + raise ValueError("Switchyard routing does not support audio") + role: Final = message.get("role") + content: Final = text_content(message.get("content")) + if role == "tool": + return { + "role": "tool", + "content": ( + { + "type": "tool_result", + "tool_call_id": message["tool_call_id"], + "content": content, + "is_error": message.get("is_error"), + }, + ), + } + if role not in ("user", "system", "developer", "assistant"): + raise ValueError("Switchyard received an unsupported message role") + calls: Final = tuple(tool_call(call) for call in sequence(message.get("tool_calls") or ())) + return {"role": role, "content": (*content, *calls)} + + +class ChatConverter(Protocol): + def get_structured_messages( + self, + data: dict[str, object], # mutable-ok: matches the request converters' existing dict parameter + *, + keep_tool_errors: bool = False, + ) -> Sequence[object] | None: ... + + +def _instantiate(converter: Callable[[], ChatConverter]) -> ChatConverter: + return converter() + + +def chat_messages( + call_type: str | None, + model: str, + messages: Sequence[Mapping[str, object]] | None, + input: object, + request_kwargs: Mapping[str, object], +) -> tuple[Mapping[str, object], ...] | None: + from litellm.llms import load_guardrail_translation_mappings + + converter: Final = ( + load_guardrail_translation_mappings().get(CallTypes(call_type)) if call_type in _CALL_TYPES else None + ) + if converter is None: + return None + request: Final = { + **request_kwargs, + "model": model, + **({"messages": [dict(message) for message in messages]} if messages is not None else {}), + **({"input": input} if input is not None else {}), + } + return tuple(map(mapping, _instantiate(converter).get_structured_messages(request, keep_tool_errors=True) or ())) + + +def request_to_switchyard(messages: Sequence[Mapping[str, object]]) -> Mapping[str, object]: + if not messages: + raise ValueError("Switchyard requires at least one message") + return { + "model": "auto", + "messages": tuple(message_to_switchyard(message) for message in messages), + "stream": False, + } + + +def judge_message(value: object) -> Mapping[str, object]: + message: Final = mapping(value) + blocks: Final = text_content(message.get("content")) + return {"role": message["role"], "content": "".join(str(block["text"]) for block in blocks)} + + +def judge_request(request: Mapping[str, object], model: str) -> Mapping[str, object]: + output: Final = mapping(request.get("output") or {}) + sampling: Final = mapping(request.get("sampling") or {}) + options: Final = { + "max_tokens": output.get("max_output_tokens"), + "response_format": output.get("response_format"), + "temperature": sampling.get("temperature"), + "top_p": sampling.get("top_p"), + } + return { + "model": model, + "stream": False, + "messages": [ + judge_message(message) + for message in ( + *sequence(request.get("instructions") or ()), + *sequence(request.get("messages") or ()), + ) + ], + **{key: value for key, value in options.items() if value is not None}, + } + + +def judge_response(response: ModelResponse) -> LlmResponse.Agg: + from switchyard_rust.libsy import LlmResponse + + body: Final = _MAPPING.validate_json(response.model_dump_json(exclude_unset=True)) + choices: Final = tuple(mapping(choice) for choice in sequence(body["choices"])) + return LlmResponse.Agg( + { + "id": response.id, + "model": response.model, + "outputs": tuple(message_to_switchyard(choice["message"]) for choice in choices), + } + ) + + +def _decision_text(value: object) -> str: + return value if isinstance(value, str) else json.dumps(value) + + +def _decision_question(name: str, question: Mapping[str, object]) -> Mapping[str, object]: + kind: Final = mapping(question["kind"]) + if kind.get("type") != "choice": + raise ValueError("Switchyard decision routing asks choice questions only") + return { + "type": "choice", + "name": name, + "instructions": _decision_text(question.get("instructions")), + "choices": [ + { + "value": option["id"], + **({"description": _decision_text(option["description"])} if option.get("description") else {}), + } + for option in map(mapping, sequence(mapping(kind["data"])["options"])) + ], + } + + +def decision_request(request: Mapping[str, object], model: str) -> Mapping[str, object]: + questions: Final = mapping(request["questions"]) + return { + "model": model, + "input": _decision_text(request["context"]), + "questions": [_decision_question(name, mapping(question)) for name, question in questions.items()], + } + + +def _choice_value(answer: object) -> Mapping[str, object]: + if not isinstance(answer, OpenAIChoiceAnswer): + raise TypeError("The decision model refused or returned a non-choice answer") + return { + "type": "choice", + "data": { + "selected": systemone_choice_key(answer.choice), + "probabilities": {systemone_choice_key(p.value): p.probability for p in answer.probabilities}, + }, + } + + +def decision_response(response: OpenAIDecisionResponse, request: Mapping[str, object]) -> Mapping[str, object]: + names: Final = tuple(mapping(request["questions"])) + if len(response.answers) != len(names): + raise ValueError("The decision model did not answer every question") + return { + "id": None, + "model": response.model, + "answers": { + answer.name or name: { + "value": _choice_value(answer), + "provider_confidence": getattr(answer, "confidence", None), + } + for name, answer in zip(names, response.answers, strict=True) + }, + "usage": { + "input_tokens": response.usage.input_tokens, + "output_tokens": response.usage.output_tokens, + "total_tokens": response.usage.total_tokens, + }, + } diff --git a/litellm/router_utils/cooldown_callbacks.py b/litellm/router_utils/cooldown_callbacks.py index 25821703029..9e90b72ba1b 100644 --- a/litellm/router_utils/cooldown_callbacks.py +++ b/litellm/router_utils/cooldown_callbacks.py @@ -37,10 +37,10 @@ async def router_cooldown_event_callback( deployment_id, ) return - _litellm_params: Final = _deployment["litellm_params"] + _litellm_params: Final = _deployment.litellm_params temp_litellm_params = copy.deepcopy(_litellm_params) temp_litellm_params = dict(temp_litellm_params) - _model_name: Final = _deployment.get("model_name", None) or "" + _model_name: Final = _deployment.model_name or "" _api_base: Final = litellm.get_api_base(model=_model_name, optional_params=temp_litellm_params) or "" model_info: Final = _deployment["model_info"] model_id: Final = model_info.id diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 20f9e114747..068782f7a08 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -5,7 +5,7 @@ from dataclasses import dataclass, field from datetime import datetime from enum import Enum from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import httpx from pydantic import TypeAdapter @@ -798,7 +798,7 @@ async def run_async_fallback( fallback_depth: int, include_fallback_errors: bool = False, **kwargs, -) -> Any: +) -> object: """ Loops through all the fallback model groups and calls kwargs["original_function"] with the arguments and keyword arguments provided. diff --git a/litellm/router_utils/handle_error.py b/litellm/router_utils/handle_error.py index 97bf4710fca..a1e6d2b2cad 100644 --- a/litellm/router_utils/handle_error.py +++ b/litellm/router_utils/handle_error.py @@ -51,7 +51,7 @@ async def send_llm_exception_alert( # the proxy is already instrumented to send LLM API call failures return - litellm_debug_info: Final = getattr(original_exception, "litellm_debug_info", None) + litellm_debug_info: Final[str | None] = getattr(original_exception, "litellm_debug_info", None) exception_str = str(original_exception) if litellm_debug_info is not None: exception_str += litellm_debug_info diff --git a/litellm/router_utils/health_state_cache.py b/litellm/router_utils/health_state_cache.py index 1c9cb10a48b..c4cc9350cfa 100644 --- a/litellm/router_utils/health_state_cache.py +++ b/litellm/router_utils/health_state_cache.py @@ -9,6 +9,7 @@ import time from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final +from pydantic import TypeAdapter from typing_extensions import TypedDict from litellm import verbose_logger @@ -25,6 +26,9 @@ else: Span = Any +_CACHED_HEALTH_SNAPSHOT: Final = TypeAdapter(object) + + class DeploymentHealthStateValue(TypedDict): is_healthy: bool timestamp: float @@ -92,7 +96,7 @@ class DeploymentHealthCache: str(e), ) - def _extract_unhealthy_ids(self, raw: Any) -> set[str]: + def _extract_unhealthy_ids(self, raw: object) -> set[str]: """Given raw cache value, return set of non-stale unhealthy deployment IDs.""" if not raw or not isinstance(raw, dict): return set() @@ -109,7 +113,7 @@ class DeploymentHealthCache: async def async_get_unhealthy_deployment_ids(self, parent_otel_span: Span | None = None) -> set[str]: """Return set of deployment IDs currently marked unhealthy and not stale.""" try: - raw: Final = await self.cache.async_get_cache(key=self.CACHE_KEY) + raw: Final = _CACHED_HEALTH_SNAPSHOT.validate_python(await self.cache.async_get_cache(key=self.CACHE_KEY)) return self._extract_unhealthy_ids(raw) except Exception as e: verbose_logger.debug( @@ -122,7 +126,7 @@ class DeploymentHealthCache: def get_unhealthy_deployment_ids(self, parent_otel_span: Span | None = None) -> set[str]: """Sync version: return set of deployment IDs currently marked unhealthy and not stale.""" try: - raw: Final = self.cache.get_cache(key=self.CACHE_KEY) + raw: Final = _CACHED_HEALTH_SNAPSHOT.validate_python(self.cache.get_cache(key=self.CACHE_KEY)) return self._extract_unhealthy_ids(raw) except Exception as e: verbose_logger.debug( diff --git a/litellm/router_utils/model_group_pricing.py b/litellm/router_utils/model_group_pricing.py new file mode 100644 index 00000000000..afadbc3eb38 --- /dev/null +++ b/litellm/router_utils/model_group_pricing.py @@ -0,0 +1,69 @@ +import math +from collections.abc import Mapping +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info +from litellm.litellm_core_utils.llm_cost_calc.utils import get_vertex_regional_endpoint_uplift +from litellm.types.router import DeploymentPricing, PriceTier + +_RAW_PRICE_TIERS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +def _numeric_text(value: str) -> float | None: + try: + parsed: Final = float(value) + except ValueError: + return None + return parsed if math.isfinite(parsed) else None + + +def _cost_value(value: object) -> float | None: + match value: + case bool(): + return None + case int() | float(): + return float(value) + case str(): + return _numeric_text(value) + case _: + return None + + +def _cost_fields(fields: Mapping[str, object], token_uplift: float) -> Mapping[str, float]: + return { + name: cost * token_uplift if "token" in name else cost + for name, value in fields.items() + if "cost" in name and (cost := _cost_value(value)) is not None + } + + +def _price_tiers(tiered_pricing: object, token_uplift: float) -> tuple[PriceTier, ...]: + try: + return tuple( + PriceTier.model_validate({"range": tier.get("range"), "costs": _cost_fields(tier, token_uplift)}) + for tier in _RAW_PRICE_TIERS.validate_python(tiered_pricing) + ) + except ValidationError: + return () + + +def deployment_pricing( + provider: str, + deployment_id: str | None, + deployment_model: str, + resolved_model_info: Mapping[str, object], + vertex_location: str | None, +) -> DeploymentPricing | None: + """Prices from the entry billing charges this deployment at, so a cost-map variant its own pricing shadows + is never shown. A deployment that declares no pricing is billed at its resolved cost-map entry. Token prices + carry the Vertex regional-endpoint uplift billing applies at the deployment's location.""" + billed_model_info: Final = deployment_pricing_model_info(model_id=deployment_id, deployment_model=deployment_model) + model_info: Final = resolved_model_info if billed_model_info is None else billed_model_info + token_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location) + costs: Final = _cost_fields(model_info, token_uplift) + tiers: Final = _price_tiers(model_info.get("tiered_pricing"), token_uplift) + if not costs and not tiers: + return None + return DeploymentPricing(provider=provider, costs=costs, tiered_pricing=tiers) diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index 062886b44fd..9badee854fb 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -148,7 +148,7 @@ class EncryptedContentAffinityCheck(CustomLogger): return any("encrypted_content_affinity" in checks for checks in model_group_affinity_config.values()) - def _is_enabled_for_model_group(self, model_group: str) -> bool: + def is_enabled_for_model_group(self, model_group: str) -> bool: group_checks: Final = self.model_group_affinity_config.get(model_group) return self.enable_global_affinity or ( group_checks is not None and "encrypted_content_affinity" in group_checks @@ -530,7 +530,7 @@ class EncryptedContentAffinityCheck(CustomLogger): """ request_kwargs = request_kwargs or {} typed_healthy_deployments: Final = cast(list[dict[str, object]], healthy_deployments) - if not self._is_enabled_for_model_group(model): + if not self.is_enabled_for_model_group(model): return typed_healthy_deployments typed_request_kwargs: Final = _REQUEST_KWARGS_ADAPTER.validate_python(request_kwargs) target_order: Final = typed_request_kwargs.get("_target_order") diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index e5d944b3aef..2f0cbc0d474 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -4,23 +4,37 @@ Check if prompt caching is valid for a given deployment Route to previously cached model id, if valid """ -from collections.abc import Callable, Mapping, Sequence -from typing import Final, TypedDict, cast +import secrets +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final, NamedTuple, Protocol, TypedDict, cast +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly from litellm import verbose_logger from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT from litellm.integrations.anthropic_cache_control_hook import ( AllToolParamValues, AnthropicCacheControlHook, ) from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.routing_common_utils import deployment_ids_failed_in_request 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.types.llms.openai import AllMessageValues, ChatCompletionToolParam +from litellm.responses.litellm_completion_transformation.transformation import ( + RESPONSES_TOOL_TYPES_WITHOUT_CHAT_FORM, + RESPONSES_WEB_SEARCH_TOOL_TYPES, + LiteLLMCompletionResponsesConfig, + ResponsesToolChatForm, + ResponseTools, +) +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam, ResponseInputParam from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import ( get_base_model_from_metadata, @@ -28,7 +42,7 @@ from litellm.utils import ( is_prompt_caching_valid_prompt, ) -from ..prompt_caching_cache import PromptCachingCache +from ..prompt_caching_cache import PROMPT_CACHE_PIN_TTL_SECONDS, PromptCachingCache class _PromptCachingSuccessEventKwargs(TypedDict, total=False): @@ -36,7 +50,165 @@ class _PromptCachingSuccessEventKwargs(TypedDict, total=False): litellm_params: ReadOnly[Mapping[str, object] | None] -def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int: +class _ResponsesPin(NamedTuple): + cache_key: str + marked_ttl_seconds: int + cacheable_deployment_ids: tuple[str, ...] + + +class ResponsesPinsInFlight(Protocol): + def get_cache(self, key: str) -> object: ... + + def set_cache(self, key: str, value: object) -> None: ... + + def delete_cache(self, key: str) -> None: ... + + +class _ResponsesPrompt(NamedTuple): + messages: list[AllMessageValues] # mutable-ok: the prefix helpers take chat messages as a list + tools: list[AllToolParamValues] | None # mutable-ok: the prefix helpers take chat tools as a list + key_tools: list[ChatCompletionToolParam] | None # mutable-ok: the prefix helpers take chat tools as a list + + +_RESPONSES_PIN_METADATA_KEY: Final = "responses_prompt_cache_pin" +_RESPONSES_PINS_IN_FLIGHT_LIMIT: Final = 10_000 +_RESPONSES_PIN_LIFETIME_SECONDS: Final = 3600 +_RETRY_RECORDS: Final = TypeAdapter(tuple[object, ...]) +_RESPONSES_CALL_TYPES: Final = frozenset((CallTypes.responses.value, CallTypes.aresponses.value)) +_RESPONSES_TOOL_TYPES_KEPT_AS_SENT: Final = RESPONSES_TOOL_TYPES_WITHOUT_CHAT_FORM | RESPONSES_WEB_SEARCH_TOOL_TYPES + + +def new_responses_pins_in_flight() -> ResponsesPinsInFlight: + return InMemoryCache( + max_size_in_memory=_RESPONSES_PINS_IN_FLIGHT_LIMIT, default_ttl=_RESPONSES_PIN_LIFETIME_SECONDS + ) + + +def _key_chat_tool(source: Mapping[str, object], chat_tool: Mapping[str, object]) -> Mapping[str, object]: + function: Final = as_str_mapping(chat_tool.get("function")) + if source.get("type") != "function" or function is None: + return chat_tool + return {**chat_tool, "function": {key: source[key] for key in function if key in source}} + + +def _key_chat_tools( + sources: Sequence[Mapping[str, object]], forms: Sequence[ResponsesToolChatForm] +) -> Iterator[Mapping[str, object]]: + for source, form in zip(sources, forms, strict=True): + yield from (_key_chat_tool(source, chat_tool) for chat_tool in form.chat_tools) + + +def _responses_prompt(request_data: Mapping[str, object]) -> _ResponsesPrompt | None: + request_input: Final = request_data.get("input") + request_tools: Final = request_data.get("tools") + if request_data.get("previous_response_id") is not None or not isinstance(request_input, str | list): + return None + listed_tools: Final = ( + cast(ResponseTools, request_tools) # cast-ok: a bad tool raises inside the guard below and is caught + if isinstance(request_tools, list) + else None + ) + try: + tools_kept_as_sent: Final = tuple( + tool for tool in listed_tools or () if tool.get("type") in _RESPONSES_TOOL_TYPES_KEPT_AS_SENT + ) + convertible_tools: Final = ( + [tool for tool in listed_tools if tool.get("type") not in _RESPONSES_TOOL_TYPES_KEPT_AS_SENT] + if listed_tools is not None + else None + ) + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # its input type is untyped + input=cast(str | ResponseInputParam, request_input), # cast-ok: a bad item raises and is caught here + responses_api_request={"instructions": request_data.get("instructions")}, + item_cache_control=True, + only_since_compaction=True, + ) + forms: Final = LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(convertible_tools) + except Exception as error: # noqa: BLE001 # the provider judges the input, a routing preference never does + verbose_logger.debug("Prompt Caching: skipping an input with no chat form: %s", type(error).__name__) + return None + tools: Final = list(chain.from_iterable(form.chat_tools for form in forms)) + key_tools: Final = [*_key_chat_tools(convertible_tools or (), forms), *tools_kept_as_sent] + return _ResponsesPrompt( + messages=cast(list[AllMessageValues], messages), # cast-ok: the bridge's chat messages, read as plain JSON + tools=cast(list[AllToolParamValues], tools) or None, # cast-ok: the bridge's chat tools, read as plain JSON + key_tools=cast(list[ChatCompletionToolParam], key_tools) or None, # cast-ok: tools are hashed as plain JSON + ) + + +def _retry_records(metadata: Mapping[str, object]) -> tuple[object, ...]: + try: + return _RETRY_RECORDS.validate_python(metadata.get("previous_models", ())) + except ValidationError: + return () + + +def _record_deployment_ids(records: tuple[object, ...]) -> Iterator[str]: + deployment_ids: Final = ((as_str_mapping(record) or {}).get("deployment_id") for record in records) + return (deployment_id for deployment_id in deployment_ids if isinstance(deployment_id, str)) + + +def _retried_deployment_ids(metadata: Mapping[str, object]) -> frozenset[str] | None: + records: Final = _retry_records(metadata) + retry_count: Final = metadata.get("request_retry_count") + if type(retry_count) is int and retry_count > len(records): + return None + return frozenset(_record_deployment_ids(records)) + + +def _deployment_ids_tried_in_request(request_data: Mapping[str, object]) -> frozenset[str] | None: + retried: Final = tuple( + _retried_deployment_ids(as_str_mapping(request_data.get(bucket)) or {}) + for bucket in ("litellm_metadata", "metadata") + ) + known: Final = tuple(ids for ids in retried if ids is not None) + if len(known) < len(retried): + return None + return deployment_ids_failed_in_request(request_data).union(*known) + + +def _previous_response_model_id(request_data: Mapping[str, object]) -> str | None: + previous_response_id: Final = request_data.get("previous_response_id") + model_id: Final = ( + ResponsesAPIRequestUtils.get_model_id_from_response_id(previous_response_id) + if isinstance(previous_response_id, str) + else None + ) + tried: Final = _deployment_ids_tried_in_request(request_data) if model_id is not None else None + return model_id if model_id is not None and tried is not None and model_id not in tried else None + + +def _deployment_min_tokens(deployment: Mapping[str, Mapping[str, object]], fallback: int) -> int: + deployment_model: Final = deployment.get("litellm_params", {}).get("model") + return get_prompt_cache_min_tokens(model=deployment_model) if isinstance(deployment_model, str) else fallback + + +async def _cacheable_deployment_ids( + healthy_deployments: Sequence[Mapping[str, Mapping[str, object]]], + reached_minimum: int, + reaches: Callable[[int], Awaitable[bool]], +) -> tuple[str, ...]: + minimums: Final = tuple( + (str(deployment["model_info"]["id"]), _deployment_min_tokens(deployment, reached_minimum)) + for deployment in healthy_deployments + ) + reached: Final = MappingProxyType( + { + minimum: minimum <= reached_minimum or await reaches(minimum) + for minimum in frozenset(minimum for _, minimum in minimums) + } + ) + return tuple(sorted(deployment_id for deployment_id, minimum in minimums if reached[minimum])) + + +def _deployment_models(healthy_deployments: Sequence[Mapping[str, object]]) -> Iterator[str]: + models: Final = ( + (as_str_mapping(deployment.get("litellm_params")) or {}).get("model") for deployment in healthy_deployments + ) + return (model for model in models if isinstance(model, str) and model) + + +def _get_min_token_count_for_deployments(healthy_deployments: Sequence[Mapping[str, object]]) -> int: """ Returns the lowest minimum cacheable prefix across a model group. @@ -52,11 +224,7 @@ def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int losing a cache hit it had earned. The lowest can only cost a lookup that finds nothing. """ return min( - ( - get_prompt_cache_min_tokens(model=deployment["litellm_params"]["model"]) - for deployment in healthy_deployments - if deployment.get("litellm_params", {}).get("model") - ), + (get_prompt_cache_min_tokens(model=model) for model in _deployment_models(healthy_deployments)), default=DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, ) @@ -75,10 +243,20 @@ def _honors_one_hour_cache_ttl( class PromptCachingDeploymentCheck(CustomLogger): - def __init__(self, cache: DualCache, is_priority_group: Callable[[str], bool] | None = None): + def __init__( + self, + cache: DualCache, + is_priority_group: Callable[[str], bool] | None = None, + responses_pins_in_flight: ResponsesPinsInFlight | None = None, + group_has_alternatives: Callable[[str], bool] | None = None, + ): super().__init__() self.cache = cache self.is_priority_group = is_priority_group + self.group_has_alternatives: Final = group_has_alternatives + self.responses_pins_in_flight: Final = ( + responses_pins_in_flight if responses_pins_in_flight is not None else new_responses_pins_in_flight() + ) async def async_filter_deployments( self, @@ -88,43 +266,81 @@ class PromptCachingDeploymentCheck(CustomLogger): request_kwargs: dict | None = None, parent_otel_span: Span | None = None, ) -> list[dict]: - if self.is_priority_group is not None and self.is_priority_group(model): - return healthy_deployments - if request_kwargs is not None and request_kwargs.get("_target_order") is not None: - return healthy_deployments - if not healthy_deployments[1:]: - return healthy_deployments + preferred_id: Final = await self.async_preferred_deployment_id( + model, healthy_deployments, messages, request_kwargs + ) + preferred: Final = [ + deployment for deployment in healthy_deployments if deployment["model_info"]["id"] == preferred_id + ] + return preferred[:1] or healthy_deployments - request_data: Final = ( - cast( # cast-ok: request_kwargs is a bare dict; this only reads its keys - Mapping[str, object], request_kwargs - ) - if request_kwargs is not None + async def async_preferred_deployment_id( + self, + model: str, + healthy_deployments: Sequence[Mapping[str, Mapping[str, object]]], + messages: list[AllMessageValues] | None, # mutable-ok: the request's messages list, read as is + request_data: Mapping[str, object] | None, + ) -> str | None: + if self.is_priority_group is not None and self.is_priority_group(model): + return None + if request_data is not None and request_data.get("_target_order") is not None: + return None + pin_metadata: Final = ( + request_data.get("litellm_metadata") if messages is None and request_data is not None else None + ) + single_healthy: Final = not healthy_deployments[1:] + if single_healthy and not ( + isinstance(pin_metadata, dict) + and self.group_has_alternatives is not None + and self.group_has_alternatives(model) + ): + return None + previous_attempt_token: Final = (as_str_mapping(pin_metadata) or {}).get(_RESPONSES_PIN_METADATA_KEY) + previous_response_model_id: Final = ( + _previous_response_model_id(request_data) + if request_data is not None and isinstance(pin_metadata, dict) + else None + ) + if previous_response_model_id is not None: + return previous_response_model_id + responses_prompt: Final = ( + await offload_token_count(_responses_prompt)(request_data) + if request_data is not None and isinstance(pin_metadata, dict) else None ) request_tools: Final = ( - cast( # cast-ok: request_kwargs is untyped; the stand-down scan duck-types every tool it reads + responses_prompt.tools + if responses_prompt is not None + else cast( # cast-ok: request_kwargs is untyped; the stand-down scan duck-types every tool it reads list[AllToolParamValues] | None, request_data.get("tools") ) - if request_data is not None + if request_data is not None and messages is not None else None ) + request_messages: Final = responses_prompt.messages if responses_prompt is not None else messages prompt_messages: Final = ( PromptCachingCache.prepend_system_prompt( - messages, request_data.get("system") if request_data is not None else None + request_messages, request_data.get("system") if request_data is not None else None ) - if messages is not None + if request_messages is not None else None ) - if prompt_messages is not None and await offload_token_count(is_prompt_caching_valid_prompt)( - messages=prompt_messages, - model=model, - min_token_count=_get_min_token_count_for_deployments(healthy_deployments), - tools=cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes the counter duck-types - list[ChatCompletionToolParam] | None, - request_tools, - ), - ): + + async def reaches(minimum: int) -> bool: + return await offload_token_count( + is_prompt_caching_valid_prompt + )( + messages=prompt_messages, + model=model, + min_token_count=minimum, + tools=cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes the counter duck-types + list[ChatCompletionToolParam] | None, + request_tools, + ), + ) + + group_minimum: Final = _get_min_token_count_for_deployments(healthy_deployments) + if prompt_messages is not None and await reaches(group_minimum): prompt_cache: Final = PromptCachingCache( cache=self.cache, ) @@ -134,34 +350,49 @@ class PromptCachingDeploymentCheck(CustomLogger): ## be derived from the messages as they will be sent, not as they arrive here. affinity_messages: Final = AnthropicCacheControlHook.messages_with_default_injections( messages=prompt_messages, - models=( - deployment["litellm_params"]["model"] - for deployment in healthy_deployments - if isinstance(deployment.get("litellm_params"), dict) and deployment["litellm_params"].get("model") - ), + models=_deployment_models(healthy_deployments), tools=request_tools, enable_prompt_caching=( - request_kwargs.get("enable_prompt_caching") is True if request_kwargs is not None else None + request_data.get("enable_prompt_caching") is True if request_data is not None else None ), - request_kwargs=request_kwargs, + request_kwargs=request_data, ) - - model_id_dict: Final = await prompt_cache.async_get_model_id( - messages=affinity_messages, - tools=cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes - Sequence[ChatCompletionToolParam] | None, - request_tools, - ), + prefix_tools: Final = ( + responses_prompt.key_tools + if responses_prompt is not None + else cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes + Sequence[ChatCompletionToolParam] | None, request_tools + ) ) - 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: - return [deployment] + positions: Final = await PromptCachingCache.async_prefix_positions(affinity_messages, prefix_tools) + if responses_prompt is not None and isinstance(pin_metadata, dict) and positions: + if isinstance(previous_attempt_token, str): + self.responses_pins_in_flight.delete_cache(previous_attempt_token) + pin_token: Final = secrets.token_hex(16) + self.responses_pins_in_flight.set_cache( + pin_token, + _ResponsesPin( + cache_key=positions[-1].cache_key, + marked_ttl_seconds=await offload_token_count(PromptCachingCache.pin_ttl_seconds)( + affinity_messages, prefix_tools + ), + cacheable_deployment_ids=await _cacheable_deployment_ids( + healthy_deployments, + group_minimum, + reaches, + ), + ), + ) + pin_metadata[_RESPONSES_PIN_METADATA_KEY] = pin_token + if single_healthy: + return None - return healthy_deployments + model_id_dict: Final = await prompt_cache.async_get_model_id_at(positions) + tried: Final = _deployment_ids_tried_in_request(request_data or {}) if model_id_dict is not None else None + if model_id_dict is not None and tried is not None and model_id_dict["model_id"] not in tried: + return model_id_dict["model_id"] + + return None def exempt_from_message_redaction(self) -> bool: return True @@ -173,6 +404,7 @@ class PromptCachingDeploymentCheck(CustomLogger): start_time: object, end_time: object, ) -> None: + recorded_pin: Final = self._take_responses_pin(kwargs) standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: @@ -180,6 +412,10 @@ class PromptCachingDeploymentCheck(CustomLogger): call_type: Final = standard_logging_object["call_type"] + if call_type in _RESPONSES_CALL_TYPES: + await self._pin_responses_request(recorded_pin, kwargs, standard_logging_object) + return + if ( call_type != CallTypes.completion.value and call_type != CallTypes.acompletion.value @@ -241,3 +477,33 @@ class PromptCachingDeploymentCheck(CustomLogger): ) return + + def _take_responses_pin(self, kwargs: _PromptCachingSuccessEventKwargs) -> object: + litellm_params: Final = kwargs.get("litellm_params") or {} + metadata: Final = litellm_params.get("litellm_metadata") + pin_token: Final = (as_str_mapping(metadata) or {}).get(_RESPONSES_PIN_METADATA_KEY) + if not isinstance(pin_token, str): + return None + if isinstance(metadata, dict): + metadata.pop(_RESPONSES_PIN_METADATA_KEY, None) # pyright: ignore[reportUnknownMemberType] # the caller's metadata dict + recorded_pin: Final = self.responses_pins_in_flight.get_cache(pin_token) + self.responses_pins_in_flight.delete_cache(pin_token) + return recorded_pin + + async def _pin_responses_request( + self, + recorded_pin: object, + kwargs: _PromptCachingSuccessEventKwargs, + standard_logging_object: StandardLoggingPayload, + ) -> None: + litellm_params: Final = kwargs.get("litellm_params") or {} + model_id: Final = standard_logging_object["model_id"] + if not isinstance(recorded_pin, _ResponsesPin) or model_id not in recorded_pin.cacheable_deployment_ids: + return + await PromptCachingCache(cache=self.cache).async_set_pin( + model_id, + recorded_pin.cache_key, + recorded_pin.marked_ttl_seconds + if _honors_one_hour_cache_ttl(standard_logging_object, litellm_params) + else PROMPT_CACHE_PIN_TTL_SECONDS, + ) diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index 1d91255e28c..e01d744dfb6 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -432,18 +432,18 @@ class PromptCachingCache: if not positions: return - ttl: Final = ( + await self.async_set_pin( + model_id, + positions[-1].cache_key, await offload_token_count(PromptCachingCache.pin_ttl_seconds)(messages, tools) if honors_one_hour_ttl - else PROMPT_CACHE_PIN_TTL_SECONDS + else PROMPT_CACHE_PIN_TTL_SECONDS, ) - self.cache.in_memory_cache.delete_cache(positions[-1].cache_key) + + async def async_set_pin(self, model_id: str, cache_key: str, ttl: int) -> None: + self.cache.in_memory_cache.delete_cache(cache_key) with service_target(_PROMPT_CACHE_PINS_TARGET): - await self.cache.async_set_cache( - positions[-1].cache_key, - PromptCachingCacheValue(model_id=model_id), - ttl=ttl, - ) + await self.cache.async_set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=ttl) async def async_get_model_id( self, @@ -455,7 +455,10 @@ class PromptCachingCache: same way the provider cache does, so a breakpoint that moved forward since the last turn still lands on the deployment whose cache holds the earlier prefix. """ - cache_keys: Final = _lookback_keys(await PromptCachingCache.async_prefix_positions(messages, tools)) + return await self.async_get_model_id_at(await PromptCachingCache.async_prefix_positions(messages, tools)) + + async def async_get_model_id_at(self, positions: tuple[PrefixPosition, ...]) -> PromptCachingCacheValue | None: + cache_keys: Final = _lookback_keys(positions) if not cache_keys: return None diff --git a/litellm/skills/main.py b/litellm/skills/main.py index 6f57bd1bc96..3262a05ca60 100644 --- a/litellm/skills/main.py +++ b/litellm/skills/main.py @@ -5,7 +5,7 @@ Provides create, list, get, and delete operations for skills import asyncio import contextvars -from collections.abc import Coroutine, Mapping +from collections.abc import Coroutine, Mapping, Sequence from functools import partial from typing import Any, Final @@ -14,6 +14,7 @@ import httpx import litellm from litellm.constants import request_timeout from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import is_str_keyed_dict from litellm.llms.base_llm.skills.transformation import BaseSkillsAPIConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.llms.anthropic_skills import ( @@ -45,15 +46,15 @@ def _get_user_api_key_auth_from_kwargs(kwargs: Mapping[str, object]) -> Any | No def _get_skill_request_metadata( kwargs: Mapping[str, object], - extra_body: dict[str, Any] | None, -) -> dict[str, Any] | None: - if extra_body and isinstance(extra_body.get("metadata"), dict): - return extra_body["metadata"] + extra_body: Mapping[str, object] | None, +) -> dict[str, object] | None: + extra_body_metadata: Final = extra_body.get("metadata") if extra_body else None + if is_str_keyed_dict(extra_body_metadata): + return extra_body_metadata metadata: Final = kwargs.get("metadata") - if isinstance(metadata, dict) and isinstance(metadata.get("requester_metadata"), dict): - return metadata["requester_metadata"] - return None + requester_metadata: Final = metadata.get("requester_metadata") if is_str_keyed_dict(metadata) else None + return requester_metadata if is_str_keyed_dict(requester_metadata) else None def _get_litellm_skills_handler(): @@ -70,7 +71,7 @@ def _get_litellm_skills_handler(): @client async def acreate_skill( - files: list[Any] | None = None, + files: Sequence[object] | None = None, display_title: str | None = None, extra_headers: dict[str, str] | None = None, extra_query: Mapping[str, object] | None = None, @@ -133,7 +134,7 @@ async def acreate_skill( @client def create_skill( - files: list[Any] | None = None, + files: Sequence[object] | None = None, display_title: str | None = None, extra_headers: dict[str, str] | None = None, extra_query: Mapping[str, object] | None = None, diff --git a/litellm/text_completions/streaming_iterator.py b/litellm/text_completions/streaming_iterator.py index 2b5b4d52274..daf65b77a85 100644 --- a/litellm/text_completions/streaming_iterator.py +++ b/litellm/text_completions/streaming_iterator.py @@ -1,6 +1,7 @@ from typing import Final from litellm.constants import VLLM_CHOICE_FIELDS +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact from litellm.types.utils import Choices, ModelResponse, ModelResponseStream, TextChoices, TextCompletionResponse @@ -43,9 +44,10 @@ class TextCompletionStreamWrapper: if chunk["choices"]: chat_choice: Final = chunk["choices"][0] - text_choices: Final = TextChoices( - **{field: chat_choice[field] for field in VLLM_CHOICE_FIELDS if field in chat_choice} + choice_fields: Final = decoded_json_object( + {field: chat_choice[field] for field in VLLM_CHOICE_FIELDS if field in chat_choice} ) + text_choices: Final = TextChoices(index=0, **choice_fields) delta: Final = chat_choice["delta"] text_choices["text"] = delta["content"] text_choices["reasoning_content"] = delta.get("reasoning_content") diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 718a651bdee..e311f2ee5f9 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -1050,6 +1050,14 @@ GuardrailNumRetries: TypeAlias = Annotated[ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch update guardrails + mcp_block_behavior: Literal["tool_result", "raise"] | None = Field( + default=None, + description=( + "How to handle guardrail blocks when executing MCP tools for an LLM request. " + "'raise' propagates the guardrail error to the client. 'tool_result' or None " + "returns the block as a tool result and lets the LLM continue." + ), + ) api_key: str | None = Field(default=None, description="API key for the guardrail service") api_base: str | None = Field(default=None, description="Base URL for the guardrail service API") @@ -1162,12 +1170,13 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up description=( "For /v1/realtime sessions: 'warn' speaks the violation message and continues; " "'end_session' speaks the message and closes the connection. " + "Bedrock Nova Sonic cannot cancel a spoken turn's answer, so a blocked spoken turn always ends its session. " "For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning." ), ) realtime_violation_message: str | None = Field( default=None, - description="The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set.", + description="The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to the guardrail's block reason if not set.", ) # Model Armor params diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 9f63d27f5f3..c04430322c4 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -207,6 +207,7 @@ class DeploymentOptions: @dataclass(frozen=True, slots=True, kw_only=True) class SpecializedRouterOptions: + switchyard_router_config: Mapping[str, object] | None = None auto_router_config_path: str | None = None auto_router_config: str | None = None auto_router_default_model: str | None = None diff --git a/litellm/types/llms/anthropic_skills.py b/litellm/types/llms/anthropic_skills.py index d11fd32f88b..690c2226a73 100644 --- a/litellm/types/llms/anthropic_skills.py +++ b/litellm/types/llms/anthropic_skills.py @@ -2,6 +2,7 @@ Type definitions for Anthropic Skills API """ +from collections.abc import Sequence from typing import Any from typing_extensions import TypedDict @@ -16,7 +17,7 @@ class CreateSkillRequest(TypedDict, total=False): display_title: str | None """Display title for the skill (optional)""" - files: list[Any] | None + files: Sequence[object] | None """Files to upload for the skill. All files must be in the same top-level directory and must include a SKILL.md file at the root.""" diff --git a/litellm/types/llms/custom_http.py b/litellm/types/llms/custom_http.py index 47f80c52845..4220a67bddb 100644 --- a/litellm/types/llms/custom_http.py +++ b/litellm/types/llms/custom_http.py @@ -34,6 +34,7 @@ class httpxSpecialProvider(str, Enum): ROICalculator = "roi_calculator" Sandbox = "sandbox" ModelCostMap = "model_cost_map" + CodexModelCatalog = "codex_model_catalog" PasswordBreachCheck = "password_breach_check" ASGI = "asgi" AgentHarness = "agent_harness" diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index c9739eee91d..378899b667d 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -957,6 +957,7 @@ class ChatCompletionToolMessage(TypedDict): role: Literal["tool"] content: str | Iterable[ToolMessageContentPart] # writable-ok: Pydantic warns on ReadOnly TypedDict fields tool_call_id: str + is_error: NotRequired[bool] # writable-ok: Pydantic warns on ReadOnly TypedDict fields class ChatCompletionFunctionMessage(TypedDict): diff --git a/litellm/types/router.py b/litellm/types/router.py index a89cd68df8c..2c0fa6717b9 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -659,6 +659,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): tag_regex: list[str] | None = None # auto-router params + switchyard_router_config: Mapping[str, object] | None = None auto_router_config_path: str | None = None auto_router_config: str | None = None auto_router_default_model: str | None = None @@ -1041,6 +1042,21 @@ def _resolved_annotations(model_class: type[object]) -> Mapping[str, object]: return MappingProxyType(get_type_hints(model_class)) +class PriceTier(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + range: tuple[float, float] + costs: Mapping[str, float] + + +class DeploymentPricing(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + provider: str + costs: Mapping[str, float] + tiered_pricing: tuple[PriceTier, ...] = () + + class ModelGroupInfo(LiteLLMBaseModel): model_group: str providers: list[str] @@ -1071,6 +1087,7 @@ class ModelGroupInfo(LiteLLMBaseModel): configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None description: str | None = None supported_endpoints: tuple[str, ...] | None = None + deployment_pricing: tuple[DeploymentPricing, ...] = () def __init__(self, **data) -> None: for field_name, field_type in _resolved_annotations(self.__class__).items(): @@ -1445,6 +1462,7 @@ class PreRoutingStrategy(Protocol): messages: list[dict[str, Any]] | None = None, input: "str | list[Any] | None" = None, specific_deployment: bool | None = False, + call_type: str | None = None, ) -> "PreRoutingHookResponse | None": ... diff --git a/litellm/types/tag_management.py b/litellm/types/tag_management.py index e59b02af71a..2958c07ce95 100644 --- a/litellm/types/tag_management.py +++ b/litellm/types/tag_management.py @@ -27,6 +27,7 @@ class TagNewRequest(TagBase): class TagUpdateRequest(TagBase): + spend: float | None = None budget_id: str | None = None # Budget fields - if provided, will update the budget max_budget: float | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5a14dc95e7c..3ddeb118c33 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -217,6 +217,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_native_streaming: bool | None supports_native_structured_output: bool | None supports_bedrock_native_structured_output: ReadOnly[bool | None] + supports_bedrock_native_structured_output_by_default: ReadOnly[bool | None] supports_parallel_function_calling: bool | None supports_web_search: bool | None supports_reasoning: bool | None @@ -475,6 +476,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_image_512: ReadOnly[float | None] output_cost_per_image_1024: ReadOnly[float | None] output_cost_per_image_1536: ReadOnly[float | None] + output_cost_per_image_1K: ReadOnly[float | None] + output_cost_per_image_2K: ReadOnly[float | None] ocr_cost_per_page: float | None # for OCR models ocr_cost_per_page_batches: ReadOnly[float | None] ocr_cost_per_credit: float | None # for OCR models priced by credit @@ -3381,7 +3384,10 @@ class StandardLoggingRoutingDecision(TypedDict, total=False): """Per-request provenance for a pre-routing strategy (auto-router) decision.""" router_model_name: str - router_type: ReadOnly[Literal["complexity", "adaptive", "quality", "semantic"]] + router_type: ReadOnly[Literal["complexity", "adaptive", "quality", "semantic", "switchyard"]] + switchyard_algorithm: ReadOnly[str] + switchyard_evidence: ReadOnly[object] + switchyard_fallback_reason: ReadOnly[str] router_config_id: ReadOnly[str] router_config_updated_at: ReadOnly[str] router_config_fingerprint: ReadOnly[str] @@ -3442,10 +3448,12 @@ class StandardLoggingKubernetesPodRouting(TypedDict): # so a redacted row stays explainable. `test_every_routing_decision_field_is_classified` # fails if a field is added to the record without being placed in one set or the other. PROMPT_QUOTING_ROUTING_DECISION_FIELDS: frozenset[str] = frozenset( - {"signals", "matched_keyword", "escalation_keyword", "classifier_crux"} + {"signals", "matched_keyword", "escalation_keyword", "classifier_crux", "switchyard_evidence"} ) DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset( { + "switchyard_algorithm", + "switchyard_fallback_reason", "router_model_name", "router_type", "router_config_id", @@ -4145,6 +4153,8 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_image_512: float | None = None output_cost_per_image_1024: float | None = None output_cost_per_image_1536: float | None = None + output_cost_per_image_1K: float | None = None + output_cost_per_image_2K: float | None = None input_cost_per_pixel: float | None = None output_cost_per_pixel: float | None = None @@ -4303,6 +4313,7 @@ DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset( "supports_regex_lookaround", "supports_assistant_prefill", "supports_bedrock_native_structured_output", + "supports_bedrock_native_structured_output_by_default", "supports_per_turn_control", "supports_pdf_input", "supports_continuation", diff --git a/litellm/types/videos/main.py b/litellm/types/videos/main.py index 4d6f8fc6a30..b02633accb1 100644 --- a/litellm/types/videos/main.py +++ b/litellm/types/videos/main.py @@ -1,4 +1,5 @@ import builtins +from collections.abc import Sequence from typing import Any, Literal from openai.types.audio.transcription_create_params import FileTypes @@ -76,6 +77,19 @@ class VideoResponse(LiteLLMBaseModel): return self.dict() +class VideoImageURL(TypedDict): + url: ReadOnly[str] + + +class VideoInputReference(TypedDict): + type: ReadOnly[Literal["image_url"]] + image_url: ReadOnly[VideoImageURL] + + +class VideoFrameImage(VideoInputReference): + frame_type: ReadOnly[Literal["first_frame", "last_frame"]] + + class VideoCreateOptionalRequestParams(TypedDict, total=False): """ TypedDict for Optional parameters supported by OpenAI's video creation API. @@ -88,6 +102,12 @@ class VideoCreateOptionalRequestParams(TypedDict, total=False): parameters: dict[str, object] | None # Provider-specific parameters block passed directly to the API model: str | None resolution: ReadOnly[str | None] + duration: ReadOnly[int] + aspect_ratio: ReadOnly[str] + frame_images: ReadOnly[Sequence[VideoFrameImage]] + input_references: ReadOnly[Sequence[VideoInputReference]] + generate_audio: ReadOnly[bool] + seed: ReadOnly[int] seconds: str | None size: str | None characters: list[dict[str, str]] | None diff --git a/litellm/utils.py b/litellm/utils.py index bbfbbd33230..32385de64bf 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -90,6 +90,8 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import ( ) from litellm.litellm_core_utils.core_helpers import ( bind_budget_reservation_to_callbacks, + has_attributes, + iter_items, normalize_drop_params, unbind_budget_reservation_from_callbacks, ) @@ -4167,6 +4169,13 @@ def _registration_base( return _RegistrationWrite(key=key, entry=dict(fallback.entry)) +def native_structured_output_support(model_info: Mapping[str, object]) -> bool | None: + if model_info.get("supports_bedrock_native_structured_output_by_default"): + return model_info.get("supports_bedrock_native_structured_output") is not False + declared: Final = model_info.get("supports_native_structured_output") + return declared if isinstance(declared, bool) else None + + def _base_fields_replaced_by(value: Mapping[str, object]) -> frozenset[str]: return ( frozenset({"explicit_cache_read_input_token_cost"}) @@ -4499,8 +4508,9 @@ def get_optional_params_transcription( passed_params.pop("OPENAI_TRANSCRIPTION_PARAMS") passed_params.pop("model") - custom_llm_provider = passed_params.pop("custom_llm_provider") - drop_params = normalize_drop_params(passed_params.pop("drop_params")) + passed_params.pop("custom_llm_provider") + passed_params.pop("drop_params") + normalized_drop_params: Final = normalize_drop_params(drop_params) special_params: Final[Mapping[str, object]] = passed_params.pop("kwargs") from litellm.litellm_core_utils.audio_utils.transcription_streaming import ( transcription_stream_flag, @@ -4511,7 +4521,7 @@ def get_optional_params_transcription( passed_params[k] = v sends_through_openai_sdk: Final = custom_llm_provider in ("openai", "azure") - drops_unsupported_params: Final = drop_params is True or litellm.drop_params is True + drops_unsupported_params: Final = normalized_drop_params is True or litellm.drop_params is True transcription_params: Final = ( OPENAI_TRANSCRIPTION_PARAMS if sends_through_openai_sdk or drops_unsupported_params @@ -4563,7 +4573,7 @@ def get_optional_params_transcription( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=drop_params if drop_params is not None else False, + drop_params=normalized_drop_params if normalized_drop_params is not None else False, ) elif custom_llm_provider == "bedrock": from litellm.llms.bedrock.audio_transcription import ( @@ -4572,7 +4582,7 @@ def get_optional_params_transcription( ) match map_openai_transcription_params( - non_default_params, drop_params=drop_params is True or litellm.drop_params is True + non_default_params, drop_params=normalized_drop_params is True or litellm.drop_params is True ): case UnsupportedTranscriptionParams(message=message): raise UnsupportedParamsError(message=message, llm_provider="bedrock") @@ -4585,7 +4595,7 @@ def get_optional_params_transcription( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=drop_params if drop_params is not None else False, + drop_params=normalized_drop_params if normalized_drop_params is not None else False, ) optional_params = add_provider_specific_params_to_optional_params( @@ -6555,26 +6565,41 @@ def count_characters(text: str) -> int: _count_characters = count_characters +class _ResponsesOutputCarrier(Protocol): + @property + def output(self) -> object: ... + + +class _ResponsesContentCarrier(Protocol): + @property + def content(self) -> object: ... + + +class _ResponsesTextCarrier(Protocol): + @property + def text(self) -> str: ... + + def get_response_string(response_obj: ModelResponse | ModelResponseStream) -> str: # Handle Responses API streaming events if hasattr(response_obj, "type") and hasattr(response_obj, "response"): # This is a Responses API streaming event (e.g., ResponseCreatedEvent, ResponseCompletedEvent) # Extract text from the response object's output if available - responses_api_response: Final = getattr(response_obj, "response", None) - if responses_api_response and hasattr(responses_api_response, "output"): + responses_api_response: Final[object] = getattr(response_obj, "response", None) + if responses_api_response and has_attributes(responses_api_response, _ResponsesOutputCarrier, "output"): output_list: Final = responses_api_response.output # Use list accumulation to avoid O(n^2) string concatenation: # repeatedly doing `response_str += part` copies the full string each time # because Python strings are immutable, so total work grows with n^2. response_output_parts: Final[list[str]] = [] - for output_item in output_list: + for output_item in iter_items(output_list): # Handle output items with content array - if hasattr(output_item, "content"): - for content_part in output_item.content: - if hasattr(content_part, "text"): + if has_attributes(output_item, _ResponsesContentCarrier, "content"): + for content_part in iter_items(output_item.content): + if has_attributes(content_part, _ResponsesTextCarrier, "text"): response_output_parts.append(content_part.text) # Handle output items with direct text field - elif hasattr(output_item, "text"): + elif has_attributes(output_item, _ResponsesTextCarrier, "text"): response_output_parts.append(output_item.text) return "".join(response_output_parts) @@ -7559,6 +7584,8 @@ def get_model_info_helper( output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None), output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None), output_cost_per_image=_model_info.get("output_cost_per_image", None), + output_cost_per_image_1K=_model_info.get("output_cost_per_image_1K", None), + output_cost_per_image_2K=_model_info.get("output_cost_per_image_2K", None), output_cost_per_pixel=_model_info.get("output_cost_per_pixel", None), output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None), output_cost_per_image_token_batches=_model_info.get("output_cost_per_image_token_batches", None), @@ -7587,10 +7614,13 @@ def get_model_info_helper( supports_pdf_input=_model_info.get("supports_pdf_input", None), supports_embedding_image_input=_model_info.get("supports_embedding_image_input", None), supports_native_streaming=_model_info.get("supports_native_streaming", None), - supports_native_structured_output=_model_info.get("supports_native_structured_output", None), + supports_native_structured_output=native_structured_output_support(_model_info), supports_bedrock_native_structured_output=_model_info.get( "supports_bedrock_native_structured_output", None ), + supports_bedrock_native_structured_output_by_default=_model_info.get( + "supports_bedrock_native_structured_output_by_default", None + ), supports_web_search=_model_info.get("supports_web_search", None), supports_url_context=_model_info.get("supports_url_context", None), supports_continuation=_model_info.get("supports_continuation", None), @@ -10679,8 +10709,8 @@ class ProviderConfigManager: # given model is on it; None keeps the chat-completions bridge. return litellm.BedrockOpenAIResponsesConfig.for_model(model) elif litellm.LlmProviders.BEDROCK_MANTLE == provider: - # Both decisions are data-driven from the model's price-map entry, with - # no model-name logic. Capability (can it serve Responses?) comes from + # Capability and wire path are data-driven from the model's price-map + # entry, with no model-name logic. Capability (can it serve Responses?) comes from # mantle_supports_responses (supported_endpoints / mode); # chat-only models (gpt-oss safeguard, nvidia, ...) return None and keep # the chat-completions emulation (responses/main.py "config is None"). @@ -10688,6 +10718,7 @@ class ProviderConfigManager: # use_openai_responses_path flag: gpt-5.x and gemma-4-* on # /openai/v1/responses, everything else (incl. gpt-oss) on # /v1/responses. + from litellm.llms.bedrock.common_utils import bedrock_model_encrypts_agent_messages from litellm.llms.bedrock_mantle.common_utils import ( mantle_base_segment, mantle_supports_responses, @@ -10696,7 +10727,8 @@ class ProviderConfigManager: if not model or not mantle_supports_responses(model, litellm.model_cost): return None return litellm.BedrockMantleResponsesAPIConfig( - use_openai_path=mantle_base_segment(model, litellm.model_cost) == "openai/v1" + use_openai_path=mantle_base_segment(model, litellm.model_cost) == "openai/v1", + encrypted_agent_messages=bedrock_model_encrypts_agent_messages(model), ) return None diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 598b56f4a79..b83c9ef48b1 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -437,7 +437,7 @@ class VectorStoreRegistry: # cast to VectorStoreConfig litellm_vector_store_config = LiteLLM_VectorStoreConfig(**vector_store_config) vector_store_name = litellm_vector_store_config.get("vector_store_name") - vector_store_litellm_params: dict[str, Any] = litellm_vector_store_config.get("litellm_params") or {} + vector_store_litellm_params = litellm_vector_store_config.get("litellm_params") or {} vector_store_id = vector_store_litellm_params.get("vector_store_id") if vector_store_id is None: diff --git a/litellm/videos/utils.py b/litellm/videos/utils.py index b0a78691a70..fa6296739e7 100644 --- a/litellm/videos/utils.py +++ b/litellm/videos/utils.py @@ -1,6 +1,8 @@ -from typing import Any, Final, cast +from collections.abc import Mapping +from typing import Final, cast import litellm +from litellm.llms.base_llm.videos.async_task_utils import OPENROUTER_VIDEO_FIELDS from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.types.videos.main import VideoCreateOptionalRequestParams from litellm.utils import filter_out_litellm_params @@ -47,7 +49,7 @@ class VideoGenerationRequestUtils: @staticmethod def get_requested_video_generation_optional_param( - params: dict[str, Any], + params: Mapping[str, object], ) -> VideoCreateOptionalRequestParams: """ Filter parameters to only include those defined in VideoCreateOptionalRequestParams. @@ -60,9 +62,8 @@ class VideoGenerationRequestUtils: """ params = dict(params or {}) - raw_kwargs = params.get("kwargs", {}) - if not isinstance(raw_kwargs, dict): - raw_kwargs = {} + nested_kwargs: Final = params.get("kwargs", {}) + raw_kwargs: Final[dict[str, object]] = nested_kwargs if isinstance(nested_kwargs, dict) else {} kwargs_extra_body: Final = raw_kwargs.pop("extra_body", None) top_level_extra_body: Final = params.get("extra_body") @@ -92,7 +93,9 @@ class VideoGenerationRequestUtils: merged_extra_body = filter_out_litellm_params(kwargs=merged_extra_body) if merged_extra_body: optional_params["extra_body"] = merged_extra_body - optional_params.update(merged_extra_body) + optional_params.update( + {key: value for key, value in merged_extra_body.items() if key not in OPENROUTER_VIDEO_FIELDS} + ) optional_params.pop("timeout", None) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index aa358f4f6f3..77244a32f6f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1961,6 +1961,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2041,6 +2042,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2116,6 +2118,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2242,7 +2245,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2287,6 +2291,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2331,7 +2336,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2376,6 +2382,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2420,7 +2427,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2465,6 +2473,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2504,7 +2513,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2549,6 +2559,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2593,7 +2604,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2633,6 +2645,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2715,7 +2728,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -2759,7 +2773,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -2803,7 +2818,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -2847,7 +2863,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -2891,7 +2908,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -3023,6 +3041,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -3068,6 +3087,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -3113,6 +3133,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -3158,6 +3179,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -3203,6 +3225,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -12097,16 +12120,19 @@ }, "azure/tts-1": { "deprecation_date": "2026-12-15", + "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "input_cost_per_character": 1.5e-05, "litellm_provider": "azure", - "mode": "audio_speech" + "mode": "audio_speech", + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule" }, "azure/tts-1-hd": { "deprecation_date": "2026-12-15", + "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "input_cost_per_character": 3.3e-05, "litellm_provider": "azure", "mode": "audio_speech", - "source": "https://azure.microsoft.com/en-us/pricing/details/azure-openai/" + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule" }, "azure/us/gpt-4.1-2025-04-14": { "deprecation_date": "2027-04-14", @@ -12859,11 +12885,13 @@ "supports_vision": true }, "azure/whisper-1": { - "deprecation_date": "2026-12-15", + "deprecation_date": "2027-02-26", + "deprecation_source": "https://management.azure.com/subscriptions/a88257a5-3323-4ef9-95bf-5681a46c9624/providers/Microsoft.CognitiveServices/locations/eastus/models?api-version=2024-10-01", "input_cost_per_second": 0.0001, "litellm_provider": "azure", "mode": "audio_transcription", - "output_cost_per_second": 0.0001 + "output_cost_per_second": 0.0001, + "source": "https://management.azure.com/subscriptions/a88257a5-3323-4ef9-95bf-5681a46c9624/providers/Microsoft.CognitiveServices/locations/eastus/models?api-version=2024-10-01" }, "azure_ai/Cohere-embed-v3-english": { "deprecation_date": "2027-02-15", @@ -13837,12 +13865,12 @@ "azure_ai/jais-30b-chat": { "deprecation_date": "2026-01-30", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", - "input_cost_per_token": 0.0032, + "input_cost_per_token": 3.2e-06, "litellm_provider": "azure_ai", "max_input_tokens": 8192, "max_output_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.00971, + "output_cost_per_token": 9.71e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", "supports_pdf_input": false }, @@ -31630,7 +31658,7 @@ "input_cost_per_audio_token_batches": 2.5e-07 }, "gemini/gemini-embedding-2-preview": { - "deprecation_date": "2026-08-10", + "deprecation_date": "2026-11-06", "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", "input_cost_per_audio_token": 6.5e-06, "input_cost_per_audio_token_batches": 3.25e-06, @@ -31961,7 +31989,8 @@ "rpm": 1000.0, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", + "successor_model": "gemini/gemini-nano-banana-2.1", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -32516,7 +32545,7 @@ "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, "rpm": 2000, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -32562,6 +32591,7 @@ "input_cost_per_token_flex": 2.5e-07, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, + "successor_model": "gemini/gemini-3.6-flash", "supports_audio_input": true, "supports_computer_use": true }, @@ -32578,7 +32608,7 @@ "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, "rpm": 2000, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -32625,7 +32655,8 @@ "output_cost_per_token_batches": 4.5e-06, "input_cost_per_token_flex": 7.5e-07, "output_cost_per_token_flex": 4.5e-06, - "cache_read_input_token_cost_flex": 8e-08 + "cache_read_input_token_cost_flex": 8e-08, + "successor_model": "gemini/gemini-3.6-flash" }, "gemini/gemini-3.6-flash": { "prompt_cache_min_tokens": 4096, @@ -32705,7 +32736,7 @@ "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, "rpm": 2000, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -32748,7 +32779,8 @@ }, "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014 + "google_maps_grounding_cost_per_query": 0.014, + "successor_model": "gemini/gemini-3.8-flash" }, "gemini/gemini-3.8-flash": { "prompt_cache_min_tokens": 4096, @@ -33610,7 +33642,7 @@ "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 0, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_modalities": [ "text" ], @@ -33625,7 +33657,8 @@ "supports_system_messages": false, "supports_vision": false, "supports_web_search": false, - "output_cost_per_image": 0.08 + "output_cost_per_image": 0.08, + "successor_model": "gemini/lyria-3.5" }, "gemini/veo-3.1-fast-generate-preview": { "deprecation_date": "2026-10-22", @@ -44089,6 +44122,7 @@ "supports_vision": true }, "mistral/mistral-large-2512": { + "deprecation_date": "2027-04-30", "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 5e-07, "litellm_provider": "mistral", @@ -44164,6 +44198,7 @@ "supports_vision": true }, "mistral/mistral-medium-3-5": { + "deprecation_date": "2027-04-30", "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 1.5e-06, "litellm_provider": "mistral", @@ -55051,14 +55086,16 @@ "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference" }, "wandb/deepseek-ai/DeepSeek-V3.1": { + "cache_read_input_token_cost": 2e-07, "supports_reasoning": true, "max_input_tokens": 128000, "max_output_tokens": 128000, - "input_cost_per_token": 5.5e-07, - "output_cost_per_token": 1.65e-06, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 1.5e-06, "litellm_provider": "wandb", "mode": "chat", - "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference" + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", + "supports_prompt_caching": true }, "wandb/meta-llama/Llama-3.3-70B-Instruct": { "max_input_tokens": 128000, @@ -55130,8 +55167,8 @@ "watsonx/core42/jais-13b-chat": { "max_input_tokens": 8192, "max_output_tokens": 8192, - "input_cost_per_token": 0.0005, - "output_cost_per_token": 0.002, + "input_cost_per_token": 1.8e-06, + "output_cost_per_token": 1.8e-06, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -64183,6 +64220,39 @@ "supports_tool_choice": true, "supports_vision": false }, + "tencent/hy4": { + "cache_read_input_token_cost": 4.2e-08, + "comment": "TokenHub bills Hy4 at 50% of list until 2026-12-31 23:59:59 Beijing time; these are the list rates.", + "input_cost_per_token": 8.34e-07, + "input_cost_per_token_cache_hit": 4.2e-08, + "litellm_provider": "tencent", + "max_input_tokens": 960000, + "max_output_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.501e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true + }, + "tencent/hy4-preview": { + "cache_read_input_token_cost": 4.2e-08, + "comment": "TokenHub bills Hy4 preview at 50% of list until 2026-12-31 23:59:59 Beijing time; these are the list rates.", + "input_cost_per_token": 8.34e-07, + "input_cost_per_token_cache_hit": 4.2e-08, + "litellm_provider": "tencent", + "max_input_tokens": 960000, + "max_output_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.501e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true + }, "tencent/minimax-m3": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 6e-08, @@ -64726,6 +64796,7 @@ "source": "https://docs.mistral.ai/models/model-cards/mistral-moderation-26-03" }, "mistral/voxtral-mini-2602": { + "deprecation_date": "2027-04-30", "input_cost_per_second": 5e-05, "litellm_provider": "mistral", "mode": "audio_transcription", @@ -64742,6 +64813,7 @@ "supports_audio_input": true }, "mistral/voxtral-mini-transcribe-realtime-2602": { + "deprecation_date": "2027-04-30", "input_cost_per_second": 0.0001, "litellm_provider": "mistral", "mode": "audio_transcription", @@ -64758,6 +64830,7 @@ "supports_audio_input": true }, "mistral/voxtral-mini-tts-2603": { + "deprecation_date": "2027-04-30", "input_cost_per_character": 1.6e-05, "litellm_provider": "mistral", "mode": "audio_speech", @@ -66431,13 +66504,13 @@ "max_output_tokens": 262144, "max_input_tokens": 262000, "input_cost_per_token": 1.3e-07, - "output_cost_per_token": 2.8e-07, - "cache_read_input_token_cost": 7e-08, + "output_cost_per_token": 2.7e-07, + "cache_read_input_token_cost": 5e-08, "supports_prompt_caching": true, "litellm_provider": "wandb", "mode": "chat", "supports_vision": false, - "source": "https://wandb.ai/site/pricing/tokens/" + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference" }, "wandb/deepseek-ai/DeepSeek-V4-Pro": { "supports_reasoning": true, @@ -66514,14 +66587,14 @@ "supports_reasoning": true, "max_output_tokens": 262144, "max_input_tokens": 262000, - "input_cost_per_token": 7.1e-07, - "output_cost_per_token": 3.5e-06, - "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 3e-06, + "cache_read_input_token_cost": 1.2e-07, "supports_prompt_caching": true, "litellm_provider": "wandb", "mode": "chat", "supports_vision": true, - "source": "https://wandb.ai/site/pricing/tokens/" + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference" }, "wandb/moonshotai/Kimi-K2.6": { "supports_reasoning": true, @@ -66587,15 +66660,17 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/Qwen/Qwen3.6-35B-A3B": { + "cache_read_input_token_cost": 3.5e-08, "supports_reasoning": true, "max_output_tokens": 262144, "max_input_tokens": 262000, - "input_cost_per_token": 2.5e-07, - "output_cost_per_token": 1.25e-06, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 8.9e-07, "litellm_provider": "wandb", "mode": "chat", "supports_vision": true, - "source": "https://wandb.ai/site/pricing/tokens/" + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", + "supports_prompt_caching": true }, "wandb/Qwen/Qwen3.6-27B": { "supports_reasoning": true, @@ -66649,18 +66724,18 @@ "litellm_provider": "wandb", "mode": "chat", "supports_reasoning": true, - "input_cost_per_token": 1e-07, + "input_cost_per_token": 6e-08, "output_cost_per_token": 1.5e-07, - "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost": 1.5e-08, "max_input_tokens": 131000, "supports_prompt_caching": true, - "source": "https://wandb.ai/site/pricing/tokens/" + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference" }, "wandb/zai-org/GLM-5.2": { "supports_reasoning": true, "max_output_tokens": 262144, "max_input_tokens": 262000, - "input_cost_per_token": 7.6e-07, + "input_cost_per_token": 6.2e-07, "output_cost_per_token": 2.42e-06, "cache_read_input_token_cost": 1.4e-07, "supports_prompt_caching": true, @@ -67955,7 +68030,9 @@ "input_cost_per_image": 0.06, "litellm_provider": "xai", "mode": "image_generation", - "source": "https://docs.x.ai/docs/models", + "output_cost_per_image_1K": 0.04, + "output_cost_per_image_2K": 0.06, + "source": "https://docs.x.ai/developers/migration/imagine-image-quality-nov-2", "supported_endpoints": [ "/v1/images/generations" ], @@ -70849,6 +70926,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token_above_272k_tokens": 1e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, "source": "https://docs.perplexity.ai/docs/agent-api/models" }, "perplexity/openai/gpt-5.4": { @@ -74215,14 +74293,14 @@ ] }, "wandb/deepseek-ai/DeepSeek-V4.1-Flash": { - "cache_read_input_token_cost": 3e-08, - "input_cost_per_token": 2e-07, + "cache_read_input_token_cost": 2e-09, + "input_cost_per_token": 1.3e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "max_output_tokens": 1048576, "mode": "chat", "output_cost_per_token": 6.5e-07, - "source": "https://wandb.ai/site/pricing/tokens/", + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_prompt_caching": true, "supports_reasoning": true, "supports_vision": true @@ -74243,13 +74321,13 @@ "supports_vision": true }, "wandb/zai-org/GLM-5.3-Flash": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 1.5e-07, + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 1.2e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "mode": "chat", - "output_cost_per_token": 5e-07, - "source": "https://docs.wandb.ai/inference/models.md", + "output_cost_per_token": 3.8e-07, + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -75544,7 +75622,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -75584,7 +75663,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -75624,6 +75704,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -75697,6 +75778,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -76995,13 +77077,13 @@ "supports_vision": true }, "azure/whisper": { - "deprecation_date": "2026-12-15", - "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", + "deprecation_date": "2027-02-26", + "deprecation_source": "https://management.azure.com/subscriptions/a88257a5-3323-4ef9-95bf-5681a46c9624/providers/Microsoft.CognitiveServices/locations/eastus/models?api-version=2024-10-01", "input_cost_per_second": 0.0001, "litellm_provider": "azure", "mode": "audio_transcription", "output_cost_per_second": 0.0001, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://management.azure.com/subscriptions/a88257a5-3323-4ef9-95bf-5681a46c9624/providers/Microsoft.CognitiveServices/locations/eastus/models?api-version=2024-10-01" }, "gemini/gemini-nano-banana-2.1": { "input_cost_per_token": 1.5e-06, @@ -77412,6 +77494,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -77469,6 +77552,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -77646,6 +77730,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -77698,6 +77783,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -77745,6 +77831,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -77858,6 +77945,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output_by_default": true, "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, @@ -78605,13 +78693,13 @@ "supports_vision": true }, "wandb/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2.5e-07, + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 1.5e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "max_output_tokens": 943718, "mode": "chat", - "output_cost_per_token": 1.5e-06, + "output_cost_per_token": 5e-07, "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_function_calling": true, "supports_prompt_caching": true, @@ -78907,12 +78995,12 @@ "supports_vision": true }, "wandb/reflection/Beam-1.0-501B-A23B": { - "cache_read_input_token_cost": 1e-07, - "input_cost_per_token": 1e-06, + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 3e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "mode": "chat", - "output_cost_per_token": 3e-06, + "output_cost_per_token": 1e-06, "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_function_calling": true, "supports_prompt_caching": true, @@ -78923,11 +79011,11 @@ }, "wandb/zai-org/GLM-5.3": { "cache_read_input_token_cost": 1.9e-07, - "input_cost_per_token": 1.2e-06, + "input_cost_per_token": 7.5e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "mode": "chat", - "output_cost_per_token": 4e-06, + "output_cost_per_token": 3e-06, "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_function_calling": true, "supports_prompt_caching": true, @@ -83064,13 +83152,16 @@ "mode": "chat", "max_input_tokens": 1000000, "max_output_tokens": 393216, - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 8e-07, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-1-flash-1017" + "source": "https://empiriolabs.ai/models/deepseek-v4-1-flash-1017", + "cache_read_input_token_cost": 2e-08, + "supports_prompt_caching": true, + "supports_function_calling": true }, "empiriolabs/deepseek-v4-flash-0731": { "litellm_provider": "empiriolabs", @@ -83082,7 +83173,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731" + "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-flash-0731:variant1": { "litellm_provider": "empiriolabs", @@ -83094,7 +83186,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731" + "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro": { "litellm_provider": "empiriolabs", @@ -83106,7 +83199,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro-0813": { "litellm_provider": "empiriolabs", @@ -83118,7 +83212,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro-0813" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro-0813", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro:variant1": { "litellm_provider": "empiriolabs", @@ -83136,7 +83231,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro:variant2": { "litellm_provider": "empiriolabs", @@ -83156,7 +83252,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro:variant3": { "litellm_provider": "empiriolabs", @@ -83168,7 +83265,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/fugu-max": { "litellm_provider": "empiriolabs", @@ -83366,7 +83464,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/glm-5-2" + "source": "https://empiriolabs.ai/models/glm-5-2", + "supports_function_calling": true }, "empiriolabs/glm-5-2:variant2": { "litellm_provider": "empiriolabs", @@ -83384,7 +83483,8 @@ "search_context_size_high": 0.013 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/glm-5-2" + "source": "https://empiriolabs.ai/models/glm-5-2", + "supports_function_calling": true }, "empiriolabs/glm-5-3": { "litellm_provider": "empiriolabs", @@ -83438,7 +83538,8 @@ "supports_reasoning": true, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-6" + "source": "https://empiriolabs.ai/models/kimi-k2-6", + "supports_function_calling": true }, "empiriolabs/kimi-k2-7-code": { "litellm_provider": "empiriolabs", @@ -83457,7 +83558,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-7-code" + "source": "https://empiriolabs.ai/models/kimi-k2-7-code", + "supports_function_calling": true }, "empiriolabs/kimi-k2-7-code-highspeed": { "litellm_provider": "empiriolabs", @@ -83476,7 +83578,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-7-code-highspeed" + "source": "https://empiriolabs.ai/models/kimi-k2-7-code-highspeed", + "supports_function_calling": true }, "empiriolabs/kimi-k2-7-code:variant1": { "litellm_provider": "empiriolabs", @@ -83497,7 +83600,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-7-code" + "source": "https://empiriolabs.ai/models/kimi-k2-7-code", + "supports_function_calling": true }, "empiriolabs/kimi-k3": { "litellm_provider": "empiriolabs", @@ -83516,7 +83620,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k3" + "source": "https://empiriolabs.ai/models/kimi-k3", + "supports_function_calling": true }, "empiriolabs/kimi-k3-1-pro": { "litellm_provider": "empiriolabs", @@ -83558,7 +83663,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/mimo-v2-6-flash" + "source": "https://empiriolabs.ai/models/mimo-v2-6-flash", + "supports_function_calling": true }, "empiriolabs/mimo-v2-6-pro": { "litellm_provider": "empiriolabs", @@ -83578,7 +83684,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/mimo-v2-6-pro" + "source": "https://empiriolabs.ai/models/mimo-v2-6-pro", + "supports_function_calling": true }, "empiriolabs/mimo-v2-6-pro-ultraspeed": { "litellm_provider": "empiriolabs", @@ -83592,7 +83699,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/mimo-v2-6-pro-ultraspeed" + "source": "https://empiriolabs.ai/models/mimo-v2-6-pro-ultraspeed", + "supports_function_calling": true }, "empiriolabs/minimax-m2-7": { "litellm_provider": "empiriolabs", @@ -83707,7 +83815,8 @@ "search_context_size_high": 0.14 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/mistral-small-4" + "source": "https://empiriolabs.ai/models/mistral-small-4", + "supports_function_calling": true }, "empiriolabs/muse-glimmer-30b": { "litellm_provider": "empiriolabs", @@ -83722,7 +83831,8 @@ "supports_reasoning": true, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/muse-glimmer-30b" + "source": "https://empiriolabs.ai/models/muse-glimmer-30b", + "supports_function_calling": true }, "empiriolabs/muse-spark-1-1": { "litellm_provider": "empiriolabs", @@ -83788,7 +83898,8 @@ "supports_system_messages": true, "supports_vision": true, "supports_prompt_caching": true, - "source": "https://empiriolabs.ai/models/nova-lite-1-0" + "source": "https://empiriolabs.ai/models/nova-lite-1-0", + "supports_function_calling": true }, "empiriolabs/nova-lite-2": { "litellm_provider": "empiriolabs", @@ -83802,7 +83913,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_prompt_caching": true, - "source": "https://empiriolabs.ai/models/nova-lite-2" + "source": "https://empiriolabs.ai/models/nova-lite-2", + "supports_function_calling": true }, "empiriolabs/nova-micro-1-0": { "litellm_provider": "empiriolabs", @@ -83814,7 +83926,8 @@ "cache_read_input_token_cost": 2.24e-08, "supports_system_messages": true, "supports_prompt_caching": true, - "source": "https://empiriolabs.ai/models/nova-micro-1-0" + "source": "https://empiriolabs.ai/models/nova-micro-1-0", + "supports_function_calling": true }, "empiriolabs/nova-pro-1-0": { "litellm_provider": "empiriolabs", @@ -83825,7 +83938,8 @@ "output_cost_per_token": 9.6e-06, "supports_system_messages": true, "supports_vision": true, - "source": "https://empiriolabs.ai/models/nova-pro-1-0" + "source": "https://empiriolabs.ai/models/nova-pro-1-0", + "supports_function_calling": true }, "empiriolabs/perplexity-advanced-deep-research": { "litellm_provider": "empiriolabs", @@ -83925,7 +84039,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b" + "source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-27b": { "litellm_provider": "empiriolabs", @@ -83946,7 +84061,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-27b" + "source": "https://empiriolabs.ai/models/qwen3-5-27b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-35b-a3b": { "litellm_provider": "empiriolabs", @@ -83967,7 +84083,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b" + "source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-397b-a17b": { "litellm_provider": "empiriolabs", @@ -83988,7 +84105,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b" + "source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-4b": { "litellm_provider": "empiriolabs", @@ -84039,7 +84157,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-flash" + "source": "https://empiriolabs.ai/models/qwen3-5-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-5-flash:variant1": { "litellm_provider": "empiriolabs", @@ -84062,7 +84181,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-flash" + "source": "https://empiriolabs.ai/models/qwen3-5-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-5-omni-flash": { "litellm_provider": "empiriolabs", @@ -84080,7 +84200,8 @@ "search_context_size_high": 0.015 }, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/qwen3-5-omni-flash" + "source": "https://empiriolabs.ai/models/qwen3-5-omni-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-5-omni-plus": { "litellm_provider": "empiriolabs", @@ -84099,7 +84220,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/qwen3-5-omni-plus" + "source": "https://empiriolabs.ai/models/qwen3-5-omni-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-5-plus": { "litellm_provider": "empiriolabs", @@ -84108,8 +84230,8 @@ "max_output_tokens": 65536, "input_cost_per_token": 3.6e-07, "output_cost_per_token": 2.21e-06, - "input_cost_per_token_above_256k_tokens": 1.08e-06, - "output_cost_per_token_above_256k_tokens": 6.62e-06, + "input_cost_per_token_above_256k_tokens": 4.5e-07, + "output_cost_per_token_above_256k_tokens": 2.76e-06, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, @@ -84120,7 +84242,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-plus" + "source": "https://empiriolabs.ai/models/qwen3-5-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-5-plus:variant1": { "litellm_provider": "empiriolabs", @@ -84143,7 +84266,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-plus" + "source": "https://empiriolabs.ai/models/qwen3-5-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-6-27b": { "litellm_provider": "empiriolabs", @@ -84162,7 +84286,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-27b" + "source": "https://empiriolabs.ai/models/qwen3-6-27b", + "supports_function_calling": true }, "empiriolabs/qwen3-6-35b-a3b": { "litellm_provider": "empiriolabs", @@ -84176,7 +84301,8 @@ "supports_reasoning": true, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-35b-a3b" + "source": "https://empiriolabs.ai/models/qwen3-6-35b-a3b", + "supports_function_calling": true }, "empiriolabs/qwen3-6-flash": { "litellm_provider": "empiriolabs", @@ -84197,7 +84323,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-flash" + "source": "https://empiriolabs.ai/models/qwen3-6-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-6-flash:variant1": { "litellm_provider": "empiriolabs", @@ -84218,7 +84345,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-flash" + "source": "https://empiriolabs.ai/models/qwen3-6-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-6-plus": { "litellm_provider": "empiriolabs", @@ -84239,7 +84367,8 @@ "search_context_size_high": 0.026 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-plus" + "source": "https://empiriolabs.ai/models/qwen3-6-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-6-plus:variant1": { "litellm_provider": "empiriolabs", @@ -84260,7 +84389,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-plus" + "source": "https://empiriolabs.ai/models/qwen3-6-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-7-flash": { "litellm_provider": "empiriolabs", @@ -84287,7 +84417,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-flash" + "source": "https://empiriolabs.ai/models/qwen3-7-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-7-flash:variant1": { "litellm_provider": "empiriolabs", @@ -84314,7 +84445,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-flash" + "source": "https://empiriolabs.ai/models/qwen3-7-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-7-max": { "litellm_provider": "empiriolabs", @@ -84332,7 +84464,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-max" + "source": "https://empiriolabs.ai/models/qwen3-7-max", + "supports_function_calling": true }, "empiriolabs/qwen3-7-max:variant1": { "litellm_provider": "empiriolabs", @@ -84350,7 +84483,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-max" + "source": "https://empiriolabs.ai/models/qwen3-7-max", + "supports_function_calling": true }, "empiriolabs/qwen3-7-plus": { "litellm_provider": "empiriolabs", @@ -84371,7 +84505,8 @@ "search_context_size_high": 0.03 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-plus" + "source": "https://empiriolabs.ai/models/qwen3-7-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-7-plus:variant1": { "litellm_provider": "empiriolabs", @@ -84395,7 +84530,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-plus" + "source": "https://empiriolabs.ai/models/qwen3-7-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-8-27b": { "litellm_provider": "empiriolabs", @@ -84430,7 +84566,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-flash" + "source": "https://empiriolabs.ai/models/qwen3-8-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-8-max": { "litellm_provider": "empiriolabs", @@ -84449,7 +84586,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-max" + "source": "https://empiriolabs.ai/models/qwen3-8-max", + "supports_function_calling": true }, "empiriolabs/qwen3-8-max-0902": { "litellm_provider": "empiriolabs", @@ -84468,7 +84606,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-max-0902" + "source": "https://empiriolabs.ai/models/qwen3-8-max-0902", + "supports_function_calling": true }, "empiriolabs/qwen3-8-max:variant1": { "litellm_provider": "empiriolabs", @@ -84487,7 +84626,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-max" + "source": "https://empiriolabs.ai/models/qwen3-8-max", + "supports_function_calling": true }, "empiriolabs/qwen3-8-omni-flash": { "litellm_provider": "empiriolabs", @@ -84507,7 +84647,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/qwen3-8-omni-flash" + "source": "https://empiriolabs.ai/models/qwen3-8-omni-flash", + "supports_function_calling": true }, "empiriolabs/seed-2-0-code": { "litellm_provider": "empiriolabs", @@ -84521,7 +84662,8 @@ "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, - "source": "https://empiriolabs.ai/models/seed-2-0-code" + "source": "https://empiriolabs.ai/models/seed-2-0-code", + "supports_function_calling": true }, "empiriolabs/seed-2-0-lite": { "litellm_provider": "empiriolabs", @@ -84536,7 +84678,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-0-lite" + "source": "https://empiriolabs.ai/models/seed-2-0-lite", + "supports_function_calling": true }, "empiriolabs/seed-2-0-mini": { "litellm_provider": "empiriolabs", @@ -84551,7 +84694,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-0-mini" + "source": "https://empiriolabs.ai/models/seed-2-0-mini", + "supports_function_calling": true }, "empiriolabs/seed-2-0-pro": { "litellm_provider": "empiriolabs", @@ -84566,7 +84710,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-0-pro" + "source": "https://empiriolabs.ai/models/seed-2-0-pro", + "supports_function_calling": true }, "empiriolabs/seed-2-1-turbo": { "litellm_provider": "empiriolabs", @@ -84579,7 +84724,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-1-turbo" + "source": "https://empiriolabs.ai/models/seed-2-1-turbo", + "supports_function_calling": true }, "empiriolabs/skylark-embedding-vision": { "litellm_provider": "empiriolabs", @@ -86235,7 +86381,7 @@ "input_cost_per_token": 2e-06, "litellm_provider": "baseten", "max_input_tokens": 1000000, - "max_output_tokens": 262144, + "max_output_tokens": 1048576, "mode": "chat", "output_cost_per_token": 1e-05, "source": "https://inference.baseten.co/v1/models", @@ -86570,5 +86716,170 @@ "supports_tool_choice": true, "supports_vision": true, "token_count_framing": 100.0 + }, + "moonshotai.kimi-k3.1": { + "cache_creation_input_token_cost": 5.5e-06, + "cache_creation_input_token_cost_flex": 2.75e-06, + "cache_creation_input_token_cost_priority": 9.625e-06, + "cache_read_input_token_cost": 2.2e-07, + "cache_read_input_token_cost_flex": 1.1e-07, + "cache_read_input_token_cost_priority": 3.85e-07, + "input_cost_per_token": 4.4e-06, + "input_cost_per_token_flex": 2.2e-06, + "input_cost_per_token_priority": 7.7e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 868928, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.2e-05, + "output_cost_per_token_flex": 1.1e-05, + "output_cost_per_token_priority": 3.85e-05, + "reasoning_effort_levels": [ + "none", + "low", + "medium", + "high", + "xhigh", + "max" + ], + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k3-1-highspeed.md", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_none_reasoning_effort": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_regex_lookaround": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "us.moonshotai.kimi-k3.1": { + "cache_creation_input_token_cost": 5.5e-06, + "cache_creation_input_token_cost_flex": 2.75e-06, + "cache_creation_input_token_cost_priority": 9.625e-06, + "cache_read_input_token_cost": 2.2e-07, + "cache_read_input_token_cost_flex": 1.1e-07, + "cache_read_input_token_cost_priority": 3.85e-07, + "input_cost_per_token": 4.4e-06, + "input_cost_per_token_flex": 2.2e-06, + "input_cost_per_token_priority": 7.7e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 868928, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.2e-05, + "output_cost_per_token_flex": 1.1e-05, + "output_cost_per_token_priority": 3.85e-05, + "reasoning_effort_levels": [ + "none", + "low", + "medium", + "high", + "xhigh", + "max" + ], + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k3-1-highspeed.md", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_none_reasoning_effort": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_regex_lookaround": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.moonshotai.kimi-k3.1": { + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_flex": 2.5e-06, + "cache_creation_input_token_cost_priority": 8.75e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_flex": 1e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_flex": 2e-06, + "input_cost_per_token_priority": 7e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 868928, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-05, + "output_cost_per_token_flex": 1e-05, + "output_cost_per_token_priority": 3.5e-05, + "reasoning_effort_levels": [ + "none", + "low", + "medium", + "high", + "xhigh", + "max" + ], + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k3-1-highspeed.md", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_none_reasoning_effort": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_regex_lookaround": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "in.moonshotai.kimi-k3.1": { + "cache_creation_input_token_cost": 5.5e-06, + "cache_creation_input_token_cost_flex": 2.75e-06, + "cache_creation_input_token_cost_priority": 9.625e-06, + "cache_read_input_token_cost": 2.2e-07, + "cache_read_input_token_cost_flex": 1.1e-07, + "cache_read_input_token_cost_priority": 3.85e-07, + "input_cost_per_token": 4.4e-06, + "input_cost_per_token_flex": 2.2e-06, + "input_cost_per_token_priority": 7.7e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 868928, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.2e-05, + "output_cost_per_token_flex": 1.1e-05, + "output_cost_per_token_priority": 3.85e-05, + "reasoning_effort_levels": [ + "none", + "low", + "medium", + "high", + "xhigh", + "max" + ], + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k3-1-highspeed.md", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_none_reasoning_effort": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_regex_lookaround": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "wandb/deepseek-ai/DeepSeek-V4.1-Flash-1017": { + "cache_read_input_token_cost": 2e-09, + "input_cost_per_token": 1.3e-07, + "litellm_provider": "wandb", + "max_input_tokens": 1049000, + "max_output_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6.5e-07, + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": true } } diff --git a/pyproject.toml b/pyproject.toml index c53fbb4d888..02834f604d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.109.0" +version = "1.110.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.16" @@ -54,6 +54,7 @@ Documentation = "https://docs.litellm.ai" # sub-packages and the opentelemetry trio move in lockstep, and grpcio is # supply-chain-pinned to a vetted, aged release. [project.optional-dependencies] +switchyard = ["nemo-switchyard==0.4.0a1"] proxy = [ "gunicorn>=23.0.0,<24.0", "uvicorn>=0.33.0,<1.0", @@ -84,8 +85,8 @@ proxy = [ "mcp>=2.3.0,<3", "httpx2>=2.10.0,<3", "pydantic>=2.12.0,<3", - "litellm-proxy-extras==0.4.112", - "litellm-enterprise==0.1.81", + "litellm-proxy-extras==0.4.114", + "litellm-enterprise==0.1.83", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", @@ -268,6 +269,7 @@ proxy-dev = [ "opentelemetry-instrumentation-fastapi==0.66b1", "azure-identity==1.25.2", "a2a-sdk==1.1.0", + "nemo-switchyard==0.4.0a1", "backports-zstd==1.8.0; python_version < '3.14'", "httptools>=0.7.1,<0.8", "semver>=3.0.4,<4", @@ -406,7 +408,7 @@ litellm-enterprise = { workspace = true } members = ["enterprise", "litellm-proxy-extras", "packaging/litellm-proxy"] [tool.commitizen] -version = "1.109.0" +version = "1.110.0" version_files = [ "pyproject.toml:^version", ] diff --git a/ruff.toml b/ruff.toml index 715e5551fad..037b15dbfe9 100644 --- a/ruff.toml +++ b/ruff.toml @@ -9,7 +9,7 @@ lint.extend-select = [ "PLW0133", "PYI030", "PYI041", "PYI064", "RET501", "RUF010", "RUF022", "RUF023", "RUF051", "S113", "SIM114", "SIM118", "SIM201", "SIM211", "SIM222", "TC005", "UP006", "UP007", "UP008", "UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045", - "C404", "C419", + "C404", "C419", "TID253", ] # RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip # `# noqa` directives that protect rules enforced elsewhere. List those codes as external @@ -35,6 +35,9 @@ format.exclude = ["**/enterprise/**"] lint.exclude = ["litellm/__init__.py", "litellm/proxy/example_config_yaml/*", "tests/*"] +[lint.flake8-tidy-imports] +banned-module-level-imports = ["prisma"] + [lint.per-file-ignores] "litellm/main.py" = ["F401"] "litellm/types/caching.py" = ["F401"] diff --git a/tests/e2e/ui/fixtures/migratedPages.ts b/tests/e2e/ui/fixtures/migratedPages.ts index 6fa8d0c2587..e49168ebb4a 100644 --- a/tests/e2e/ui/fixtures/migratedPages.ts +++ b/tests/e2e/ui/fixtures/migratedPages.ts @@ -79,7 +79,7 @@ export const MIGRATED_E2E_PAGES: Readonly> = { group: "Tools", content: { role: "heading", name: "Tool Policies" }, }, - skills: { segment: "skills", linkName: "Skills", content: { role: "heading", name: "Skills" } }, + skills: { segment: "skills", linkName: "Skills Hub", content: { role: "heading", name: "Skills Hub" } }, caching: { segment: "caching", linkName: "Response Cache", content: { role: "tab", name: "Cache Settings" } }, "cost-tracking": { segment: "cost-tracking", diff --git a/tests/e2e/ui/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts index 3943526aebf..af261044c84 100644 --- a/tests/e2e/ui/helpers/navigation.ts +++ b/tests/e2e/ui/helpers/navigation.ts @@ -79,7 +79,7 @@ export async function setAriaHidden(page: PlaywrightPage, hidden: boolean): Prom } await page.keyboard.press("Escape"); await expect(panel).toBeHidden(); - const aria = page.getByRole("banner").getByRole("button", { name: "Aria Agent", exact: true }); + const aria = sidebarLink(page, "Aria Agent"); if (premium && hidden) { await expect(aria).toBeHidden(); } else { diff --git a/tests/e2e/ui/tests/logs/logsPagination.spec.ts b/tests/e2e/ui/tests/logs/logsPagination.spec.ts index a611f760e8b..b81708d1cd1 100644 --- a/tests/e2e/ui/tests/logs/logsPagination.spec.ts +++ b/tests/e2e/ui/tests/logs/logsPagination.spec.ts @@ -1,6 +1,12 @@ import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; -import { navigateToPage, dismissFeedbackPopup, setAriaHidden } from "../../helpers/navigation"; +import { + navigateToPage, + dismissFeedbackPopup, + expectUiRoute, + setAriaHidden, + sidebarLink, +} from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; import { CHAT_MODEL_A, @@ -51,12 +57,18 @@ async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): await page.getByRole("option", { name: size, exact: true }).click(); } -test.describe("Logs pagination with Aria in the header", () => { +async function openAriaFromSidebar(page: PlaywrightPage): Promise { + await sidebarLink(page, "Aria Agent").click(); + await expectUiRoute(page, "aria"); + await expect(page.getByRole("region", { name: "Aria Agent", exact: true })).toBeVisible(); +} + +test.describe("Logs pagination with Aria in the sidebar", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); test("keeps pagination reachable and Aria available at 1440px", async ({ page, request }) => { await page.setViewportSize({ width: 1440, height: 900 }); - const alias = `e2e-aria-header-${uniqueSuffix()}`; + const alias = `e2e-aria-sidebar-${uniqueSuffix()}`; const key = await createVirtualKey(request, { key_alias: alias }); try { const ids = await Promise.all( @@ -76,20 +88,13 @@ test.describe("Logs pagination with Aria in the header", () => { await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); await visibleTestId(page, "pagination-next").click(); await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 26-26 of 26"); - - const aria = page.getByRole("banner").getByRole("button", { name: "Aria Agent", exact: true }); - await aria.click(); - await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); - await page.keyboard.press("Escape"); - await expect(aria).toBeFocused(); await visibleTestId(page, "pagination-first").click(); await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); + + await openAriaFromSidebar(page); for (const destination of [Page.NewUsage, Page.LlmPlayground, Page.ApiKeys]) { await navigateToPage(page, destination); - await aria.click(); - await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); - await page.keyboard.press("Escape"); - await expect(aria).toBeFocused(); + await openAriaFromSidebar(page); } } finally { await deleteVirtualKey(request, key.token); diff --git a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts index 517ec64c7df..18a685f052d 100644 --- a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts +++ b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts @@ -39,21 +39,16 @@ test.describe("AI Hub (internal admin view)", () => { }); }); - test("AI Hub tab list renders Model Hub, Agent Hub, MCP Hub and Skill Hub", async ({ page }) => { + test("AI Hub tab list renders Model Hub, Agent Hub, MCP Hub and Skills Hub", async ({ page }) => { await navigateToPage(page, Page.ModelHubTable); - // The tab strip lives in the main view; check each tab is present and clickable. - // (The "Claude Code Plugin Marketplace" tab from the manual-QA checklist was - // renamed to "Skill Hub" — verify the current label here so the test stays - // in sync with the UI.) - // // Note: unlike the public /ui/model_hub_table view (test below), the admin // ModelHubTable renders all four tabs unconditionally — there are no `&&` // guards around Agent Hub or MCP Hub in the source // (ModelHubTable.tsx ~L436-439). Asserting all four here is intentional: // this pins the manual-QA contract that the AI Hub tab strip exposes // exactly these labels regardless of seeded agent/MCP data. - for (const tabName of ["Model Hub", "Agent Hub", "MCP Hub", "Skill Hub"]) { + for (const tabName of ["Model Hub", "Agent Hub", "MCP Hub", "Skills Hub"]) { const tab = page.getByRole("tab", { name: tabName }); await expect(tab, `${tabName} tab should be present`).toBeVisible({ timeout: 5_000 }); await tab.click(); @@ -75,7 +70,7 @@ test.describe("Public model hub (/ui/model_hub_table)", () => { await expect(page.getByRole("tab")).toHaveCount(2); await page.getByRole("tab", { name: "Skills Hub", exact: true }).click(); await expect( - page.getByRole("textbox", { name: "Search skills", exact: true }), + page.getByRole("textbox", { name: "Search skills, tags, or people…", exact: true }), ).toBeVisible(); }); @@ -105,7 +100,7 @@ test.describe("Public model hub (/ui/model_hub_table)", () => { await dismissFeedbackPopup(page); await page.getByRole("checkbox", { name: domain, exact: true }).check(); await page - .getByRole("textbox", { name: "Search skills", exact: true }) + .getByRole("textbox", { name: "Search skills, tags, or people…", exact: true }) .fill(skillName); await page.getByRole("link").filter({ hasText: skillName }).click(); diff --git a/tests/integration/compatibility/test_codex_model_catalog.py b/tests/integration/compatibility/test_codex_model_catalog.py index 6aeeead0d90..a513088d082 100644 --- a/tests/integration/compatibility/test_codex_model_catalog.py +++ b/tests/integration/compatibility/test_codex_model_catalog.py @@ -36,7 +36,7 @@ pytestmark: Final = pytest.mark.shares_worker_state T = TypeVar("T") -CLIENT_VERSION: Final = "0.159.3" +CLIENT_VERSION: Final = "0.166.0" CONSISTENT_READS: Final = 6 CONVERGENCE_SECONDS: Final = 60 RELOAD_MARGIN_SECONDS: Final = 5.0 diff --git a/tests/integration/configuration/test_codex_catalog_boot.py b/tests/integration/configuration/test_codex_catalog_boot.py index 08c60debfcf..6564827e2c6 100644 --- a/tests/integration/configuration/test_codex_catalog_boot.py +++ b/tests/integration/configuration/test_codex_catalog_boot.py @@ -18,7 +18,7 @@ from tests.integration._support.process import owned_proxy, owned_proxy_process pytestmark: Final = pytest.mark.shares_worker_state -CLIENT_VERSION: Final = "0.159.3" +CLIENT_VERSION: Final = "0.166.0" FRESH_CONNECTION: Final = {"Connection": "close"} CONSISTENT_READS: Final = 8 CONVERGENCE_SECONDS: Final = 60 diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 8f11b3ce2f1..ffdb6752bf2 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -1603,32 +1603,44 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { }, "azure/gpt-6-luna": { "cache_read_input_token_cost": 1.3e-08, + "cache_read_input_token_cost_flex": 6.5e-09, "cache_creation_input_token_cost": 1.625e-07, + "cache_creation_input_token_cost_flex": 8.125e-08, "cache_creation_input_token_cost_above_272k_tokens": 3.25e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.625e-07, "input_cost_per_token": 1.3e-07, + "input_cost_per_token_flex": 6.5e-08, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.5e-07, + "output_cost_per_token_flex": 3.25e-07, "supports_function_calling": True, "supports_none_reasoning_effort": True, "input_cost_per_token_above_272k_tokens": 2.6e-07, + "input_cost_per_token_above_272k_tokens_flex": 1.3e-07, "output_cost_per_token_above_272k_tokens": 1.3e-06, + "output_cost_per_token_above_272k_tokens_flex": 4.875e-07, "cache_read_input_token_cost_above_272k_tokens": 2.6e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 1.3e-08, }, "azure/gpt-6-astra": { "cache_read_input_token_cost": 1.25e-06, + "cache_read_input_token_cost_flex": 6.25e-07, "cache_creation_input_token_cost": 1.5625e-05, + "cache_creation_input_token_cost_flex": 7.8125e-06, "cache_creation_input_token_cost_above_272k_tokens": 3.125e-05, "input_cost_per_token": 1.25e-05, + "input_cost_per_token_flex": 6.25e-06, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.25e-05, + "output_cost_per_token_flex": 3.125e-05, "supports_function_calling": True, "supports_none_reasoning_effort": True, "input_cost_per_token_above_272k_tokens": 2.5e-05, @@ -1637,20 +1649,28 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { }, "azure/gpt-6.1-sol": { "cache_read_input_token_cost": 2.4e-07, + "cache_read_input_token_cost_flex": 6e-08, "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_flex": 1.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 3e-06, "input_cost_per_token": 2.4e-06, + "input_cost_per_token_flex": 1.2e-06, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.2e-05, + "output_cost_per_token_flex": 6e-06, "supports_function_calling": True, "supports_none_reasoning_effort": True, "input_cost_per_token_above_272k_tokens": 4.8e-06, + "input_cost_per_token_above_272k_tokens_flex": 2.4e-06, "output_cost_per_token_above_272k_tokens": 2.4e-05, + "output_cost_per_token_above_272k_tokens_flex": 9e-06, "cache_read_input_token_cost_above_272k_tokens": 4.8e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 1.2e-07, }, } diff --git a/tests/integration/management/test_tag_budget_update.py b/tests/integration/management/test_tag_budget_update.py new file mode 100644 index 00000000000..c8ce1bdc8f4 --- /dev/null +++ b/tests/integration/management/test_tag_budget_update.py @@ -0,0 +1,20 @@ +from typing import Final +from uuid import uuid4 + +from tests.integration._support.client import Gateway, object_value, string_value + + +def test_tag_budget_can_be_updated_with_its_existing_budget_id(gateway: Gateway) -> None: + name: Final = f"integration-tag-{uuid4().hex}" + gateway.post("/tag/new", {"name": name, "max_budget": 10}) + with gateway.scenario() as scenario: + scenario.cleanups.callback(gateway.post, "/tag/delete", {"name": name}) + before: Final = object_value(gateway.post("/tag/info", {"names": [name]})[name]) + budget_id: Final = string_value(object_value(before["litellm_budget_table"])["budget_id"]) + scenario.cleanups.callback(gateway.post, "/budget/delete", {"id": budget_id}) + response: Final = gateway.request( + "POST", "/tag/update", {"name": name, "budget_id": budget_id, "max_budget": 20} + ) + assert response.status_code == 200, response.text + after: Final = object_value(gateway.post("/tag/info", {"names": [name]})[name]) + assert object_value(after["litellm_budget_table"])["max_budget"] == 20, after diff --git a/tests/integration/observability/test_otel_team_request_metadata.py b/tests/integration/observability/test_otel_team_request_metadata.py new file mode 100644 index 00000000000..b5caf71c14d --- /dev/null +++ b/tests/integration/observability/test_otel_team_request_metadata.py @@ -0,0 +1,227 @@ +"""A team's or key's own New Relic destination gets the operator-allowlisted request metadata as +``litellm.metadata.*``, whether or not ``otel`` is listed in ``litellm_settings.callbacks`` (LIT-9783).""" + +from __future__ import annotations + +import json +import uuid +from collections.abc import Iterator, Mapping +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal + +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value +from integration._support.otlp_sink import ( + ConnectSink, + Span, + owned_connect_sink, + recorded_requests, + recorded_spans, + spans_for_trace, +) +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +OPERATION: Final = "gen_ai.operation.name" +SERVER_KIND: Final = 2 +PROMOTED: Final = ("litellm.metadata.correlation_id", "litellm.metadata.correlation.id") +ALLOWLIST: Final = ( + "requester_metadata.correlation_id", + "requester_metadata.correlation.id", + "requester_metadata.headers.authorization", + "requester_metadata.headers.x-api-key", + "requester_metadata.headers.cookie", + "requester_metadata.headers.proxy-authorization", + "user_api_key_user_id", + "user_api_key_alias", + "requester_ip_address", +) + + +def _chat(_request: Request) -> Reply: + return Reply( + body=( + b'{"id":"chatcmpl-meta","object":"chat.completion","created":1,"model":"gpt-4o-mini",' + b'"choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],' + b'"usage":{"prompt_tokens":7,"completion_tokens":2,"total_tokens":9}}' + ) + ) + + +def _config(directory: Path, callbacks: tuple[str, ...]) -> Path: + loaded: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = { + **loaded, + "litellm_settings": {**loaded.get("litellm_settings", {}), "callbacks": list(callbacks)}, + "general_settings": {**loaded.get("general_settings", {}), "user_api_key_cache_ttl": 2}, + } + path: Final = directory / "team-request-metadata.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def provider() -> Iterator[Wire]: + with wire_server(_chat) as wire: + yield wire + + +@pytest.fixture(scope="module") +def newrelic(tmp_path_factory: pytest.TempPathFactory) -> Iterator[ConnectSink]: + with owned_connect_sink(tmp_path_factory.mktemp("newrelic-request-metadata")) as sink: + yield sink + + +@pytest.fixture(scope="module", params=[(), ("otel",)], ids=["team-only", "with-proxy-otel"]) +def proxy( + request: pytest.FixtureRequest, newrelic: ConnectSink, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("team-request-metadata") + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + directory, + { + "LITELLM_OTEL_V2": "1", + "OTEL_BSP_SCHEDULE_DELAY": "300", + "LITELLM_OTEL_BAGGAGE_METADATA_KEYS": ",".join(ALLOWLIST), + "HTTPS_PROXY": newrelic.proxy_url, + "NO_PROXY": "127.0.0.1,localhost", + "REQUESTS_CA_BUNDLE": newrelic.ca_pem, + }, + config=_config(directory, request.param), + remove_environment=("OTEL_ENDPOINT", "OTEL_EXPORTER", "OTEL_EXPORTER_OTLP_ENDPOINT"), + workers=2, + ) as owned, + ): + yield owned.gateway + + +def _server_and_model_spans(spans: tuple[Span, ...]) -> tuple[Span, ...]: + return tuple(span for span in spans if span["kind"] == SERVER_KIND or OPERATION in span["attributes"]) + + +def _team_callback_key(proxy: Gateway, scenario: Scenario, destination: Mapping[str, JsonValue]) -> str: + team: Final = scenario.team() + attached: Final = proxy.request("POST", f"/team/{team}/callback", dict(destination)) + assert attached.status_code == 200, attached.text + return scenario.key(team_id=team) + + +@pytest.mark.parametrize("level", ["team", "key"]) +def test_tenant_newrelic_spans_carry_allowlisted_request_metadata( + level: Literal["team", "key"], proxy: Gateway, provider: Wire, newrelic: ConnectSink +) -> None: + with proxy.scenario() as scenario: + model: Final = scenario.model(api_base=provider.url + "/v1") + destination: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "callback_name": "newrelic", + "callback_type": "success_and_failure", + "callback_vars": {"newrelic_api_key": f"nr-{level}-{uuid.uuid4().hex}", "newrelic_region": "us"}, + } + ) + key: Final = ( + _team_callback_key(proxy, scenario, destination) + if level == "team" + else scenario.key(metadata={"logging": [dict(destination)]}) + ) + since: Final = recorded_spans(newrelic.control_url)[0] + trace_id: Final = uuid.uuid4().hex + correlation: Final = f"corr-{trace_id}" + response: Final = proxy.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"correlation_id": correlation, "correlation.id": f"{correlation}-dot"}, + }, + key=key, + headers={"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"}, + ) + assert response.status_code == 200, response.text + + spans: Final = _server_and_model_spans( + eventually( + lambda: spans_for_trace(recorded_spans(newrelic.control_url, since)[1], trace_id), + lambda found: len(_server_and_model_spans(found)) >= 2, + seconds=40, + ) + ) + + assert sorted(span["kind"] == SERVER_KIND for span in spans) == [False, True], [span["name"] for span in spans] + assert [{name: span["attributes"].get(name) for name in PROMOTED} for span in spans] == [ + {PROMOTED[0]: correlation, PROMOTED[1]: f"{correlation}-dot"} + ] * len(spans), [(span["name"], sorted(span["attributes"])) for span in spans] + + +def _export_credentials(newrelic: ConnectSink, span_ids: frozenset[str]) -> frozenset[str]: + return frozenset( + str(object_value(request.get("headers") or {}).get("api-key")) + for request in recorded_requests(newrelic.control_url) + if span_ids & {span_id for span_id in (request.get("span_ids") or ()) if isinstance(span_id, str)} + ) + + +def test_tenant_newrelic_spans_never_carry_proxy_or_caller_credentials( + proxy: Gateway, provider: Wire, newrelic: ConnectSink +) -> None: + with proxy.scenario() as scenario: + model: Final = scenario.model(api_base=provider.url + "/v1") + team_newrelic_key: Final = f"nr-team-{uuid.uuid4().hex}" + key: Final = _team_callback_key( + proxy, + scenario, + { + "callback_name": "newrelic", + "callback_type": "success_and_failure", + "callback_vars": {"newrelic_api_key": team_newrelic_key, "newrelic_region": "us"}, + }, + ) + planted_newrelic_key: Final = f"nr-planted-{uuid.uuid4().hex}" + credential_headers: Final = MappingProxyType( + { + "cookie": f"session=cookie-{uuid.uuid4().hex}", + "x-api-key": f"xapikey-{uuid.uuid4().hex}", + "proxy-authorization": f"Basic proxyauth-{uuid.uuid4().hex}", + } + ) + since: Final = recorded_spans(newrelic.control_url)[0] + trace_id: Final = uuid.uuid4().hex + response: Final = proxy.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "correlation_id": f"corr-{trace_id}", + "callback_vars": {"newrelic_api_key": planted_newrelic_key}, + }, + }, + key=key, + headers={"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01", **credential_headers}, + ) + assert response.status_code == 200, response.text + traced: Final = eventually( + lambda: spans_for_trace(recorded_spans(newrelic.control_url, since)[1], trace_id), + lambda found: len(_server_and_model_spans(found)) >= 2, + seconds=40, + ) + + exported: Final = json.dumps([[span["attributes"], span["resource"]] for span in traced]) + secrets: Final = { + "master key": proxy.key, + "virtual key": key, + "team New Relic key": team_newrelic_key, + **credential_headers, + } + assert [name for name, secret in secrets.items() if secret in exported] == [], [ + sorted(span["attributes"]) for span in traced + ] + assert _export_credentials(newrelic, frozenset(span["span_id"] for span in traced)) == {team_newrelic_key} diff --git a/tests/integration/providers/test_bedrock_batch_files_wire.py b/tests/integration/providers/test_bedrock_batch_files_wire.py index 823bcda7e06..162fc5aac27 100644 --- a/tests/integration/providers/test_bedrock_batch_files_wire.py +++ b/tests/integration/providers/test_bedrock_batch_files_wire.py @@ -88,10 +88,17 @@ _SCHEMA: Final = { @pytest.mark.parametrize( - "backend", ["us.anthropic.claude-haiku-5-5", "us.anthropic.claude-opus-4-8"], ids=["haiku_5_5", "opus_4_8"] + ("backend", "native_output"), + [ + ("us.anthropic.claude-sonnet-5-5", False), + ("us.anthropic.claude-haiku-5-5", True), + ("us.anthropic.claude-opus-4-7", False), + ("us.anthropic.claude-opus-4-8", True), + ], + ids=["opt_in_native_output_model", "native_output_model", "opus_4_7", "opus_4_8"], ) -def test_json_schema_batch_record_with_thinking_uploads_the_json_tool_and_no_native_output( - gateway: Gateway, backend: str +def test_json_schema_batch_record_with_thinking_uploads_native_output_only_on_models_that_always_take_it( + gateway: Gateway, backend: str, native_output: bool ) -> None: record: Final = { "custom_id": "schema-record", @@ -125,7 +132,16 @@ def test_json_schema_batch_record_with_thinking_uploads_the_json_tool_and_no_nat (upload,) = wire.drain() (stored,) = tuple(JSON_OBJECT.validate_json(line) for line in upload.body.decode().splitlines() if line.strip()) model_input: Final = object_value(stored["modelInput"]) - assert "format" not in object_value(model_input.get("output_config", {})), model_input - (tool,) = (object_value(entry) for entry in list_value(model_input["tools"])) + output_format: Final = object_value(model_input.get("output_config", {})).get("format") + tools: Final = tuple(object_value(entry) for entry in list_value(model_input.get("tools", []))) + if native_output: + assert model_input["output_config"] == {"format": {"type": "json_schema", "schema": _SCHEMA}}, model_input + assert model_input["anthropic_beta"] == ["structured-outputs-2025-11-13"], model_input + assert "tools" not in model_input, model_input + assert model_input["thinking"] == {"type": "adaptive"}, model_input + assert not tools, model_input + return + assert output_format is None, model_input + (tool,) = tools assert tool["name"] == _JSON_TOOL, model_input assert object_value(tool["input_schema"])["properties"] == _SCHEMA["properties"], model_input diff --git a/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py b/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py index 0788b9ac85a..159097c80ee 100644 --- a/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py +++ b/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py @@ -121,3 +121,28 @@ async def test_async_openai_sdk_responses_stream_is_served_by_the_native_respons request: Final = _native_request(wire) assert _body(request) == {"model": GPT, "input": _prompt(marker), "stream": True}, request.body assert _spend_row(completed.response.id) == _success_row(model) + + +def test_codex_agent_message_with_encrypted_content_reaches_the_native_responses_route_intact( + gateway: Gateway, +) -> None: + marker: Final = uuid.uuid4().hex + agent_message: Final[JsonValue] = { + "type": "agent_message", + "author": "/root/reviewer", + "recipient": "/root", + "content": [ + {"type": "input_text", "text": "Worker message:\n"}, + {"type": "encrypted_content", "encrypted_content": f"gAAAAABsynthetic-agent-ciphertext-{marker}=="}, + ], + } + user_turn: Final[JsonValue] = {"type": "message", "role": "user", "content": _prompt(marker)} + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire) + response: Final = gateway.request( + "POST", + "/v1/responses", + {"model": model, "input": [agent_message, user_turn], "cache": {"no-cache": True}}, + ) + assert response.status_code == 200, response.text + assert _body(_native_request(wire))["input"] == [agent_message, user_turn] diff --git a/tests/integration/providers/test_bedrock_mantle_codex_input_wire.py b/tests/integration/providers/test_bedrock_mantle_codex_input_wire.py index aa66e82475b..1383910f08c 100644 --- a/tests/integration/providers/test_bedrock_mantle_codex_input_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_codex_input_wire.py @@ -88,6 +88,35 @@ def test_codex_agent_message_context_compaction_and_local_shell_call_reach_mantl assert JSON_OBJECT.validate_json(forwarded[0].body)["input"] == expected_input, forwarded[0].body +def test_codex_agent_message_with_encrypted_content_reaches_mantle_intact(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + agent_message: Final[JsonValue] = { + "type": "agent_message", + "author": "/root/reviewer", + "recipient": "/root", + "content": [ + {"type": "input_text", "text": "Worker message:\n"}, + {"type": "encrypted_content", "encrypted_content": f"gAAAAABsynthetic-agent-ciphertext-{marker}=="}, + ], + } + expected_input: Final = [agent_message, user_turn(f"next turn {marker}")] + + def mantle_peer(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/openai/v1/responses", request.target + assert JSON_OBJECT.validate_json(request.body)["input"] == expected_input, request.body + return Reply(body=RESPONSE) + + with wire_server(mantle_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=MODEL, api_key=TOKEN, api_base=wire.url, aws_region_name="us-east-2") + response: Final = gateway.request( + "POST", "/v1/responses", {"model": model, "input": expected_input, "store": False} + ) + assert response.status_code == 200, response.text + forwarded: Final = wire.drain() + assert len(forwarded) == 1, forwarded + assert JSON_OBJECT.validate_json(forwarded[0].body)["input"] == expected_input, forwarded[0].body + + SHELL_TOOL: Final[JsonValue] = { "type": "function", "name": "shell", diff --git a/tests/integration/providers/test_bedrock_mantle_responses_wire.py b/tests/integration/providers/test_bedrock_mantle_responses_wire.py index 3bd83b5019b..84340132fbc 100644 --- a/tests/integration/providers/test_bedrock_mantle_responses_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_responses_wire.py @@ -44,7 +44,7 @@ def _codex_history(marker: str) -> list[JsonValue]: { "type": "agent_message", "id": "msg_agent", - "content": [{"type": "text", "text": "sub-agent said "}, {"type": "text", "encrypted_content": "hello"}], + "content": [{"type": "input_text", "text": "sub-agent said "}, {"type": "input_text", "text": "hello"}], }, {"type": "context_compaction", "id": "cmp_1", "encrypted_content": "compacted-history"}, { diff --git a/tests/integration/providers/test_bedrock_native_structured_output_wire.py b/tests/integration/providers/test_bedrock_native_structured_output_wire.py index 2502bd90bd6..ceb40879b9f 100644 --- a/tests/integration/providers/test_bedrock_native_structured_output_wire.py +++ b/tests/integration/providers/test_bedrock_native_structured_output_wire.py @@ -18,9 +18,10 @@ from pydantic import JsonValue, TypeAdapter _HAIKU: Final = "us.anthropic.claude-haiku-5-5" _FABLE: Final = "us.anthropic.claude-fable-5" _SONNET: Final = "us.anthropic.claude-sonnet-4-6" -_OPUS: Final = "us.anthropic.claude-opus-5-5" -_OPT_IN: Final = _OPUS -_UNFORCEABLE_TOOL: Final = _OPUS +_NEITHER_KEY: Final = "us.anthropic.claude-opus-4-7" +_OPT_IN: Final = _NEITHER_KEY +_BEDROCK_KEY_ONLY: Final = "us.anthropic.claude-sonnet-5-5" +_UNFORCEABLE_TOOL: Final = "anthropic.claude-opus-5-5" _OPUS_4_8: Final = "us.anthropic.claude-opus-4-8" _DEEPSEEK: Final = "deepseek.v3.2" _AWS: Final[dict[str, JsonValue]] = { @@ -622,11 +623,9 @@ def test_messages_thinking_on_a_bedrock_converse_native_model_gets_native_output @pytest.mark.parametrize("route", ["converse", "invoke"]) -@pytest.mark.parametrize( - ("backend", "thinks_by_default"), [(_OPUS_4_8, False), (_HAIKU, True)], ids=["opus_4_8", "haiku_5_5"] -) -def test_bedrock_native_model_with_no_thinking_param_gets_native_output_only_when_it_thinks_by_default( - gateway: Gateway, route: str, backend: str, thinks_by_default: bool +@pytest.mark.parametrize("backend", [_OPUS_4_8, _HAIKU], ids=["opus_4_8", "haiku_5_5"]) +def test_bedrock_native_model_with_no_thinking_param_gets_native_output( + gateway: Gateway, route: str, backend: str ) -> None: marker, marker_hex = _marker() with wire_server(_bedrock) as wire, gateway.scenario() as scenario: @@ -639,11 +638,9 @@ def test_bedrock_native_model_with_no_thinking_param_gets_native_output_only_whe assert not message.get("tool_calls"), response.text _, body = _one_call(wire, marker_hex) invoke: Final = route == "invoke" - if thinks_by_default: - _assert_native_without_tools(body, invoke, _SCHEMA) - return - assert _native_schema(body, invoke) is None, body - assert _tools(body, invoke)[_JSON_TOOL]["properties"] == _SCHEMA["properties"], body + _assert_native_without_tools(body, invoke, _SCHEMA) + if invoke: + assert body["anthropic_beta"] == ["structured-outputs-2025-11-13"], body @pytest.mark.parametrize("route", ["converse", "invoke"]) @@ -683,11 +680,16 @@ def test_a_model_that_cannot_force_the_json_tool_retries_a_rejected_native_outpu @pytest.mark.parametrize( ("backend", "model_info", "native"), [ - (_OPUS, None, False), + (_NEITHER_KEY, None, False), (_OPT_IN, {"supports_bedrock_native_structured_output": True}, True), (_OPUS_4_8, {"supports_bedrock_native_structured_output": False}, False), + ( + _OPUS_4_8, + {"supports_native_structured_output": True, "supports_bedrock_native_structured_output": False}, + True, + ), ], - ids=["neither_key", "deployment_opt_in", "deployment_opt_out"], + ids=["neither_key", "deployment_opt_in", "deployment_opt_out", "explicit_native_with_bedrock_key_off"], ) def test_thinking_request_follows_the_bedrock_native_key_the_model_resolves_to( gateway: Gateway, backend: str, model_info: Mapping[str, JsonValue] | None, native: bool @@ -897,7 +899,7 @@ def test_messages_on_bedrock_invoke_follows_the_bedrock_native_key_without_think ) -> None: marker, marker_hex = _marker() with wire_server(_bedrock) as wire, gateway.scenario() as scenario: - model: Final = _deployment(scenario, wire, "invoke", _HAIKU) + model: Final = _deployment(scenario, wire, "invoke", _BEDROCK_KEY_ONLY) response: Final = gateway.request( "POST", "/v1/messages", @@ -945,9 +947,9 @@ def test_malformed_schema_on_a_native_model_reaches_bedrock_once_without_a_proxy ("route", "backend", "thinking", "schema", "native"), [ ("converse", _SONNET, False, _NUMERIC_FORMATS_SCHEMA, True), - ("invoke", _HAIKU, True, _NUMERIC_FORMATS_SCHEMA, True), + ("invoke", _BEDROCK_KEY_ONLY, True, _NUMERIC_FORMATS_SCHEMA, True), ("converse", _SONNET, False, _WIDE_SCHEMA, True), - ("converse", _HAIKU, True, _WIDE_SCHEMA, False), + ("converse", _BEDROCK_KEY_ONLY, True, _WIDE_SCHEMA, False), ], ids=["numeric_formats_old_key", "numeric_formats_new_key_invoke", "wide_old_key", "wide_new_key"], ) @@ -1046,7 +1048,7 @@ def test_bedrock_rejecting_opt_in_native_output_gets_one_retry_without_it( ids=["other_400", "native_model"], ) @pytest.mark.parametrize("route", ["converse", "invoke"]) -def test_bedrock_400_is_not_retried_unless_it_rejects_opt_in_native_output( +def test_bedrock_400_is_not_retried_unless_it_rejects_native_output( gateway: Gateway, route: str, backend: str, rejection: str ) -> None: marker, marker_hex = _marker() @@ -1057,10 +1059,19 @@ def test_bedrock_400_is_not_retried_unless_it_rejects_opt_in_native_output( "/v1/chat/completions", _chat_body(model, marker, thinking=_ADAPTIVE, response_format=_response_format(_SCHEMA)), ) - assert response.status_code == 400, response.text - assert rejection.split(": ", 1)[1] in response.text, response.text - _, body = _one_call(wire, marker_hex) - assert _native_schema(body, route == "invoke") == _SCHEMA, body + invoke: Final = route == "invoke" + if rejection == _MAX_TOKENS_400: + assert response.status_code == 400, response.text + assert rejection.split(": ", 1)[1] in response.text, response.text + _, body = _one_call(wire, marker_hex) + _assert_native_without_tools(body, invoke, _SCHEMA) + return + assert _content_json(_chat_message(response)) == _payload(marker_hex), response.text + rejected, retried = (_JSON.validate_json(request.body) for request in wire.drain()) + _assert_native_without_tools(rejected, invoke, _SCHEMA) + assert _native_schema(retried, invoke) is None, retried + assert _tools(retried, invoke)[_JSON_TOOL]["properties"] == _SCHEMA["properties"], retried + assert marker_hex in json.dumps(retried), retried def _always_rejecting(message: str) -> Callable[[Request], Reply]: @@ -1075,7 +1086,7 @@ def _always_rejecting(message: str) -> Callable[[Request], Reply]: [(_OPUS_4_8, _MAX_TOKENS_400), (_SONNET, _FABLE_REJECTION)], ids=["other_400", "native_model"], ) -def test_messages_400_on_bedrock_invoke_reaches_the_caller_after_one_call_unless_it_rejects_opt_in_native_output( +def test_messages_400_on_bedrock_invoke_reaches_the_caller_after_one_call_unless_it_rejects_native_output( gateway: Gateway, backend: str, rejection: str ) -> None: marker, marker_hex = _marker() @@ -1088,7 +1099,17 @@ def test_messages_400_on_bedrock_invoke_reaches_the_caller_after_one_call_unless ) assert response.status_code == 400, response.text assert rejection.split(": ", 1)[1] in response.text, response.text - _one_call(wire, marker_hex) + if rejection == _MAX_TOKENS_400: + _, body = _one_call(wire, marker_hex) + _assert_native_without_tools(body, True, _SCHEMA) + return + rejected, retried = (_JSON.validate_json(request.body) for request in wire.drain()) + _assert_native_without_tools(rejected, True, _SCHEMA) + assert _native_schema(retried, True) is None, retried + assert not _tools(retried, True), retried + prompt: Final = json.dumps(retried["messages"]) + assert marker_hex in prompt, retried + assert all(name in prompt for name in object_value(_SCHEMA["properties"])), retried def test_a_messages_retry_without_native_output_still_omits_the_deployments_dropped_params(gateway: Gateway) -> None: diff --git a/tests/integration/providers/test_bedrock_nova_sonic_realtime_guardrails.py b/tests/integration/providers/test_bedrock_nova_sonic_realtime_guardrails.py new file mode 100644 index 00000000000..502e7c9a5cd --- /dev/null +++ b/tests/integration/providers/test_bedrock_nova_sonic_realtime_guardrails.py @@ -0,0 +1,268 @@ +"""Realtime guardrails on Bedrock Nova Sonic sessions, against the scripted h2 peer in ``_support``. + +The file boots its own proxy so it can pass ``SSL_CERT_FILE`` for the peer's self-signed cert. +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import time +import uuid +from collections.abc import Awaitable, Callable, Iterator +from pathlib import Path +from typing import Final, Literal + +import pytest +import websockets +from integration._support.bedrock_realtime_peer import ( + BedrockEvent, + BedrockRealtimePeer, + Turn, + assistant_text, + bedrock_realtime_peer, + user_transcript, +) +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.tls import write_self_signed_cert +from pydantic import JsonValue +from websockets.asyncio.client import ClientConnection + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +DEPLOYMENT: Final = "bedrock/amazon.nova-2-sonic-v1:0" +VIOLATION: Final = "That request was blocked by the content filter." +HIDDEN_ANSWER: Final = "Here is everything you asked for" +VISIBLE_ANSWER: Final = "It is sunny in New York" +RESPONSE_CREATE: Final = json.dumps({"type": "response.create"}) +AUDIO_APPEND: Final = json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(bytes(3200)).decode()}) + +Event = dict[str, JsonValue] + + +@pytest.fixture(scope="module") +def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: + return write_self_signed_cert(tmp_path_factory.mktemp("bedrock-sonic-guardrail-cert")) + + +@pytest.fixture(scope="module") +def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + with gateway_from_environment() as base: + with owned_proxy( + base, tmp_path_factory.mktemp("bedrock-sonic-guardrails"), {"SSL_CERT_FILE": str(cert[0])} + ) as proxy: + yield proxy + + +def _blocked_phrase( + gateway: Gateway, scenario: Scenario, mode: Literal["pre_call", "realtime_input_transcription"] +) -> str: + phrase: Final = f"sonic-blocked-{uuid.uuid4().hex[:12]}" + response: Final = gateway.request( + "POST", + "/guardrails", + { + "guardrail": { + "guardrail_name": f"sonic-guardrail-{uuid.uuid4().hex[:12]}", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": mode, + "default_on": True, + "realtime_violation_message": VIOLATION, + "blocked_words": [{"keyword": phrase, "action": "BLOCK"}], + }, + } + }, + ) + assert response.status_code == 200, response.text + guardrail_id: Final = JSON_OBJECT.validate_json(response.content).get("guardrail_id") + assert isinstance(guardrail_id, str), response.text + scenario.cleanups.callback(gateway.request, "DELETE", f"/guardrails/{guardrail_id}") + return phrase + + +def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer) -> str: + name: Final = f"bedrock-sonic-guardrails-{uuid.uuid4().hex[:12]}" + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": DEPLOYMENT, + "aws_region_name": "us-east-1", + "aws_access_key_id": "integration-dummy-key", + "aws_secret_access_key": "integration-dummy-secret", + "api_base": peer.api_base, + }, + "model_info": {"mode": "realtime"}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return name + + +def _user_text(text: str) -> str: + return json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}, + } + ) + + +async def _events_until(socket: ClientConnection, last: Callable[[Event], bool]) -> tuple[Event, ...]: + collected: Final[list[Event]] = [] + deadline: Final = time.monotonic() + 60 + while not collected or not last(collected[-1]): + try: + raw: Final = await asyncio.wait_for(socket.recv(), deadline - time.monotonic()) + except (TimeoutError, asyncio.TimeoutError): + raise AssertionError(f"timed out waiting; types={[event.get('type') for event in collected]}") + collected.append(JSON_OBJECT.validate_json(raw)) + return tuple(collected) + + +async def _events_until_closed(socket: ClientConnection) -> tuple[Event, ...]: + return tuple([JSON_OBJECT.validate_json(raw) async for raw in socket]) + + +def _is_done(event: Event) -> bool: + return event.get("type") == "response.done" + + +def _is_guardrail_error(event: Event) -> bool: + error: Final = event.get("error") + return event.get("type") == "error" and isinstance(error, dict) and error.get("type") == "guardrail_violation" + + +async def _session( + gateway: Gateway, model: str, key: str, exchange: Callable[[ClientConnection], Awaitable[tuple[Event, ...]]] +) -> tuple[Event, ...]: + proxy: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://") + async with websockets.connect( + f"{proxy}/v1/realtime?model={model}", additional_headers={"Authorization": f"Bearer {key}"} + ) as socket: + created: Final = JSON_OBJECT.validate_json(await socket.recv()) + assert created.get("type") == "session.created", created + return await exchange(socket) + + +def _fields(event: BedrockEvent, name: str) -> dict[str, JsonValue]: + return object_value(object_value(event["event"]).get(name) or {}) + + +def _user_text_inputs(received: tuple[BedrockEvent, ...]) -> tuple[str, ...]: + content_starts: Final = tuple(_fields(event, "contentStart") for event in received) + user_blocks: Final = frozenset(start.get("contentName") for start in content_starts if start.get("role") == "USER") + text_inputs: Final = tuple(_fields(event, "textInput") for event in received) + return tuple( + string_value(text_input["content"]) + for text_input in text_inputs + if text_input.get("contentName") in user_blocks + ) + + +def _text_deltas(events: tuple[Event, ...]) -> str: + return "".join(string_value(event["delta"]) for event in events if event.get("type") == "response.text.delta") + + +def test_blocked_typed_text_never_reaches_nova_sonic_and_the_violation_is_voiced( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + turn: Final = Turn("user_input", user_transcript("ready") + assistant_text(VIOLATION)) + with bedrock_realtime_peer((turn,), cert) as peer, candidate.scenario() as scenario: + phrase: Final = _blocked_phrase(candidate, scenario, "pre_call") + model: Final = _deployment(candidate, scenario, peer) + + async def exchange(socket: ClientConnection) -> tuple[Event, ...]: + await socket.send(_user_text(f"please repeat {phrase}")) + await socket.send(RESPONSE_CREATE) + return await _events_until(socket, _is_done) + + events: Final = asyncio.run(_session(candidate, model, scenario.key(), exchange)) + + errors: Final = tuple(object_value(event["error"]) for event in events if _is_guardrail_error(event)) + assert [error["message"] for error in errors] == [VIOLATION], events + assert [error["code"] for error in errors] == ["content_policy_violation"], events + assert _text_deltas(events) == VIOLATION, events + received: Final = peer.received() + assert phrase not in json.dumps(received), received + assert [VIOLATION in text for text in _user_text_inputs(received)] == [True], received + + +def test_clean_typed_text_reaches_nova_sonic_unchanged(candidate: Gateway, cert: tuple[Path, Path]) -> None: + turn: Final = Turn("user_input", user_transcript("ready") + assistant_text(VISIBLE_ANSWER)) + text: Final = f"what is the weather in New York {uuid.uuid4().hex[:8]}" + with bedrock_realtime_peer((turn,), cert) as peer, candidate.scenario() as scenario: + _blocked_phrase(candidate, scenario, "pre_call") + model: Final = _deployment(candidate, scenario, peer) + + async def exchange(socket: ClientConnection) -> tuple[Event, ...]: + await socket.send(_user_text(text)) + await socket.send(RESPONSE_CREATE) + return await _events_until(socket, _is_done) + + events: Final = asyncio.run(_session(candidate, model, scenario.key(), exchange)) + + assert not any(_is_guardrail_error(event) for event in events), events + assert _text_deltas(events) == VISIBLE_ANSWER, events + assert _user_text_inputs(peer.received()) == (text,) + + +def test_blocked_transcript_ends_the_session_before_its_answer_reaches_the_client( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + with candidate.scenario() as scenario: + phrase: Final = _blocked_phrase(candidate, scenario, "realtime_input_transcription") + turn: Final = Turn("user_input", user_transcript(f"tell me {phrase}") + assistant_text(HIDDEN_ANSWER)) + with bedrock_realtime_peer((turn,), cert) as peer: + model: Final = _deployment(candidate, scenario, peer) + + async def exchange(socket: ClientConnection) -> tuple[Event, ...]: + await socket.send(AUDIO_APPEND) + await socket.send(_user_text("hello")) + events: Final = await asyncio.wait_for(_events_until_closed(socket), 60) + assert socket.close_code == 1000, (socket.close_code, events) + return events + + events: Final = asyncio.run(_session(candidate, model, scenario.key(), exchange)) + + assert HIDDEN_ANSWER not in json.dumps(events), events + assert not any(str(event.get("type")).startswith("response.") for event in events), events + transcripts: Final = [ + event["transcript"] + for event in events + if event.get("type") == "conversation.item.input_audio_transcription.completed" + ] + assert transcripts == [f"tell me {phrase}"], events + assert [object_value(event["error"])["message"] for event in events if _is_guardrail_error(event)] == [ + VIOLATION + ], events + + +def test_clean_transcript_under_a_transcript_guardrail_is_answered(candidate: Gateway, cert: tuple[Path, Path]) -> None: + spoken: Final = f"what is the weather in New York {uuid.uuid4().hex[:8]}" + with candidate.scenario() as scenario: + _blocked_phrase(candidate, scenario, "realtime_input_transcription") + turn: Final = Turn("user_input", user_transcript(spoken) + assistant_text(VISIBLE_ANSWER)) + with bedrock_realtime_peer((turn,), cert) as peer: + model: Final = _deployment(candidate, scenario, peer) + + async def exchange(socket: ClientConnection) -> tuple[Event, ...]: + await socket.send(AUDIO_APPEND) + await socket.send(_user_text("hello")) + return await _events_until(socket, _is_done) + + events: Final = asyncio.run(_session(candidate, model, scenario.key(), exchange)) + + assert not any(_is_guardrail_error(event) for event in events), events + assert _text_deltas(events) == VISIBLE_ANSWER, events diff --git a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py index 2712d6e30cd..89d6ff50d53 100644 --- a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py +++ b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py @@ -1591,13 +1591,6 @@ def _assert_voice_session_ungated(session: Session, observed: tuple[dict[str, Js assert _session_updates(observed) == (VOICE_UPDATE,), _session_updates(observed) -@pytest.mark.xfail( - strict=True, - reason=( - "the realtime route hands the transcript guardrail only the request's guardrails list, so the key and " - "team metadata that opts out of a default-on guardrail never reaches it" - ), -) def test_key_opted_out_of_the_transcript_guardrail_is_not_gated_by_the_prompt_guardrail( guardrail_proxy: OwnedProxy, ) -> None: @@ -1609,13 +1602,6 @@ def test_key_opted_out_of_the_transcript_guardrail_is_not_gated_by_the_prompt_gu _assert_voice_session_ungated(session, _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id)) -@pytest.mark.xfail( - strict=True, - reason=( - "the realtime route hands the transcript guardrail only the request's guardrails list, so the key and " - "team metadata that opts out of a default-on guardrail never reaches it" - ), -) def test_team_opted_out_of_the_transcript_guardrail_is_not_gated(guardrail_proxy: OwnedProxy) -> None: with guardrail_proxy.gateway.scenario() as scenario: handle: Final = _scripted(scenario, _voice_scenario(BLOCKED_TRANSCRIPT)) diff --git a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py index 19928ecd6e0..95e07593b2a 100644 --- a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -48,7 +48,7 @@ _THINKING_OFF_TOOL_MODEL: Final = "claude-opus-4-8" _UNMAPPED_MODEL: Final = "claude-scripted-9" _BEDROCK_INVOKE_MODEL: Final = "anthropic.claude-sonnet-5" _BEDROCK_TOOL_PATH_MODEL: Final = "anthropic.claude-opus-4-8" -_BEDROCK_JSON_TOOL_MODEL: Final = "us.anthropic.claude-opus-4-8" +_BEDROCK_JSON_TOOL_MODEL: Final = "us.anthropic.claude-opus-4-7" _THINKING_SIGNATURE: Final = "c2NyaXB0ZWQtc2lnbmF0dXJl" _PROJECT: Final = "scripted-project" _LOCATION: Final = "us-east5" diff --git a/tests/integration/routing/test_prompt_cache_pin_ttl.py b/tests/integration/routing/test_prompt_cache_pin_ttl.py index 733f7859ff0..96121402def 100644 --- a/tests/integration/routing/test_prompt_cache_pin_ttl.py +++ b/tests/integration/routing/test_prompt_cache_pin_ttl.py @@ -387,7 +387,7 @@ def _responses_id(rig: _Rig, body: Mapping[str, JsonValue], *, stream: bool) -> @pytest.mark.parametrize("stream", [False, True], ids=["non-stream", "stream"]) -def test_responses_api_one_hour_breakpoint_answers_and_writes_no_pin(rig: _Rig, stream: bool) -> None: +def test_responses_api_one_hour_breakpoint_pins_the_served_deployment_for_an_hour(rig: _Rig, stream: bool) -> None: before: Final = _before(rig) marker: Final = new_marker() user: Final = f"pin-responses-{marker}" @@ -404,9 +404,12 @@ def test_responses_api_one_hour_breakpoint_answers_and_writes_no_pin(rig: _Rig, ], } assert _responses_id(rig, body, stream=stream).startswith("resp_") - assert rig.served_by(marker) in {_LEFT, _RIGHT} + served: Final = rig.served_by(marker) + assert served in {_LEFT, _RIGHT} end_user_row_lands(user) - held(lambda: pins_since(rig.redis, before), lambda found: not found, holding=3, seconds=20) + _, pin = one_new_pin(rig.redis, before) + assert pin.model_id == served, pin + assert lives_one_hour(pin), pin def _ask_surface(rig: _Rig, surface: str, session: str, marker: str) -> str: diff --git a/tests/integration/routing/test_prompt_caching_affinity_tool_call_marks.py b/tests/integration/routing/test_prompt_caching_affinity_tool_call_marks.py index 9d4bf62adf5..38129013294 100644 --- a/tests/integration/routing/test_prompt_caching_affinity_tool_call_marks.py +++ b/tests/integration/routing/test_prompt_caching_affinity_tool_call_marks.py @@ -2,7 +2,7 @@ from pathlib import Path from typing import Final import pytest -from integration._support.client import Gateway, eventually +from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows, scratch_database from integration._support.process import owned_proxy_process from integration._support.wire import Wire, wire_server @@ -24,6 +24,7 @@ from pydantic import JsonValue _MODEL: Final = "affinity-claude" _FOLLOW_UPS: Final = 24 +_PINS_ROUTE: Final = "/memory-usage-in-mem-cache-items" _LONG_SYSTEM: Final = "Answer from the tool results below. " + "lorem ipsum " * 1500 @@ -91,3 +92,68 @@ def test_router_affinity_is_lost_when_auto_caching_stands_down_for_tool_call_mar assert [status for status, _, _ in follow_ups] == [200] * _FOLLOW_UPS, [text for _, _, text in follow_ups] assert {first_marker, *follow_up_markers} == served["left"] | served["right"] assert all(served[side] & set(follow_up_markers) for side in served), served + + +def _responses_first_turn(marker: str) -> list[JsonValue]: + return [ + {"role": "user", "content": "What is the weather in London?"}, + {"type": "function_call", "call_id": "call_london", "name": "lookup_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_london", "output": "sunny"}, + {"role": "user", "content": final_text(marker)}, + ] + + +def _responses_body(session: str, request_input: list[JsonValue]) -> dict[str, JsonValue]: + return { + "model": _MODEL, + "instructions": f"{_LONG_SYSTEM} session {session}", + "tools": [{"type": "function", "name": "lookup_weather", "parameters": {"type": "object", "properties": {}}}], + "input": request_input, + } + + +def _pins(gateway: Gateway) -> dict[str, JsonValue]: + return object_value(object_value(gateway.request("GET", _PINS_ROUTE).json())["prompt_caching_pin_cache"]) + + +@pytest.mark.timeout(240) +def test_responses_thread_holds_its_deployment_when_the_bridge_logs_chat_messages( + gateway: Gateway, tmp_path: Path +) -> None: + session: Final = new_marker() + first_marker: Final = new_marker() + follow_up_markers: Final = tuple(new_marker() for _ in range(_FOLLOW_UPS)) + first_turn: Final = _responses_first_turn(first_marker) + with scratch_database() as database_url, wire_server(anthropic_peer) as left, wire_server(anthropic_peer) as right: + config: Final = owned_config( + tmp_path, + [ + {**anthropic_deployment(_MODEL, left.url), "model_info": {"id": f"responses-left-{session}"}}, + {**anthropic_deployment(_MODEL, right.url), "model_info": {"id": f"responses-right-{session}"}}, + ], + router_settings={"optional_pre_call_checks": ["prompt_caching"]}, + ) + with ( + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": database_url}, + config=config, + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=1, + ) as owned, + owned.gateway.scenario() as scenario, + ): + key: Final = scenario.key(metadata={"enable_prompt_caching": True}) + first: Final = owned.gateway.request("POST", "/v1/responses", _responses_body(session, first_turn), key=key) + assert first.status_code == 200, first.text + eventually(lambda: _pins(owned.gateway), bool, seconds=30) + follow_ups: Final = tuple( + owned.gateway.request( + "POST", "/v1/responses", _responses_body(session, _follow_up(first_turn, marker)), key=key + ) + for marker in follow_up_markers + ) + served: Final = {"left": _served(left), "right": _served(right)} + assert [response.status_code for response in follow_ups] == [200] * _FOLLOW_UPS, follow_ups[0].text + assert {first_marker, *follow_up_markers} in tuple(served.values()), served diff --git a/tests/integration/spend/test_redis_ttl_preserving_token_increment.py b/tests/integration/spend/test_redis_ttl_preserving_token_increment.py index 8b7a21f73d5..daf76bf4740 100644 --- a/tests/integration/spend/test_redis_ttl_preserving_token_increment.py +++ b/tests/integration/spend/test_redis_ttl_preserving_token_increment.py @@ -11,7 +11,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler, ) from litellm.proxy.utils import InternalUsageCache -from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.redis.stores.rate_limits import CounterIncrement, RateLimitScope @pytest.mark.asyncio @@ -22,11 +22,12 @@ async def test_async_increment_tokens_with_ttl_preservation() -> None: handler: Final = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(DualCache(redis_cache=redis_cache)) ) - assert handler.token_increment_script is not None suffix: Final = uuid.uuid4().hex[:8] - key_with_ttl: Final = f"{{test_ttl}}:with_ttl:{suffix}" - key_without_ttl: Final = f"{{test_ttl}}:without_ttl:{suffix}" + scope_with_ttl: Final = RateLimitScope("test_ttl", f"with_ttl:{suffix}") + scope_without_ttl: Final = RateLimitScope("test_ttl", f"without_ttl:{suffix}") + key_with_ttl: Final = f"{{test_ttl:with_ttl:{suffix}}}:tokens" + key_without_ttl: Final = f"{{test_ttl:without_ttl:{suffix}}}:tokens" try: await redis_cache.async_delete_cache(key_with_ttl) @@ -34,8 +35,8 @@ async def test_async_increment_tokens_with_ttl_preservation() -> None: await handler.async_increment_tokens_with_ttl_preservation( pipeline_operations=[ - RedisPipelineIncrementOperation(key=key_with_ttl, increment_value=10.0, ttl=60), - RedisPipelineIncrementOperation(key=key_without_ttl, increment_value=5.0, ttl=None), + CounterIncrement(scope_with_ttl.counter("tokens"), 10.0, 60), + CounterIncrement(scope_without_ttl.counter("tokens"), 5.0, None), ] ) @@ -50,8 +51,8 @@ async def test_async_increment_tokens_with_ttl_preservation() -> None: await handler.async_increment_tokens_with_ttl_preservation( pipeline_operations=[ - RedisPipelineIncrementOperation(key=key_with_ttl, increment_value=15.0, ttl=60), - RedisPipelineIncrementOperation(key=key_without_ttl, increment_value=7.0, ttl=None), + CounterIncrement(scope_with_ttl.counter("tokens"), 15.0, 60), + CounterIncrement(scope_without_ttl.counter("tokens"), 7.0, None), ] ) diff --git a/tests/integration/spend/test_tag_spend_reset.py b/tests/integration/spend/test_tag_spend_reset.py new file mode 100644 index 00000000000..2104539f5d7 --- /dev/null +++ b/tests/integration/spend/test_tag_spend_reset.py @@ -0,0 +1,75 @@ +import json +import uuid +from typing import Final + +import httpx + +from tests.integration._support.client import Gateway, eventually +from tests.integration._support.database import read_rows +from tests.integration._support.provider import PROVIDER_URL, SharedProvider +from tests.integration._support.wire import Reply + + +def test_admin_can_reset_tag_spend_and_resume_budget_blocked_requests( + gateway: Gateway, provider: SharedProvider +) -> None: + tag: Final = f"tag-reset-{uuid.uuid4().hex}" + completion: Final = Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + with gateway.scenario() as scenario: + model: Final = scenario.model( + api_base=f"{PROVIDER_URL}/v1", input_cost_per_token=0.001, output_cost_per_token=0.002 + ) + gateway.post("/tag/new", {"name": tag, "max_budget": 0.0001, "budget_duration": "30d"}) + + def delete_tag() -> None: + gateway.post("/tag/delete", {"name": tag}) + + scenario.cleanups.callback(delete_tag) + + def tagged_request() -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"tag reset {uuid.uuid4().hex}"}], + "metadata": {"tags": [tag]}, + }, + ) + + provider.expect(completion) + first: Final = tagged_request() + assert first.status_code == 200, first.text + assert len(provider.received()) == 1 + eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_TagTable" WHERE tag_name=%s', (tag,)), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0.0001, + seconds=70, + ) + blocked: Final = tagged_request() + assert blocked.status_code == 422, blocked.text + assert blocked.json()["error"]["type"] == "budget_exceeded", blocked.text + assert f"Tag={tag}" in blocked.text, blocked.text + assert provider.received() == () + reset: Final = gateway.request("POST", "/tag/update", {"name": tag, "spend": 0}) + assert reset.status_code == 200, reset.text + provider.expect(completion) + try: + resumed: Final = tagged_request() + assert resumed.status_code == 200, resumed.text + assert len(provider.received()) == 1 + finally: + provider.replies.clear() diff --git a/tests/integration/streaming/test_responses_top_level_error_events.py b/tests/integration/streaming/test_responses_top_level_error_events.py index 63bdadceebc..2cbc4570dec 100644 --- a/tests/integration/streaming/test_responses_top_level_error_events.py +++ b/tests/integration/streaming/test_responses_top_level_error_events.py @@ -22,12 +22,14 @@ from integration._support.responses_stream import ( data_frames, delta, error_body, + error_event, frame, function_tools, healthy_stream, messages_body, openai_error_event, openai_error_stream, + response_object, serve, spend_row, sse_events, @@ -277,6 +279,38 @@ def test_native_responses_stream_fails_with_the_provider_code_and_message(gatewa _assert_failure_logged(response.headers["x-litellm-call-id"], "429", RATE_LIMIT_MESSAGE) +_FLEX_MESSAGE: Final = "Flex processing is temporarily unavailable. Please try again later or use standard processing." +_FLEX_ERROR: Final = {"type": "resource_unavailable", "code": "flex_unavailable", "message": _FLEX_MESSAGE} + + +def _flex_outage(identity: str, kind: str) -> tuple[bytes, ...]: + failure: Final = ( + error_event(_FLEX_ERROR) + if kind == "error" + else { + "type": "response.failed", + "sequence_number": 1, + "response": response_object(identity, "failed", error=_FLEX_ERROR), + } + ) + return (frame(created(identity)), frame(failure)) + + +@pytest.mark.parametrize("kind", ["error", "response.failed"]) +def test_native_responses_stream_flex_outage_fails_as_a_rate_limit(gateway: Gateway, kind: str) -> None: + response, error = _responses_frames(gateway, _flex_outage(_identity(), kind)) + assert error["code"] == "rate_limit_exceeded" and _FLEX_MESSAGE in string_value(error["message"]), error + _assert_failure_logged(response.headers["x-litellm-call-id"], "429", _FLEX_MESSAGE) + + +@pytest.mark.parametrize("kind", ["error", "response.failed"]) +def test_chat_stream_flex_outage_is_a_throttling_error_like_the_non_streamed_429(gateway: Gateway, kind: str) -> None: + response, _ = _chat(gateway, _flex_outage(_identity(), kind)) + assert response.status_code == 429, response.text + assert_provider_typed_rate_limit(error_body(response), _FLEX_MESSAGE) + _assert_failure_logged(response.headers["x-litellm-call-id"], "429", _FLEX_MESSAGE) + + def test_native_responses_stream_error_after_text_fails_with_the_provider_code(gateway: Gateway) -> None: identity: Final = _identity() stream: Final = ( diff --git a/tests/unit/a2a_protocol/test_card_resolver.py b/tests/unit/a2a_protocol/test_card_resolver.py index 4cfa8dab45b..cdf651bd259 100644 --- a/tests/unit/a2a_protocol/test_card_resolver.py +++ b/tests/unit/a2a_protocol/test_card_resolver.py @@ -14,6 +14,7 @@ import pytest from litellm.a2a_protocol.card_resolver import ( LiteLLMA2ACardResolver, fix_agent_card_url, + get_agent_card_url, is_localhost_or_internal_url, normalize_agent_card_interfaces, resolve_a2a_protocol_version, @@ -117,6 +118,37 @@ def test_fix_agent_card_url_updates_interface_when_top_level_is_localhost(): assert result.supported_interfaces[0].url == "https://my-public-agent.example.com/" +@pytest.mark.parametrize( + ("card", "expected"), + [ + (SimpleNamespace(url="https://agent.example.com/"), "https://agent.example.com/"), + ( + SimpleNamespace(url="", supported_interfaces=[SimpleNamespace(url="https://interface.example.com/")]), + "https://interface.example.com/", + ), + ( + SimpleNamespace(supported_interfaces=[SimpleNamespace(url="https://interface.example.com/")]), + "https://interface.example.com/", + ), + (SimpleNamespace(url=""), None), + (SimpleNamespace(url=None, supported_interfaces=[]), None), + (SimpleNamespace(), None), + ], + ids=[ + "top_level_url", + "empty_url_falls_to_interface", + "interfaces_only", + "empty_url_without_interfaces", + "no_interfaces", + "bare", + ], +) +def test_get_agent_card_url_prefers_the_top_level_url_then_the_first_interface( + card: SimpleNamespace, expected: str | None +): + assert get_agent_card_url(card) == expected + + def test_normalize_agent_card_interfaces_downgrades_miscased_interfaces_to_the_0_3_dialect(): pb2 = pytest.importorskip("a2a.types.a2a_pb2") diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index e9a075af507..bd14b505323 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -259,6 +259,29 @@ def test_cache_rejects_unsupported_type() -> None: assert str(raised.value).split("Supported types: ")[1].split(", ") == [member.value for member in LiteLLMCacheType] +class _PingableInMemoryCache(InMemoryCache): + async def ping(self) -> bool: + return True + + async def delete_cache_keys(self, keys: list[str]) -> None: + for key in keys: + self.delete_cache(key) + + +@pytest.mark.asyncio +async def test_cache_ping_and_delete_cache_keys_reach_a_backend_that_supports_them() -> None: + backend: Final = _PingableInMemoryCache() + cache: Final = Cache(type="local", _backend=backend) + backend.set_cache("kept", "kept value") + backend.set_cache("dropped", "dropped value") + + assert await cache.ping() is True + await cache.delete_cache_keys(["dropped"]) + + assert backend.get_cache("dropped") is None + assert backend.get_cache("kept") == "kept value" + + def test_semantic_cache_scope_rejects_unknown_value(): with pytest.raises(ValueError, match="'team' is not a valid SemanticCacheScope"): _semantic_cache(semantic_cache_scope="team") diff --git a/tests/unit/caching/test_redis_semantic_cache.py b/tests/unit/caching/test_redis_semantic_cache.py index 1b2fbb4fd3d..9b6f6de0d7d 100644 --- a/tests/unit/caching/test_redis_semantic_cache.py +++ b/tests/unit/caching/test_redis_semantic_cache.py @@ -2,6 +2,7 @@ from collections.abc import Iterator from contextlib import contextmanager from importlib import import_module import sys +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -300,6 +301,28 @@ def test_redis_semantic_cache_matches_bytes_cache_key(): ) +_STORED_RESPONSE: Final = {"content": "Paris", "finish_reason": None, "cached": True} + + +@pytest.mark.parametrize( + ("cached_response", "expected"), + [ + (str(_STORED_RESPONSE), _STORED_RESPONSE), + (str(_STORED_RESPONSE).encode("utf-8"), _STORED_RESPONSE), + ("{'content': 'Paris'", None), + ], + ids=["python_repr", "python_repr_bytes", "unparseable"], +) +def test_redis_semantic_cache_reads_back_a_value_stored_as_its_python_repr( + cached_response: str | bytes, expected: dict[str, object] | None +): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + + assert redis_semantic_cache._get_cache_logic(cached_response) == expected + + def test_redis_semantic_cache_rejects_pre_isolation_unscoped_hit(): """Pre-isolation entries with no cache-key field cannot be safely reassigned to a caller's scope and are treated as misses.""" diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 146a6a3ff1c..2f236596960 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -4,7 +4,7 @@ import importlib import os import shutil import tempfile -from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator +from collections.abc import AsyncIterator, Callable, Container, Coroutine, Iterable, Iterator from dataclasses import dataclass, field from pathlib import Path from types import MappingProxyType @@ -496,7 +496,7 @@ def declare_mid_conversation_system(monkeypatch: pytest.MonkeyPatch) -> Callable return declare -NativeStructuredOutputKey: TypeAlias = Literal["legacy", "bedrock"] | None +NativeStructuredOutputKey: TypeAlias = Literal["legacy", "bedrock", "default"] | None @pytest.fixture @@ -506,7 +506,8 @@ def set_native_structured_output( def set_flags(model_cost_key: str, native_key: NativeStructuredOutputKey) -> None: entry: Final = litellm.model_cost[model_cost_key] monkeypatch.setitem(entry, "supports_native_structured_output", native_key == "legacy") - monkeypatch.setitem(entry, "supports_bedrock_native_structured_output", native_key == "bedrock") + monkeypatch.setitem(entry, "supports_bedrock_native_structured_output", native_key in ("bedrock", "default")) + monkeypatch.setitem(entry, "supports_bedrock_native_structured_output_by_default", native_key == "default") litellm.get_model_info.cache_clear() return set_flags @@ -595,3 +596,9 @@ def pytest_sessionfinish() -> None: _close_handler_if_needed(getattr(litellm, name, None)) _run_coroutine_if_needed(close_litellm_async_clients()) enable_socket() + + +def in_operator_error(value: Container[object]) -> str: + with pytest.raises(TypeError) as native: + _ = "key" in value + return str(native.value) diff --git a/tests/unit/enterprise/integrations/test_prometheus_unit_tests.py b/tests/unit/enterprise/integrations/test_prometheus_unit_tests.py index 13fef7a0945..ebe022072dc 100644 --- a/tests/unit/enterprise/integrations/test_prometheus_unit_tests.py +++ b/tests/unit/enterprise/integrations/test_prometheus_unit_tests.py @@ -23,7 +23,7 @@ from litellm import Router from litellm.caching.caching import DualCache from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback -from litellm.types.router import ModelInfo +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.types.utils import BudgetConfig, GenericBudgetConfigType @@ -243,11 +243,11 @@ async def test_router_cooldown_event_callback(): """ # Mock routing context mock_ctx = MagicMock() - mock_deployment = { - "litellm_params": {"model": "gpt-5-mini"}, - "model_name": "gpt-5-mini", - "model_info": ModelInfo(id="test-model-id"), - } + mock_deployment = Deployment( + model_name="gpt-5-mini", + litellm_params=LiteLLM_Params(model="gpt-5-mini"), + model_info=ModelInfo(id="test-model-id"), + ) mock_ctx.registry.return_value.deployment.return_value = mock_deployment # Create a real PrometheusLogger instance @@ -303,11 +303,11 @@ async def test_router_cooldown_event_callback_no_prometheus(): """ # Mock routing context mock_ctx = MagicMock() - mock_deployment = { - "litellm_params": {"model": "gpt-5-mini"}, - "model_name": "gpt-5-mini", - "model_info": ModelInfo(id="test-model-id"), - } + mock_deployment = Deployment( + model_name="gpt-5-mini", + litellm_params=LiteLLM_Params(model="gpt-5-mini"), + model_info=ModelInfo(id="test-model-id"), + ) mock_ctx.registry.return_value.deployment.return_value = mock_deployment await router_cooldown_event_callback( diff --git a/tests/unit/integrations/bitbucket/test_bitbucket_prompt_manager.py b/tests/unit/integrations/bitbucket/test_bitbucket_prompt_manager.py index a1a88653ee6..b90f563496e 100644 --- a/tests/unit/integrations/bitbucket/test_bitbucket_prompt_manager.py +++ b/tests/unit/integrations/bitbucket/test_bitbucket_prompt_manager.py @@ -2,6 +2,7 @@ import json import re from unittest.mock import MagicMock, patch +import httpx import pytest @@ -10,6 +11,7 @@ from litellm.integrations.bitbucket.bitbucket_prompt_manager import ( BitBucketPromptManager, BitBucketPromptTemplate, ) +from litellm.llms.custom_httpx.http_handler import HTTPHandler def test_bitbucket_prompt_template_creation(): @@ -502,3 +504,19 @@ rate: 0.5""" assert parsed["disabled"] is False assert parsed["count"] == 42 assert parsed["rate"] == 0.5 + + +def test_bitbucket_client_reads_repository_info_and_branches_from_the_api() -> None: + def bitbucket_api(request: httpx.Request) -> httpx.Response: + if request.url.path == "/2.0/repositories/test-workspace/test-repo": + return httpx.Response(200, json={"full_name": "test-workspace/test-repo", "is_private": True}) + if request.url.path == "/2.0/repositories/test-workspace/test-repo/refs/branches": + return httpx.Response(200, json={"values": [{"name": "main"}, {"name": "dev"}]}) + return httpx.Response(404) + + client = BitBucketClient({"workspace": "test-workspace", "repository": "test-repo", "access_token": "test-token"}) + client.http_handler = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(bitbucket_api))) + + assert client.get_repository_info() == {"full_name": "test-workspace/test-repo", "is_private": True} + assert [branch["name"] for branch in client.get_branches()] == ["main", "dev"] + assert client.test_connection() is True diff --git a/tests/unit/integrations/focus/test_focus_database.py b/tests/unit/integrations/focus/test_focus_database.py index 06240eac387..e5d046efb78 100644 --- a/tests/unit/integrations/focus/test_focus_database.py +++ b/tests/unit/integrations/focus/test_focus_database.py @@ -109,3 +109,19 @@ async def test_should_build_frame_from_rows_recovered_for_double_hashed_keys(mon assert result["api_key_alias"].to_list() == ["joined", "batch-worker"] assert result["team_id"].to_list() == ["team-0", "team-1"] + + +@pytest.mark.asyncio +async def test_table_info_returns_the_daily_user_spend_columns(monkeypatch: pytest.MonkeyPatch): + columns = [ + {"column_name": "id", "data_type": "text", "is_nullable": "NO"}, + {"column_name": "spend", "data_type": "double precision", "is_nullable": "YES"}, + ] + db, query_mock = _setup_db(monkeypatch, columns) + + info = await db.get_table_info() + + assert info == {"columns": columns, "table_name": "LiteLLM_DailyUserSpend"} + query_text, *params = query_mock.await_args.args + assert "WHERE table_name = 'LiteLLM_DailyUserSpend'" in query_text + assert params == [] diff --git a/tests/unit/integrations/otel/test_otel_v2_logger.py b/tests/unit/integrations/otel/test_otel_v2_logger.py index ac1f5557bf6..017eb745d90 100644 --- a/tests/unit/integrations/otel/test_otel_v2_logger.py +++ b/tests/unit/integrations/otel/test_otel_v2_logger.py @@ -1793,6 +1793,54 @@ def test_pre_call_hook_seeds_baggage_onto_server_and_child_spans(): ) +@pytest.mark.parametrize("guardrails_only", [False, True]) +def test_proxy_pre_call_hook_seeds_request_metadata_without_a_listed_otel_callback(monkeypatch, guardrails_only): + """A V2 logger that is registered but not in ``litellm.callbacks`` (a team-only destination) still gets the + caller's allowlisted ``requester_metadata`` from the proxy pre-call hook, keeps the model auth seeded, and a + ``guardrails_only`` scan seeds nothing.""" + import litellm + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server + from litellm.proxy.utils import ProxyLogging + + cfg: Final = OpenTelemetryV2Config(exporter="in_memory", baggage_metadata_keys=["requester_metadata.trace_id"]) + exporter: Final = InMemorySpanExporter() + logger: Final = OpenTelemetryV2(config=cfg, tracer_provider=providers.build_tracer_provider(cfg, exporter=exporter)) + monkeypatch.setattr(proxy_server, "open_telemetry_logger", logger) + monkeypatch.setattr(litellm, "callbacks", []) + server: Final = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + data: Final = { + "model": "resolved-deployment", + "metadata": {"requester_metadata": {"trace_id": "abc", "other": "x"}}, + } + + async def _flow(): + logger.seed_request_identity(_Auth(), model="client-alias") + await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=_Auth(), data=data, call_type="completion", guardrails_only=guardrails_only + ) + await logger.async_service_success_hook(payload=_ServicePayload("redis", "set"), parent_otel_span=server) + + with trace.use_span(server, end_on_exit=False): + asyncio.run(_flow()) + server.end() + + spans: Final = {s.name: s for s in exporter.get_finished_spans()} + seeded: Final = None if guardrails_only else "abc" + observed: Final = { + name: ( + spans[name].attributes.get(f"{LiteLLM.METADATA_PREFIX}trace_id"), + spans[name].attributes.get("gen_ai.request.model"), + ) + for name in (LITELLM_PROXY_REQUEST_SPAN_NAME, "redis set") + } + assert observed == { + LITELLM_PROXY_REQUEST_SPAN_NAME: (seeded, "client-alias"), + "redis set": (seeded, "client-alias"), + }, observed + assert not any(any(k.endswith("other") for k in s.attributes) for s in spans.values()), sorted(spans) + + def test_pre_call_hook_promotes_nested_request_metadata_key(): """``baggage_metadata_keys: [requester_metadata.trace_id]`` reads the caller's ``metadata.trace_id`` (snapshotted by the proxy under ``requester_metadata``) diff --git a/tests/unit/integrations/test_custom_guardrail_post_call_once.py b/tests/unit/integrations/test_custom_guardrail_post_call_once.py new file mode 100644 index 00000000000..fb823156929 --- /dev/null +++ b/tests/unit/integrations/test_custom_guardrail_post_call_once.py @@ -0,0 +1,566 @@ +import asyncio +import json +from collections.abc import AsyncIterable, Callable, Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal +from unittest.mock import AsyncMock, patch + +import httpx +import pytest +import respx +from fastapi import HTTPException, Request, Response +from fastapi.responses import StreamingResponse +from mcp.types import Tool as MCPTool +from pydantic import BaseModel, TypeAdapter + +import litellm +from litellm.caching.caching import Cache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import ProxyLogging +from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler +from litellm.router import Router +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CallTypes, Choices, GenericGuardrailAPIInputs, LLMResponseTypes, Message, ModelResponse + +pytestmark = pytest.mark.respx(assert_all_called=False) + +REPLY: Final = "hi from the model" +REWRITTEN: Final = "rewritten by the proxy" +MASKED: Final = "masked by G" + +Route = Literal["acompletion", "aresponses", "anthropic_messages"] + +_OPENAI_API_BASE: Final = "https://api.openai.com/v1" +_FAILING_API_BASE: Final = "https://primary.invalid/v1" +_MCP_API_BASE: Final = "https://mcp-rounds.invalid/v1" + +_CHAT_REPLY: Final[Mapping[str, object]] = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 0, + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": REPLY}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +_CHAT_SSE: Final = ( + "".join( + f"data: {json.dumps(chunk)}\n\n" + for chunk in ( + { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 0, + "model": "gpt-4o", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": REPLY}, "finish_reason": None}], + }, + { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 0, + "model": "gpt-4o", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, + ) + ) + + "data: [DONE]\n\n" +) +_RESPONSES_REPLY: Final[Mapping[str, object]] = { + "id": "resp_1", + "object": "response", + "created_at": 0, + "status": "completed", + "model": "gpt-4o", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": REPLY, "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, +} +_RESPONSES_SSE: Final = "".join( + f"data: {json.dumps(event)}\n\n" + for event in ( + {"type": "response.created", "sequence_number": 0, "response": {**_RESPONSES_REPLY, "output": []}}, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": REPLY, + }, + {"type": "response.completed", "sequence_number": 2, "response": _RESPONSES_REPLY}, + ) +) +_TOOL_CALL_ROUND: Final[Mapping[str, object]] = { + **_RESPONSES_REPLY, + "id": "resp_tool_call", + "output": [ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_echo", + "name": "echo", + "arguments": '{"value": "PING"}', + "status": "completed", + } + ], +} +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) + + +class _CountedGuardrail(CustomGuardrail): + checked: tuple[str, ...] = () + + +class ApplyGuardrailG(_CountedGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + self.checked = (*self.checked, " ".join(inputs.get("texts") or ())) + return inputs + + +def _reply_text(response: LLMResponseTypes) -> str: + match response: + case ModelResponse(choices=[Choices(message=Message(content=str() as text))]): + return text + case _: + return str(response) + + +class OwnHookG(_CountedGuardrail): + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + ) -> None: + self.checked = (*self.checked, _reply_text(response)) + + +class InheritedApplyGuardrailG(ApplyGuardrailG): + pass + + +class MaskingG(ApplyGuardrailG): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + checked: Final = await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + return {**checked, "texts": [MASKED for _ in checked.get("texts") or ()]} + + +class BlockingG(ApplyGuardrailG): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + raise HTTPException(status_code=400, detail="blocked by G") + + +class RewritingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "texts": [REWRITTEN for _ in inputs.get("texts") or ()]} + + +class ToolCallBlockingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + if input_type == "response" and inputs.get("tool_calls"): + raise HTTPException(status_code=400, detail="tool calls are not allowed") + return inputs + + +def _guardrail_g( + style: type[_CountedGuardrail] = ApplyGuardrailG, default_on: bool = False, run_in_parallel: bool = False +) -> _CountedGuardrail: + return style( + guardrail_name="G", + event_hook=GuardrailEventHooks.post_call, + default_on=default_on, + run_in_parallel=run_in_parallel, + ) + + +def _known_gap(action_item: str) -> pytest.MarkDecorator: + return pytest.mark.xfail( + raises=AssertionError, + strict=True, + reason=f"broken before #40154, out of scope for LIT-10481, tracked as RCA {action_item}", + ) + + +def _deployment(model_name: str, guardrails: tuple[str, ...], api_base: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-test", + "api_base": api_base, + **({"guardrails": list(guardrails)} if guardrails else {}), + }, + } + + +def _router(primary: tuple[str, ...], fallback: tuple[str, ...] | None = None, cached: bool = False) -> Router: + if fallback is None: + return Router(model_list=[_deployment("gpt", primary, _OPENAI_API_BASE)], cache_responses=cached) + return Router( + model_list=[ + _deployment("gpt", primary, _FAILING_API_BASE), + _deployment("gpt-backup", fallback, _OPENAI_API_BASE), + ], + fallbacks=[{"gpt": ["gpt-backup"]}], + num_retries=0, + ) + + +def _json_or_sse(json_reply: Mapping[str, object], sse_reply: str) -> Callable[[httpx.Request], httpx.Response]: + def reply(request: httpx.Request) -> httpx.Response: + if _JSON_OBJECT.validate_json(request.content).get("stream"): + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=sse_reply.encode()) + return httpx.Response(200, json=dict(json_reply)) + + return reply + + +async def _drain(stream: AsyncIterable[object]) -> str: + return "".join([str(chunk) async for chunk in stream]) + + +@pytest.fixture(autouse=True) +def provider(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> Iterator[respx.MockRouter]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + ProxyLogging._callback_capabilities_cache.clear() # pyright: ignore[reportPrivateUsage] # keyed on id(), which a collected guardrail from an earlier test can reuse + for endpoint in ("chat/completions", "responses"): + respx_mock.post(f"{_FAILING_API_BASE}/{endpoint}").mock( + return_value=httpx.Response(500, json={"error": {"message": "primary down"}}) + ) + respx_mock.post(f"{_OPENAI_API_BASE}/chat/completions").mock(side_effect=_json_or_sse(_CHAT_REPLY, _CHAT_SSE)) + respx_mock.post(f"{_OPENAI_API_BASE}/responses").mock(side_effect=_json_or_sse(_RESPONSES_REPLY, _RESPONSES_SSE)) + yield respx_mock + litellm.in_memory_llm_clients_cache.flush_cache() + + +@dataclass(frozen=True, slots=True) +class ProxyCase: + route: Route = "acompletion" + stream: bool = False + primary: tuple[str, ...] = ("G",) + fallback: tuple[str, ...] | None = None + request: tuple[str, ...] = () + default_on: bool = False + style: type[_CountedGuardrail] = ApplyGuardrailG + cached: bool = False + client_sees: str = REPLY + checks: tuple[str, ...] = (REPLY,) + + +PROXY_CASES: Final = ( + pytest.param(ProxyCase(), id="A-chat"), + pytest.param(ProxyCase(route="aresponses"), id="A-responses"), + pytest.param(ProxyCase(route="anthropic_messages"), id="A-messages"), + pytest.param(ProxyCase(request=("G",)), id="A-chat-model-and-request"), + pytest.param(ProxyCase(style=MaskingG, client_sees=MASKED), id="A-chat-g-masks-its-own-output"), + pytest.param(ProxyCase(stream=True), id="B-chat-stream"), + pytest.param(ProxyCase(route="aresponses", stream=True), id="B-responses-stream"), + pytest.param(ProxyCase(route="anthropic_messages", stream=True), id="B-messages-stream"), + pytest.param(ProxyCase(primary=(), fallback=("G",)), id="E-only-on-fallback"), + pytest.param( + ProxyCase(request=("rewriter",), client_sees=REWRITTEN, checks=(REPLY, REWRITTEN)), + id="G-rewritten-after-check", + ), + pytest.param(ProxyCase(fallback=("G",)), id="H-on-both-deployments"), + pytest.param(ProxyCase(primary=(), request=("G",)), id="I-request-chat"), + pytest.param(ProxyCase(route="aresponses", primary=(), request=("G",)), id="I-request-responses"), + pytest.param(ProxyCase(style=OwnHookG), id="J-A-chat-own-hook", marks=_known_gap("AI-5")), + pytest.param(ProxyCase(primary=(), fallback=("G",), style=OwnHookG), id="E-only-on-fallback-own-hook"), + pytest.param(ProxyCase(primary=(), request=("G",), style=OwnHookG), id="I-request-chat-own-hook"), + pytest.param( + ProxyCase(primary=(), request=("G",), style=InheritedApplyGuardrailG), + id="K-I-request-chat-inherited-apply-guardrail", + marks=_known_gap("AI-2"), + ), + pytest.param(ProxyCase(primary=("other",), default_on=True), id="L-default-on"), +) + +_ROUTE_PATHS: Final[Mapping[Route, str]] = MappingProxyType( + { + "acompletion": "/v1/chat/completions", + "aresponses": "/v1/responses", + "anthropic_messages": "/v1/messages", + } +) +_ROUTE_BODIES: Final[Mapping[Route, Mapping[str, object]]] = MappingProxyType( + { + "acompletion": {"messages": [{"role": "user", "content": "hello"}]}, + "aresponses": {"input": "hello"}, + "anthropic_messages": {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 32}, + } +) + + +async def _proxy_reply(case: ProxyCase) -> str: + from litellm.proxy.proxy_server import ProxyConfig, select_data_generator + + path: Final = _ROUTE_PATHS[case.route] + router: Final = _router(case.primary, case.fallback, case.cached) + response: Final = await ProxyBaseLLMRequestProcessing( + data={ + "model": "gpt", + **_ROUTE_BODIES[case.route], + **({"stream": True} if case.stream else {}), + **({"guardrails": list(case.request)} if case.request else {}), + } + ).base_process_llm_request( + request=Request({"type": "http", "method": "POST", "path": path, "headers": [], "query_string": b""}), + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", request_route=path), + route_type=case.route, + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + general_settings={}, + proxy_config=ProxyConfig(), + select_data_generator=select_data_generator, + llm_router=router, + is_streaming_request=case.stream, + ) + text: Final = ( + await _drain(response.body_iterator) + if isinstance(response, StreamingResponse) + else json.dumps(response.model_dump() if isinstance(response, BaseModel) else response) + ) + await asyncio.sleep(0) + return text + + +@pytest.mark.parametrize("case", PROXY_CASES) +async def test_a_proxy_request_checks_each_distinct_response_once( + monkeypatch: pytest.MonkeyPatch, case: ProxyCase +) -> None: + guardrail: Final = _guardrail_g(case.style, default_on=case.default_on) + rewriter: Final = RewritingGuardrail(guardrail_name="rewriter", event_hook=GuardrailEventHooks.post_call) + monkeypatch.setattr(litellm, "callbacks", [rewriter, guardrail]) + + reply: Final = await _proxy_reply(case) + + assert case.client_sees in reply + assert sorted(guardrail.checked) == sorted(case.checks) + + +@pytest.mark.parametrize("route", ["acompletion", "aresponses"]) +async def test_a_cache_hit_is_checked_like_a_fresh_response( + monkeypatch: pytest.MonkeyPatch, provider: respx.MockRouter, route: Route +) -> None: + guardrail: Final = _guardrail_g() + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr(litellm, "cache", Cache(type="local")) + + case: Final = ProxyCase(route=route, cached=True) + replies: Final = (await _proxy_reply(case), await _proxy_reply(case)) + + assert provider.calls.call_count == 1, "the second request must be served from the cache" + assert all(REPLY in reply for reply in replies) + assert guardrail.checked == (REPLY, REPLY) + + +async def test_a_blocking_verdict_blocks_the_proxy_response_after_one_check(monkeypatch: pytest.MonkeyPatch) -> None: + guardrail: Final = _guardrail_g(BlockingG) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + with pytest.raises(HTTPException, check=lambda e: e.detail == "blocked by G"): + await _proxy_reply(ProxyCase()) + + assert guardrail.checked == (REPLY,) + + +SdkCaller = Literal["sdk-call", "sdk-router"] + + +@dataclass(frozen=True, slots=True) +class SdkCase: + caller: SdkCaller = "sdk-call" + stream: bool = False + style: type[_CountedGuardrail] = ApplyGuardrailG + + +SDK_CASES: Final = ( + pytest.param(SdkCase(), id="C-sdk-call"), + pytest.param(SdkCase(caller="sdk-router"), id="C-sdk-router"), + pytest.param(SdkCase(style=OwnHookG), id="C-sdk-call-own-hook"), + pytest.param(SdkCase(stream=True), id="D-sdk-call-stream", marks=_known_gap("AI-4")), +) + + +async def _sdk_reply(case: SdkCase) -> str: + messages: Final = [{"role": "user", "content": "hello"}] + response: Final = await ( + _router(("G",)).acompletion(model="gpt", messages=messages, stream=case.stream) + if case.caller == "sdk-router" + else litellm.acompletion( + model="openai/gpt-4o", api_key="sk-test", guardrails=["G"], messages=messages, stream=case.stream + ) + ) + text: Final = await _drain(response) if isinstance(response, AsyncIterable) else str(response) + await asyncio.sleep(0) + return text + + +@pytest.mark.parametrize("case", SDK_CASES) +async def test_an_sdk_call_checks_its_response_once(monkeypatch: pytest.MonkeyPatch, case: SdkCase) -> None: + guardrail: Final = _guardrail_g(case.style) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + reply: Final = await _sdk_reply(case) + + assert REPLY in reply + assert guardrail.checked == (REPLY,) + + +async def test_mcp_auto_execute_checks_the_tool_call_round_before_running_its_tools( + monkeypatch: pytest.MonkeyPatch, provider: respx.MockRouter +) -> None: + monkeypatch.setattr( + litellm, "callbacks", [ToolCallBlockingGuardrail(guardrail_name="G", event_hook=GuardrailEventHooks.post_call)] + ) + provider.post(f"{_MCP_API_BASE}/responses").mock( + side_effect=[httpx.Response(200, json=_TOOL_CALL_ROUND), httpx.Response(200, json=_RESPONSES_REPLY)] + ) + echo_tool: Final = MCPTool.model_validate( + {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, + by_name=False, + ) + run_tools: Final = AsyncMock(return_value=[{"tool_call_id": "call_echo", "result": "PING", "name": "echo"}]) + + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_without_openai_transform", + AsyncMock(return_value=([echo_tool], {"echo": "demo"})), + ), + patch.object(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", run_tools), + pytest.raises(HTTPException, check=lambda e: e.detail == "tool calls are not allowed"), + ): + await litellm.aresponses( + model="openai/gpt-4o", + input="Call echo with PING", + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + api_key="sk-test", + api_base=_MCP_API_BASE, + guardrails=["G"], + proxy_server_request={"url": "http://localhost:4000/v1/responses", "method": "POST"}, + litellm_metadata={"user_api_key_request_route": "/v1/responses"}, + ) + + run_tools.assert_not_awaited() + + +def _same_response(reply: ModelResponse) -> ModelResponse: + return reply + + +def _model_renamed(reply: ModelResponse) -> ModelResponse: + return reply.model_copy(update={"model": "public-alias"}) + + +def _rewritten(_: ModelResponse) -> ModelResponse: + return ModelResponse(choices=[{"message": {"role": "assistant", "content": REWRITTEN}}]) + + +PROXY_LOOP_CASES: Final = ( + pytest.param(_same_response, (REPLY,), id="A-same-response"), + pytest.param(_model_renamed, (REPLY,), id="A-model-renamed"), + pytest.param(_rewritten, (REPLY, REWRITTEN), id="G-rewritten"), +) + + +@pytest.mark.parametrize("parallel", [False, True], ids=["sequential", "parallel"]) +@pytest.mark.parametrize(("proxy_sees", "checks"), PROXY_LOOP_CASES) +async def test_each_proxy_post_call_loop_skips_only_output_the_deployment_checked( + monkeypatch: pytest.MonkeyPatch, + parallel: bool, + proxy_sees: Callable[[ModelResponse], ModelResponse], + checks: tuple[str, ...], +) -> None: + guardrail: Final = _guardrail_g(run_in_parallel=parallel) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + data: Final = { + "model": "gpt", + "proxy_server_request": {"url": "http://localhost:4000/v1/chat/completions"}, + "metadata": {"guardrails": ["G"]}, + } + reply: Final = ModelResponse(model="provider-model", choices=[{"message": {"role": "assistant", "content": REPLY}}]) + + await guardrail.async_post_call_success_deployment_hook( + request_data={**data, "guardrails": ["G"]}, response=reply, call_type=CallTypes.acompletion + ) + await ProxyLogging(user_api_key_cache=UserApiKeyCache()).post_call_success_hook( + data=data, response=proxy_sees(reply), user_api_key_dict=UserAPIKeyAuth() + ) + + assert sorted(guardrail.checked) == sorted(checks) + + +def _ocr_reply(markdown: str) -> OCRResponse: + return OCRResponse.model_validate( + {"pages": [{"index": 0, "markdown": markdown, "images": [], "dimensions": None}], "model": "mistral-ocr-latest"} + ) + + +@pytest.mark.parametrize( + ("proxy_sees", "checks"), + [ + pytest.param(REPLY, (REPLY,), id="A-ocr-same-pages"), + pytest.param(REWRITTEN, (REPLY, REWRITTEN), id="G-ocr-rewritten-pages"), + ], +) +async def test_an_ocr_response_is_skipped_only_when_its_pages_were_checked( + monkeypatch: pytest.MonkeyPatch, proxy_sees: str, checks: tuple[str, ...] +) -> None: + guardrail: Final = _guardrail_g() + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + data: Final = { + "model": "ocr", + "document": {"type": "document_url", "document_url": "https://example.com/doc.pdf"}, + "proxy_server_request": {"url": "http://localhost:4000/v1/ocr"}, + "metadata": {"guardrails": ["G"], "user_api_key_request_route": "/v1/ocr"}, + } + + await guardrail.async_post_call_success_deployment_hook( + request_data={**data, "guardrails": ["G"]}, response=_ocr_reply(REPLY), call_type=CallTypes.aocr + ) + await ProxyLogging(user_api_key_cache=UserApiKeyCache()).post_call_success_hook( + data=data, response=_ocr_reply(proxy_sees), user_api_key_dict=UserAPIKeyAuth(request_route="/v1/ocr") + ) + + assert sorted(guardrail.checked) == sorted(checks) diff --git a/tests/unit/integrations/test_email_alerting.py b/tests/unit/integrations/test_email_alerting.py new file mode 100644 index 00000000000..35de2b112c6 --- /dev/null +++ b/tests/unit/integrations/test_email_alerting.py @@ -0,0 +1,33 @@ +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock + +import pytest + +from litellm.integrations.email_alerting import get_all_team_member_emails +from litellm.proxy import proxy_server + + +@pytest.mark.asyncio +async def test_team_member_emails_are_the_set_emails_of_the_members_with_a_user_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + team: Final = SimpleNamespace( + members_with_roles=[{"user_id": "u-1", "role": "admin"}, {"role": "user"}, {"user_id": "u-2", "role": "user"}] + ) + query_raw: Final = AsyncMock( + return_value=[{"user_email": "a@example.com"}, {"user_email": None}, {}, {"user_email": "b@example.com"}] + ) + prisma_client: Final = SimpleNamespace( + db=SimpleNamespace( + litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=team)), + query_raw=query_raw, + ) + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + + emails: Final = await get_all_team_member_emails("team-1") + + assert emails == ["a@example.com", "b@example.com"] + assert query_raw.await_args is not None + assert query_raw.await_args.args[1] == ["u-1", "u-2"] diff --git a/tests/unit/integrations/test_galileo.py b/tests/unit/integrations/test_galileo.py index d067b31a651..80fc858331e 100644 --- a/tests/unit/integrations/test_galileo.py +++ b/tests/unit/integrations/test_galileo.py @@ -984,3 +984,32 @@ def test_galileo_output_str_of_a_chat_response_whose_message_is_not_a_litellm_me ) assert output == expected_output + + +GALILEO_MESSAGES: Final = [{"role": "user", "content": "weather in SF?"}] +GALILEO_TOOLS: Final = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}] +GALILEO_FUNCTIONS: Final = [{"name": "get_time", "parameters": {}}] + + +@pytest.mark.parametrize( + ("optional_params", "expected_input"), + [ + ( + {"functions": GALILEO_FUNCTIONS, "tools": GALILEO_TOOLS, "temperature": 0.1}, + {"messages": GALILEO_MESSAGES, "functions": GALILEO_FUNCTIONS, "tools": GALILEO_TOOLS}, + ), + (None, {"messages": GALILEO_MESSAGES}), + ], +) +def test_galileo_error_span_input_carries_the_request_functions_and_tools( + optional_params: dict[str, object] | None, expected_input: dict[str, object] +) -> None: + _, output_text, logged_input = GalileoObserve()._get_galileo_input_output_content( + {"messages": GALILEO_MESSAGES, "optional_params": optional_params}, + None, + level="ERROR", + status_message="upstream 500", + ) + + assert output_text == "upstream 500" + assert logged_input == expected_input diff --git a/tests/unit/integrations/test_humanloop.py b/tests/unit/integrations/test_humanloop.py index 469dcef9fc9..daa87ee6874 100644 --- a/tests/unit/integrations/test_humanloop.py +++ b/tests/unit/integrations/test_humanloop.py @@ -1,15 +1,18 @@ import asyncio import importlib import os -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from typing import Final +import httpx import pytest import pytest_asyncio import litellm from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE +from litellm.integrations import humanloop from litellm.integrations.humanloop import HumanLoopPromptManager +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -30,6 +33,58 @@ def test_compile_prompt(): assert compiled_prompt[0]["content"] == "You are John. Answer questions as this person. Do not break character." +def _serve_humanloop_prompt(monkeypatch: pytest.MonkeyPatch, prompt_body: Mapping[str, object]) -> None: + def humanloop_api(request: httpx.Request) -> httpx.Response: + if request.url.path == "/v5/prompts/pr_1" and request.headers.get("X-Api-Key") == "hl-key": + return httpx.Response(200, json=prompt_body) + return httpx.Response(404, json={"detail": "not found"}) + + monkeypatch.setattr( + humanloop, + "get_httpx_client", + lambda: HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(humanloop_api))), + ) + + +@pytest.mark.parametrize( + ("template", "expected_template"), + [ + ({"role": "system", "content": "You are {{person}}."}, [{"role": "system", "content": "You are {{person}}."}]), + ( + [{"role": "system", "content": "Be brief."}, {"role": "user", "content": "{{question}}"}], + [{"role": "system", "content": "Be brief."}, {"role": "user", "content": "{{question}}"}], + ), + ], +) +def test_humanloop_prompt_from_api_lists_its_template_and_keeps_only_openai_params( + monkeypatch: pytest.MonkeyPatch, template: object, expected_template: list[dict[str, str]] +) -> None: + _serve_humanloop_prompt( + monkeypatch, + { + "id": "pr_1", + "name": "support-bot", + "template": template, + "model": "gpt-4o", + "temperature": 0.3, + "max_tokens": 64, + }, + ) + + prompt: Final = HumanLoopPromptManager()._get_prompt_from_id_api("pr_1", "hl-key") + + assert prompt["prompt_template"] == expected_template + assert prompt["model"] == "gpt-4o" + assert prompt["optional_params"] == {"temperature": 0.3, "max_tokens": 64} + + +def test_humanloop_prompt_from_api_rejects_a_template_that_is_not_messages(monkeypatch: pytest.MonkeyPatch) -> None: + _serve_humanloop_prompt(monkeypatch, {"id": "pr_1", "template": "plain text", "model": "gpt-4o"}) + + with pytest.raises(ValueError, match="Invalid template message type: "): + HumanLoopPromptManager()._get_prompt_from_id_api("pr_1", "hl-key") + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/integrations/test_lago.py b/tests/unit/integrations/test_lago.py new file mode 100644 index 00000000000..21e866fa747 --- /dev/null +++ b/tests/unit/integrations/test_lago.py @@ -0,0 +1,55 @@ +import json +from datetime import datetime, timezone +from typing import Final + +import httpx +import pytest + +from litellm.integrations.lago import LagoLogger +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.types.utils import ModelResponse, Usage + +CALL_TIME: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + + +@pytest.fixture +def lago_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LAGO_API_KEY", "lago-key") + monkeypatch.setenv("LAGO_API_BASE", "https://lago.example") + monkeypatch.setenv("LAGO_API_EVENT_CODE", "llm-usage") + monkeypatch.delenv("LAGO_API_CHARGE_BY", raising=False) + + +def test_lago_event_bills_the_response_token_usage_to_the_end_user(lago_env: None) -> None: + sent: Final[list[httpx.Request]] = [] + + def lago(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={}) + + logger: Final = LagoLogger() + logger.sync_http_handler = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(lago))) + response: Final = ModelResponse(model="gpt-4o", usage=Usage(prompt_tokens=12, completion_tokens=5, total_tokens=17)) + + logger.log_success_event( + kwargs={ + "model": "gpt-4o", + "response_cost": 0.002, + "litellm_params": {"metadata": {}, "proxy_server_request": {"body": {"user": "end-user-1"}}}, + }, + response_obj=response, + start_time=CALL_TIME, + end_time=CALL_TIME, + ) + + assert [str(request.url) for request in sent] == ["https://lago.example/api/v1/events"] + event: Final = json.loads(sent[0].content)["event"] + assert event["external_subscription_id"] == "end-user-1" + assert event["code"] == "llm-usage" + assert event["properties"] == { + "model": "gpt-4o", + "response_cost": 0.002, + "prompt_tokens": 12, + "completion_tokens": 5, + "total_tokens": 17, + } diff --git a/tests/unit/integrations/test_langfuse_otel.py b/tests/unit/integrations/test_langfuse_otel.py index 71087b4de51..aa9a51a93ee 100644 --- a/tests/unit/integrations/test_langfuse_otel.py +++ b/tests/unit/integrations/test_langfuse_otel.py @@ -22,6 +22,7 @@ from opentelemetry.trace import Span import litellm from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, REDACTED_BY_LITELLM from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger +from litellm.integrations.langfuse.langfuse_otel_attributes import LangfuseLLMObsOTELAttributes from litellm.integrations.opentelemetry import LITELLM_PROXY_REQUEST_SPAN_NAME, OpenTelemetry, OpenTelemetryConfig from litellm.integrations.otel.model.destination import OtelDestination from litellm.integrations.otel.plumbing.context import set_request_destinations @@ -1636,3 +1637,31 @@ def test_a_request_without_a_langfuse_destination_still_exports_through_the_lega assert _spans_the_legacy_logger_exports((_LANGFUSE_DESTINATION.model_copy(update={"callback_name": "arize"}),)) == ( logged ) + + +OBSERVATION_MESSAGES: Final = [{"role": "user", "content": "weather in SF?"}] +OBSERVATION_TOOLS: Final = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}] +OBSERVATION_FUNCTIONS: Final = [{"name": "get_time", "parameters": {}}] + + +@pytest.mark.parametrize( + ("request_kwargs", "expected_input"), + [ + ( + { + "messages": OBSERVATION_MESSAGES, + "optional_params": {"tools": OBSERVATION_TOOLS, "functions": OBSERVATION_FUNCTIONS, "top_p": 0.5}, + }, + {"messages": OBSERVATION_MESSAGES, "functions": OBSERVATION_FUNCTIONS, "tools": OBSERVATION_TOOLS}, + ), + ({"messages": OBSERVATION_MESSAGES}, {"messages": OBSERVATION_MESSAGES}), + ], +) +def test_observation_input_carries_the_request_functions_and_tools( + request_kwargs: Mapping[str, object], expected_input: Mapping[str, object] +) -> None: + span: Final = _TRACER.start_span("litellm_request") + + LangfuseLLMObsOTELAttributes.set_messages(span, request_kwargs) + + assert json.loads(str(span.attributes["langfuse.observation.input"])) == expected_input diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index 393905da6a2..c4fd7ed51c7 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -7569,6 +7569,57 @@ class TestOpentelemetryUnitTests: f"{SpanAttributes.LLM_COMPLETIONS.value}.1.function_call.name": "get_news", } + def test_chat_response_tool_calls_are_recorded_per_call_on_the_span(self): + from litellm.types.utils import ModelResponse + + response: Final = ModelResponse( + id="chatcmpl-tools", + model="gpt-4o", + choices=[ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + }, + { + "id": "call_time", + "type": "function", + "function": {"name": "get_time", "arguments": "{}"}, + }, + ], + }, + } + ], + ) + span: Final = TracerProvider().get_tracer("otel-tool-call-tests").start_span("litellm_request") + + OpenTelemetry().set_attributes( + span=span, + kwargs={ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "weather and time in SF?"}], + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "openai"}, + "standard_logging_object": {"id": "call-tools", "call_type": "completion", "metadata": {}}, + }, + response_obj=response, + ) + + completions: Final = SpanAttributes.LLM_COMPLETIONS.value + assert {key: value for key, value in span.attributes.items() if ".function_call." in key} == { + f"{completions}.0.function_call.name": "get_weather", + f"{completions}.0.function_call.arguments": '{"city": "SF"}', + f"{completions}.1.function_call.name": "get_time", + f"{completions}.1.function_call.arguments": "{}", + } + @pytest.mark.asyncio async def test_opentelemetry_integration(self): """ diff --git a/tests/unit/integrations/test_prompt_layer.py b/tests/unit/integrations/test_prompt_layer.py new file mode 100644 index 00000000000..4ebe2ce2fdb --- /dev/null +++ b/tests/unit/integrations/test_prompt_layer.py @@ -0,0 +1,75 @@ +import json +from datetime import datetime, timezone +from typing import Final + +import httpx +import pytest + +import litellm +from litellm.integrations.prompt_layer import PromptLayerLogger +from litellm.llms.custom_httpx.http_handler import HTTPHandler + +CALL_TIME: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) +REQUEST_KWARGS: Final = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "optional_params": {"temperature": 0.2}, + "litellm_params": {"metadata": {"pl_tags": ["prod"], "team": "search"}}, +} + + +def _log_to_promptlayer( + monkeypatch: pytest.MonkeyPatch, track_request_body: object +) -> tuple[list[tuple[str, object]], list[str]]: + posted: Final[list[tuple[str, object]]] = [] + logged: Final[list[str]] = [] + + def promptlayer(request: httpx.Request) -> httpx.Response: + posted.append((request.url.path, json.loads(request.content))) + if request.url.path == "/rest/track-request": + return httpx.Response(200, json=track_request_body) + return httpx.Response(200, json={"success": True}) + + monkeypatch.setenv("PROMPTLAYER_API_KEY", "pl-key") + monkeypatch.setattr( + litellm, "module_level_client", HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(promptlayer))) + ) + PromptLayerLogger().log_event( + kwargs=REQUEST_KWARGS, + response_obj={"id": "chatcmpl-1"}, + start_time=CALL_TIME, + end_time=CALL_TIME, + print_verbose=logged.append, + ) + return posted, logged + + +def test_prompt_layer_tracks_metadata_under_the_request_id_it_got_back(monkeypatch: pytest.MonkeyPatch) -> None: + posted, _ = _log_to_promptlayer(monkeypatch, {"success": True, "request_id": 42}) + + assert [path for path, _ in posted] == ["/rest/track-request", "/rest/track-metadata"] + track_request: Final = posted[0][1] + assert isinstance(track_request, dict) + assert track_request["tags"] == ["prod"] + assert track_request["kwargs"] == { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "temperature": 0.2, + } + assert posted[1][1] == {"request_id": 42, "api_key": "pl-key", "metadata": {"team": "search"}} + + +@pytest.mark.parametrize( + ("track_request_body", "logged_error"), + [ + ({"success": False, "request_id": 42}, "Promptlayer did not successfully log the response!"), + (["success", "request_id"], "'list' object has no attribute 'get'"), + ], +) +def test_prompt_layer_skips_metadata_and_logs_the_error_when_tracking_did_not_succeed( + monkeypatch: pytest.MonkeyPatch, track_request_body: object, logged_error: str +) -> None: + posted, logged = _log_to_promptlayer(monkeypatch, track_request_body) + + assert [path for path, _ in posted] == ["/rest/track-request"] + assert logged_error in logged[-1] diff --git a/tests/unit/integrations/test_wandb.py b/tests/unit/integrations/test_wandb.py index 9341c002378..dfddf683ad1 100644 --- a/tests/unit/integrations/test_wandb.py +++ b/tests/unit/integrations/test_wandb.py @@ -1,3 +1,8 @@ +import importlib.util +import sys +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import ModuleType, SimpleNamespace from typing import Final from unittest.mock import MagicMock @@ -21,3 +26,65 @@ def test_wandb_logging(): pass except Exception as e: print(e) + + +@dataclass(frozen=True) +class _TraceResult: + inputs: Mapping[str, object] + outputs: Mapping[str, object] + + +@dataclass(frozen=True) +class _TraceSpan: + name: str + attributes: Mapping[str, object] + start_time_ms: int + end_time_ms: int + span_kind: str + results: Sequence[_TraceResult] + + +@dataclass(frozen=True) +class _TraceTree: + root_span: _TraceSpan + model_dict: Mapping[str, object] + + +def _weights_biases_with_stub_trace_tree(monkeypatch: pytest.MonkeyPatch) -> ModuleType: + data_types: Final = ModuleType("wandb.sdk.data_types") + data_types.trace_tree = SimpleNamespace( + Result=_TraceResult, Span=_TraceSpan, WBTraceTree=_TraceTree, SpanKind=SimpleNamespace(LLM="LLM") + ) + for name in ("wandb", "wandb.sdk"): + monkeypatch.setitem(sys.modules, name, ModuleType(name)) + monkeypatch.setitem(sys.modules, "wandb.sdk.data_types", data_types) + spec: Final = importlib.util.find_spec("litellm.integrations.weights_biases") + assert spec is not None and spec.loader is not None + module: Final = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +def test_wandb_chat_completion_trace_lists_each_request_message_with_its_role(monkeypatch: pytest.MonkeyPatch) -> None: + weights_biases: Final = _weights_biases_with_stub_trace_tree(monkeypatch) + request: Final = { + "model": "gpt-4o", + "messages": [{"role": "system", "content": "Be brief."}, {"role": "user", "content": "Hi"}], + } + response: Final = { + "object": "chat.completion", + "model": "gpt-4o", + "created": 1767225600, + "choices": [{"message": {"role": "assistant", "content": "Hello!"}}], + } + + trace: Final = weights_biases.OpenAIRequestResponseResolver()(request, response, 1.5) + + assert trace.root_span.results == [ + _TraceResult( + inputs={"request": "\n\n**system**: Be brief.\n\n\n**user**: Hi\n"}, + outputs={"response": "\n\n**assistant**: Hello!\n"}, + ) + ] + assert trace.root_span.end_time_ms - trace.root_span.start_time_ms == 1500 diff --git a/tests/unit/integrations/test_weave_otel.py b/tests/unit/integrations/test_weave_otel.py index e8f06c00e41..a3a022d514b 100644 --- a/tests/unit/integrations/test_weave_otel.py +++ b/tests/unit/integrations/test_weave_otel.py @@ -1,13 +1,20 @@ +import json import os +from collections.abc import Mapping +from typing import Final from unittest.mock import MagicMock, patch import pytest +from opentelemetry.sdk.trace import TracerProvider +from litellm.integrations._types.open_inference import SpanAttributes as OpenInferenceSpanAttributes from litellm.integrations.weave.weave_otel import ( + WeaveLLMObsOTELAttributes, _set_weave_specific_attributes, get_weave_otel_config, ) from litellm.types.integrations.weave_otel import WeaveOtelConfig, WeaveSpanAttributes +from litellm.types.utils import ModelResponse def test_get_weave_otel_config(): @@ -154,3 +161,47 @@ def test_set_weave_specific_attributes_thread_id_and_is_turn(): mock_safe_set.assert_any_call( mock_span, WeaveSpanAttributes.IS_TURN.value, True ) + + +_TRACER: Final = TracerProvider().get_tracer("weave-otel-tests") +WEAVE_MESSAGES: Final = [{"role": "user", "content": "weather in SF?"}] +WEAVE_TOOLS: Final = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}] +WEAVE_FUNCTIONS: Final = [{"name": "get_time", "parameters": {}}] + + +@pytest.mark.parametrize( + ("request_kwargs", "expected_input"), + [ + ( + { + "messages": WEAVE_MESSAGES, + "optional_params": {"tools": WEAVE_TOOLS, "functions": WEAVE_FUNCTIONS, "top_p": 0.5}, + }, + {"messages": WEAVE_MESSAGES, "functions": WEAVE_FUNCTIONS, "tools": WEAVE_TOOLS}, + ), + ({"messages": WEAVE_MESSAGES, "optional_params": None}, {"messages": WEAVE_MESSAGES}), + ], +) +def test_weave_input_value_carries_the_request_functions_and_tools( + request_kwargs: Mapping[str, object], expected_input: Mapping[str, object] +) -> None: + span: Final = _TRACER.start_span("litellm_request") + + WeaveLLMObsOTELAttributes.set_messages(span, request_kwargs) + + assert json.loads(str(span.attributes[OpenInferenceSpanAttributes.INPUT_VALUE])) == expected_input + + +def test_weave_output_value_is_the_dumped_model_response() -> None: + span: Final = _TRACER.start_span("litellm_request") + response: Final = ModelResponse( + id="chatcmpl-1", + model="gpt-4o", + choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "Sunny."}}], + ) + + _set_weave_specific_attributes(span, {"model": "gpt-4o"}, response) + + output: Final = json.loads(str(span.attributes[OpenInferenceSpanAttributes.OUTPUT_VALUE])) + assert output["id"] == "chatcmpl-1" + assert output["choices"][0]["message"]["content"] == "Sunny." diff --git a/tests/unit/interactions/test_gemini_interactions_transformation.py b/tests/unit/interactions/test_gemini_interactions_transformation.py index ebc99e8e88b..2f167254b0f 100644 --- a/tests/unit/interactions/test_gemini_interactions_transformation.py +++ b/tests/unit/interactions/test_gemini_interactions_transformation.py @@ -598,3 +598,14 @@ class TestTransformRequestSchemaCoalescing: ) assert len(optional_params["response_format"]) == 1 assert body_retry["response_format"] == body["response_format"] + + def test_image_config_that_is_not_a_mapping_is_rejected(self, config): + with pytest.raises(TypeError, match=r"^'str' object is not a mapping$"): + config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="draw a sunset", + optional_params={"generation_config": {"image_config": "1:1"}}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) diff --git a/tests/unit/interactions/test_litellm_responses_bridge.py b/tests/unit/interactions/test_litellm_responses_bridge.py index 9330ae44ad8..a40b8ca0365 100644 --- a/tests/unit/interactions/test_litellm_responses_bridge.py +++ b/tests/unit/interactions/test_litellm_responses_bridge.py @@ -6,7 +6,7 @@ the litellm_responses bridge provider, which calls litellm.responses() internall """ from collections.abc import Mapping -from datetime import datetime +from datetime import datetime, timezone from typing import Final import pytest @@ -144,6 +144,26 @@ class TestBridgeResponseTransformation: assert datetime.fromisoformat(interaction.created).timestamp() == 1700000000 assert interaction.updated == interaction.created + @pytest.mark.parametrize( + ("created_at", "expected_created"), + [ + (datetime(2026, 1, 2, 3, 4, 5, tzinfo=timezone.utc), "2026-01-02T03:04:05+00:00"), + (None, None), + ], + ids=["datetime", "missing"], + ) + def test_responses_response_created_at_that_is_not_a_timestamp( + self, created_at: datetime | None, expected_created: str | None + ): + reply: Final = _responses_api_response("completed", None).model_copy(update={"created_at": created_at}) + + interaction: Final = LiteLLMResponsesInteractionsConfig.transform_responses_response_to_interactions_response( + reply, None + ) + + assert interaction.created == expected_created + assert interaction.updated == expected_created + def test_responses_response_with_a_status_that_is_not_text_is_refused(self): reply: Final = _responses_api_response("completed", None).model_copy(update={"status": 7}) diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3964f472bd9..fd7ec0d6bca 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -21,6 +21,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( calculate_cache_writing_cost, generic_cost_per_token, get_billed_token_rates, + get_regional_uplift_multiplier, get_token_type_cost_breakdown, ) from litellm.llms.gemini.image_generation.cost_calculator import ( @@ -2781,6 +2782,45 @@ def test_fedramp_uplift_defaults_to_one(model_info: ModelInfo, data_residency: s assert get_fedramp_uplift_multiplier(model_info, data_residency) == 1.0 +@pytest.mark.parametrize( + "model_info, data_residency", + [({f"regional_processing_uplift_multiplier_{r.value}": "1.25"}, r.value.upper()) for r in DataResidency], +) +def test_regional_uplift_reads_the_multiplier_named_by_each_residency(model_info: ModelInfo, data_residency: str): + assert get_regional_uplift_multiplier(model_info, data_residency) == 1.25 + + +@pytest.mark.parametrize( + "model_info", + [{"regional_processing_uplift_multiplier_us": value} for value in (1.25, "1.25", b"1.25", bytearray(b"1.25"))], +) +def test_regional_uplift_accepts_every_value_float_converts(model_info: ModelInfo): + assert get_regional_uplift_multiplier(model_info, "us") == 1.25 + + +class _IndexOnly: + def __index__(self) -> int: + return 2 + + +@pytest.mark.parametrize("model_info", [{"regional_processing_uplift_multiplier_us": _IndexOnly()}]) +def test_regional_uplift_accepts_an_index_only_value(model_info: ModelInfo): + assert get_regional_uplift_multiplier(model_info, "us") == 2.0 + + +@pytest.mark.parametrize( + "model_info, data_residency", + [ + ({"regional_processing_uplift_multiplier_eu": "not-a-number"}, "eu"), + ({"regional_processing_uplift_multiplier_eu": [1.25]}, "eu"), + ({"regional_processing_uplift_multiplier_apac": 1.25}, "apac"), + ({"regional_processing_uplift_multiplier_eu": 1.25}, None), + ], +) +def test_regional_uplift_defaults_to_one(model_info: ModelInfo, data_residency: str | None): + assert get_regional_uplift_multiplier(model_info, data_residency) == 1.0 + + def test_service_tier_suffixes_constant_in_sync_with_enum(): from litellm.litellm_core_utils.llm_cost_calc.utils import _SERVICE_TIER_SUFFIXES from litellm.types.utils import ServiceTier diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index 5d6d14a4ca8..02476545cae 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -4,7 +4,7 @@ import datetime import json import logging import threading -from collections.abc import Iterator, Mapping +from collections.abc import Container, Iterator, Mapping from typing import Final import httpx @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_or_create_metadata_bucket, get_parent_otel_span_from_kwargs, get_provider_response_headers_from_hidden_params, + json_contains, map_finish_reason, normalize_drop_params, preserve_upstream_non_openai_attributes, @@ -33,6 +34,7 @@ from litellm.litellm_core_utils.core_helpers import ( redact_nested_match_and_regex_keys, response_cost_in_hidden_params, set_provider_response_headers_in_hidden_params, + subscript, unbind_budget_reservation_from_callbacks, without, ) @@ -855,6 +857,40 @@ def test_billed_usage_cost_keeps_the_upstream_charge_after_litellm_wrote_its_bil assert billed_usage_cost(unreported, "openai") is None +def test_subscript_reads_keys_and_indexes_and_rejects_values_without_item_access() -> None: + assert subscript({"role": "user"}, "role") == "user" + assert subscript(["system", "user"], 1) == "user" + with pytest.raises(TypeError, match="'int' object is not subscriptable"): + subscript(5, "role") + + +class _IterOnly: + def __iter__(self) -> Iterator[object]: + return iter(("error",)) + + +class _GetItemOnly: + def __getitem__(self, index: int) -> object: + return ("error",)[index] + + +@pytest.mark.parametrize( + "value", + [{"error": 1}, {"data": 1}, ["error"], "an error", "", _IterOnly(), _GetItemOnly()], +) +def test_json_contains_answers_what_the_in_operator_answers(value: Container[object]) -> None: + assert json_contains(value, "error") is ("error" in value) + + +@pytest.mark.parametrize("value", [5, 1.5, None, True, object()]) +def test_json_contains_raises_the_in_operators_own_error(value: Container[object]) -> None: + with pytest.raises(TypeError) as native: + _ = "error" in value + with pytest.raises(TypeError) as ours: + json_contains(value, "error") + assert str(ours.value) == str(native.value) + + def test_without_drops_the_named_keys_and_keeps_the_rest_in_order() -> None: source: Final = {"a": 1, "b": 2, "c": 3, "d": 4} diff --git a/tests/unit/litellm_core_utils/test_dd_tracing.py b/tests/unit/litellm_core_utils/test_dd_tracing.py index 251f1842e6e..e626d327989 100644 --- a/tests/unit/litellm_core_utils/test_dd_tracing.py +++ b/tests/unit/litellm_core_utils/test_dd_tracing.py @@ -1,10 +1,14 @@ import json +from dataclasses import dataclass +from typing import Final from unittest.mock import MagicMock, patch import pytest +from litellm.litellm_core_utils import dd_tracing from litellm.litellm_core_utils.dd_tracing import ( + set_active_span_tag, should_use_dd_profiler, _should_use_dd_tracer, ) @@ -97,3 +101,33 @@ def test_should_use_dd_profiler(): mock_get_secret.return_value = False assert should_use_dd_profiler() is False mock_get_secret.assert_called_once_with("USE_DDPROFILER", False) + + +@dataclass(frozen=True) +class _CurrentSpanTracer: + span: object + + def current_span(self) -> object: + return self.span + + +class _SetTagOnlySpan: + def __init__(self) -> None: + self.tags: dict[str, str] = {} + + def set_tag(self, key: str, value: str) -> None: + self.tags[key] = value + + +def test_set_active_span_tag_uses_set_tag_on_a_span_without_set_tag_str(monkeypatch: pytest.MonkeyPatch) -> None: + span: Final = _SetTagOnlySpan() + monkeypatch.setattr(dd_tracing, "tracer", _CurrentSpanTracer(span)) + + assert set_active_span_tag("litellm.call_id", "call-1") is True + assert span.tags == {"litellm.call_id": "call-1"} + + +def test_set_active_span_tag_reports_false_for_a_span_that_takes_no_tags(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(dd_tracing, "tracer", _CurrentSpanTracer(object())) + + assert set_active_span_tag("litellm.call_id", "call-1") is False diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index e8ea0c04f6f..58fd03e0cbc 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -2701,6 +2701,55 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests( dummy_logger.log_stream_event.assert_not_called() +def test_sync_success_callbacks_run_when_litellm_params_is_a_params_model( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.integrations.custom_logger import CustomLogger + from litellm.types.router import GenericLiteLLMParams + + class DummyLogger(CustomLogger): + pass + + logging_obj.stream = False + logging_obj.call_type = "completion" + logging_obj.model_call_details["litellm_params"] = GenericLiteLLMParams() + dummy_logger = DummyLogger() + dummy_logger.log_success_event = MagicMock() + logging_obj.dynamic_success_callbacks = [dummy_logger] + logging_obj.dynamic_async_success_callbacks = [] + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + + logging_obj.success_handler( + result=ModelResponse( + id="resp-123", + model="gpt-4o-mini", + choices=[{"message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop", "index": 0}], + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + ) + ) + + dummy_logger.log_success_event.assert_called_once() + + +@pytest.mark.parametrize(("extra_params", "expected"), [({}, True), ({"no-log": True}, False)]) +def test_should_run_callback_reads_no_log_from_a_params_model( + logging_obj: LitellmLogging, + monkeypatch: pytest.MonkeyPatch, + extra_params: dict[str, object], + expected: bool, +) -> None: + from litellm.types.router import GenericLiteLLMParams + + monkeypatch.setattr(litellm, "global_disable_no_log_param", False) + + should_run: Final = logging_obj.should_run_callback( + callback="langfuse", litellm_params=GenericLiteLLMParams(**extra_params), event_hook="success_handler" + ) + + assert should_run is expected + + @pytest.mark.asyncio @pytest.mark.parametrize("different_headers", [False, True]) async def test_completed_success_preserves_generic_logger_configuration( diff --git a/tests/unit/litellm_core_utils/test_realtime_guardrails.py b/tests/unit/litellm_core_utils/test_realtime_guardrails.py new file mode 100644 index 00000000000..26ab2887883 --- /dev/null +++ b/tests/unit/litellm_core_utils/test_realtime_guardrails.py @@ -0,0 +1,481 @@ +import json +import sys +from typing import Final, Literal + +import pytest +from fastapi import HTTPException + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import CustomCodeGuardrail +from litellm.litellm_core_utils.realtime_guardrails import ( + RealtimeGuardrailBlock, + RealtimeGuardrailRewrite, + RealtimeMessageText, + realtime_guarded_item, + realtime_guardrail_verdict, + rewritten_item, +) +from litellm.types.guardrails import BlockedWord, ContentFilterAction, ContentFilterPattern, GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + + +class RaisingGuardrail(CustomGuardrail): + def __init__(self, error: Exception, realtime_violation_message: str | None = None) -> None: + super().__init__( + guardrail_name="raising", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + realtime_violation_message=realtime_violation_message, + ) + self.error = error + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + raise self.error + + +class MaskingGuardrail(CustomGuardrail): + def __init__(self, word: str) -> None: + super().__init__(guardrail_name=f"mask-{word}", event_hook=GuardrailEventHooks.pre_call, default_on=True) + self.word = word + self.screened: list[tuple[str, ...]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + texts: Final = inputs.get("texts", []) + self.screened.append(tuple(texts)) + return {**inputs, "texts": [text.replace(self.word, "[MASKED]") for text in texts]} + + +class LastTextRedactingGuardrail(MaskingGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + last: Final = inputs.get("texts", [""])[-1] + if self.word not in last: + return inputs + return {**inputs, "texts": [last.replace(self.word, "[MASKED]")]} + + +def _item_create(item: object) -> dict: + return {"type": "conversation.item.create", "item": item} + + +@pytest.mark.parametrize("role", ["user", "assistant", "system"]) +def test_only_the_typed_text_of_a_message_is_screened_whatever_its_role(role: str) -> None: + client_event: Final = _item_create( + { + "type": "message", + "role": role, + "content": [{"type": "input_audio", "audio": "AAAA"}, {"type": "input_text", "text": "hello"}], + } + ) + + assert realtime_guarded_item(client_event) == RealtimeMessageText(role=role, texts=("hello",)) + + +@pytest.mark.parametrize("part_type", ["output_text", "text"]) +def test_assistant_text_written_by_the_client_is_screened(part_type: str) -> None: + client_event: Final = _item_create( + {"type": "message", "role": "assistant", "content": [{"type": part_type, "text": "hello"}]} + ) + + assert realtime_guarded_item(client_event) == RealtimeMessageText(role="assistant", texts=("hello",)) + + +@pytest.mark.parametrize( + "client_event", + [ + pytest.param( + _item_create({"type": "message", "role": "user", "content": [{"type": "input_audio", "audio": "AAAA"}]}), + id="audio-only user message", + ), + pytest.param( + _item_create({"type": "function_call_output", "call_id": "call_1", "output": ""}), id="empty tool output" + ), + pytest.param(_item_create("hello"), id="item that is not an object"), + ], +) +def test_items_without_client_text_are_not_screened(client_event: dict) -> None: + assert realtime_guarded_item(client_event) is None + + +@pytest.mark.asyncio +async def test_a_plain_string_block_reason_becomes_the_violation_message() -> None: + guardrail: Final = RaisingGuardrail(HTTPException(status_code=400, detail="hello is not allowed")) + + block: Final = await realtime_guardrail_verdict( + ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) + ) + + assert isinstance(block, RealtimeGuardrailBlock) + assert block.message == "hello is not allowed" + + +@pytest.mark.asyncio +async def test_a_configured_violation_message_replaces_the_block_reason() -> None: + guardrail: Final = RaisingGuardrail(HTTPException(status_code=400, detail={"error": "keyword 'hello' detected"})) + configured: Final = RaisingGuardrail( + HTTPException(status_code=400, detail="hello is not allowed"), realtime_violation_message="Please rephrase" + ) + + block: Final = await realtime_guardrail_verdict( + ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) + ) + configured_block: Final = await realtime_guardrail_verdict( + ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (configured,) + ) + + assert isinstance(block, RealtimeGuardrailBlock) + assert block.message == "keyword 'hello' detected" + assert isinstance(configured_block, RealtimeGuardrailBlock) + assert configured_block.message == "Please rephrase" + + +@pytest.mark.asyncio +async def test_masked_texts_come_back_as_a_rewrite_and_reach_later_guardrails_masked() -> None: + first: Final = MaskingGuardrail("hunter2") + second: Final = MaskingGuardrail("swordfish") + + verdict: Final = await realtime_guardrail_verdict( + ("my hunter2", "and swordfish", "hi"), "user", (GuardrailEventHooks.pre_call,), {}, None, (first, second) + ) + + assert verdict == RealtimeGuardrailRewrite(texts=("my [MASKED]", "and [MASKED]", "hi")) + assert sorted(second.screened[:2]) == [("my [MASKED] and swordfish hi",), ("my [MASKED]and swordfishhi",)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("parts", "screened"), + [ + pytest.param( + ("Compare this photo ", "with that one"), + ("Compare this photo with that one",), + id="a space at the boundary", + ), + pytest.param( + ("Compare this photo", "with that one"), + ("Compare this photowith that one", "Compare this photo with that one"), + id="no space at the boundary", + ), + ], +) +async def test_a_clean_multi_part_item_sends_each_guardrail_only_the_joins_it_needs_one_text_per_call( + parts: tuple[str, ...], screened: tuple[str, ...] +) -> None: + guardrail: Final = MaskingGuardrail("hunter2") + + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) + ) + + assert verdict is None + assert sorted(guardrail.screened) == sorted((text,) for text in screened) + + +@pytest.mark.asyncio +async def test_texts_a_guardrail_leaves_unchanged_need_no_rewrite() -> None: + verdict: Final = await realtime_guardrail_verdict( + ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail("hunter2"),) + ) + + assert verdict is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("parts", "phrase"), + [ + pytest.param(("tell me the secret", "plan"), "secret plan", id="split at a word break"), + pytest.param(("tell me about secret-", "project"), "secret-project", id="split mid-word"), + ], +) +async def test_a_phrase_split_across_typed_text_parts_is_still_blocked(parts: tuple[str, str], phrase: str) -> None: + client_event: Final = _item_create( + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": part} for part in parts]} + ) + guarded_item: Final = realtime_guarded_item(client_event) + assert isinstance(guarded_item, RealtimeMessageText) + + phrase_filter: Final = ContentFilterGuardrail( + guardrail_name="phrase-filter", + blocked_words=[BlockedWord(keyword=phrase, action=ContentFilterAction.BLOCK)], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + verdict: Final = await realtime_guardrail_verdict( + guarded_item.texts, guarded_item.role, (GuardrailEventHooks.pre_call,), {}, None, (phrase_filter,) + ) + + assert isinstance(verdict, RealtimeGuardrailBlock) + + +@pytest.mark.asyncio +async def test_a_mask_inside_one_part_keeps_every_part_in_place() -> None: + first_photo: Final = {"type": "input_image", "image_url": "data:image/png;base64,AAA"} + second_photo: Final = {"type": "input_image", "image_url": "data:image/png;base64,BBB"} + client_event: Final = _item_create( + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Compare this photo"}, + first_photo, + {"type": "input_text", "text": "with hunter2's photo"}, + second_photo, + ], + } + ) + guarded_item: Final = realtime_guarded_item(client_event) + assert isinstance(guarded_item, RealtimeMessageText) + + verdict: Final = await realtime_guardrail_verdict( + guarded_item.texts, "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail("hunter2"),) + ) + + assert isinstance(verdict, RealtimeGuardrailRewrite) + assert rewritten_item(client_event, verdict.texts)["content"] == [ + {"type": "input_text", "text": "Compare this photo"}, + first_photo, + {"type": "input_text", "text": "with [MASKED]'s photo"}, + second_photo, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("parts", "word", "masked"), + [ + pytest.param(("my secret", "plan is here"), "secret plan", "my [MASKED] is here", id="split at a word break"), + pytest.param(("my code is hun", "ter2 ok"), "hunter2", "my code is [MASKED] ok", id="split mid-word"), + ], +) +async def test_a_mask_across_parts_replaces_them_with_the_masked_text_the_model_would_have_read( + parts: tuple[str, ...], word: str, masked: str +) -> None: + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail(word),) + ) + + assert verdict == RealtimeGuardrailRewrite(texts=(masked,)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "parts", + [ + pytest.param(("abc", "123-45-6789 hunter2 rocks"), id="masks inside one part"), + pytest.param(("abc", "123-45-6789 hunter2", "rocks"), id="a mask inside a part and one across parts"), + ], +) +async def test_no_mask_one_text_part_gets_is_lost_when_its_joined_text_masks_differently( + parts: tuple[str, ...], +) -> None: + content_filter: Final = ContentFilterGuardrail( + guardrail_name="ssn-and-keyword-mask", + patterns=[ + ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.MASK) + ], + blocked_words=[BlockedWord(keyword="hunter2 rocks", action=ContentFilterAction.MASK)], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (content_filter,) + ) + + assert isinstance(verdict, RealtimeGuardrailRewrite) + forwarded: Final = " ".join(verdict.texts) + assert "123-45-6789" not in forwarded + assert "hunter2 rocks" not in forwarded + assert forwarded.startswith("abc") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("parts", "rewritten"), + [ + pytest.param(("my code is", "hunter2"), ("my code is", "[MASKED]"), id="mask inside one part"), + pytest.param(("my code is hun", "ter2"), ("my code is [MASKED]",), id="mask across parts"), + ], +) +async def test_a_guardrail_that_answers_with_one_text_still_masks_parts_joined_without_a_space( + parts: tuple[str, ...], rewritten: tuple[str, ...] +) -> None: + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (LastTextRedactingGuardrail("hunter2"),) + ) + + assert verdict == RealtimeGuardrailRewrite(texts=rewritten) + + +@pytest.mark.asyncio +async def test_a_mask_in_a_message_with_more_text_parts_than_the_recursion_limit_still_comes_back() -> None: + parts: Final = (*(f"part {index}." for index in range(sys.getrecursionlimit())), "my code is hunter2.") + + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail("hunter2"),) + ) + + assert verdict == RealtimeGuardrailRewrite(texts=(*parts[:-1], "my code is [MASKED].")) + + +class GuardrailLogRecordingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="log-recorder", event_hook=GuardrailEventHooks.pre_call, default_on=True) + self.guardrail_log_lengths: list[int] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + self.guardrail_log_lengths.append( + len(request_data["litellm_metadata"].get("standard_logging_guardrail_information", [])) + ) + return inputs + + +@pytest.mark.asyncio +async def test_a_guardrail_sees_only_the_current_turns_guardrail_log_however_long_the_session_runs() -> None: + session_request_data: Final = {"litellm_metadata": {"user_api_key_team_id": "team-a"}} + content_filter: Final = ContentFilterGuardrail( + guardrail_name="phrase-filter", + blocked_words=[BlockedWord(keyword="hunter2", action=ContentFilterAction.BLOCK)], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + recorder: Final = GuardrailLogRecordingGuardrail() + + for turn in ("hello", "how are you", "goodbye"): + await realtime_guardrail_verdict( + (turn,), "user", (GuardrailEventHooks.pre_call,), session_request_data, None, (content_filter, recorder) + ) + + assert recorder.guardrail_log_lengths == [1, 1, 1] + assert session_request_data == {"litellm_metadata": {"user_api_key_team_id": "team-a"}} + + +class ScopedGuardrail(RaisingGuardrail): + def __init__( + self, *, skip_system_message: bool | None, skip_tool_message: bool | None, scan_only_tool_results: bool | None + ) -> None: + super().__init__(HTTPException(status_code=400, detail="blocked")) + self.skip_system_message_in_guardrail = skip_system_message + self.skip_tool_message_in_guardrail = skip_tool_message + self.scan_only_tool_results = scan_only_tool_results + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "screened_roles"), + [ + pytest.param({}, ("user", "assistant", "system", "tool"), id="no scope"), + pytest.param({"skip_system_message": True}, ("user", "assistant", "tool"), id="skip system"), + pytest.param({"skip_tool_message": True}, ("user", "assistant", "system"), id="skip tool"), + pytest.param({"scan_only_tool_results": True}, ("tool",), id="only tool results"), + ], +) +async def test_a_guardrail_screens_only_the_roles_its_scope_allows( + scope: dict[str, bool], screened_roles: tuple[str, ...] +) -> None: + guardrail: Final = ScopedGuardrail( + skip_system_message=scope.get("skip_system_message"), + skip_tool_message=scope.get("skip_tool_message"), + scan_only_tool_results=scope.get("scan_only_tool_results"), + ) + + verdicts: Final = { + role: await realtime_guardrail_verdict( + ("hello",), role, (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) + ) + for role in ("user", "assistant", "system", "tool") + } + + assert tuple(role for role, verdict in verdicts.items() if verdict is not None) == screened_roles + + +def test_a_rewrite_puts_the_masked_text_in_place_of_the_typed_text_parts() -> None: + audio: Final = {"type": "input_audio", "audio": "AAAA"} + client_event: Final = _item_create( + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "a"}, audio, {"type": "input_text", "text": "b"}], + } + ) + + assert rewritten_item(client_event, ("[MASKED] b",)) == { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "[MASKED] b"}, audio], + } + + +def test_a_rewrite_replaces_the_tool_output() -> None: + client_event: Final = _item_create({"type": "function_call_output", "call_id": "call_1", "output": "hunter2"}) + + assert rewritten_item(client_event, ("[MASKED]",)) == { + "type": "function_call_output", + "call_id": "call_1", + "output": "[MASKED]", + } + + +@pytest.mark.asyncio +async def test_a_guardrail_that_crashes_is_not_treated_as_a_block() -> None: + guardrail: Final = RaisingGuardrail(RuntimeError("guardrail backend down")) + + with pytest.raises(RuntimeError, match="guardrail backend down"): + await realtime_guardrail_verdict(("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,)) + + +_TEAM_POLICY_CODE: Final = ( + "def apply_guardrail(inputs, request_data, input_type):\n" + ' if request_data["team_id"] == "team-strict":\n' + ' return block("blocked for this team")\n' + " return allow()\n" +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_id", "blocked"), + [pytest.param("team-strict", True, id="policy team"), pytest.param("team-other", False, id="other team")], +) +async def test_a_team_scoped_guardrail_sees_the_session_team(team_id: str, blocked: bool) -> None: + team_policy = CustomCodeGuardrail( + custom_code=_TEAM_POLICY_CODE, guardrail_name="team-policy", event_hook="pre_call", default_on=True + ) + + verdict = await realtime_guardrail_verdict( + ("hello",), + "user", + (GuardrailEventHooks.pre_call,), + {"litellm_metadata": {"user_api_key_team_id": team_id}}, + None, + (team_policy,), + ) + + assert isinstance(verdict, RealtimeGuardrailBlock) is blocked diff --git a/tests/unit/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py index eec1dfcf459..7f205b0bd00 100644 --- a/tests/unit/litellm_core_utils/test_realtime_streaming.py +++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py @@ -1056,6 +1056,117 @@ async def test_transcription_captured_in_backend_to_client(): assert logging_obj.model_call_details["messages"] == streaming.input_messages +@pytest.mark.asyncio +@pytest.mark.parametrize("transcript", [None, 123, ["a"], {"text": "a"}]) +async def test_a_non_string_transcript_does_not_end_the_session(transcript: object) -> None: + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + transcript_event: Final = json.dumps( + {"type": "conversation.item.input_audio_transcription.completed", "transcript": transcript, "item_id": "i"} + ) + later_event: Final = json.dumps({"type": "response.done", "response": {"id": "resp_1"}}) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[transcript_event.encode(), later_event.encode(), ConnectionClosed(None, None)] + ) + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.model_call_details = {} + logging_obj.dispatch_success_handlers = AsyncMock() + + await RealTimeStreaming(client_ws, backend_ws, logging_obj).backend_to_client_send_messages() + + sent_types: Final = [json.loads(call.args[0])["type"] for call in client_ws.send_text.call_args_list] + assert sent_types == ["conversation.item.input_audio_transcription.completed", "response.done"] + backend_ws.send.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transcript", [None, 123, ["a"]]) +async def test_a_non_string_transcript_fails_the_session_when_a_transcript_guardrail_would_screen_it( + transcript: object, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + transcript_event: Final = json.dumps( + {"type": "conversation.item.input_audio_transcription.completed", "transcript": transcript, "item_id": "i"} + ) + later_event: Final = json.dumps({"type": "response.done", "response": {"id": "resp_1"}}) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[transcript_event.encode(), later_event.encode(), ConnectionClosed(None, None)] + ) + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.model_call_details = {} + logging_obj.dispatch_success_handlers = AsyncMock() + + close: Final = await RealTimeStreaming(client_ws, backend_ws, logging_obj).backend_to_client_send_messages() + + assert close.code == 1011 + sent_types: Final = [json.loads(call.args[0])["type"] for call in client_ws.send_text.call_args_list] + assert sent_types == ["conversation.item.input_audio_transcription.completed"] + backend_ws.send.assert_not_called() + + +def _native_lifecycle_transcription_guardrail() -> CustomGuardrail: + class _NativeLifecycleGuardrail(CustomGuardrail): + use_native_lifecycle_hooks = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + return _NativeLifecycleGuardrail( + guardrail_name="test_native_guard", + event_hook=GuardrailEventHooks.realtime_input_transcription, + default_on=True, + ) + + +def _tool_results_only_transcription_guardrail() -> CustomGuardrail: + class _ToolResultsOnlyGuardrail(CustomGuardrail): + scan_only_tool_results = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + return _ToolResultsOnlyGuardrail( + guardrail_name="test_tool_results_guard", + event_hook=GuardrailEventHooks.realtime_input_transcription, + default_on=True, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "guardrail_factory", [_native_lifecycle_transcription_guardrail, _tool_results_only_transcription_guardrail] +) +async def test_a_non_string_transcript_keeps_the_session_when_no_guardrail_would_screen_the_users_turn( + guardrail_factory: Callable[[], CustomGuardrail], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "callbacks", [guardrail_factory()]) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + transcript_event: Final = json.dumps( + {"type": "conversation.item.input_audio_transcription.completed", "transcript": None, "item_id": "i"} + ) + later_event: Final = json.dumps({"type": "response.done", "response": {"id": "resp_1"}}) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[transcript_event.encode(), later_event.encode(), ConnectionClosed(None, None)] + ) + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.model_call_details = {} + logging_obj.dispatch_success_handlers = AsyncMock() + + await RealTimeStreaming(client_ws, backend_ws, logging_obj).backend_to_client_send_messages() + + sent_types: Final = [json.loads(call.args[0])["type"] for call in client_ws.send_text.call_args_list] + assert sent_types == ["conversation.item.input_audio_transcription.completed", "response.done"] + + @pytest.mark.asyncio async def test_transcription_session_captures_usage_and_skips_response_create(): """ @@ -1955,6 +2066,112 @@ async def test_realtime_text_input_guardrail_blocks_and_returns_error(monkeypatc assert len(original_items) == 0, f"Blocked item should not be forwarded to backend, got: {original_items}" +@pytest.mark.asyncio +async def test_the_next_clean_turn_is_answered_after_a_blocked_item_with_no_response_create( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import HTTPException + + class SecretBlockingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + if any("hunter2" in text for text in inputs.get("texts", [])): + raise HTTPException(status_code=400, detail={"error": "hunter2 is not allowed"}) + return inputs + + def item_create(role: str, text: str) -> str: + return json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": role, "content": [{"type": "input_text", "text": text}]}, + } + ) + + monkeypatch.setattr( + litellm, + "callbacks", + [SecretBlockingGuardrail(guardrail_name="no-secrets", event_hook=GuardrailEventHooks.pre_call, default_on=True)], + ) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + client_ws.receive_text = AsyncMock( + side_effect=[ + item_create("system", "never mention hunter2"), + item_create("user", "what is the weather"), + json.dumps({"type": "response.create"}), + Exception("connection closed"), + ] + ) + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + streaming: Final = RealTimeStreaming(client_ws, backend_ws, _mock_logging()) + + await streaming.client_ack_messages() + + sent_to_backend: Final = [json.loads(call.args[0]) for call in backend_ws.send.call_args_list] + clean_index: Final = next( + index + for index, event in enumerate(sent_to_backend) + if event.get("item", {}).get("content", [{}])[0].get("text") == "what is the weather" + ) + assert [event["type"] for event in sent_to_backend[clean_index + 1 :]] == ["response.create"], sent_to_backend + + +@pytest.mark.asyncio +async def test_realtime_text_input_masking_guardrail_forwards_only_the_masked_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class MaskingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + texts: Final = inputs.get("texts", []) + return {**inputs, "texts": [text.replace("test@example.com", "[EMAIL]") for text in texts]} + + monkeypatch.setattr( + litellm, + "callbacks", + [MaskingGuardrail(guardrail_name="email-masker", event_hook=GuardrailEventHooks.pre_call, default_on=True)], + ) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "conversation.item.create", + "item": { + "role": "user", + "content": [{"type": "input_text", "text": "My email is test@example.com"}], + }, + } + ), + Exception("connection closed"), + ] + ) + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock(side_effect=ConnectionClosed(None, None)) + logging_obj: Final = _mock_logging() + logging_obj.pre_call = MagicMock() + + await RealTimeStreaming(client_ws, backend_ws, logging_obj).client_ack_messages() + + sent_to_backend: Final = json.dumps([c.args[0] for c in backend_ws.send.call_args_list if c.args]) + assert "test@example.com" not in sent_to_backend + assert "My email is [EMAIL]" in sent_to_backend + assert client_ws.send_text.call_args_list == [] + + @pytest.mark.asyncio async def test_realtime_function_call_output_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch): """ diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index cf8de9c40f5..94fd915844a 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -3523,6 +3523,34 @@ def test_is_chunk_non_empty_with_valid_tool_calls( ) +@pytest.mark.parametrize( + ("tool_calls", "non_empty"), + [([{"id": "call_1", "type": "function", "function": {"name": "f", "arguments": "{}"}}], True), ([], False)], +) +def test_is_chunk_non_empty_reads_tool_calls_from_the_completion_obj( + initialized_custom_stream_wrapper: CustomStreamWrapper, tool_calls: list[dict[str, object]], non_empty: bool +): + assert ( + initialized_custom_stream_wrapper.is_chunk_non_empty( + completion_obj={"tool_calls": tool_calls}, + model_response=ModelResponseStream(), + response_obj={}, + ) + is non_empty + ) + + +def test_is_chunk_non_empty_rejects_completion_obj_tool_calls_without_a_length( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + with pytest.raises(TypeError, match="object of type 'generator' has no len"): + initialized_custom_stream_wrapper.is_chunk_non_empty( + completion_obj={"tool_calls": (tool_call for tool_call in ())}, + model_response=ModelResponseStream(), + response_obj={}, + ) + + def _make_chunk(content: Optional[str]) -> ModelResponseStream: return ModelResponseStream( id="test", @@ -4655,6 +4683,25 @@ def test_dispatch_triton_stream( assert initialized_custom_stream_wrapper.received_finish_reason == "stop" +@pytest.mark.parametrize( + "chunk", + [ + 'data: {"text_output": "triton text", "is_finished": true, "stop_reason": "stop", "generated_token_count": 2}', + b'data: {"text_output": "triton text", "is_finished": true, "stop_reason": "stop", "generated_token_count": 2}', + ], +) +def test_handle_triton_stream_parses_sse_text_output_lines( + initialized_custom_stream_wrapper: CustomStreamWrapper, chunk: str | bytes +): + assert initialized_custom_stream_wrapper.handle_triton_stream(chunk) == { + "text": "triton text", + "is_finished": True, + "finish_reason": "stop", + "prompt_tokens": 0, + "completion_tokens": 2, + } + + def test_dispatch_text_completion_openai_with_usage( initialized_custom_stream_wrapper: CustomStreamWrapper, ): diff --git a/tests/unit/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py b/tests/unit/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py index 01eff7f6bbf..1a10acd51ad 100644 --- a/tests/unit/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py +++ b/tests/unit/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py @@ -1,9 +1,16 @@ """Tests for litellm/llms/a2a/chat/guardrail_translation/handler.py.""" import json +from typing import Final, Literal +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey +from litellm.types.agents import LiteLLMSendMessageResponse +from litellm.types.utils import GenericGuardrailAPIInputs def _text_event(text: str) -> str: @@ -37,3 +44,53 @@ class TestA2AGuardrailHandlerStreamingScanKey: def test_json_items_that_are_not_objects_are_ignored(self): key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hi"), "[1, 2]", "5"]) assert key.texts == ("hi",) + + +SSN_REPLY: Final = {"kind": "message", "role": "agent", "parts": [{"kind": "text", "text": "SSN 123-45-6789"}]} + + +class _MaskingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Logging | None = None, + ) -> GenericGuardrailAPIInputs: + return {"texts": [text.replace("123-45-6789", "[SSN]") for text in inputs.get("texts", [])]} + + +class _BlockingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Logging | None = None, + ) -> GenericGuardrailAPIInputs: + raise ValueError(f"blocked {inputs.get('texts')}") + + +class TestA2AGuardrailHandlerOutputResponse: + async def test_dict_response_gets_the_guardrailed_text_in_place(self): + response: Final = {"jsonrpc": "2.0", "id": "req-1", "result": json.loads(json.dumps(SSN_REPLY))} + + processed: Final = await A2AGuardrailHandler().process_output_response( + response, _MaskingGuardrail(guardrail_name="mask-ssn") + ) + + assert processed is response + assert response["result"]["parts"][0]["text"] == "SSN [SSN]" + + async def test_pydantic_response_text_reaches_the_guardrail(self): + response: Final = LiteLLMSendMessageResponse(id="req-1", result=json.loads(json.dumps(SSN_REPLY))) + + with pytest.raises(ValueError, match=r"blocked \['SSN 123-45-6789'\]"): + await A2AGuardrailHandler().process_output_response(response, _BlockingGuardrail(guardrail_name="block")) + + async def test_response_of_an_unknown_type_is_returned_without_running_the_guardrail(self): + processed: Final = await A2AGuardrailHandler().process_output_response( + "SSN 123-45-6789", _BlockingGuardrail(guardrail_name="block") + ) + + assert processed == "SSN 123-45-6789" diff --git a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 8affe2e136e..b6229dbaec0 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -3458,3 +3458,41 @@ def test_reading_a_dropped_budget_for_a_guardrail_logs_no_drop_warning(caplog: p assert messages == [{"role": "user", "content": "hi"}] assert not [record for record in caplog.records if "Dropping" in record.getMessage()] + + +FAILED_TOOL_RESULT_REQUEST: Final = { + "model": "claude-sonnet", + "max_tokens": 16, + "messages": [ + {"role": "user", "content": "Rename the helper"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "edit0", "name": "Edit", "input": {"file_path": "a.py"}}], + }, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "edit0", "content": "String not found", "is_error": True} + ], + }, + ], +} + + +@pytest.mark.parametrize("keep_tool_errors", [True, False]) +def test_structured_messages_carry_tool_failures_only_when_asked(keep_tool_errors: bool) -> None: + messages: Final = AnthropicMessagesHandler().get_structured_messages( + FAILED_TOOL_RESULT_REQUEST, keep_tool_errors=keep_tool_errors + ) + assert messages is not None + (tool_message,) = [message for message in messages if message["role"] == "tool"] + assert tool_message.get("is_error") is (True if keep_tool_errors else None) + + +def test_request_sent_to_a_model_never_carries_tool_failure_marker() -> None: + from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter + + request, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + anthropic_message_request=FAILED_TOOL_RESULT_REQUEST # pyright: ignore[reportArgumentType] # a JSON request in the Anthropic shape + ) + assert "is_error" not in json.dumps(request["messages"], default=str) diff --git a/tests/unit/llms/anthropic/messages/test_stream_assembly.py b/tests/unit/llms/anthropic/messages/test_stream_assembly.py index f8cd00ac402..81b72bfbf06 100644 --- a/tests/unit/llms/anthropic/messages/test_stream_assembly.py +++ b/tests/unit/llms/anthropic/messages/test_stream_assembly.py @@ -7,6 +7,7 @@ from litellm.llms.anthropic.messages.stream_assembly import ( anthropic_message_delta_fields, anthropic_sse_error_events, assemble_anthropic_sse_stream, + collapse_pure_text_chunks, has_unreadable_json_event, is_anthropic_sse_stream, is_raw_sse_stream, @@ -379,3 +380,24 @@ def test_assembly_through_an_upstream_error_reads_past_the_error_frame(error_fra chunks: Final = (*ANTHROPIC_STREAM[:4], error_frame) assert assemble_anthropic_sse_stream(chunks, through_upstream_error=True) is not None + + +def test_collapsing_a_text_run_rejects_a_delta_whose_text_is_not_a_string(): + chunks: Final = ( + _frame( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}}, + ), + _frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": 7}}, + ), + _frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + ) + + with pytest.raises(TypeError, match="sequence item 1: expected str instance, int found"): + collapse_pure_text_chunks(chunks) diff --git a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index 51ba2c34cd7..b14500bea1e 100644 --- a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -1,6 +1,7 @@ import base64 import json from collections.abc import Mapping +from io import BytesIO from typing import Final import httpx @@ -144,6 +145,27 @@ def test_flux2_image_edit_rejects_too_many_references(model: str, reference_imag ) +def _flux2_reference_request(image: object) -> dict[str, object]: + request, _ = AzureFoundryFlux2ImageEditConfig().transform_image_edit_request( + model="FLUX.2-pro", + prompt="Add a red hat", + image=image, + image_edit_optional_request_params={}, + litellm_params={}, + headers={}, + ) + return request + + +def test_flux2_image_edit_reads_a_file_like_reference_image(): + assert _flux2_reference_request(BytesIO(b"png-bytes"))["input_image"] == base64.b64encode(b"png-bytes").decode() + + +def test_flux2_image_edit_rejects_a_reference_image_it_cannot_read(): + with pytest.raises(ValueError, match="Unsupported image type: "): + _flux2_reference_request("cat.png") + + @pytest.mark.parametrize("dimensions", ({"size": "2048x1024"}, {"width": 2048, "height": 1024}, {"width": "2048", "height": "1024"})) @pytest.mark.usefixtures("local_model_cost_map") def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[str, int | str]): diff --git a/tests/unit/llms/base_llm/responses/test_codex_compat.py b/tests/unit/llms/base_llm/responses/test_codex_compat.py index be811b17a82..2e66e986aa3 100644 --- a/tests/unit/llms/base_llm/responses/test_codex_compat.py +++ b/tests/unit/llms/base_llm/responses/test_codex_compat.py @@ -7,11 +7,14 @@ smoke test passes and hides the problem entirely. """ import json +from functools import partial import pytest from litellm.llms.base_llm.responses.codex_compat import normalize_codex_input_items +normalize = partial(normalize_codex_input_items, keep_encrypted_agent_messages=True) + USER = {"role": "user", "content": "hi"} @@ -22,7 +25,7 @@ class TestAgentMessage: "role": "assistant", "content": [{"type": "output_text", "text": "prior turn"}], } - out, types = normalize_codex_input_items([item, USER]) + out, types = normalize([item, USER]) assert types == ("agent_message",) assert out[0] == { "type": "message", @@ -30,40 +33,55 @@ class TestAgentMessage: "content": ({"type": "output_text", "text": "prior turn"},), } - def test_encrypted_content_slot_is_used_as_text(self): + @pytest.mark.parametrize( + "opaque_block", + [{"type": "encrypted_content", "encrypted_content": "ciphertext"}, {"encrypted_content": "ciphertext"}], + ) + def test_encrypted_item_reaches_a_decrypting_backend_intact(self, opaque_block): + item = { + "type": "agent_message", + "id": "amsg_1", + "author": "/root/reviewer", + "recipient": "/root", + "content": [{"type": "input_text", "text": "Worker message:\n"}, opaque_block], + } + out, types = normalize([item, USER]) + assert out == [item, USER] + assert types == () + + def test_encrypted_content_slot_is_used_as_text_when_the_backend_never_encrypts(self): """Codex puts the plaintext payload there when the model issued no encrypted args.""" item = {"type": "agent_message", "content": [{"encrypted_content": "plain"}]} - out, _ = normalize_codex_input_items([item, USER]) + out, types = normalize([item, USER], keep_encrypted_agent_messages=False) assert out[0]["content"] == ({"type": "output_text", "text": "plain"},) + assert types == ("agent_message",) def test_non_list_content_yields_no_text_and_drops_the_item(self): - out, types = normalize_codex_input_items([{"type": "agent_message", "content": "not a list"}, USER]) + out, types = normalize([{"type": "agent_message", "content": "not a list"}, USER]) assert out == [USER] assert types == ("agent_message",) def test_textless_item_is_dropped(self): - out, types = normalize_codex_input_items([{"type": "agent_message", "content": []}, USER]) + out, types = normalize([{"type": "agent_message", "content": []}, USER]) assert out == [USER] assert types == ("agent_message",) class TestContextCompaction: def test_becomes_compaction(self): - out, types = normalize_codex_input_items([{"type": "context_compaction", "encrypted_content": "abc"}, USER]) + out, types = normalize([{"type": "context_compaction", "encrypted_content": "abc"}, USER]) assert out[0] == {"type": "compaction", "encrypted_content": "abc"} assert types == ("context_compaction",) @pytest.mark.parametrize("bad", [{}, {"encrypted_content": ""}, {"encrypted_content": 7}]) def test_without_usable_content_is_dropped(self, bad): - out, _ = normalize_codex_input_items([{"type": "context_compaction", **bad}, USER]) + out, _ = normalize([{"type": "context_compaction", **bad}, USER]) assert out == [USER] class TestLocalShellCall: def test_becomes_the_function_call_its_output_pairs_with(self): - out, types = normalize_codex_input_items( - [{"type": "local_shell_call", "call_id": "c1", "action": {"command": ["ls"]}}, USER] - ) + out, types = normalize([{"type": "local_shell_call", "call_id": "c1", "action": {"command": ["ls"]}}, USER]) assert out[0] == { "type": "function_call", "call_id": "c1", @@ -73,21 +91,21 @@ class TestLocalShellCall: assert types == ("local_shell_call",) def test_missing_action_yields_empty_arguments(self): - out, _ = normalize_codex_input_items([{"type": "local_shell_call", "call_id": "c1"}, USER]) + out, _ = normalize([{"type": "local_shell_call", "call_id": "c1"}, USER]) assert out[0]["arguments"] == "{}" def test_without_call_id_is_dropped(self): - out, _ = normalize_codex_input_items([{"type": "local_shell_call"}, USER]) + out, _ = normalize([{"type": "local_shell_call"}, USER]) assert out == [USER] class TestPassthroughAndShape: def test_string_input_untouched(self): - assert normalize_codex_input_items("just a prompt") == ("just a prompt", ()) + assert normalize("just a prompt") == ("just a prompt", ()) def test_unrelated_items_untouched_and_no_types_reported(self): items = [USER, {"type": "message", "role": "assistant", "content": []}] - out, types = normalize_codex_input_items(items) + out, types = normalize(items) assert out == items assert types == () @@ -96,7 +114,7 @@ class TestPassthroughAndShape: per-item return protocol, so it is dropped. Other non-mapping entries pass through untouched. This matches the behaviour before the normalizer moved out of the bedrock_mantle config.""" - out, types = normalize_codex_input_items(["a string", 42, None, USER]) + out, types = normalize(["a string", 42, None, USER]) assert out == ["a string", 42, USER] assert types == () @@ -106,11 +124,11 @@ class TestPassthroughAndShape: {"type": "agent_message", "content": [{"text": "x"}]}, {"type": "local_shell_call", "call_id": "c2"}, ] - _, types = normalize_codex_input_items(items) + _, types = normalize(items) assert types == ("agent_message", "local_shell_call") def test_returns_a_list_not_a_tuple(self): """The input->messages conversion downstream narrows on isinstance(input, list); a tuple silently yields zero messages and the provider rejects the request.""" - out, _ = normalize_codex_input_items([{"type": "agent_message", "content": [{"text": "x"}]}, USER]) + out, _ = normalize([{"type": "agent_message", "content": [{"text": "x"}]}, USER]) assert isinstance(out, list) diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py index 053574d7dd6..4f17f8b10e1 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py @@ -570,6 +570,7 @@ def _invoke_params_for_schema( ("us.anthropic.claude-opus-4-8", "bedrock", {"thinking": {"type": "adaptive"}}, True), ("us.anthropic.claude-opus-4-8", "bedrock", {"reasoning_effort": "low"}, True), ("us.anthropic.claude-haiku-5-5", "bedrock", {"thinking": {"type": "disabled"}}, False), + ("us.anthropic.claude-haiku-5-5", "default", {"thinking": {"type": "disabled"}}, True), ], ids=[ "legacy-key", @@ -578,6 +579,7 @@ def _invoke_params_for_schema( "bedrock-key-adaptive", "bedrock-key-effort", "bedrock-key-thinking-disabled", + "default-key-thinking-disabled", ], ) def test_bedrock_native_output_key_applies_only_when_the_model_will_think( @@ -2047,9 +2049,9 @@ def _invoke_chat_params( ) @pytest.mark.parametrize("model", ["us.anthropic.claude-opus-4-8", "global.anthropic.claude-sonnet-5"]) def test_invoke_uses_opt_in_native_output_when_the_opt_in_drove_native_output( - local_model_cost_map, model, thinking_params, monkeypatch + set_native_structured_output, model, thinking_params ): - monkeypatch.setitem(litellm.model_cost[model], "supports_native_structured_output", False) + set_native_structured_output(model, "bedrock") optional_params: Final = _invoke_chat_params( model, None, {"response_format": _CITY_RESPONSE_FORMAT, **thinking_params} ) @@ -2084,8 +2086,12 @@ def test_invoke_uses_opt_in_native_output_is_false_when_native_output_did_not_co ) -def test_invoke_uses_default_native_output_for_schema_retry(local_model_cost_map: None) -> None: - model: Final = "us.anthropic.claude-opus-5" +@pytest.mark.parametrize( + "model", + ["us.anthropic.claude-opus-5", "us.anthropic.claude-sonnet-5"], + ids=["opus-5", "sonnet-5"], +) +def test_invoke_uses_default_native_output_for_schema_retry(local_model_cost_map: None, model: str) -> None: optional_params: Final = _invoke_chat_params( model, None, {"response_format": _CITY_RESPONSE_FORMAT, "thinking": {"type": "disabled"}} ) diff --git a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py index 7d1b502df2b..d43cc771185 100644 --- a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -599,6 +599,139 @@ async def test_native_output_retry_is_bounded_by_the_request_timeout_not_the_str assert _sent_paths(respx_mock) == ["converse-stream", "converse"] +# Bedrock's 403 for a role without the streaming action, seen live on us.anthropic.claude-sonnet-5 on 2026-11-06 +_STREAM_PERMISSION_DENIAL: Final = json.dumps( + { + "Message": "User: arn:aws:sts::123456789012:assumed-role/invoke-only/qa is not authorized to perform: " + "bedrock:InvokeModelWithResponseStream on resource: arn:aws:bedrock:us-east-1:123456789012:inference-profile/" + "us.anthropic.claude-opus-5-5 because no identity-based policy allows the " + "bedrock:InvokeModelWithResponseStream action" + } +).encode() + + +def _denies_streaming_then_answers(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/converse-stream"): + return httpx.Response(403, content=_STREAM_PERMISSION_DENIAL, request=request) + return httpx.Response(200, content=CONVERSE_ANSWER, headers={"content-type": "application/json"}, request=request) + + +@pytest.mark.asyncio +async def test_a_denied_native_output_stream_is_sent_again_buffered_with_native_output( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post(url__regex=_BEDROCK_RUNTIME_URL).mock(side_effect=_denies_streaming_then_answers) + router: Final = _bedrock_router( + "bedrock/converse/us.anthropic.claude-opus-5-5", timeout=60, stream_timeout=STREAM_TIMEOUT_SECONDS + ) + + response: Final = await router.acompletion( + model="bedrock", + messages=[{"role": "user", "content": "hi"}], + stream=True, + thinking={"type": "disabled"}, + response_format={"type": "json_schema", "json_schema": {"name": "answer", "schema": _ANSWER_SCHEMA}}, + ) + + assert "".join([chunk.choices[0].delta.content or "" async for chunk in response]) == "hi" + assert _sent_paths(respx_mock) == ["converse-stream", "converse"] + retry_body: Final = json.loads(respx_mock.calls[1].request.content) + assert "textFormat" in retry_body.get("outputConfig", {}) and "toolConfig" not in retry_body + + +_INVOKE_ANSWER: Final = json.dumps( + { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-opus-5-5", + "content": [{"type": "text", "text": '{"answer": "hi"}'}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, + } +).encode() + + +def _denies_invoke_streaming_then_answers(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/invoke-with-response-stream"): + return httpx.Response(403, content=_STREAM_PERMISSION_DENIAL, request=request) + return httpx.Response(200, content=_INVOKE_ANSWER, headers={"content-type": "application/json"}, request=request) + + +@pytest.mark.asyncio +async def test_a_denied_native_output_stream_on_the_invoke_route_is_sent_again_buffered_to_invoke( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post(url__regex=_BEDROCK_RUNTIME_URL).mock(side_effect=_denies_invoke_streaming_then_answers) + router: Final = _bedrock_router( + "bedrock/invoke/us.anthropic.claude-opus-5-5", timeout=60, stream_timeout=STREAM_TIMEOUT_SECONDS + ) + + response: Final = await router.acompletion( + model="bedrock", + messages=[{"role": "user", "content": "hi"}], + stream=True, + thinking={"type": "disabled"}, + response_format={"type": "json_schema", "json_schema": {"name": "answer", "schema": _ANSWER_SCHEMA}}, + ) + + assert "".join([chunk.choices[0].delta.content or "" async for chunk in response]) == '{"answer": "hi"}' + assert _sent_paths(respx_mock) == ["invoke-with-response-stream", "invoke"] + retry_body: Final = json.loads(respx_mock.calls[1].request.content) + assert "format" in retry_body.get("output_config", {}) and "tools" not in retry_body + + +_INVOKE_JSON_TOOL_ANSWER: Final = json.dumps( + { + "id": "msg_2", + "type": "message", + "role": "assistant", + "model": "claude-opus-5-5", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "json_tool_call", "input": {"answer": "hi"}}], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, + } +).encode() + + +def _denies_invoke_streaming_and_rejects_native_output(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/invoke-with-response-stream"): + return httpx.Response(403, content=_STREAM_PERMISSION_DENIAL, request=request) + if "format" in json.loads(request.content).get("output_config", {}): + return httpx.Response(400, content=_NATIVE_OUTPUT_REJECTION, request=request) + return httpx.Response( + 200, content=_INVOKE_JSON_TOOL_ANSWER, headers={"content-type": "application/json"}, request=request + ) + + +@pytest.mark.asyncio +async def test_a_rejected_schema_after_a_denied_invoke_stream_falls_back_to_the_json_tool_still_buffered( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post(url__regex=_BEDROCK_RUNTIME_URL).mock(side_effect=_denies_invoke_streaming_and_rejects_native_output) + router: Final = _bedrock_router( + "bedrock/invoke/us.anthropic.claude-opus-5-5", timeout=60, stream_timeout=STREAM_TIMEOUT_SECONDS + ) + + response: Final = await router.acompletion( + model="bedrock", + messages=[{"role": "user", "content": "hi"}], + stream=True, + thinking={"type": "disabled"}, + response_format={"type": "json_schema", "json_schema": {"name": "answer", "schema": _ANSWER_SCHEMA}}, + ) + + assert "".join([chunk.choices[0].delta.content or "" async for chunk in response]) == '{"answer": "hi"}' + assert _sent_paths(respx_mock) == ["invoke-with-response-stream", "invoke", "invoke"] + json_tool_body: Final = json.loads(respx_mock.calls[2].request.content) + assert "format" not in json_tool_body.get("output_config", {}) and "tools" in json_tool_body + + @pytest.mark.asyncio async def test_failed_native_output_retry_reports_the_request_timeout_it_was_bounded_by( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index b37d6977c10..4dcf312737d 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -5225,9 +5225,9 @@ def test_a_model_that_cannot_force_the_json_tool_gets_native_output_without_thin ("model", "thinking_params", "native"), [ ("us.anthropic.claude-opus-5-5", None, True), - ("us.anthropic.claude-opus-5-5", {"thinking": {"type": "adaptive"}}, False), + ("us.anthropic.claude-opus-5-5", {"thinking": {"type": "adaptive"}}, True), ("global.anthropic.claude-fable-5-1", None, True), - ("global.anthropic.claude-fable-5-1", {"reasoning_effort": "low"}, False), + ("global.anthropic.claude-fable-5-1", {"reasoning_effort": "low"}, True), ("us-gov.anthropic.claude-opus-5-5", None, False), ("us-gov.anthropic.claude-fable-5-1", None, False), ("us-gov-west-1/anthropic.claude-opus-5-5", None, False), @@ -5235,7 +5235,7 @@ def test_a_model_that_cannot_force_the_json_tool_gets_native_output_without_thin ("us-east-1/anthropic.claude-opus-5-5", None, True), ], ) -def test_models_whose_thinking_stops_under_native_output_keep_the_json_tool_only_when_thinking( +def test_native_output_follows_model_cost_map_for_bedrock_profiles( model: str, thinking_params: Mapping[str, object] | None, native: bool ) -> None: _, request = _converse_request_for_json_schema(model, False, thinking_params) @@ -5248,6 +5248,17 @@ def test_models_whose_thinking_stops_under_native_output_keep_the_json_tool_only @pytest.mark.parametrize( ("model", "thinking_params", "native"), [ + ("us.anthropic.claude-opus-5-5", {"thinking": {"type": "disabled"}}, True), + ("global.anthropic.claude-fable-5-1", {"thinking": {"type": "adaptive"}}, True), + ("global.anthropic.claude-fable-5-1", {"thinking": {"type": "disabled"}}, True), + ("us.anthropic.claude-sonnet-5", None, True), + ("us.anthropic.claude-sonnet-5", {"thinking": {"type": "adaptive"}}, True), + ("us.anthropic.claude-sonnet-5", {"thinking": {"type": "disabled"}}, True), + ("us.anthropic.claude-haiku-5-5", None, True), + ("us.anthropic.claude-haiku-5-5", {"thinking": {"type": "adaptive"}}, True), + ("us.anthropic.claude-haiku-5-5", {"thinking": {"type": "disabled"}}, True), + ("us-gov.anthropic.claude-sonnet-5", None, False), + ("us-gov.anthropic.claude-haiku-5-5", None, False), ("us.anthropic.claude-opus-5", None, True), ("us.anthropic.claude-opus-5", {"thinking": {"type": "adaptive"}}, True), ("us.anthropic.claude-opus-5", {"thinking": {"type": "disabled"}}, True), @@ -5258,6 +5269,17 @@ def test_models_whose_thinking_stops_under_native_output_keep_the_json_tool_only ("us-gov.anthropic.claude-opus-4-8", None, False), ], ids=[ + "opus-5-5-disabled", + "fable-5-1-adaptive", + "fable-5-1-disabled", + "sonnet-5-unset", + "sonnet-5-adaptive", + "sonnet-5-disabled", + "haiku-5-5-unset", + "haiku-5-5-adaptive", + "haiku-5-5-disabled", + "sonnet-5-us-gov", + "haiku-5-5-us-gov", "opus-5-unset", "opus-5-adaptive", "opus-5-disabled", @@ -5268,7 +5290,7 @@ def test_models_whose_thinking_stops_under_native_output_keep_the_json_tool_only "opus-4-8-us-gov", ], ) -def test_opus_5_and_4_8_cost_map_defaults_use_native_output_across_thinking_modes( +def test_model_cost_native_output_defaults_hold_across_thinking_modes( model: str, thinking_params: Mapping[str, object] | None, native: bool ) -> None: _, request = _converse_request_for_json_schema(model, False, thinking_params) @@ -5282,6 +5304,30 @@ def test_opus_5_and_4_8_cost_map_defaults_use_native_output_across_thinking_mode assert [tool["toolSpec"]["name"] for tool in request["toolConfig"]["tools"]] == ["json_tool_call"] +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + ("model", "thinking_params"), + [ + ("us.anthropic.claude-opus-5-5", {"thinking": {"type": "adaptive"}}), + ("global.anthropic.claude-fable-5-1", {"thinking": {"type": "adaptive"}}), + ("us.anthropic.claude-sonnet-5", {"thinking": {"type": "disabled"}}), + ("us.anthropic.claude-haiku-5-5", None), + ("us.anthropic.claude-opus-5", None), + ], +) +@pytest.mark.parametrize("stream", [False, True]) +def test_deployment_opt_out_keeps_the_json_tool_on_models_that_default_to_native_output( + model: str, thinking_params: Mapping[str, object] | None, stream: bool +) -> None: + optional_params, request = _converse_request_for_json_schema( + model, stream, thinking_params, model_info={"supports_bedrock_native_structured_output": False} + ) + + assert "outputConfig" not in request + assert [tool["toolSpec"]["name"] for tool in request["toolConfig"]["tools"]] == ["json_tool_call"] + assert optional_params.get("fake_stream", False) is stream + + def _converse_capital_answer(request: httpx.Request) -> httpx.Response: return httpx.Response( 200, @@ -5295,9 +5341,7 @@ def _converse_capital_answer(request: httpx.Request) -> httpx.Response: @pytest.mark.usefixtures("local_model_cost_map") -@pytest.mark.parametrize( - "model", ["us.anthropic.claude-opus-5-5", "global.anthropic.claude-sonnet-5-5", "global.anthropic.claude-fable-5-1"] -) +@pytest.mark.parametrize("model", ["global.anthropic.claude-sonnet-5-5"]) def test_streamed_json_schema_on_a_model_that_cannot_force_a_tool_still_calls_converse_without_streaming( model: str, ) -> None: @@ -5610,7 +5654,9 @@ def test_schema_past_a_native_output_size_limit_keeps_the_tool_path( @pytest.mark.parametrize( - ("native_key", "native_past_limit"), [("legacy", True), ("bedrock", False)], ids=["legacy-key", "bedrock-key"] + ("native_key", "native_past_limit"), + [("legacy", True), ("default", True), ("bedrock", False)], + ids=["legacy-key", "default-key", "bedrock-key"], ) def test_field_count_limit_sends_the_json_tool_only_on_the_bedrock_native_key( set_native_structured_output, native_key, native_past_limit @@ -10335,7 +10381,6 @@ def test_converse_reserves_the_tool_cache_point_within_the_cap(tool_cache_params assert request["toolConfig"]["tools"][-1] == cache_point - @pytest.mark.parametrize( "unsent_tool_params", [ @@ -10385,6 +10430,7 @@ def test_converse_leaves_the_cap_to_markers_when_a_marked_tool_is_not_sent( ] assert cache_point not in request.get("toolConfig", {}).get("tools", []) + def _converse_tool_turn(*result_parts: dict[str, object]) -> list[dict[str, object]]: return [ { @@ -10506,7 +10552,9 @@ def test_converse_translates_prompt_cache_breakpoint_on_an_unmapped_arn_only_whe messages=[ { "role": "user", - "content": [{"type": "text", "text": "Hello", "prompt_cache_breakpoint": {"mode": "explicit", **ttl}}], + "content": [ + {"type": "text", "text": "Hello", "prompt_cache_breakpoint": {"mode": "explicit", **ttl}} + ], } ], optional_params={}, @@ -10518,7 +10566,6 @@ def test_converse_translates_prompt_cache_breakpoint_on_an_unmapped_arn_only_whe assert user_content(unflagged_arn, {"ttl": "30m"}) == [{"text": "Hello"}] - def test_converse_prompt_cache_breakpoint_follows_the_arn_flag_over_its_base_model(local_model_cost_map): from litellm import Router @@ -10545,7 +10592,9 @@ def test_converse_prompt_cache_breakpoint_follows_the_arn_flag_over_its_base_mod messages=[ { "role": "user", - "content": [{"type": "text", "text": "Hello", "prompt_cache_breakpoint": {"mode": "explicit", **ttl}}], + "content": [ + {"type": "text", "text": "Hello", "prompt_cache_breakpoint": {"mode": "explicit", **ttl}} + ], } ], optional_params={}, @@ -10633,7 +10682,9 @@ def test_converse_prompt_cache_breakpoint_on_guarded_user_text_keeps_its_cache_p def test_converse_guarded_user_text_keeps_no_cache_point_on_a_model_not_known_to_take_one(local_model_cost_map): request: Final = litellm.AmazonConverseConfig()._transform_request( model="arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/unmapped1", - messages=[{"role": "user", "content": [{"type": "text", "text": "Hello", "cache_control": {"type": "ephemeral"}}]}], + messages=[ + {"role": "user", "content": [{"type": "text", "text": "Hello", "cache_control": {"type": "ephemeral"}}]} + ], optional_params={"guardrailConfig": {"guardrailIdentifier": "gr-1", "guardrailVersion": "DRAFT"}}, litellm_params={}, ) @@ -11638,9 +11689,9 @@ _OUTPUT_CONFIG_FORMAT: Final = { ) @pytest.mark.parametrize("model", ["us.anthropic.claude-opus-4-8", "global.anthropic.claude-sonnet-5"]) def test_uses_opt_in_native_output_when_the_opt_in_drove_native_output( - local_model_cost_map, monkeypatch, model, request_params + set_native_structured_output, model, request_params ): - monkeypatch.setitem(litellm.model_cost[model], "supports_native_structured_output", False) + set_native_structured_output(model, "bedrock") optional_params: Final = _converse_params(model, None, request_params) assert AmazonConverseConfig().uses_opt_in_native_output(model, optional_params, {"model_info": {"id": "d1"}}) @@ -11672,9 +11723,9 @@ def test_uses_opt_in_native_output_is_false_when_native_output_did_not_come_from def test_uses_opt_in_native_output_reads_the_base_model_of_an_application_inference_profile( - local_model_cost_map, monkeypatch + set_native_structured_output, ): - monkeypatch.setitem(litellm.model_cost["us.anthropic.claude-opus-4-8"], "supports_native_structured_output", False) + set_native_structured_output("us.anthropic.claude-opus-4-8", "bedrock") profile: Final = "arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/profile-id" litellm_params: Final = {"base_model": "bedrock/us.anthropic.claude-opus-4-8"} optional_params: Final = _converse_params( @@ -11768,8 +11819,9 @@ def test_a_refused_native_schema_is_sent_again_only_when_the_retry_drops_native_ @pytest.mark.respx(assert_all_called=True) +@pytest.mark.parametrize("model", ["us.anthropic.claude-sonnet-5", "us.anthropic.claude-opus-5"]) def test_a_default_native_schema_rejection_retries_with_the_json_tool( - respx_mock: respx.Router, local_model_cost_map: None + respx_mock: respx.Router, local_model_cost_map: None, model: str ) -> None: route: Final = respx_mock.post(url__regex=r"/converse$").mock( return_value=httpx.Response(400, json={"message": _BEDROCK_SCHEMA_REFUSAL}) @@ -11777,7 +11829,7 @@ def test_a_default_native_schema_rejection_retries_with_the_json_tool( with pytest.raises(litellm.BadRequestError, match=re.escape(_BEDROCK_SCHEMA_REFUSAL)): completion( - model="bedrock/converse/us.anthropic.claude-opus-5", + model=f"bedrock/converse/{model}", messages=[{"role": "user", "content": "What is the capital of France?"}], response_format=CAPITAL_RESPONSE_FORMAT, thinking={"type": "disabled"}, @@ -11938,7 +11990,9 @@ def test_converse_sends_no_tools_when_drop_params_drops_every_hosted_mcp_tool(mo def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: return tuple( - r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() + r.getMessage() + for r in caplog.records + if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() ) @@ -12102,9 +12156,15 @@ def test_converse_sends_only_the_header_betas_when_header_and_body_both_carry_be def _five_cache_hint_conversation() -> list[dict[str, object]]: return [ - {"role": "system", "content": [{"type": "text", "text": "system prompt", "cache_control": {"type": "ephemeral"}}]}, + { + "role": "system", + "content": [{"type": "text", "text": "system prompt", "cache_control": {"type": "ephemeral"}}], + }, {"role": "user", "content": [{"type": "text", "text": "turn 1", "cache_control": {"type": "ephemeral"}}]}, - {"role": "assistant", "content": [{"type": "text", "text": "answer 1", "cache_control": {"type": "ephemeral"}}]}, + { + "role": "assistant", + "content": [{"type": "text", "text": "answer 1", "cache_control": {"type": "ephemeral"}}], + }, {"role": "user", "content": [{"type": "text", "text": "turn 2", "cache_control": {"type": "ephemeral"}}]}, {"role": "assistant", "content": "answer 2"}, {"role": "user", "content": [{"type": "text", "text": "turn 3", "cache_control": {"type": "ephemeral"}}]}, @@ -12121,7 +12181,10 @@ def _five_cache_hint_conversation() -> list[dict[str, object]]: ], ) def test_converse_claude_request_over_the_cache_control_cap_drops_the_oldest_hints_only_under_drop_params( - drop_params: bool | str, expected_system_cache_points: int, monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None + drop_params: bool | str, + expected_system_cache_points: int, + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, ) -> None: monkeypatch.setattr(litellm, "drop_params", False) messages: Final = _five_cache_hint_conversation() diff --git a/tests/unit/llms/bedrock/chat/test_native_output_fallback.py b/tests/unit/llms/bedrock/chat/test_native_output_fallback.py index 81832d7d8bf..c87451d9577 100644 --- a/tests/unit/llms/bedrock/chat/test_native_output_fallback.py +++ b/tests/unit/llms/bedrock/chat/test_native_output_fallback.py @@ -149,6 +149,122 @@ def test_a_rejection_reaches_the_caller_when_the_retry_would_send_native_output_ assert bedrock.sent == (None,) +# Bedrock's 403 for a role without the streaming action, seen live on us.anthropic.claude-sonnet-5 on 2026-11-06 +_STREAM_PERMISSION_DENIAL: Final = BedrockError( + status_code=403, + message='{"Message":"User: arn:aws:sts::123456789012:assumed-role/invoke-only/qa is not authorized to perform: ' + "bedrock:InvokeModelWithResponseStream on resource: arn:aws:bedrock:us-west-2:123456789012:inference-profile/" + "us.anthropic.claude-sonnet-5 because no identity-based policy allows the " + 'bedrock:InvokeModelWithResponseStream action"}', +) +_INVOKE_PERMISSION_DENIAL: Final = BedrockError( + status_code=403, + message='{"Message":"User: arn:aws:sts::123456789012:assumed-role/no-bedrock/qa is not authorized to perform: ' + 'bedrock:InvokeModel because no identity-based policy allows the bedrock:InvokeModel action"}', +) +_BUFFERED_NATIVE_STREAM: Final = NativeOutputFallback( + optional_params={"outputConfig": "native", "fake_stream": True}, litellm_params={} +) + + +def _buffered(fallback: NativeOutputFallback) -> NativeOutputFallback: + return NativeOutputFallback( + optional_params={**fallback.optional_params, "fake_stream": True}, litellm_params=fallback.litellm_params + ) + + +def _call_with_buffered_stream_fallback( + bedrock: ScriptedBedrock, fallback: NativeOutputFallback | None, buffered_stream_fallback: NativeOutputFallback +) -> object: + return _resolve( + call_with_native_output_fallback( + bedrock, + None if fallback is None else lambda: fallback, + _MODEL, + buffered_stream_fallback=lambda: buffered_stream_fallback, + ) + ) + + +@_MODES +def test_a_denied_native_stream_is_sent_again_once_as_a_buffered_native_request( + asynchronous: bool, caplog: pytest.LogCaptureFixture +) -> None: + bedrock: Final = ScriptedBedrock(errors=(_STREAM_PERMISSION_DENIAL, None), asynchronous=asynchronous) + + with caplog.at_level("WARNING", logger="LiteLLM"): + response: Final = _call_with_buffered_stream_fallback(bedrock, _FALLBACK, _BUFFERED_NATIVE_STREAM) + + assert isinstance(response, ModelResponse) and response.id == "answer-2" + assert bedrock.sent == (None, _BUFFERED_NATIVE_STREAM) + assert "bedrock:InvokeModelWithResponseStream" in caplog.text and _MODEL in caplog.text + + +@_MODES +def test_a_buffered_native_request_bedrock_rejects_is_sent_again_buffered_with_the_fallback_params( + asynchronous: bool, +) -> None: + bedrock: Final = ScriptedBedrock( + errors=(_STREAM_PERMISSION_DENIAL, _GRAMMAR_REJECTION, None), asynchronous=asynchronous + ) + + response: Final = _call_with_buffered_stream_fallback(bedrock, _FALLBACK, _BUFFERED_NATIVE_STREAM) + + assert isinstance(response, ModelResponse) and response.id == "answer-3" + assert bedrock.sent == (None, _BUFFERED_NATIVE_STREAM, _buffered(_FALLBACK)) + + +@_MODES +def test_strict_tools_bedrock_rejects_after_a_buffered_native_retry_are_dropped_on_the_buffered_json_tool_retry( + asynchronous: bool, +) -> None: + bedrock: Final = ScriptedBedrock( + errors=(_STREAM_PERMISSION_DENIAL, _GRAMMAR_REJECTION, _GRAMMAR_REJECTION, None), asynchronous=asynchronous + ) + + response: Final = _resolve( + call_with_native_output_fallback( + bedrock, + lambda: _NATIVE_RETRY_WITH_STRICT_TOOLS, + "us.anthropic.claude-sonnet-4-6", + strict_tools_fallback=lambda: _STRICT_FALLBACK, + buffered_stream_fallback=lambda: _BUFFERED_NATIVE_STREAM, + ) + ) + + assert isinstance(response, ModelResponse) and response.id == "answer-4" + assert bedrock.sent == ( + None, + _BUFFERED_NATIVE_STREAM, + _buffered(_NATIVE_RETRY_WITH_STRICT_TOOLS), + _buffered(_COMPOSED_FALLBACK), + ) + + +@_MODES +@pytest.mark.parametrize( + ("error", "buffered_stream_fallback"), + [ + ( + _STREAM_PERMISSION_DENIAL, + NativeOutputFallback(optional_params={"outputConfig": "native"}, litellm_params={}), + ), + (_INVOKE_PERMISSION_DENIAL, _BUFFERED_NATIVE_STREAM), + ], + ids=["retry_would_stream_again", "invoke_denied"], +) +def test_a_permission_denial_reaches_the_caller_unless_the_retry_buffers_the_stream( + asynchronous: bool, error: BedrockError, buffered_stream_fallback: NativeOutputFallback +) -> None: + bedrock: Final = ScriptedBedrock(errors=(error,), asynchronous=asynchronous) + + with pytest.raises(BedrockError) as raised: + _call_with_buffered_stream_fallback(bedrock, _FALLBACK, buffered_stream_fallback) + + assert raised.value is error + assert bedrock.sent == (None,) + + @pytest.mark.parametrize( ("error", "advises_opting_out"), [(_REJECTION, True), (_SCHEMA_REFUSAL, False)], diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index 93bda5b444d..fdd6d5e61bc 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -21,7 +21,7 @@ from botocore.credentials import Credentials import litellm from litellm.constants import DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET -from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.bedrock.common_utils import BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES, BedrockError from litellm.llms.bedrock.files.transformation import BedrockFilesConfig, BedrockJsonlFilesTransformation from tests.unit.proxy.conftest import httpx_transport @@ -2069,6 +2069,76 @@ class TestBedrockBatchAnthropicRowParams: assert ("format" in model_input.get("output_config", {})) is native_output assert (tool_names == {"json_tool_call"}) is not native_output + @pytest.mark.parametrize( + ("model", "native_over_the_limit"), + [(MODEL, True), ("bedrock/us.anthropic.claude-opus-5-5", False)], + ids=["already_native_model", "native_by_default_model"], + ) + @pytest.mark.parametrize("over_the_limit", [False, True], ids=["at_the_limit", "over_the_limit"]) + def test_only_native_by_default_models_send_json_schemas_wider_than_the_limit_through_the_json_tool( + self, local_model_cost_map, model, native_over_the_limit, over_the_limit + ): + native_output: Final = native_over_the_limit or not over_the_limit + field_names: Final = tuple( + f"field_{index:03d}" for index in range(BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES + over_the_limit) + ) + wide_schema: Final = { + "type": "object", + "properties": {name: {"type": "string"} for name in field_names}, + "required": list(field_names), + "additionalProperties": False, + } + model_input: Final = self._transform( + "/v1/chat/completions", + { + "messages": [{"role": "user", "content": "Describe a ceramic coffee mug"}], + "response_format": {"type": "json_schema", "json_schema": {"name": "wide", "schema": wide_schema}}, + }, + model=model, + ) + tool_names: Final = frozenset(tool["name"] for tool in model_input.get("tools", [])) + + assert ("format" in model_input.get("output_config", {})) is native_output + assert (tool_names == {"json_tool_call"}) is not native_output + + def test_a_deployment_opt_out_sends_uploaded_json_schema_rows_through_the_json_tool(self, local_model_cost_map): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + record: Final = { + "custom_id": "row-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": self.MODEL, + "messages": [{"role": "user", "content": "Weather in Paris?"}], + "response_format": { + "type": "json_schema", + "json_schema": {"name": "weather", "schema": self.PARAMETERS}, + }, + }, + } + request: Final = BedrockFilesConfig().transform_create_file_request( + model=self.MODEL, + create_file_data={ + "file": ("batch.jsonl", json.dumps(record).encode(), "application/jsonl"), + "purpose": "batch", + }, + optional_params={ + "aws_access_key_id": "test-key-id", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + }, + litellm_params={ + "s3_bucket_name": "litellm-batch-bucket", + "litellm_metadata": {"model_info": {"supports_bedrock_native_structured_output": False}}, + }, + ) + assert isinstance(request, dict) + model_input: Final = json.loads(request["data"].splitlines()[0])["modelInput"] + + assert "format" not in model_input.get("output_config", {}) + assert [tool["name"] for tool in model_input["tools"]] == ["json_tool_call"] + def test_provider_native_params_still_pass_through(self): model_input = self._transform( "/v1/chat/completions", diff --git a/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py b/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py index 6ca70324162..e3aa01a609d 100644 --- a/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py +++ b/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py @@ -1,4 +1,17 @@ +import re +from datetime import datetime +from typing import Final + +import httpx +import pytest + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.image_edit.stability_transformation import BedrockStabilityImageEditConfig +from litellm.types.utils import ImageResponse +from tests.unit.conftest import in_operator_error + +INPAINT_MODEL: Final = "stability.stable-image-inpaint-v1:0" def test_transform_image_edit_request_returns_the_json_body_and_no_files(): @@ -12,3 +25,56 @@ def test_transform_image_edit_request_returns_the_json_body_and_no_files(): ) assert request == ({"output_format": "png", "prompt": "add a red hat", "image": "iVBORy1ieXRlcw=="}, {}) + + +def _edit_response(body: object) -> ImageResponse: + return BedrockStabilityImageEditConfig().transform_image_edit_response( + model=INPAINT_MODEL, + raw_response=httpx.Response(200, json=body), + logging_obj=Logging( + model=INPAINT_MODEL, + messages=[], + stream=False, + call_type="image_edit", + start_time=datetime(2026, 1, 1), + litellm_call_id="stability-edit-call", + function_id="stability-edit-function", + ), + ) + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_transform_image_edit_response_returns_every_non_empty_image_in_order(): + response = _edit_response({"images": ["aW1hZ2Ux", "", "aW1hZ2Uy"], "finish_reasons": [None], "seeds": [7]}) + + assert [(image.b64_json, image.url) for image in response.data] == [("aW1hZ2Ux", None), ("aW1hZ2Uy", None)] + + +@pytest.mark.parametrize( + ("body", "status_code", "message"), + [ + ({"errors": ["prompt was flagged"]}, 200, "Bedrock Stability error: ['prompt was flagged']"), + ( + {"images": [], "finish_reasons": ["Filter reason: prompt"]}, + 400, + "Bedrock Stability error: Filter reason: prompt", + ), + ], +) +def test_transform_image_edit_response_raises_the_error_bedrock_reports(body: object, status_code: int, message: str): + with pytest.raises(BedrockError) as error: + _edit_response(body) + + assert (error.value.status_code, error.value.message) == (status_code, message) + + +@pytest.mark.parametrize( + ("body", "message"), + [ + (5, in_operator_error(5)), + ({"finish_reasons": 5}, "'int' object is not subscriptable"), + ], +) +def test_transform_image_edit_response_rejects_a_scalar_where_bedrock_sends_a_container(body: object, message: str): + with pytest.raises(TypeError, match=re.escape(message)): + _edit_response(body) diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index cf2f25faaa9..1221667cb51 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -10,23 +10,38 @@ from typing import Final, Literal from unittest.mock import MagicMock import pytest +from fastapi import HTTPException from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict, Unpack import litellm from litellm.constants import ( BEDROCK_REALTIME_SDK_SUPPORTED_RANGE, + REALTIME_SESSION_ENDED_ABNORMALLY_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, WEBSOCKET_CLOSE_REASON_MAX_BYTES, ) +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.realtime_errors import websocket_close_reason +from litellm.litellm_core_utils.realtime_guardrails import BLOCKED_TOOL_OUTPUT from litellm.llms.bedrock.common_utils import BedrockError -from litellm.llms.bedrock.realtime.handler import BedrockRealtime, SessionEndSignals -from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.llms.bedrock.realtime.handler import ( + GUARDRAIL_UNAVAILABLE_MESSAGE, + BedrockRealtime, + BedrockRealtimeGuardrails, + SessionEndSignals, +) from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.realtime_api import main as realtime_main -from litellm.integrations.custom_logger import CustomLogger -from litellm.types.utils import StandardLoggingPayload, StandardLoggingPayloadErrorInformation +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import ( + GenericGuardrailAPIInputs, + StandardLoggingPayload, + StandardLoggingPayloadErrorInformation, +) from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault @@ -223,10 +238,19 @@ class DisconnectingClientWS: return self._messages.pop(0) raise RuntimeError("client disconnected") - async def send_text(self, message): + async def send_text(self, message: str) -> None: self.sent_to_client.append(message) +class GuardrailErrorUndeliveredClientWS(DisconnectingClientWS): + """Client that disconnects just as the proxy tries to tell it about a guardrail violation""" + + async def send_text(self, message: str) -> None: + if "guardrail_violation" in message: + raise RuntimeError("client disconnected") + await super().send_text(message) + + class ClosableClientWS: def __init__(self): self.closed = False @@ -1052,7 +1076,16 @@ class TestBedrockRealtimeProviderFailurePropagation: assert other_provider.calls == 0 @pytest.mark.asyncio - async def test_readiness_failure_on_an_input_first_start_replays_the_input(self, stub_aws_sdk_client): + @pytest.mark.parametrize( + "guardrail_awaits_its_backend", [False, True], ids=["no guardrail", "a guardrail that awaits its backend"] + ) + async def test_readiness_failure_on_an_input_first_start_replays_the_input( + self, stub_aws_sdk_client, monkeypatch, guardrail_awaits_its_backend: bool + ): + guardrail: Final = AwaitingGuardrail( + "hunter2", guardrail_name="remote", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail] if guardrail_awaits_its_backend else []) handler = BedrockRealtime() websocket = RealtimeClientWS([self.USER_TEXT]) healthy_stream = SessionBoundBedrockStream() @@ -1687,5 +1720,614 @@ class TestBedrockRealtimeSdkImportErrors: assert BEDROCK_REALTIME_SDK_SUPPORTED_RANGE in close_reason +class SessionEndOptions(TypedDict, total=False): + on_violation: ReadOnly[str] + end_session_after_n_fails: ReadOnly[int] + + +class BlockedWordGuardrailOptions(SessionEndOptions, total=False): + guardrail_name: ReadOnly[str] + event_hook: ReadOnly[GuardrailEventHooks] + default_on: ReadOnly[bool] + realtime_violation_message: ReadOnly[str] + + +class BlockedWordGuardrail(CustomGuardrail): + def __init__(self, blocked_word: str, **kwargs: Unpack[BlockedWordGuardrailOptions]) -> None: + super().__init__(**kwargs) + self.blocked_word = blocked_word + self.checked: list[tuple[tuple[str, ...], object]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + self.checked.append((tuple(inputs["texts"]), request_data.get("user_api_key_dict"))) + if any(self.blocked_word in text for text in inputs["texts"]): + raise HTTPException(status_code=400, detail={"error": f"{self.blocked_word} is not allowed"}) + return inputs + + +class MaskingGuardrail(BlockedWordGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "texts": [text.replace(self.blocked_word, "[MASKED]") for text in inputs["texts"]]} + + +class AwaitingGuardrail(BlockedWordGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + await asyncio.sleep(0.05) + return await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + + +class SystemSkippingGuardrail(BlockedWordGuardrail): + skip_system_message_in_guardrail: Final = True + + +class CrashingGuardrail(BlockedWordGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + raise RuntimeError("guardrail backend down") + + +class TestBedrockRealtimeGuardrails: + """Realtime guardrails must screen Nova Sonic sessions the way they screen OpenAI realtime sessions""" + + MODEL: Final = "amazon.nova-sonic-v1:0" + VIOLATION: Final = "I can't help with that." + BLOCKED_TEXT: Final = "my secret is hunter2" + CLEAN_TEXT: Final = "what is the weather like" + RESPONSE_CREATE: Final = json.dumps({"type": "response.create"}) + + @staticmethod + def _guardrail( + event_hook: GuardrailEventHooks = GuardrailEventHooks.pre_call, + default_on: bool = True, + **kwargs: Unpack[SessionEndOptions], + ) -> BlockedWordGuardrail: + return BlockedWordGuardrail( + "hunter2", + guardrail_name="no_secrets", + event_hook=event_hook, + default_on=default_on, + realtime_violation_message=TestBedrockRealtimeGuardrails.VIOLATION, + **kwargs, + ) + + @staticmethod + def _user_text(text: str, role: str = "user") -> str: + return json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": role, "content": [{"type": "input_text", "text": text}]}, + } + ) + + @staticmethod + def _bedrock_events(stream: FakeBedrockStream) -> list[dict]: + return [json.loads(chunk.value.bytes_.decode("utf-8"))["event"] for chunk in stream.input_stream.sent] + + @classmethod + def _text_inputs(cls, stream: FakeBedrockStream) -> list[str]: + events: Final = cls._bedrock_events(stream) + user_contents: Final = { + event["contentStart"]["contentName"] + for event in events + if event.get("contentStart", {}).get("role") == "USER" + } + return [ + event["textInput"]["content"] + for event in events + if event.get("textInput", {}).get("contentName") in user_contents + ] + + @staticmethod + def _client_events(client_ws: DisconnectingClientWS | RealtimeClientWS) -> list[dict]: + return [json.loads(message) for message in client_ws.sent_to_client] + + @classmethod + def _guardrail_errors(cls, client_ws: DisconnectingClientWS | RealtimeClientWS) -> list[dict]: + return [ + event["error"] + for event in cls._client_events(client_ws) + if event["type"] == "error" and event["error"]["type"] == "guardrail_violation" + ] + + async def _forward_client( + self, + messages: list[str], + guardrail: BlockedWordGuardrail, + client_ws_type: type[DisconnectingClientWS] | type[ConnectedClientWS] = DisconnectingClientWS, + ) -> tuple[FakeBedrockStream, DisconnectingClientWS | ConnectedClientWS]: + stream: Final = FakeBedrockStream() + client_ws: Final = client_ws_type(messages) + await asyncio.wait_for( + BedrockRealtime()._forward_client_to_bedrock( + client_ws, + stream, + BedrockRealtimeConfig(), + self.MODEL, + {}, + signals=SessionEndSignals(), + guardrails=BedrockRealtimeGuardrails(callbacks=lambda: (guardrail,)), + ), + timeout=1, + ) + return stream, client_ws + + @staticmethod + def _user_transcript(text: str) -> tuple[str, ...]: + return ( + json.dumps( + { + "event": { + "contentStart": { + "role": "USER", + "type": "TEXT", + "additionalModelFields": json.dumps({"generationStage": "FINAL"}), + } + } + } + ), + json.dumps({"event": {"textOutput": {"content": text}}}), + json.dumps({"event": {"contentEnd": {"stopReason": "PARTIAL_TURN"}}}), + ) + + @staticmethod + def _assistant_text(text: str) -> tuple[str, ...]: + return ( + json.dumps({"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}), + json.dumps({"event": {"textOutput": {"content": text}}}), + json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN"}}}), + ) + + async def _forward_bedrock( + self, + payloads: tuple[str, ...], + guardrail: BlockedWordGuardrail, + client_audio_streamed: bool = True, + client_ws_type: type[ConnectedClientWS] | type[GuardrailErrorUndeliveredClientWS] = ConnectedClientWS, + ) -> tuple[ConnectedClientWS | GuardrailErrorUndeliveredClientWS, list[object]]: + config: Final = BedrockRealtimeConfig() + config.client_audio_streamed = client_audio_streamed + client_ws: Final = client_ws_type([]) + logged: Final = [ + event + async for event in BedrockRealtime()._forward_bedrock_to_client( + ScriptedBedrockStream(payloads), + client_ws, + config, + self.MODEL, + FakeLogging(), + {}, + signals=_closing_input(), + guardrails=BedrockRealtimeGuardrails(callbacks=lambda: (guardrail,)), + ) + ] + return client_ws, logged + + @pytest.mark.asyncio + async def test_blocked_typed_message_never_reaches_bedrock_and_the_violation_is_voiced( + self, stub_aws_models: None + ) -> None: + guardrail: Final = self._guardrail() + + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT), self.RESPONSE_CREATE], guardrail + ) + clean_stream, _ = await self._forward_client( + [self._user_text(self.CLEAN_TEXT), self.RESPONSE_CREATE], guardrail + ) + + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert [f'"{self.VIOLATION}"' in text for text in self._text_inputs(stream)] == [True] + assert self._guardrail_errors(client_ws) == [ + {"type": "guardrail_violation", "message": self.VIOLATION, "code": "content_policy_violation"} + ] + assert guardrail.checked == [((self.BLOCKED_TEXT,), None), ((self.CLEAN_TEXT,), None)] + audio_inputs = [sum("audioInput" in event for event in self._bedrock_events(s)) for s in (stream, clean_stream)] + assert audio_inputs[0] == audio_inputs[1] > 0, "the client's own response.create must not start a 2nd turn" + + @pytest.mark.asyncio + @pytest.mark.parametrize("role", ["assistant", "system"]) + async def test_blocked_typed_message_is_screened_whatever_its_role(self, stub_aws_models: None, role: str) -> None: + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT, role=role)], self._guardrail() + ) + + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + + @pytest.mark.asyncio + async def test_a_system_message_nova_receives_as_a_user_turn_is_screened_by_a_system_skipping_guardrail( + self, stub_aws_models: None + ) -> None: + guardrail: Final = SystemSkippingGuardrail( + "hunter2", + guardrail_name="no_secrets", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + realtime_violation_message=self.VIOLATION, + ) + + stream, client_ws = await self._forward_client([self._user_text(self.BLOCKED_TEXT, role="system")], guardrail) + + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + + @pytest.mark.asyncio + async def test_the_next_clean_turn_is_answered_after_a_blocked_item_with_no_response_create( + self, stub_aws_models: None + ) -> None: + stream, _ = await self._forward_client( + [ + self._user_text(self.BLOCKED_TEXT, role="system"), + self._user_text(self.CLEAN_TEXT), + self.RESPONSE_CREATE, + ], + self._guardrail(), + ) + + events: Final = self._bedrock_events(stream) + clean_text_index: Final = next( + i for i, event in enumerate(events) if event.get("textInput", {}).get("content") == self.CLEAN_TEXT + ) + assert any("audioInput" in event for event in events[clean_text_index:]), "the client's response.create runs" + + @pytest.mark.asyncio + async def test_client_message_the_guardrail_cannot_parse_never_reaches_bedrock( + self, stub_aws_models: None + ) -> None: + lone_surrogate_text: Final = self._user_text(f"{self.BLOCKED_TEXT} \ud800") + + stream, _ = await self._forward_client( + [lone_surrogate_text, self._user_text(self.CLEAN_TEXT)], self._guardrail() + ) + + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert self._text_inputs(stream) == [self.CLEAN_TEXT] + + @pytest.mark.asyncio + async def test_clean_typed_message_reaches_bedrock_unchanged(self, stub_aws_models: None) -> None: + guardrail: Final = self._guardrail() + + stream, client_ws = await self._forward_client([self._user_text(self.CLEAN_TEXT)], guardrail) + + assert self._text_inputs(stream) == [self.CLEAN_TEXT] + assert client_ws.sent_to_client == [] + assert guardrail.checked == [((self.CLEAN_TEXT,), None)] + + @pytest.mark.asyncio + async def test_masked_typed_message_reaches_bedrock_masked_and_is_answered(self, stub_aws_models: None) -> None: + guardrail: Final = MaskingGuardrail( + "hunter2", guardrail_name="mask_secrets", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT), self.RESPONSE_CREATE], guardrail + ) + + events: Final = self._bedrock_events(stream) + assert "hunter2" not in json.dumps(events) + assert self._text_inputs(stream) == ["my secret is [MASKED]"] + assert client_ws.sent_to_client == [] + assert any("audioInput" in event for event in events), "the client's response.create runs" + + @pytest.mark.asyncio + async def test_masked_tool_output_reaches_bedrock_masked(self, stub_aws_models: None) -> None: + guardrail: Final = MaskingGuardrail( + "hunter2", guardrail_name="mask_secrets", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + tool_output: Final = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "function_call_output", "call_id": "tool_1", "output": self.BLOCKED_TEXT}, + } + ) + + stream, client_ws = await self._forward_client([tool_output], guardrail) + + events: Final = self._bedrock_events(stream) + assert [event["toolResult"]["content"] for event in events if "toolResult" in event] == [ + "my secret is [MASKED]" + ] + assert client_ws.sent_to_client == [] + + @pytest.mark.asyncio + async def test_the_block_reason_is_spoken_when_no_message_is_configured(self, stub_aws_models: None) -> None: + guardrail: Final = BlockedWordGuardrail( + "hunter2", guardrail_name="no_secrets", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT), self.RESPONSE_CREATE], guardrail + ) + + assert self.BLOCKED_TEXT not in json.dumps(self._bedrock_events(stream)) + assert ['"hunter2 is not allowed"' in text for text in self._text_inputs(stream)] == [True] + assert [error["message"] for error in self._guardrail_errors(client_ws)] == ["hunter2 is not allowed"] + + @pytest.mark.asyncio + async def test_guardrail_that_crashes_ends_the_session_as_an_outage_without_forwarding_the_text( + self, stub_aws_models: None + ) -> None: + guardrail: Final = CrashingGuardrail( + "hunter2", guardrail_name="flaky", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + stream: Final = SessionBoundBedrockStream() + client_ws: Final = ConnectedClientWS([self._user_text(self.CLEAN_TEXT), self.RESPONSE_CREATE]) + logging_obj: Final = FakeLogging() + + async def open_stream() -> SessionBoundBedrockStream: + return stream + + await asyncio.wait_for( + BedrockRealtime()._run_session( + client_ws, + open_stream, + self.MODEL, + logging_obj, + None, + guardrails=BedrockRealtimeGuardrails(callbacks=lambda: (guardrail,)), + ), + timeout=2, + ) + + assert self.CLEAN_TEXT not in json.dumps(self._bedrock_events(stream)) + assert [event["error"] for event in self._client_events(client_ws) if event["type"] == "error"] == [ + {"type": "guardrail_error", "message": GUARDRAIL_UNAVAILABLE_MESSAGE} + ] + assert client_ws.closed + assert logging_obj.model_call_details.get(REALTIME_SESSION_ENDED_ABNORMALLY_KEY) is True + + @pytest.mark.asyncio + async def test_blocked_tool_output_is_replaced_and_answered_by_the_model_itself( + self, stub_aws_models: None + ) -> None: + tool_output: Final = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "function_call_output", "call_id": "tool_1", "output": self.BLOCKED_TEXT}, + } + ) + + stream, client_ws = await self._forward_client([tool_output, self.RESPONSE_CREATE], self._guardrail()) + + events: Final = self._bedrock_events(stream) + assert "hunter2" not in json.dumps(events) + tool_result_index: Final = next(i for i, event in enumerate(events) if "toolResult" in event) + assert events[tool_result_index]["toolResult"]["content"] == BLOCKED_TOOL_OUTPUT + tool_starts: Final = [ + event["contentStart"] for event in events if event.get("contentStart", {}).get("type") == "TOOL" + ] + assert [start["toolResultInputConfiguration"]["toolUseId"] for start in tool_starts] == ["tool_1"] + assert self._text_inputs(stream) == [] + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + assert any("audioInput" in event for event in events[tool_result_index:]), "the client's response.create runs" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("event_hook", "default_on"), + [ + (GuardrailEventHooks.pre_call, False), + (GuardrailEventHooks.post_call, True), + (GuardrailEventHooks.realtime_input_transcription, True), + ], + ids=["not-default-on", "post-call-hook", "transcription-hook"], + ) + async def test_guardrail_that_does_not_apply_to_typed_text_changes_nothing( + self, stub_aws_models: None, event_hook: GuardrailEventHooks, default_on: bool + ) -> None: + guardrail: Final = self._guardrail(event_hook=event_hook, default_on=default_on) + + stream, client_ws = await self._forward_client([self._user_text(self.BLOCKED_TEXT)], guardrail) + + assert self._text_inputs(stream) == [self.BLOCKED_TEXT] + assert client_ws.sent_to_client == [] + assert guardrail.checked == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("guardrail_kwargs", "blocked_messages"), + [({"on_violation": "end_session"}, 1), ({"end_session_after_n_fails": 2}, 2)], + ids=["on-violation-end-session", "after-two-fails"], + ) + async def test_end_session_guardrail_closes_the_session( + self, stub_aws_models: None, guardrail_kwargs: SessionEndOptions, blocked_messages: int + ) -> None: + messages: Final = [self._user_text(self.BLOCKED_TEXT)] * blocked_messages + [self._user_text(self.CLEAN_TEXT)] + + stream, client_ws = await self._forward_client( + messages, self._guardrail(**guardrail_kwargs), client_ws_type=ConnectedClientWS + ) + + assert client_ws.closed + assert len(self._guardrail_errors(client_ws)) == blocked_messages + assert self.CLEAN_TEXT not in self._text_inputs(stream) + assert sum(self.VIOLATION in text for text in self._text_inputs(stream)) == blocked_messages - 1 + assert stream.input_stream.closed + + @pytest.mark.asyncio + @pytest.mark.parametrize("guardrail_kwargs", [{}, {"on_violation": "end_session"}], ids=["default", "end-session"]) + async def test_blocked_transcript_ends_the_session_before_its_answer_reaches_the_client( + self, guardrail_kwargs: SessionEndOptions + ) -> None: + payloads: Final = ( + *self._user_transcript(self.BLOCKED_TEXT), + *self._assistant_text("Your secret is safe with me"), + *self._user_transcript(self.CLEAN_TEXT), + *self._assistant_text("It is sunny"), + ) + guardrail: Final = self._guardrail( + event_hook=GuardrailEventHooks.realtime_input_transcription, **guardrail_kwargs + ) + + client_ws, _ = await self._forward_bedrock(payloads, guardrail) + + client_events: Final = self._client_events(client_ws) + assert client_ws.closed + assert not any(event["type"].startswith("response.") for event in client_events) + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + assert guardrail.checked == [((self.BLOCKED_TEXT,), None)] + + @pytest.mark.asyncio + async def test_blocked_transcript_answer_never_reaches_the_client_when_the_previous_answer_closes_late( + self, + ) -> None: + earlier_answer: Final = self._assistant_text("Earlier answer") + payloads: Final = ( + *earlier_answer[:2], + *self._user_transcript(self.BLOCKED_TEXT), + earlier_answer[2], + *self._assistant_text("Your secret is safe with me"), + ) + + client_ws, _ = await self._forward_bedrock( + payloads, self._guardrail(event_hook=GuardrailEventHooks.realtime_input_transcription) + ) + + client_events: Final = self._client_events(client_ws) + assert "Your secret is safe with me" not in json.dumps(client_events) + assert [event["delta"] for event in client_events if event["type"] == "response.text.delta"] == [ + "Earlier answer" + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("event_hook", "default_on", "client_audio_streamed"), + [ + (GuardrailEventHooks.realtime_input_transcription, True, False), + (GuardrailEventHooks.pre_call, True, True), + (GuardrailEventHooks.realtime_input_transcription, False, True), + ], + ids=["typed-text-session", "pre-call-hook", "not-default-on"], + ) + async def test_transcript_guardrail_that_does_not_apply_changes_nothing( + self, event_hook: GuardrailEventHooks, default_on: bool, client_audio_streamed: bool + ) -> None: + guardrail: Final = self._guardrail(event_hook=event_hook, default_on=default_on) + + client_ws, _ = await self._forward_bedrock( + (*self._user_transcript(self.BLOCKED_TEXT), *self._assistant_text("Your secret is safe with me")), + guardrail, + client_audio_streamed=client_audio_streamed, + ) + + client_events: Final = self._client_events(client_ws) + assert [event["delta"] for event in client_events if event["type"] == "response.text.delta"] == [ + "Your secret is safe with me" + ] + assert self._guardrail_errors(client_ws) == [] + assert guardrail.checked == [] + + @pytest.mark.asyncio + async def test_session_ends_when_the_client_leaves_before_hearing_about_a_blocked_message( + self, stub_aws_models: None + ) -> None: + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT), self._user_text(self.CLEAN_TEXT)], + self._guardrail(), + client_ws_type=GuardrailErrorUndeliveredClientWS, + ) + + assert self._text_inputs(stream) == [] + assert self._client_events(client_ws) == [] + + @pytest.mark.asyncio + async def test_answers_stop_when_the_client_leaves_before_hearing_about_a_blocked_transcript(self) -> None: + guardrail: Final = self._guardrail(event_hook=GuardrailEventHooks.realtime_input_transcription) + + _, logged = await self._forward_bedrock( + (*self._user_transcript(self.BLOCKED_TEXT), *self._assistant_text("Your secret is safe with me")), + guardrail, + client_ws_type=GuardrailErrorUndeliveredClientWS, + ) + + assert not any(event["type"].startswith("response.") for event in logged) + + @pytest.mark.asyncio + async def test_proxy_request_guardrails_and_key_reach_the_bedrock_session( + self, stub_aws_sdk_client: dict[str, object], monkeypatch: pytest.MonkeyPatch + ) -> None: + class ProxyLogging(FakeLogging): + def update_from_kwargs(self, **kwargs: object) -> None: + pass + + guardrail: Final = self._guardrail(default_on=False) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + stream: Final = ScriptedBedrockStream([], receiver_type=DrainedThenOpenBedrockReceiver) + stub_aws_sdk_client["streams"] = [stream] + user_api_key_dict: Final = object() + client_ws: Final = DisconnectingClientWS([self._user_text(self.BLOCKED_TEXT)]) + + await asyncio.wait_for( + realtime_main._arealtime.__wrapped__( + model=f"bedrock/{self.MODEL}", + websocket=client_ws, + litellm_logging_obj=ProxyLogging(), + aws_region_name="us-east-1", + aws_access_key_id="k", + aws_secret_access_key="s", + user_api_key_dict=user_api_key_dict, + metadata={"guardrails": ["no_secrets"]}, + ), + timeout=1, + ) + + assert guardrail.checked == [((self.BLOCKED_TEXT,), user_api_key_dict)] + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + + @pytest.mark.asyncio + async def test_a_key_that_opted_out_of_global_guardrails_is_not_screened( + self, stub_aws_sdk_client: dict[str, object], monkeypatch: pytest.MonkeyPatch + ) -> None: + class ProxyLogging(FakeLogging): + def update_from_kwargs(self, **kwargs: object) -> None: + pass + + guardrail: Final = self._guardrail() + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + stream: Final = ScriptedBedrockStream([], receiver_type=DrainedThenOpenBedrockReceiver) + stub_aws_sdk_client["streams"] = [stream] + client_ws: Final = DisconnectingClientWS([self._user_text(self.BLOCKED_TEXT)]) + + await asyncio.wait_for( + realtime_main._arealtime.__wrapped__( + model=f"bedrock/{self.MODEL}", + websocket=client_ws, + litellm_logging_obj=ProxyLogging(), + aws_region_name="us-east-1", + aws_access_key_id="k", + aws_secret_access_key="s", + metadata={"user_api_key_metadata": {"disable_global_guardrails": True}}, + ), + timeout=1, + ) + + assert guardrail.checked == [] + assert self._guardrail_errors(client_ws) == [] + assert self.BLOCKED_TEXT in json.dumps(self._bedrock_events(stream)) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/unit/llms/bedrock/responses/test_bedrock_openai_responses.py b/tests/unit/llms/bedrock/responses/test_bedrock_openai_responses.py index a1466cfdae4..8844eae97bb 100644 --- a/tests/unit/llms/bedrock/responses/test_bedrock_openai_responses.py +++ b/tests/unit/llms/bedrock/responses/test_bedrock_openai_responses.py @@ -296,6 +296,28 @@ class TestCodexHistoryNormalization: "user", ] + def test_agent_message_with_encrypted_content_reaches_the_endpoint_intact(self): + """Only the endpoint can decrypt the block for the recipient model; as assistant text it is gibberish.""" + agent_message = { + "type": "agent_message", + "author": "/root/reviewer", + "recipient": "/root", + "content": [ + {"type": "input_text", "text": "Worker message:\n"}, + {"type": "encrypted_content", "encrypted_content": "gAAAAABpZ3ciphertext=="}, + ], + } + cfg = BedrockOpenAIResponsesConfig.for_model(MODEL) + assert cfg is not None + body = cfg.transform_responses_api_request( + model=MODEL, + input=[json.loads(json.dumps(agent_message)), {"role": "user", "content": "carry on"}], + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert body["input"] == [agent_message, {"role": "user", "content": "carry on"}] + def test_a_first_turn_request_is_untouched(self): """The rejected types are history items, so turn one exercises none of this.""" original = [{"role": "user", "content": "first turn"}] diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 0c1c4f5475a..7cc4629e497 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -1,11 +1,10 @@ - import asyncio, importlib, litellm, pytest import respx from httpx import Response from typing import Final -from litellm.llms.bedrock.common_utils import( +from litellm.llms.bedrock.common_utils import ( AmazonBedrockGlobalConfig, BedrockModelInfo, NativeOutputFallback, @@ -54,9 +53,7 @@ def test_bedrock_response_stream_shape_lazy_loads_once(): import litellm.llms.bedrock.common_utils as mod sentinel = MagicMock() - with patch.object( - mod, "_load_bedrock_response_stream_shape", return_value=sentinel - ) as mock_load: + with patch.object(mod, "_load_bedrock_response_stream_shape", return_value=sentinel) as mock_load: assert mod.get_bedrock_response_stream_shape() is sentinel assert mod.get_bedrock_response_stream_shape() is sentinel mock_load.assert_called_once() @@ -103,9 +100,7 @@ def test_bedrock_response_stream_shape_is_structure_shape(): from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape loaded_shape = get_bedrock_response_stream_shape() - assert ( - loaded_shape is not None - ), "get_bedrock_response_stream_shape() is None — botocore may not be installed" + assert loaded_shape is not None, "get_bedrock_response_stream_shape() is None — botocore may not be installed" shape: StructureShape = loaded_shape assert isinstance(shape, StructureShape) assert shape.name == "ResponseStream" @@ -170,9 +165,7 @@ def test_deepseek_cris(): Test that DeepSeek models with cross-region inference prefix use converse route """ bedrock_model_info = BedrockModelInfo - bedrock_route = bedrock_model_info.get_bedrock_route( - model="bedrock/us.deepseek.r1-v1:0" - ) + bedrock_route = bedrock_model_info.get_bedrock_route(model="bedrock/us.deepseek.r1-v1:0") assert bedrock_route == "converse" @@ -228,7 +221,10 @@ def test_bare_opaque_arn_with_no_invoke_provider_routes_to_converse(arn_segment: model: Final = f"bedrock/arn:aws:bedrock:us-east-1:123412341234:{arn_segment}" assert BaseAWSLLM.get_bedrock_invoke_provider(model) is None assert BedrockModelInfo.get_bedrock_route(model) == "converse" - assert BedrockModelInfo.get_bedrock_route(f"bedrock/llama/arn:aws:bedrock:us-east-1:123412341234:{arn_segment}") == "invoke" + assert ( + BedrockModelInfo.get_bedrock_route(f"bedrock/llama/arn:aws:bedrock:us-east-1:123412341234:{arn_segment}") + == "invoke" + ) def test_govcloud_cross_region_inference_prefix(): @@ -238,27 +234,19 @@ def test_govcloud_cross_region_inference_prefix(): bedrock_model_info = BedrockModelInfo # Test us-gov prefix is stripped correctly for Claude models - base_model = bedrock_model_info.get_base_model( - model="bedrock/us-gov.anthropic.claude-haiku-4-5-20251001-v1:0" - ) + base_model = bedrock_model_info.get_base_model(model="bedrock/us-gov.anthropic.claude-haiku-4-5-20251001-v1:0") assert base_model == "anthropic.claude-haiku-4-5-20251001-v1:0" # Test us-gov prefix is stripped correctly for different Claude versions - base_model = bedrock_model_info.get_base_model( - model="bedrock/us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0" - ) + base_model = bedrock_model_info.get_base_model(model="bedrock/us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0") assert base_model == "anthropic.claude-sonnet-4-5-20250929-v1:0" # Test us-gov prefix is stripped correctly for Haiku models - base_model = bedrock_model_info.get_base_model( - model="bedrock/us-gov.anthropic.claude-3-haiku-20240307-v1:0" - ) + base_model = bedrock_model_info.get_base_model(model="bedrock/us-gov.anthropic.claude-3-haiku-20240307-v1:0") assert base_model == "anthropic.claude-3-haiku-20240307-v1:0" # Test us-gov prefix is stripped correctly for Meta models - base_model = bedrock_model_info.get_base_model( - model="bedrock/us-gov.meta.llama3-8b-instruct-v1:0" - ) + base_model = bedrock_model_info.get_base_model(model="bedrock/us-gov.meta.llama3-8b-instruct-v1:0") assert base_model == "meta.llama3-8b-instruct-v1:0" @@ -272,23 +260,14 @@ def test_context_window_suffix_stripped_for_cost_lookup(): """ from litellm.llms.bedrock.common_utils import get_bedrock_base_model - assert ( - get_bedrock_base_model("us.anthropic.claude-opus-4-6-v1[1m]") - == "anthropic.claude-opus-4-6-v1" - ) - assert ( - get_bedrock_base_model("us.anthropic.claude-sonnet-4-6[1m]") - == "anthropic.claude-sonnet-4-6" - ) + assert get_bedrock_base_model("us.anthropic.claude-opus-4-6-v1[1m]") == "anthropic.claude-opus-4-6-v1" + assert get_bedrock_base_model("us.anthropic.claude-sonnet-4-6[1m]") == "anthropic.claude-sonnet-4-6" assert ( get_bedrock_base_model("global.anthropic.claude-opus-4-5-20251101-v1:0[1m]") == "anthropic.claude-opus-4-5-20251101-v1:0" ) # Ensure models without suffix are unaffected - assert ( - get_bedrock_base_model("us.anthropic.claude-opus-4-6-v1") - == "anthropic.claude-opus-4-6-v1" - ) + assert get_bedrock_base_model("us.anthropic.claude-opus-4-6-v1") == "anthropic.claude-opus-4-6-v1" # Ensure :51k throughput suffix still works assert ( get_bedrock_base_model("anthropic.claude-3-5-sonnet-20241022-v2:0:51k") @@ -358,9 +337,7 @@ def test_output_config_effort_normalization_uses_model_info_ceiling(monkeypatch) ("us.anthropic.claude-opus-4-7", "xhigh"), ], ) -def test_bundled_bedrock_opus_model_info_declares_output_config_effort_ceiling( - model, expected_ceiling -): +def test_bundled_bedrock_opus_model_info_declares_output_config_effort_ceiling(model, expected_ceiling): from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap model_info = GetModelCostMap.load_local_model_cost_map()[model] @@ -379,54 +356,24 @@ def test_route_prefix_matched_as_path_segment_not_substring(): or a ``/`` boundary. """ # The bedrock_mantle/ provider prefix must NOT be read as the mantle/ route. - assert ( - BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.5") != "mantle" - ) - assert ( - BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.4") == "converse" - ) - assert ( - BedrockModelInfo._explicit_mantle_route("bedrock_mantle/openai.gpt-5.5") - is False - ) + assert BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.5") != "mantle" + assert BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.4") == "converse" + assert BedrockModelInfo._explicit_mantle_route("bedrock_mantle/openai.gpt-5.5") is False # A genuine mantle route still resolves, via the startswith branch... - assert ( - BedrockModelInfo.get_bedrock_route("mantle/anthropic.claude-mythos-preview") - == "mantle" - ) + assert BedrockModelInfo.get_bedrock_route("mantle/anthropic.claude-mythos-preview") == "mantle" # ...and via the mid-path "/mantle/" branch (after the bedrock/ provider prefix). - assert ( - BedrockModelInfo.get_bedrock_route( - "bedrock/mantle/anthropic.claude-mythos-preview" - ) - == "mantle" - ) + assert BedrockModelInfo.get_bedrock_route("bedrock/mantle/anthropic.claude-mythos-preview") == "mantle" def test_model_has_route_prefix_exercises_both_branches(): """``_model_has_route_prefix`` matches on ``startswith`` or a ``/`` boundary only.""" # startswith branch - assert ( - BedrockModelInfo._model_has_route_prefix( - "mantle/anthropic.claude-mythos-preview", "mantle/" - ) - is True - ) + assert BedrockModelInfo._model_has_route_prefix("mantle/anthropic.claude-mythos-preview", "mantle/") is True # f"/{prefix}" boundary branch - assert ( - BedrockModelInfo._model_has_route_prefix( - "bedrock/mantle/anthropic.claude-mythos-preview", "mantle/" - ) - is True - ) + assert BedrockModelInfo._model_has_route_prefix("bedrock/mantle/anthropic.claude-mythos-preview", "mantle/") is True # neither branch: the token only appears glued to another segment - assert ( - BedrockModelInfo._model_has_route_prefix( - "bedrock_mantle/openai.gpt-5.5", "mantle/" - ) - is False - ) + assert BedrockModelInfo._model_has_route_prefix("bedrock_mantle/openai.gpt-5.5", "mantle/") is False @pytest.mark.parametrize( @@ -476,16 +423,10 @@ def test_explicit_invoke_route_does_not_match_async_invoke(): """ async_invoke_model = "async_invoke/twelvelabs.marengo-embed-2-7-v1:0" assert BedrockModelInfo._explicit_invoke_route(async_invoke_model) is False - assert ( - BedrockModelInfo._explicit_invoke_route(f"bedrock/{async_invoke_model}") - is False - ) + assert BedrockModelInfo._explicit_invoke_route(f"bedrock/{async_invoke_model}") is False # ...while async_invoke/ is still detected as its own route. assert BedrockModelInfo._explicit_async_invoke_route(async_invoke_model) is True - assert ( - BedrockModelInfo._explicit_async_invoke_route(f"bedrock/{async_invoke_model}") - is True - ) + assert BedrockModelInfo._explicit_async_invoke_route(f"bedrock/{async_invoke_model}") is True def test_capability_lookups_fall_back_to_base_model_when_regional_entry_lacks_field(monkeypatch): @@ -605,7 +546,6 @@ def test_deployment_model_info_breakpoint_flag_covers_an_unmapped_arn(local_mode assert bedrock_model_accepts_cache_points(unflagged_arn) is True - _CONVERSE_REPLY: Final = { "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, "stopReason": "end_turn", @@ -638,7 +578,11 @@ _MARKED_TOOL_CONVERSATION: Final = ( } ], }, - {"role": "tool", "tool_call_id": "call_1", "content": [{"type": "text", "text": "Sunny", "cache_control": _CACHE_MARK}]}, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [{"type": "text", "text": "Sunny", "cache_control": _CACHE_MARK}], + }, {"role": "tool", "tool_call_id": "call_1_1", "content": "Cloudy"}, {"role": "assistant", "content": "Paris is sunny, Rome is cloudy.", "cache_control": _CACHE_MARK}, {"role": "user", "content": "Thanks"}, @@ -711,6 +655,7 @@ async def test_a_bedrock_deployment_sends_cache_points_by_its_own_flag_not_a_sib "flagless": not declared, } + async def test_a_routed_bedrock_deployment_sends_cache_points_when_register_model_opted_its_model_in( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None ) -> None: @@ -738,7 +683,9 @@ async def test_a_routed_bedrock_deployment_sends_cache_points_when_register_mode num_retries=0, ) - await router.acompletion(model="flagless", messages=[{"role": "system", "content": "stable"}, {"role": "user", "content": "hi"}]) + await router.acompletion( + model="flagless", messages=[{"role": "system", "content": "stable"}, {"role": "user", "content": "hi"}] + ) assert b"cachePoint" in converse.calls.last.request.content @@ -783,10 +730,13 @@ def test_a_deployments_prompt_caching_flag_written_as_a_string_or_number_decides ) assert ( - bedrock_model_accepts_cache_points(model, model_info={"id": "qa-deployment", "supports_prompt_caching": declared}) + bedrock_model_accepts_cache_points( + model, model_info={"id": "qa-deployment", "supports_prompt_caching": declared} + ) is expected ) + def test_merge_bedrock_aws_request_params_strips_caller_identity_when_deployment_has_static_credentials(): from litellm.llms.bedrock.common_utils import merge_bedrock_aws_request_params @@ -1491,6 +1441,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="session") def event_loop(): try: @@ -1500,6 +1451,7 @@ def event_loop(): yield loop loop.close() + @pytest.fixture(scope="function") def setup_and_teardown(event_loop): import litellm @@ -1533,6 +1485,7 @@ def setup_and_teardown(event_loop): if pending: event_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "set_verbose": getattr(litellm, "set_verbose", False), @@ -1547,6 +1500,7 @@ _SCALAR_DEFAULTS = { "cohere_key": getattr(litellm, "cohere_key", None), } + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestStripBedrockRoutingPrefix: """Tests for strip_bedrock_routing_prefix function.""" @@ -1577,6 +1531,7 @@ class TestStripBedrockRoutingPrefix: == "anthropic.claude-3-sonnet-20240229-v1:0" ) + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestStripBedrockThroughputSuffix: """Tests for strip_bedrock_throughput_suffix function.""" @@ -1604,6 +1559,7 @@ class TestStripBedrockThroughputSuffix: def test_strip_throughput_suffix(self, input_model, expected): assert strip_bedrock_throughput_suffix(input_model) == expected + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestExtractModelNameFromBedrockArn: """Tests for extract_model_name_from_bedrock_arn function.""" @@ -1624,6 +1580,7 @@ class TestExtractModelNameFromBedrockArn: arn = "ARN:aws:bedrock:us-east-1:123456789012:model/my-model" assert extract_model_name_from_bedrock_arn(arn) == "my-model" + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestGetBedrockCrossRegionInferenceRegions: """Tests for get_bedrock_cross_region_inference_regions function.""" @@ -1639,6 +1596,7 @@ class TestGetBedrockCrossRegionInferenceRegions: regions = get_bedrock_cross_region_inference_regions() assert isinstance(regions, list) + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestGetBedrockBaseModel: """Tests for get_bedrock_base_model function.""" @@ -1702,6 +1660,7 @@ class TestGetBedrockBaseModel: """Test that throughput tier suffixes like :51k are stripped. Issue #19113.""" assert get_bedrock_base_model(input_model) == expected + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBedrockModelInfoWrappers: """Tests that BedrockModelInfo methods correctly wrap standalone functions.""" @@ -1723,6 +1682,7 @@ class TestBedrockModelInfoWrappers: model = "bedrock/converse/claude-3" assert BedrockModelInfo.get_non_litellm_routing_model_name(model) == strip_bedrock_routing_prefix(model) + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBedrockTokenCounter: """Tests for BedrockTokenCounter class.""" @@ -1759,12 +1719,14 @@ class TestBedrockTokenCounter: ) assert result is None + @pytest.fixture def _pr4_bedrock_env(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("AWS_ACCESS_KEY_ID", "pr4-test-aws-access-key") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "pr4-test-aws-secret-key") monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + @pytest.fixture(scope="module") def _use_local_model_cost_map() -> Iterator[None]: with pytest.MonkeyPatch.context() as monkeypatch: @@ -1772,6 +1734,7 @@ def _use_local_model_cost_map() -> Iterator[None]: importlib.reload(litellm) yield + @pytest.mark.usefixtures( "_pr4_bedrock_env", "_use_local_model_cost_map", @@ -2348,8 +2311,12 @@ def test_bedrock_capability_model_reads_base_model_only_for_application_profile_ assert bedrock_capability_model(model, base_model) == expected -_FORMAT_REJECTION: Final = "The model returned the following errors: output_config.format: Extra inputs are not permitted" -_COMPILE_REJECTION: Final = "The model returned the following errors: compilation exceeded limits: Schema is too complex." +_FORMAT_REJECTION: Final = ( + "The model returned the following errors: output_config.format: Extra inputs are not permitted" +) +_COMPILE_REJECTION: Final = ( + "The model returned the following errors: compilation exceeded limits: Schema is too complex." +) _GRAMMAR_REJECTION: Final = ( "The model returned the following errors: The compiled grammar is too large, which would cause performance issues. " "Simplify your tool schemas or reduce the number of strict tools." @@ -2394,10 +2361,10 @@ def test_is_bedrock_native_output_rejection_matches_only_the_400s_that_reject_na ("model", "model_info", "expected"), [ ("us.anthropic.claude-opus-4-8", None, False), - ("global.anthropic.claude-sonnet-5", None, True), + ("global.anthropic.claude-sonnet-5", None, False), ("us-gov.anthropic.claude-opus-4-8", None, False), ("us.anthropic.claude-opus-4-8", {"supports_bedrock_native_structured_output": False}, False), - ("us.anthropic.claude-opus-5-5", {"supports_bedrock_native_structured_output": True}, True), + ("us.anthropic.claude-opus-5-5", {"supports_bedrock_native_structured_output": True}, False), ("us.anthropic.claude-sonnet-4-6", {"supports_bedrock_native_structured_output": True}, False), ], ids=["opus_4_8", "sonnet_5", "govcloud_opus_4_8", "deployment_opt_out", "deployment_opt_in", "native_model"], @@ -2490,7 +2457,10 @@ def test_bedrock_strict_tools_fallback_drops_only_strict_and_keeps_everything_el assert fallback == NativeOutputFallback( optional_params={ "tools": [ - {**strict_tool, "function": {"name": "lookup", "description": "Look up", "parameters": {"type": "object"}}}, + { + **strict_tool, + "function": {"name": "lookup", "description": "Look up", "parameters": {"type": "object"}}, + }, plain_tool, anthropic_tool, {"toolSpec": {"name": "spec", "inputSchema": {"json": {"type": "object"}}}}, diff --git a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index f8f18852543..9e9e46af26c 100644 --- a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -24,6 +24,7 @@ from litellm.llms.bedrock_mantle.responses.transformation import ( ) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager LOOKALIKE_MANTLE_HOSTS = ( "https://bedrock-mantle.us-east-1.api.aws.internal.example.com", @@ -902,11 +903,10 @@ class TestBedrockMantleCodexInputItemNormalization: on the Codex history item types agent_message, context_compaction, and local_shell_call (verified against bedrock-mantle.us-east-1.api.aws with openai.gpt-5.6-sol), so the config must rewrite them into supported - equivalents. agent_message is what every Codex multi-agent v2 session sends, - and its encrypted_content slot carries the verbatim plaintext payload when - the upstream model never issued encrypted args, so that slot must be - preserved, not dropped. Mantle also rejects assistant messages with - input_text content, so the rewrite must use output_text.""" + equivalents. agent_message is what every Codex multi-agent v2 session sends. + One carrying an encrypted_content block is ciphertext only Mantle can decrypt, + so it is forwarded intact. Mantle also rejects assistant messages with + input_text content, so the plaintext rewrite must use output_text.""" _USER_MESSAGE = { "type": "message", @@ -914,10 +914,11 @@ class TestBedrockMantleCodexInputItemNormalization: "content": [{"type": "input_text", "text": "Continue."}], } - def _transform(self, input): - cfg = BedrockMantleResponsesAPIConfig() + def _transform(self, input, model="openai.gpt-5.6-sol"): + cfg = ProviderConfigManager.get_provider_responses_api_config(LlmProviders.BEDROCK_MANTLE, model) + assert isinstance(cfg, BedrockMantleResponsesAPIConfig) return cfg.transform_responses_api_request( - model="openai.gpt-5.6-sol", + model=model, input=input, response_api_optional_request_params={}, litellm_params=GenericLiteLLMParams(), @@ -946,31 +947,48 @@ class TestBedrockMantleCodexInputItemNormalization: }, ] - def test_agent_message_encrypted_content_payload_is_preserved(self): + def test_agent_message_with_encrypted_content_reaches_mantle_intact_and_unlogged(self, caplog): + """Only Mantle can decrypt the block for the recipient model; as assistant text it is gibberish.""" + agent_message = { + "type": "agent_message", + "author": "/root/reviewer", + "recipient": "/root", + "content": [ + {"type": "input_text", "text": "Worker message:\n"}, + {"type": "encrypted_content", "encrypted_content": "gAAAAABpZ3ciphertext=="}, + ], + } + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + body = self._transform(input=[copy.deepcopy(agent_message), self._USER_MESSAGE]) + assert body["input"] == [agent_message, self._USER_MESSAGE] + assert not [record for record in caplog.records if "rewrote Codex input item" in record.getMessage()] + + @pytest.mark.parametrize("model", ["openai.gpt-oss-120b", "google.gemma-4-31b", "xai.grok-4.3"]) + def test_models_that_never_lock_worker_messages_keep_the_text_rewrite(self, model): + """Codex puts plain text in the locked slot for these models, which Mantle rejects unless rewritten.""" body = self._transform( input=[ { "type": "agent_message", "author": "/root", - "recipient": "/root/arithmetic", + "recipient": "/root/codeword", "content": [ {"type": "input_text", "text": "Message Type: NEW_TASK\nPayload:\n"}, - {"type": "encrypted_content", "encrypted_content": "Answer the question 'what is 2+2'."}, + {"type": "encrypted_content", "encrypted_content": "Reply with ZEPHYR-1."}, ], }, self._USER_MESSAGE, - ] + ], + model=model, ) - assert body["input"][0] == { - "type": "message", - "role": "assistant", - "content": ( - { - "type": "output_text", - "text": "Message Type: NEW_TASK\nPayload:\nAnswer the question 'what is 2+2'.", - }, - ), - } + assert body["input"] == [ + { + "type": "message", + "role": "assistant", + "content": ({"type": "output_text", "text": "Message Type: NEW_TASK\nPayload:\nReply with ZEPHYR-1."},), + }, + self._USER_MESSAGE, + ] def test_agent_message_without_any_text_is_dropped(self): body = self._transform( diff --git a/tests/unit/llms/cohere/test_cohere_common_utils.py b/tests/unit/llms/cohere/test_cohere_common_utils.py new file mode 100644 index 00000000000..1d17824a181 --- /dev/null +++ b/tests/unit/llms/cohere/test_cohere_common_utils.py @@ -0,0 +1,18 @@ +from typing import Final + +import pytest + +from litellm.llms.cohere.common_utils import ModelResponseIterator + + +def test_stream_parses_object_chunks_and_rejects_a_chunk_that_is_not_an_object() -> None: + stream: Final = ModelResponseIterator( + streaming_response=iter(['{"text": "Hi", "index": 1}', "[1, 2]"]), + sync_stream=True, + ) + + first: Final = next(stream) + + assert (first["text"], first["index"], first["is_finished"]) == ("Hi", 1, False) + with pytest.raises(AttributeError, match=r"^'list' object has no attribute 'get'$"): + next(stream) diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index dd0e5762845..ba8fb3bf107 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -4253,6 +4253,7 @@ def test_video_generation_with_input_reference_keeps_file_multipart(): AZURE_AI_BASE = "https://myfoundry.services.ai.azure.com" AZURE_AI_CHAT_COMPLETIONS_URL = f"{AZURE_AI_BASE}/models/chat/completions" + def _a_tool_with_an_unsupported_field() -> dict: return { "type": "function", @@ -4260,6 +4261,7 @@ def _a_tool_with_an_unsupported_field() -> dict: "strict": True, } + A_COMPLETION = { "id": "chatcmpl-1", "object": "chat.completion", @@ -4459,6 +4461,45 @@ async def test_create_file_async_awaits_the_provider_credential_hook_instead_of_ assert client.post.call_args.kwargs["url"] == "https://files.example/v1/files" +def test_create_file_sync_sends_a_presigned_upload_with_its_own_method_url_headers_and_body(): + provider_config = Mock(spec=BaseFilesConfig) + provider_config.validate_environment.return_value = {"x-api-key": "unused-for-presigned"} + provider_config.get_complete_file_url.return_value = "https://files.example/v1/files" + provider_config.transform_create_file_request.return_value = { + "method": "PUT", + "url": "https://bucket.s3.example/batch.jsonl", + "headers": {"Authorization": "AWS4-HMAC-SHA256 signed"}, + "data": b'{"recordId": "req-1"}', + } + file_object = object() + provider_config.transform_create_file_response.return_value = file_object + upload_response = Mock(spec=httpx.Response) + client = Mock(spec=HTTPHandler) + client.put.return_value = upload_response + + result = BaseLLMHTTPHandler().create_file( + create_file_data={"file": b"{}", "purpose": "batch"}, + litellm_params={}, + provider_config=provider_config, + headers={}, + api_base=None, + api_key=None, + logging_obj=Mock(), + client=client, + timeout=30, + ) + + assert result is file_object + client.post.assert_not_called() + client.put.assert_called_once_with( + url="https://bucket.s3.example/batch.jsonl", + headers={"Authorization": "AWS4-HMAC-SHA256 signed"}, + data=b'{"recordId": "req-1"}', + timeout=30, + ) + assert provider_config.transform_create_file_response.call_args.kwargs["raw_response"] is upload_response + + @pytest.mark.asyncio async def test_create_batch_async_validates_credentials_off_the_event_loop(): provider_config = Mock(spec=BaseBatchesConfig) @@ -4890,6 +4931,92 @@ async def test_json_mode_stream_reaches_the_iterator_of_a_config_without_a_reque ] +class _PlainStreamRouteConfig(_TransformRecordingConfig): + def __init__(self): + super().__init__(transform_async=False) + + def get_complete_url(self, api_base, api_key, model, optional_params, litellm_params, stream=None): + return f"{api_base}/{'stream' if stream else 'answer'}" + + +class _StreamRouteConfig(_PlainStreamRouteConfig): + @property + def has_custom_stream_wrapper(self) -> bool: + return True + + +def _answers_only_the_non_streaming_route(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/answer"): + return httpx.Response(200, json=json.loads(request.content)) + return httpx.Response(403, json={"message": "not authorized to stream"}) + + +def _fake_streamed_completion( + acompletion: bool, + client: HTTPHandler | AsyncHTTPHandler, + fake_stream: object = True, + provider_config: BaseConfig | None = None, +): + return BaseLLMHTTPHandler().completion( + model="stub-model", + messages=[{"role": "user", "content": "hi"}], + api_base="https://llm.example/v1/chat", + custom_llm_provider="openai", + model_response=ModelResponse(), + encoding=None, + logging_obj=_streaming_logging_obj(), + optional_params={"fake_stream": fake_stream}, + timeout=10.0, + litellm_params={}, + acompletion=acompletion, + stream=True, + client=client, + provider_config=provider_config or _StreamRouteConfig(), + ) + + +async def test_a_fake_streamed_completion_is_sent_to_the_non_streaming_route(): + sync_client: Final = HTTPHandler( + client=httpx.Client(transport=httpx.MockTransport(_answers_only_the_non_streaming_route)) + ) + async_client: Final = AsyncHTTPHandler() + async_client.client = httpx.AsyncClient(transport=httpx.MockTransport(_answers_only_the_non_streaming_route)) + + sync_chunks: Final = list(_fake_streamed_completion(acompletion=False, client=sync_client)) + async_chunks: Final = [ + chunk async for chunk in await _fake_streamed_completion(acompletion=True, client=async_client) + ] + + assert (_streamed_text(sync_chunks), _streamed_text(async_chunks)) == ("sync", "sync") + + +def _answers_only_the_streaming_route(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/stream"): + return _sse_echoing_transformed_by(request) + return httpx.Response(403, json={"message": "only streaming is allowed"}) + + +@pytest.mark.parametrize("fake_stream", ["true", 1], ids=["string", "int"]) +async def test_a_fake_stream_value_other_than_true_still_streams_from_the_streaming_route(fake_stream: object): + sync_client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_answers_only_the_streaming_route))) + async_client: Final = AsyncHTTPHandler() + async_client.client = httpx.AsyncClient(transport=httpx.MockTransport(_answers_only_the_streaming_route)) + + sync_chunks: Final = list( + _fake_streamed_completion( + acompletion=False, client=sync_client, fake_stream=fake_stream, provider_config=_PlainStreamRouteConfig() + ) + ) + async_chunks: Final = [ + chunk + async for chunk in await _fake_streamed_completion( + acompletion=True, client=async_client, fake_stream=fake_stream, provider_config=_PlainStreamRouteConfig() + ) + ] + + assert (_streamed_text(sync_chunks), _streamed_text(async_chunks)) == ("sync", "sync") + + class _ImageEditRecordingConfig(BaseImageEditConfig): def __init__(self): self.transform_calls = [] diff --git a/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py b/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py index 27afbee3173..5dd8b165a0f 100644 --- a/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py +++ b/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py @@ -3,19 +3,20 @@ Unit tests for DashScope embedding transformation. """ import json +import re from typing import Final from unittest.mock import MagicMock import httpx import pytest - from litellm.llms.dashscope.common_utils import DashScopeError from litellm.llms.dashscope.embed.transformation import ( DEFAULT_API_BASE, DashScopeEmbeddingConfig, ) from litellm.types.utils import EmbeddingResponse +from tests.unit.conftest import in_operator_error WORKSPACE: Final = "https://llm-abc123.cn-beijing.maas.aliyuncs.com" @@ -171,3 +172,22 @@ def test_transform_embedding_response_error(): ) assert exc.value.status_code == 401 assert "Incorrect API key" in exc.value.message + + +def test_transform_embedding_response_rejects_a_scalar_body(): + raw = httpx.Response( + status_code=200, + content=b"5", + request=httpx.Request("POST", "https://example.com"), + ) + with pytest.raises(TypeError, match=f"^{re.escape(in_operator_error(5))}$"): + DashScopeEmbeddingConfig().transform_embedding_response( + model="text-embedding-v4", + raw_response=raw, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + api_key="sk-test", + request_data={"input": ["a"]}, + optional_params={}, + litellm_params={}, + ) diff --git a/tests/unit/llms/dashscope/test_dashscope_video_transformation.py b/tests/unit/llms/dashscope/test_dashscope_video_transformation.py index c1b4157c8ab..654722ad3e5 100644 --- a/tests/unit/llms/dashscope/test_dashscope_video_transformation.py +++ b/tests/unit/llms/dashscope/test_dashscope_video_transformation.py @@ -1434,3 +1434,17 @@ class TestDashScopeVideoEndToEndRequest: async def test_a_size_in_a_tier_the_model_lacks_never_reaches_dashscope(self): with pytest.raises(UnsupportedParamsError, match="480P"): await self._wire_request(model="dashscope/wan2.7-t2v", prompt="p", size="854x480") + + +@pytest.mark.parametrize("drop_params", [False, True]) +def test_existing_scalar_parameters_remain_accepted(drop_params: bool) -> None: + params: Final = {"seed": 7, "duration": 6, "aspect_ratio": "9:16", "generate_audio": False} + assert DashScopeVideoConfig().map_openai_params(params, "wan3.0-video", drop_params) == params + + +@pytest.mark.parametrize("field", ["frame_images", "input_references"]) +def test_first_class_video_fields_are_rejected_or_dropped(field: str) -> None: + config: Final = DashScopeVideoConfig() + with pytest.raises(UnsupportedParamsError, match=field): + config.map_openai_params({field: 1}, "video-model", False) + assert config.map_openai_params({field: 1}, "video-model", True) == config.map_openai_params({}, "video-model", True) diff --git a/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py b/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py index d94f6a53c89..c8071687b4f 100644 --- a/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py +++ b/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py @@ -4,6 +4,7 @@ Unit tests for DataForSEO Search functionality. import datetime import os +import re from unittest.mock import AsyncMock, patch import httpx @@ -15,6 +16,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome +from tests.unit.conftest import in_operator_error @pytest.mark.asyncio @@ -172,6 +174,19 @@ def test_dataforseo_malformed_item_is_rejected_without_naming_its_position(): assert "429" not in str(exc_info.value) +@pytest.mark.parametrize( + ("payload", "message"), + [ + (5, in_operator_error(5)), + ({"tasks": 5}, "object of type 'int' has no len()"), + ({"tasks": [{"status_code": 20000, "result": 5}]}, "object of type 'int' has no len()"), + ], +) +def test_dataforseo_scalar_where_the_response_needs_a_container_raises_type_error(payload: object, message: str): + with pytest.raises(TypeError, match=re.escape(message)): + _transform(payload) + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/llms/edenai/videos/test_edenai_video_transformation.py b/tests/unit/llms/edenai/videos/test_edenai_video_transformation.py index 360e4d07f24..0dd8245ee92 100644 --- a/tests/unit/llms/edenai/videos/test_edenai_video_transformation.py +++ b/tests/unit/llms/edenai/videos/test_edenai_video_transformation.py @@ -117,6 +117,9 @@ class TestCreate: assert b'name="model"\r\n\r\n' + SELLER_MODEL.encode() in request.content assert b'name="seconds"\r\n\r\n4' in request.content + def test_seed_zero_survives_drop_params(self): + assert EdenAIVideoConfig().map_openai_params({"seed": 0}, SELLER_MODEL, True) == {"seed": 0} + def test_a_reference_image_url_stays_in_the_json_body(self, eden_key, respx_mock): respx_mock.post(EDEN_VIDEOS_URL).mock(return_value=httpx.Response(200, json=_eden_video())) diff --git a/tests/unit/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/unit/llms/fal_ai/videos/test_fal_ai_video_transformation.py index d4ad22ca5a0..2505d635375 100644 --- a/tests/unit/llms/fal_ai/videos/test_fal_ai_video_transformation.py +++ b/tests/unit/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -4,6 +4,7 @@ from unittest.mock import AsyncMock, Mock import httpx import pytest +from litellm.exceptions import UnsupportedParamsError import litellm import litellm.llms.fal_ai.videos.transformation as fal_video_module @@ -24,6 +25,14 @@ H3_TEXT_MODEL = "minimax/h3/text-to-video" H3_REFERENCE_MODEL = "minimax/h3/reference-to-video" +@pytest.mark.parametrize("drop_params", [False, True]) +def test_native_video_parameters_keep_precedence(drop_params: bool) -> None: + params: Final = {"aspect_ratio": "1:1", "seed": 0, "duration": 6} + assert FalAIVideoConfig().map_openai_params({"seconds": "5", "size": "1280x720", **params}, MODEL, drop_params) == { + "resolution": "720p", **params, + } + + class TestFalAIVideoTransformation: def setup_method(self): self.config = FalAIVideoConfig() @@ -695,3 +704,11 @@ class TestFalAIVideoTransformation: ) == 5 * row["output_cost_per_second_768p"] ) + + +@pytest.mark.parametrize("field", ["frame_images", "input_references"]) +def test_first_class_video_fields_are_rejected_or_dropped(field: str) -> None: + config: Final = FalAIVideoConfig() + with pytest.raises(UnsupportedParamsError, match=field): + config.map_openai_params({field: 1}, "video-model", False) + assert config.map_openai_params({field: 1}, "video-model", True) == config.map_openai_params({}, "video-model", True) diff --git a/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py index 6538ccf1cac..fe33a26dfd7 100644 --- a/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py +++ b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py @@ -477,6 +477,84 @@ def test_transform_request_forwards_non_string_instructions_and_input_untouched( ) +CODEX_CLOCK_NAMESPACE_TOOL: Final = { + "type": "namespace", + "name": "clock", + "description": "Clock tools.", + "tools": [ + { + "type": "function", + "name": "sleep", + "description": "Sleep for some seconds.", + "parameters": {"type": "object", "properties": {"seconds": {"type": "number"}}, "required": ["seconds"]}, + } + ], +} +CODEX_ADDITIONAL_TOOLS_ITEM: Final = { + "type": "additional_tools", + "role": "developer", + "tools": [CODEX_CLOCK_NAMESPACE_TOOL], +} + + +@pytest.mark.parametrize("instructions", [None, "Answer with exactly one word."], ids=["no-instructions", "instructions"]) +def test_transform_request_moves_codex_additional_tools_item_into_tools(instructions: str | None) -> None: + user_item: Final = {"role": "user", "content": "Sleep for two seconds."} + request: Final = FireworksAIResponsesAPIConfig().transform_responses_api_request( + model="accounts/fireworks/models/glm-5p3", + input=cast(ResponseInputParam, [CODEX_ADDITIONAL_TOOLS_ITEM, user_item]), + response_api_optional_request_params={} if instructions is None else {"instructions": instructions}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert tuple(request["tools"]) == (CODEX_CLOCK_NAMESPACE_TOOL,) + assert tuple(request["input"]) == (user_item,) + assert request.get("instructions") == instructions + + +@pytest.mark.parametrize( + ("drop_params", "expected_clock_copies"), + [ + pytest.param(True, 1, id="merged once under drop_params"), + pytest.param(None, 2, id="forwarded as sent without drop_params"), + ], +) +@pytest.mark.parametrize("tools_as_tuple", [False, True], ids=["tools-list", "tools-tuple"]) +def test_transform_request_merges_codex_additional_tools_into_existing_tools_only_under_drop_params( + drop_params: bool | None, expected_clock_copies: int, tools_as_tuple: bool +) -> None: + weather_tool: Final = { + "type": "function", + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } + user_item: Final = {"role": "user", "content": "Sleep for two seconds, then check the weather in Paris."} + request: Final = FireworksAIResponsesAPIConfig().transform_responses_api_request( + model="accounts/fireworks/models/glm-5p3", + input=cast(ResponseInputParam, [CODEX_ADDITIONAL_TOOLS_ITEM, user_item, CODEX_ADDITIONAL_TOOLS_ITEM]), + response_api_optional_request_params={"tools": (weather_tool,) if tools_as_tuple else [weather_tool]}, + litellm_params=GenericLiteLLMParams(drop_params=drop_params), + headers={}, + ) + assert tuple(request["tools"]) == (weather_tool, *(CODEX_CLOCK_NAMESPACE_TOOL,) * expected_clock_copies) + assert tuple(request["input"]) == (user_item,) + + +@pytest.mark.parametrize("drop_params", [True, None], ids=["drop_params", "no-drop_params"]) +def test_transform_request_leaves_codex_additional_tools_in_input_when_tools_is_malformed( + drop_params: bool | None, +) -> None: + user_item: Final = {"role": "user", "content": "Sleep for two seconds."} + request: Final = FireworksAIResponsesAPIConfig().transform_responses_api_request( + model="accounts/fireworks/models/glm-5p3", + input=cast(ResponseInputParam, [CODEX_ADDITIONAL_TOOLS_ITEM, user_item]), + response_api_optional_request_params={"tools": {"type": "function", "name": "get_weather"}}, + litellm_params=GenericLiteLLMParams(drop_params=drop_params), + headers={}, + ) + assert tuple(request["input"]) == (CODEX_ADDITIONAL_TOOLS_ITEM, user_item) + + def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_output_items() -> None: client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b")) pydantic_input: Final = cast( diff --git a/tests/unit/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 5f2af34e513..508bd264f19 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -1,4 +1,5 @@ import contextlib +import re from typing import Final from unittest.mock import AsyncMock, patch @@ -16,6 +17,7 @@ from litellm.llms.gemini.common_utils import ( ) from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter as GeminiCountTokensHandler +from tests.unit.conftest import in_operator_error _JSON_BODY: Final = TypeAdapter(dict[str, JsonValue]) @@ -498,6 +500,18 @@ class TestGoogleAIStudioTokenCounter: # Verify the contents are unchanged assert cleaned_contents == contents_without_function_response + @pytest.mark.parametrize( + ("part", "message"), + [ + (7, in_operator_error(7)), + (frozenset({"functionResponse"}), "'frozenset' object is not subscriptable"), + ({"functionResponse": "done"}, "FunctionResponse() argument after ** must be a mapping, not str"), + ], + ) + def test_clean_contents_for_gemini_api_rejects_a_malformed_part(self, part: object, message: str) -> None: + with pytest.raises(TypeError, match=re.escape(message)): + GeminiCountTokensHandler()._clean_contents_for_gemini_api([{"role": "user", "parts": [part]}]) + def _map_image_config(raw_image_config: str) -> dict[str, object]: return map_openai_image_params_to_gemini( diff --git a/tests/unit/llms/gemini/videos/test_gemini_video_transformation.py b/tests/unit/llms/gemini/videos/test_gemini_video_transformation.py index 1ac451d17db..7177e8e9c68 100644 --- a/tests/unit/llms/gemini/videos/test_gemini_video_transformation.py +++ b/tests/unit/llms/gemini/videos/test_gemini_video_transformation.py @@ -6,9 +6,11 @@ import io import json import os from unittest.mock import MagicMock, Mock, patch +from typing import Final import httpx import pytest +from litellm.exceptions import UnsupportedParamsError from litellm.llms.gemini.videos.transformation import GeminiVideoConfig from litellm.llms.openai.cost_calculation import video_generation_cost @@ -842,3 +844,17 @@ class TestGeminiVideoCostTracking: if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +@pytest.mark.parametrize("drop_params", [False, True]) +def test_scalar_provider_parameters_are_preserved(drop_params: bool) -> None: + params: Final = {"seed": 0, "duration": 6, "aspect_ratio": "1:1", "generate_audio": False} + assert GeminiVideoConfig().map_openai_params(params, "veo-3.1-generate-preview", drop_params) == params + + +@pytest.mark.parametrize("field", ["frame_images", "input_references"]) +def test_first_class_video_fields_are_rejected_or_dropped(field: str) -> None: + config: Final = GeminiVideoConfig() + with pytest.raises(UnsupportedParamsError, match=field): + config.map_openai_params({field: 1}, "video-model", False) + assert config.map_openai_params({field: 1}, "video-model", True) == config.map_openai_params({}, "video-model", True) diff --git a/tests/unit/llms/gigachat/chat/test_gigachat_chat_streaming.py b/tests/unit/llms/gigachat/chat/test_gigachat_chat_streaming.py index 35ca93319f5..913130e5d5c 100644 --- a/tests/unit/llms/gigachat/chat/test_gigachat_chat_streaming.py +++ b/tests/unit/llms/gigachat/chat/test_gigachat_chat_streaming.py @@ -2,6 +2,10 @@ Tests for litellm.llms.gigachat.chat.streaming """ +import json +from collections.abc import AsyncIterator +from typing import Final + from litellm.llms.gigachat.chat.streaming import GigaChatModelResponseIterator @@ -85,3 +89,32 @@ class TestChunkParserUsage: assert parsed["usage"]["prompt_tokens"] == 45 assert parsed["usage"]["total_tokens"] == 52 assert parsed["usage"]["prompt_tokens_details"]["cached_tokens"] == 20 + + +def _content_chunk(text: str) -> dict[str, object]: + return {"choices": [{"delta": {"content": text}, "index": 0, "finish_reason": None}]} + + +STREAM_ITEMS: Final = ( + f"data: {json.dumps(_content_chunk('Hel'))}", + _content_chunk("lo"), + "data: {not json", + "data: [DONE]", + f"data: {json.dumps(_content_chunk('after done'))}", +) + + +class TestStreamIteration: + def test_sync_stream_reads_sse_lines_and_parsed_chunks_until_done(self): + iterator = GigaChatModelResponseIterator(streaming_response=iter(STREAM_ITEMS), sync_stream=True) + + assert [chunk["text"] for chunk in iterator] == ["Hel", "lo", ""] + + async def test_async_stream_reads_sse_lines_and_parsed_chunks_until_done(self): + async def stream() -> AsyncIterator[object]: + for item in STREAM_ITEMS: + yield item + + iterator = GigaChatModelResponseIterator(streaming_response=stream(), sync_stream=False) + + assert [chunk["text"] async for chunk in iterator] == ["Hel", "lo", ""] diff --git a/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py index 7040fbe5fcc..ef5b16606f4 100644 --- a/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py +++ b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py @@ -1,4 +1,5 @@ import json +import re from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -7,6 +8,9 @@ import pytest import respx import litellm +from litellm.llms.huggingface.embedding.handler import HuggingFaceEmbedding +from litellm.types.utils import EmbeddingResponse +from tests.unit.conftest import in_operator_error MOCK_EMBEDDING_RESPONSE = [[0.1, 0.2, 0.3, 0.4, 0.5]] @@ -227,3 +231,46 @@ async def test_huggingface_embedding_sends_the_caller_timeout( assert lookup.called assert route.calls.last.request.extensions["timeout"]["read"] == 45 assert response.data[0]["embedding"] == MOCK_EMBEDDING_RESPONSE[0] + + +async def _embedding(sync_mode: bool, **kwargs: object) -> object: + return litellm.embedding(**kwargs) if sync_mode else await litellm.aembedding(**kwargs) + + +@pytest.mark.respx(assert_all_called=False) +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_huggingface_embedding_raises_the_error_the_endpoint_returns( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.get(url__regex=r"https://huggingface\.co/.*/config\.json").respond(json={"max_position_embeddings": 512}) + respx_mock.get(HF_MODEL_LOOKUP_URL).respond(json={"pipeline_tag": "feature-extraction"}) + route: Final = respx_mock.post(HF_EMBED_URL).respond(json={"error": "Model BAAI/bge-m3 is currently loading"}) + request: Final = {"model": "huggingface/BAAI/bge-m3", "input": ["hello"], "api_key": "hf_test"} + + with pytest.raises( + litellm.APIError, match="HuggingfaceException - Model BAAI/bge-m3 is currently loading" + ) as exc_info: + await _embedding(sync_mode, **request) + + assert route.called + assert exc_info.value.status_code == 500 + + +@pytest.mark.parametrize( + ("embeddings", "message"), + [ + (5, in_operator_error(5)), + ([None], "'NoneType' object is not subscriptable"), + ({"similarities": [1]}, "cannot unpack non-iterable int object"), + ], +) +def test_process_embedding_response_rejects_a_malformed_body(embeddings: object, message: str) -> None: + with pytest.raises(TypeError, match=f"^{re.escape(message)}$"): + HuggingFaceEmbedding()._process_embedding_response( + embeddings=embeddings, + model_response=EmbeddingResponse(), + model="BAAI/bge-m3", + input=["hello"], + encoding=MagicMock(), + ) diff --git a/tests/unit/llms/litellm_proxy/skills/test_code_execution.py b/tests/unit/llms/litellm_proxy/skills/test_code_execution.py index d9c4470759f..a1cd6a4d9b2 100644 --- a/tests/unit/llms/litellm_proxy/skills/test_code_execution.py +++ b/tests/unit/llms/litellm_proxy/skills/test_code_execution.py @@ -1,3 +1,5 @@ +from typing import Final + import pytest from litellm.llms.litellm_proxy.skills.code_execution import ( @@ -104,3 +106,23 @@ class TestHandlerDefaults: handler = CodeExecutionHandler(max_iterations=falsy, sandbox_timeout=falsy) assert handler.max_iterations == DEFAULT_MAX_ITERATIONS assert handler.sandbox_timeout == DEFAULT_SANDBOX_TIMEOUT + + +class TestExecuteWithCodeExecution: + async def test_a_reply_without_tool_calls_ends_the_loop_after_one_call(self): + messages: Final = [{"role": "user", "content": "Say hi"}] + + outcome: Final = await CodeExecutionHandler().execute_with_code_execution( + model="gpt-5.6", + messages=messages, + tools=[LITELLM_CODE_EXECUTION_TOOL], + skill_files={}, + mock_response="Hi there", + ) + + response: Final = outcome["response"] + assert response is not None + assert response.choices[0].message.content == "Hi there" + assert outcome["messages"] == [*messages, {"role": "assistant", "content": "Hi there"}] + assert messages == [{"role": "user", "content": "Say hi"}] + assert (outcome["files"], outcome["execution_results"]) == ([], []) diff --git a/tests/unit/llms/minimax/videos/test_minimax_video_transformation.py b/tests/unit/llms/minimax/videos/test_minimax_video_transformation.py index 1405964134c..b512ef13d57 100644 --- a/tests/unit/llms/minimax/videos/test_minimax_video_transformation.py +++ b/tests/unit/llms/minimax/videos/test_minimax_video_transformation.py @@ -6,7 +6,8 @@ import base64 import io import json import time -from typing import Final +from collections.abc import Mapping +from typing import Final, cast from unittest.mock import Mock import httpx @@ -24,6 +25,7 @@ from litellm.llms.minimax.videos.transformation import ( _video_url_from_task, ) from litellm.types.router import GenericLiteLLMParams +from litellm.types.videos.main import VideoCreateOptionalRequestParams from litellm.types.videos.utils import ( decode_video_id_with_provider, encode_video_id_with_provider, @@ -327,6 +329,7 @@ class TestMinimaxVideoCreateRequest: assert raised.value.status_code == 400 + def test_text_inside_content_is_rejected_so_guardrails_cannot_be_bypassed(self): """Guardrails scan prompt; a text item smuggled into content would reach MiniMax unscanned while a harmless prompt passed the check.""" @@ -1316,3 +1319,106 @@ class TestMinimaxVideoInputBilling: ) assert raised.value.status_code == 400 + + +@pytest.mark.parametrize("role", ["first_frame", "last_frame", "reference_image"]) +def test_openrouter_video_images_become_provider_content(role: str) -> None: + params: Final = cast(VideoCreateOptionalRequestParams, { + "duration": 8, "aspect_ratio": "9:16", "resolution": "768P", + **( + {"input_references": [{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}}]} + if role == "reference_image" + else {"frame_images": [{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}, "frame_type": role}]} + ), + }) + config: Final = MinimaxVideoConfig() + mapped: Final = config.map_openai_params(params, "MiniMax-H3", False) + body, _, _ = config.transform_video_create_request( + model="MiniMax-H3", prompt="a cat", api_base=API_BASE, + video_create_optional_request_params=mapped, litellm_params=GenericLiteLLMParams(), headers={}, + ) + assert body["content"] == [ + {"type": "text", "text": "a cat"}, + {"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}, "role": role}, + ] + assert (body["duration"], body["ratio"]) == (8, "9:16") + assert not ({"frame_images", "input_references", "aspect_ratio"} & body.keys()) + + +@pytest.mark.parametrize("other", [ + {"content": [{"type": "image_url", "image_url": {"url": "https://example.com/raw.png"}, "role": "reference_image"}]}, + {"input_references": [{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}}]}, + {"input_reference": "https://example.com/old.png"}, + {"frame_images": [ + {"type": "image_url", "image_url": {"url": "https://example.com/first.png"}, "frame_type": "first_frame"}, + {"type": "image_url", "image_url": {"url": "https://example.com/duplicate.png"}, "frame_type": "first_frame"}, + ]}, +]) +def test_openrouter_image_conflicts_raise_or_keep_first_frame(other: object) -> None: + params: Final = cast(VideoCreateOptionalRequestParams, { + "frame_images": [{"type": "image_url", "image_url": {"url": "https://example.com/first.png"}, "frame_type": "first_frame"}], + **cast(Mapping[str, object], other), + }) + config: Final = MinimaxVideoConfig() + with pytest.raises(UnsupportedParamsError, match="drop_params"): + config.map_openai_params(params, "MiniMax-H3", False) + assert config.map_openai_params(params, "MiniMax-H3", True) == { + "content": [{"type": "image_url", "image_url": {"url": "https://example.com/first.png"}, "role": "first_frame"}], + } + + +@pytest.mark.parametrize("params", [{"seed": 0}, {"generate_audio": False}]) +def test_openrouter_unsupported_audio_and_seed_raise_or_drop(params: VideoCreateOptionalRequestParams) -> None: + config: Final = MinimaxVideoConfig() + with pytest.raises(UnsupportedParamsError, match="does not support"): + config.map_openai_params(params, "MiniMax-H3", False) + assert config.map_openai_params(params, "MiniMax-H3", True) == {} + + +def test_openrouter_duration_conflict_raises_or_prefers_duration() -> None: + config: Final = MinimaxVideoConfig() + with pytest.raises(UnsupportedParamsError, match="duration and seconds"): + config.map_openai_params({"duration": 8, "seconds": "5"}, "MiniMax-H3", False) + assert config.map_openai_params({"duration": 8, "seconds": "5"}, "MiniMax-H3", True) == {"duration": 8} + assert config.map_openai_params({"duration": 8, "seconds": "8"}, "MiniMax-H3", False) == {"duration": 8} + + +@pytest.mark.parametrize("images", [[], [{"type": "image_url", "image_url": {}}]]) +def test_openrouter_empty_and_invalid_frames(images: object) -> None: + params: Final = cast(VideoCreateOptionalRequestParams, {"frame_images": images}) + config: Final = MinimaxVideoConfig() + if images: + with pytest.raises(UnsupportedParamsError, match="Invalid video image"): + config.map_openai_params(params, "MiniMax-H3", False) + assert config.map_openai_params(params, "MiniMax-H3", True) == {} + else: + assert config.map_openai_params(params, "MiniMax-H3", False) == {"content": []} + +@pytest.mark.parametrize("references", [[], [{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}}]]) +def test_plural_and_singular_reference_conflict_raises_or_prefers_plural(references: object) -> None: + config: Final = MinimaxVideoConfig() + params: Final = cast(VideoCreateOptionalRequestParams, { + "input_references": references, "input_reference": "https://example.com/first.png", + }) + with pytest.raises(UnsupportedParamsError, match="input_references cannot be combined"): + config.map_openai_params(params, "MiniMax-H3", False) + assert config.map_openai_params(params, "MiniMax-H3", True) == { + "content": [ + {"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}, "role": "reference_image"} + ] if references else [], + } + + +def test_dropping_seconds_does_not_parse_the_discarded_value() -> None: + assert MinimaxVideoConfig().map_openai_params({"duration": 8, "seconds": "abc"}, "MiniMax-H3", True) == {"duration": 8} + + +@pytest.mark.parametrize("drop_params", [False, True]) +def test_aspect_ratio_preserves_native_ratio_precedence_over_size(drop_params: bool) -> None: + config: Final = MinimaxVideoConfig() + assert config.map_openai_params( + {"size": "1280x720", "aspect_ratio": "9:16"}, "MiniMax-H3", drop_params, + ) == {"ratio": "9:16"} + assert config.map_openai_params( + cast(VideoCreateOptionalRequestParams, {"size": "1280x720", "ratio": "9:16"}), "MiniMax-H3", drop_params, + ) == {"ratio": "9:16"} diff --git a/tests/unit/llms/mistral/test_mistral_chat_transformation.py b/tests/unit/llms/mistral/test_mistral_chat_transformation.py index 42673e17acf..696c53c0b5a 100644 --- a/tests/unit/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/unit/llms/mistral/test_mistral_chat_transformation.py @@ -740,6 +740,25 @@ class TestMistralEmptyContentHandling: assert MistralConfig._is_empty_assistant_message(message) is False +@pytest.mark.parametrize("body", [[], "ok"]) +def test_transform_response_rejects_a_body_that_is_not_a_json_object(body: object) -> None: + logging_obj: Final = MagicMock() + logging_obj.model_call_details = {} + + with pytest.raises(AttributeError, match=f"^'{type(body).__name__}' object has no attribute 'get'$"): + MistralConfig().transform_response( + model="mistral/mistral-large-latest", + raw_response=httpx.Response(200, json=body), + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={}, + messages=[{"role": "user", "content": "Hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + class TestMistralFileHandling: """Test suite for Mistral file handling functionality.""" diff --git a/tests/unit/llms/ollama/test_ollama_embedding.py b/tests/unit/llms/ollama/test_ollama_embedding.py index 97f6f8468f4..b32bb9c4415 100644 --- a/tests/unit/llms/ollama/test_ollama_embedding.py +++ b/tests/unit/llms/ollama/test_ollama_embedding.py @@ -1,4 +1,5 @@ import json +import logging from collections.abc import Mapping from types import MappingProxyType from typing import Final @@ -9,7 +10,11 @@ import pytest import respx import litellm -from litellm.llms.ollama.completion.handler import ollama_aembeddings, ollama_embeddings +from litellm.llms.ollama.completion.handler import ( + _process_ollama_embedding_response, + ollama_aembeddings, + ollama_embeddings, +) from litellm.types.utils import EmbeddingResponse @@ -121,6 +126,39 @@ def test_prompt_eval_fallback_when_missing(mock_embedding_response, mock_encodin assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] +@pytest.mark.parametrize( + ("use_encoding", "prompt_tokens", "level", "message"), + [ + (True, 5, logging.DEBUG, "Ollama response missing prompt_eval_count; estimated with encoding."), + (False, 0, logging.WARNING, "Missing prompt_eval_count and no encoding provided; defaulted to 0."), + ], +) +def test_missing_prompt_eval_count_is_reported_on_the_given_logger( + mock_encoding: MagicMock, + caplog: pytest.LogCaptureFixture, + use_encoding: bool, + prompt_tokens: int, + level: int, + message: str, +) -> None: + logger: Final = logging.getLogger("ollama-embedding-prompt-eval-test") + + with caplog.at_level(logging.DEBUG, logger=logger.name): + response: Final = _process_ollama_embedding_response( + response_json={"embeddings": [[0.1, 0.2]]}, + prompts=["hello"], + model="nomic-embed-text", + model_response=EmbeddingResponse(), + logging_obj=logger, + encoding=mock_encoding if use_encoding else None, + ) + + assert response.usage.prompt_tokens == prompt_tokens + assert [(record.levelno, record.getMessage()) for record in caplog.records if record.name == logger.name] == [ + (level, message) + ] + + @pytest.mark.parametrize( "api_base", [ diff --git a/tests/unit/llms/openai/videos/test_openai_video_transformation.py b/tests/unit/llms/openai/videos/test_openai_video_transformation.py index 89f52261cba..ec4e2ac8471 100644 --- a/tests/unit/llms/openai/videos/test_openai_video_transformation.py +++ b/tests/unit/llms/openai/videos/test_openai_video_transformation.py @@ -1,4 +1,9 @@ import io +from typing import Final + +import pytest + +from litellm.exceptions import UnsupportedParamsError from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.router import GenericLiteLLMParams @@ -115,3 +120,17 @@ def test_video_edit_request_without_file_sends_json_id_reference(): assert url == "https://api.openai.com/v1/videos/edits" assert data == {"prompt": "brighter", "video": {"id": "video_abc123"}} assert files is None + + +@pytest.mark.parametrize("drop_params", [False, True]) +def test_scalar_provider_parameters_are_preserved(drop_params: bool) -> None: + params: Final = {"seed": 0, "duration": 6, "aspect_ratio": "1:1", "generate_audio": False} + assert OpenAIVideoConfig().map_openai_params(params, "sora-2", drop_params) == params + + +@pytest.mark.parametrize("field", ["frame_images", "input_references"]) +def test_first_class_video_fields_are_rejected_or_dropped(field: str) -> None: + config: Final = OpenAIVideoConfig() + with pytest.raises(UnsupportedParamsError, match=field): + config.map_openai_params({field: 1}, "video-model", False) + assert config.map_openai_params({field: 1}, "video-model", True) == config.map_openai_params({}, "video-model", True) diff --git a/tests/unit/llms/recraft/image_edit/test_recraft_image_edit_transformation.py b/tests/unit/llms/recraft/image_edit/test_recraft_image_edit_transformation.py index 97d65935a1b..bf119b01573 100644 --- a/tests/unit/llms/recraft/image_edit/test_recraft_image_edit_transformation.py +++ b/tests/unit/llms/recraft/image_edit/test_recraft_image_edit_transformation.py @@ -170,3 +170,11 @@ class TestRecraftImageEditTransformation: ) assert "Error transforming image edit response" in str(exc_info.value) + + def test_transform_image_edit_response_rejects_a_body_that_is_not_a_json_object(self): + with pytest.raises(TypeError, match="'int' object is not subscriptable"): + self.config.transform_image_edit_response( + model=self.model, + raw_response=httpx.Response(200, json=5), + logging_obj=self.logging_obj, + ) diff --git a/tests/unit/llms/runwayml/videos/test_runway_video_transformation.py b/tests/unit/llms/runwayml/videos/test_runway_video_transformation.py index afc8e7ec4a2..6b18ef5b4b1 100644 --- a/tests/unit/llms/runwayml/videos/test_runway_video_transformation.py +++ b/tests/unit/llms/runwayml/videos/test_runway_video_transformation.py @@ -3,9 +3,11 @@ Tests for RunwayML video generation transformation. """ from unittest.mock import Mock +from typing import Final import httpx import pytest +from litellm.exceptions import UnsupportedParamsError from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.runwayml.videos.transformation import ( @@ -17,6 +19,12 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.videos.main import VideoObject +@pytest.mark.parametrize("drop_params", [False, True]) +def test_native_video_parameters_keep_precedence(drop_params: bool) -> None: + params: Final = {"seed": 0, "duration": 6, "aspect_ratio": "1:1", "generate_audio": False} + assert RunwayMLVideoConfig().map_openai_params({"seconds": "5", **params}, "gen4_turbo", drop_params) == params + + class TestRunwayMLVideoTransformation: """Test RunwayMLVideoConfig transformation class.""" @@ -372,3 +380,11 @@ class TestRunwayMLVideoTransformation: if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +@pytest.mark.parametrize("field", ["frame_images", "input_references"]) +def test_first_class_video_fields_are_rejected_or_dropped(field: str) -> None: + config: Final = RunwayMLVideoConfig() + with pytest.raises(UnsupportedParamsError, match=field): + config.map_openai_params({field: 1}, "video-model", False) + assert config.map_openai_params({field: 1}, "video-model", True) == config.map_openai_params({}, "video-model", True) diff --git a/tests/unit/llms/sap/test_sap_fetch_creds.py b/tests/unit/llms/sap/test_sap_fetch_creds.py index cc52c8991f3..34c10be026b 100644 --- a/tests/unit/llms/sap/test_sap_fetch_creds.py +++ b/tests/unit/llms/sap/test_sap_fetch_creds.py @@ -160,3 +160,15 @@ def test_credentials_without_authentication_mode(monkeypatch): # validate_credentials should raise because no authentication mode is provided with pytest.raises(ValueError, match="SAP AI Core credentials are incomplete"): sap_credentials.validate_credentials(**creds) + + +@pytest.mark.parametrize("vcap_services", [["aicore"], "aicore"]) +def test_vcap_services_that_is_not_an_object_is_rejected(monkeypatch, vcap_services): + _prep_env(monkeypatch) + monkeypatch.setenv("VCAP_SERVICES", json.dumps(vcap_services)) + + with pytest.raises( + AttributeError, + match=f"^'{type(vcap_services).__name__}' object has no attribute 'values'$", + ): + sap_credentials.fetch_credentials() diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index aec68f1fa45..0539a939bb5 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -2962,6 +2962,30 @@ def test_convert_tool_response_with_url_image(content_type: str, nested: bool) - ) +def _assistant_with_image(url: str) -> Message: + return Message( + role="assistant", + content="Here is the picture", + images=[{"type": "image_url", "image_url": {"url": url}, "index": 0}], + ) + + +@pytest.mark.parametrize( + ("message", "needs_fetch"), + [ + (_assistant_with_image("gs://litellm-test-bucket/cat"), True), + (_assistant_with_image("gs://litellm-test-bucket/cat.png"), False), + ("gs://litellm-test-bucket/cat", False), + ], +) +def test_sync_media_fetch_check_reads_pydantic_messages_and_skips_other_objects( + message: object, needs_fetch: bool +) -> None: + from litellm.llms.vertex_ai.gemini.transformation import _openai_messages_may_need_sync_media_fetch + + assert _openai_messages_may_need_sync_media_fetch([message]) is needs_fetch + + CLAUDE_THINKING_SIGNATURE: Final = "EqQBCkYIBxgCKkCfQ2x0b3VkZS1zaWduYXR1cmUtbm90LW1pbnRlZC1ieS1nZW1pbmkSDJ3lXf5sD+QVqpFQmRoM" diff --git a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py index 6ae0bbd0d72..60d7629276f 100644 --- a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py +++ b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py @@ -2,7 +2,8 @@ import base64 import json import os from datetime import datetime -from io import BytesIO +from io import BytesIO, StringIO +from pathlib import Path from typing import Dict, Final from unittest.mock import MagicMock, patch @@ -512,6 +513,30 @@ class TestVertexAIImagenImageEditTransformation: == b"test_bytearray" ) + @pytest.mark.parametrize( + ("image", "expected"), + [ + ([None, b"first image"], b"first image"), + ({"data": base64.b64encode(b"from data").decode()}, b"from data"), + ({"bytes": None, "content": b"from content"}, b"from content"), + ({"path": BytesIO(b"from path")}, b"from path"), + (StringIO("from text reader"), b"from text reader"), + ], + ) + def test_read_all_bytes_unwraps_lists_dicts_and_readers(self, image: object, expected: bytes) -> None: + assert self.config._read_all_bytes(image) == expected + + @pytest.mark.parametrize( + ("image", "refusal"), + [ + ({"path": "/etc/passwd"}, "plain string values are not accepted"), + ({"path": Path("/etc/passwd")}, "filesystem paths are not accepted"), + ], + ) + def test_read_all_bytes_refuses_to_read_a_path_given_inside_a_dict(self, image: object, refusal: str) -> None: + with pytest.raises(ValueError, match=refusal): + self.config._read_all_bytes(image) + _IMAGEN_EDIT_IMAGE_KEY = "bytesBase64Encoded" diff --git a/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py index 5c90d54ae90..ecf131f8215 100644 --- a/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -6,11 +6,12 @@ import base64 import json from collections.abc import Mapping from pathlib import Path -from typing import cast +from typing import Final, cast from unittest.mock import Mock, patch import httpx import pytest +from litellm.exceptions import UnsupportedParamsError import litellm from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -970,3 +971,17 @@ class TestImageAndParametersPassthrough: assert "parameters" not in data["parameters"] assert url.endswith(":predictLongRunning") + + +@pytest.mark.parametrize("drop_params", [False, True]) +def test_existing_scalar_parameters_remain_accepted(drop_params: bool) -> None: + params: Final = {"seed": 7, "duration": 6, "aspect_ratio": "9:16", "generate_audio": False} + assert VertexAIVideoConfig().map_openai_params(params, "veo-3.1-generate-001", drop_params) == {} + + +@pytest.mark.parametrize("field", ["frame_images", "input_references"]) +def test_first_class_video_fields_are_rejected_or_dropped(field: str) -> None: + config: Final = VertexAIVideoConfig() + with pytest.raises(UnsupportedParamsError, match=field): + config.map_openai_params({field: 1}, "video-model", False) + assert config.map_openai_params({field: 1}, "video-model", True) == config.map_openai_params({}, "video-model", True) diff --git a/tests/unit/llms/volcengine/test_volcengine_video_transformation.py b/tests/unit/llms/volcengine/test_volcengine_video_transformation.py index 1fcf49eeb13..ead17753b81 100644 --- a/tests/unit/llms/volcengine/test_volcengine_video_transformation.py +++ b/tests/unit/llms/volcengine/test_volcengine_video_transformation.py @@ -2,7 +2,8 @@ import base64 import io import json from pathlib import Path -from typing import Final +from collections.abc import Mapping +from typing import Final, cast from unittest.mock import Mock import httpx @@ -16,6 +17,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.openai.cost_calculation import video_generation_cost from litellm.llms.volcengine.videos.transformation import VolcEngineVideoConfig, VolcEngineVideoError from litellm.types.router import GenericLiteLLMParams +from litellm.types.videos.main import VideoCreateOptionalRequestParams from litellm.types.videos.utils import decode_video_id_with_provider, encode_video_id_with_provider from litellm.utils import ProviderConfigManager from litellm.videos.main import avideo_content, avideo_generation @@ -1129,3 +1131,129 @@ def test_a_video_created_through_a_routing_group_is_polled_on_the_member_that_cr hosts = [call.request.url.host for call in respx.calls] assert hosts[1::2] == hosts[0::2] + + +@pytest.mark.parametrize("role", ["first_frame", "last_frame", "reference_image"]) +def test_openrouter_video_images_become_provider_content(role: str) -> None: + params: Final = cast(VideoCreateOptionalRequestParams, { + "duration": 8, "aspect_ratio": "9:16", "resolution": "720p", "seed": 0, "generate_audio": False, + **( + {"input_references": [{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}}]} + if role == "reference_image" + else {"frame_images": [{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}, "frame_type": role}]} + ), + }) + config: Final = VolcEngineVideoConfig() + mapped: Final = config.map_openai_params(params, SEEDANCE_2_5, False) + body, _, _ = config.transform_video_create_request( + model=SEEDANCE_2_5, prompt="a cat", api_base=API_BASE, + video_create_optional_request_params=mapped, litellm_params=GenericLiteLLMParams(), headers={}, + ) + assert body["content"] == [ + {"type": "text", "text": "a cat"}, + {"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}, "role": role}, + ] + assert (body["duration"], body["ratio"]) == (8, "9:16") + assert not ({"frame_images", "input_references", "aspect_ratio"} & body.keys()) + + +@pytest.mark.parametrize("other", [ + {"content": [{"type": "image_url", "image_url": {"url": "https://example.com/raw.png"}, "role": "reference_image"}]}, + {"input_references": [{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}}]}, + {"input_reference": "https://example.com/old.png"}, + {"frame_images": [ + {"type": "image_url", "image_url": {"url": "https://example.com/first.png"}, "frame_type": "first_frame"}, + {"type": "image_url", "image_url": {"url": "https://example.com/duplicate.png"}, "frame_type": "first_frame"}, + ]}, +]) +def test_openrouter_image_conflicts_raise_or_keep_first_frame(other: object) -> None: + params: Final = cast(VideoCreateOptionalRequestParams, { + "frame_images": [{"type": "image_url", "image_url": {"url": "https://example.com/first.png"}, "frame_type": "first_frame"}], + **cast(Mapping[str, object], other), + }) + config: Final = VolcEngineVideoConfig() + with pytest.raises(UnsupportedParamsError, match="drop_params"): + config.map_openai_params(params, SEEDANCE_2_5, False) + assert config.map_openai_params(params, SEEDANCE_2_5, True) == { + "content": [{"type": "image_url", "image_url": {"url": "https://example.com/first.png"}, "role": "first_frame"}], + } + + +@pytest.mark.parametrize("roles", [("first_frame", "last_frame"), ("last_frame", "first_frame")]) +def test_large_duplicate_frame_arrays_keep_first_images_in_input_order(roles: tuple[str, str]) -> None: + params: Final = cast(VideoCreateOptionalRequestParams, { + "frame_images": [ + {"type": "image_url", "image_url": {"url": f"https://example.com/{index}.png"}, "frame_type": role} + for index, role in enumerate(roles * 5000) + ], + }) + assert VolcEngineVideoConfig().map_openai_params(params, SEEDANCE_2_5, True) == { + "content": [ + {"type": "image_url", "image_url": {"url": f"https://example.com/{index}.png"}, "role": role} + for index, role in enumerate(roles) + ], + } + + +def test_openrouter_duration_conflict_raises_or_prefers_duration() -> None: + config: Final = VolcEngineVideoConfig() + with pytest.raises(UnsupportedParamsError, match="duration and seconds"): + config.map_openai_params({"duration": 8, "seconds": "5"}, SEEDANCE_2_5, False) + assert config.map_openai_params({"duration": 8, "seconds": "5"}, SEEDANCE_2_5, True) == {"duration": 8} + assert config.map_openai_params({"duration": 8, "seconds": "8"}, SEEDANCE_2_5, False) == {"duration": 8} + + +@pytest.mark.parametrize("images", [[], [{"type": "image_url", "image_url": {}}]]) +def test_openrouter_empty_and_invalid_frames(images: object) -> None: + params: Final = cast(VideoCreateOptionalRequestParams, {"frame_images": images}) + config: Final = VolcEngineVideoConfig() + if images: + with pytest.raises(UnsupportedParamsError, match="Invalid video image"): + config.map_openai_params(params, SEEDANCE_2_5, False) + assert config.map_openai_params(params, SEEDANCE_2_5, True) == {} + else: + assert config.map_openai_params(params, SEEDANCE_2_5, False) == {"content": []} + + +def test_dropping_invalid_frames_preserves_raw_content_and_other_mappings() -> None: + content: Final = [{"type": "image_url", "image_url": {"url": "https://example.com/raw.png"}, "role": "first_frame"}] + params: Final = cast(VideoCreateOptionalRequestParams, { + "frame_images": [{"type": "image_url", "image_url": {}}], + "content": content, + "aspect_ratio": "16:9", + "duration": 8, + }) + assert VolcEngineVideoConfig().map_openai_params(params, SEEDANCE_2_5, True) == { + "content": content, + "ratio": "16:9", + "duration": 8, + } + +@pytest.mark.parametrize("references", [[], [{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}}]]) +def test_plural_and_singular_reference_conflict_raises_or_prefers_plural(references: object) -> None: + config: Final = VolcEngineVideoConfig() + params: Final = cast(VideoCreateOptionalRequestParams, { + "input_references": references, "input_reference": "https://example.com/first.png", + }) + with pytest.raises(UnsupportedParamsError, match="input_references cannot be combined"): + config.map_openai_params(params, SEEDANCE_2_5, False) + assert config.map_openai_params(params, SEEDANCE_2_5, True) == { + "content": [ + {"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}, "role": "reference_image"} + ] if references else [], + } + + +def test_dropping_seconds_does_not_parse_the_discarded_value() -> None: + assert VolcEngineVideoConfig().map_openai_params({"duration": 8, "seconds": "abc"}, SEEDANCE_2_5, True) == {"duration": 8} + + +@pytest.mark.parametrize("drop_params", [False, True]) +def test_aspect_ratio_preserves_native_ratio_precedence_over_size(drop_params: bool) -> None: + config: Final = VolcEngineVideoConfig() + assert config.map_openai_params( + {"size": "1280x720", "aspect_ratio": "9:16"}, SEEDANCE_2_5, drop_params, + ) == {"ratio": "9:16", "resolution": "720p"} + assert config.map_openai_params( + cast(VideoCreateOptionalRequestParams, {"size": "1280x720", "ratio": "9:16"}), SEEDANCE_2_5, drop_params, + ) == {"ratio": "9:16", "resolution": "720p"} diff --git a/tests/unit/llms/xai/image_generation/__init__.py b/tests/unit/llms/xai/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/xai/image_generation/test_cost_calculator.py b/tests/unit/llms/xai/image_generation/test_cost_calculator.py new file mode 100644 index 00000000000..70246534368 --- /dev/null +++ b/tests/unit/llms/xai/image_generation/test_cost_calculator.py @@ -0,0 +1,307 @@ +from datetime import datetime +from typing import Final + +import httpx +import pytest +from openai import AsyncOpenAI, OpenAI + +import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import convert_to_model_response_object +from litellm.types.utils import ImageResponse + + +def _response(ticks: object = None) -> ImageResponse: + return convert_to_model_response_object( + response_object={ + "created": 123, + "data": [{"url": "https://image.test/image.png"}], + "usage": {"cost_in_usd_ticks": ticks}, + }, + response_type="image_generation", + ) + + +@pytest.mark.parametrize("ticks", (0, 1234567890)) +@pytest.mark.parametrize("call_type", ("image_generation", "aimage_generation", "image_edit", "aimage_edit")) +def test_billed_cost_is_request_total_even_for_multiple_images(ticks: int, call_type: str) -> None: + assert ( + CostCalculatorUtils.route_image_generation_cost_calculator( + model="grok-imagine-image-2.0", + custom_llm_provider="xai", + completion_response=_response(ticks), + n=4, + call_type=call_type, + optional_params={"resolution": "2k", "quality": "medium"}, + ) + == ticks / 10_000_000_000 + ) + + +@pytest.mark.parametrize("ticks", (None, -1, True, "400000000", 1.5)) +@pytest.mark.parametrize("resolution,unit_price", ((None, 0.17), ("1k", 0.17), ("2k", 0.29))) +def test_missing_or_invalid_ticks_use_requested_resolution( + ticks: object, resolution: str | None, unit_price: float +) -> None: + params: Final = {"resolution": resolution} if resolution else {} + assert CostCalculatorUtils.route_image_generation_cost_calculator( + model="unmapped-image-fixture", + custom_llm_provider="xai", + completion_response=_response(ticks), + n=3, + optional_params={"extra_body": params}, + model_info={"output_cost_per_image_1K": 0.17, "output_cost_per_image_2K": 0.29}, + ) == pytest.approx(unit_price * 3) + + +@pytest.mark.parametrize( + "quality,call_type,resolution", + ( + ("medium", "image_generation", "1k"), + ("high", "image_generation", "2k"), + (None, "image_edit", "1k"), + ("auto", "aimage_edit", "2k"), + ("low", "image_generation", "1.5k"), + ), +) +def test_unpublished_rates_keep_legacy_fallback(quality: str | None, call_type: str, resolution: str) -> None: + assert CostCalculatorUtils.route_image_generation_cost_calculator( + model="unmapped-image-fixture", + custom_llm_provider="xai", + completion_response=_response(), + n=2, + quality=quality, + call_type=call_type, + optional_params={"resolution": resolution}, + model_info={"input_cost_per_image": 0.71, "output_cost_per_image_1K": 0.17, "output_cost_per_image_2K": 0.29}, + ) == pytest.approx(0.71 * 2) + + +@pytest.mark.parametrize("provider", ("xai", "openai")) +def test_flat_price_and_empty_response_remain_supported(provider: str) -> None: + response: Final = ImageResponse(data=[]) + assert ( + CostCalculatorUtils.route_image_generation_cost_calculator( + model="unmapped-image-fixture", + custom_llm_provider=provider, + completion_response=response, + model_info={"output_cost_per_image": 0.37}, + ) + == 0.0 + ) + assert CostCalculatorUtils.route_image_generation_cost_calculator( + model="unmapped-image-fixture", + custom_llm_provider=provider, + completion_response=response, + n=2, + model_info={"output_cost_per_image": 0.37}, + ) == pytest.approx(0.37 * 2) + + +@pytest.mark.parametrize("quality,call_type", (("low", "image_edit"), ("auto", "image_generation"), (None, None))) +def test_low_resolution_price_applies_without_flat_override(quality: str | None, call_type: str | None) -> None: + assert CostCalculatorUtils.route_image_generation_cost_calculator( + model="unmapped-image-fixture", + custom_llm_provider="xai", + completion_response=_response(), + n=2, + quality=quality, + call_type=call_type, + optional_params={"resolution": "2k"}, + model_info={"output_cost_per_image_2K": 0.29}, + ) == pytest.approx(0.29 * 2) + + +def test_deployment_flat_price_overrides_shared_resolution_price() -> None: + assert CostCalculatorUtils.route_image_generation_cost_calculator( + model="grok-imagine-image-2.0", + custom_llm_provider="xai", + completion_response=_response(), + n=2, + model_info={"output_cost_per_image": 0.83}, + ) == pytest.approx(0.83 * 2) + + +@pytest.mark.parametrize( + "price_key,resolution", + ( + ("input_cost_per_image", "1k"), + ("output_cost_per_image", "1k"), + ("output_cost_per_image_1K", "1k"), + ("output_cost_per_image_2K", "2k"), + ), +) +def test_deployment_image_price_overrides_billed_ticks(price_key: str, resolution: str) -> None: + assert CostCalculatorUtils.route_image_generation_cost_calculator( + model="grok-imagine-image-2.0", + custom_llm_provider="xai", + completion_response=_response(1234567890), + n=3, + optional_params={"resolution": resolution}, + model_info={price_key: 0.83}, + ) == pytest.approx(0.83 * 3) + + +@pytest.mark.parametrize("ticks", (None, -1, True, "600000000", 0, 600000000)) +@pytest.mark.parametrize("pixel_rate", (0.0, 0.000001)) +@pytest.mark.parametrize( + "size,pixels,n", ((None, 1024 * 1024, 3), ("1024x1024", 1024 * 1024, 1), ("512x256", 512 * 256, 2)) +) +def test_deployment_pixel_price_overrides_ticks_and_resolution_prices( + ticks: object, pixel_rate: float, size: str | None, pixels: int, n: int +) -> None: + assert CostCalculatorUtils.route_image_generation_cost_calculator( + model="grok-imagine-image-2.0", + custom_llm_provider="xai", + completion_response=_response(ticks), + size=size, + quality="low", + n=n, + model_info={"input_cost_per_pixel": pixel_rate, "output_cost_per_image_2K": 0.29}, + optional_params={"resolution": "2k"}, + call_type="image_generation", + ) == pytest.approx(pixel_rate * pixels * n) + + +@pytest.mark.parametrize("ticks", (None, 600000000)) +@pytest.mark.parametrize("price_key", ("input_cost_per_image", "output_cost_per_image")) +def test_deployment_flat_price_still_precedes_pixel_price(ticks: object, price_key: str) -> None: + assert CostCalculatorUtils.route_image_generation_cost_calculator( + model="unmapped-image-fixture", + custom_llm_provider="xai", + completion_response=_response(ticks), + n=2, + quality="low", + model_info={price_key: 0.83, "input_cost_per_pixel": 0.000001, "output_cost_per_image_2K": 0.29}, + optional_params={"resolution": "2k"}, + ) == pytest.approx(0.83 * 2) + + +@pytest.mark.asyncio +async def test_router_preserves_deployment_pixel_price_through_image_normalization() -> None: + ticks: Final = 600000000 + transport: Final = httpx.MockTransport( + lambda request: httpx.Response( + 200, + json={ + "created": 123, + "data": [{"url": "https://image.test/output.png"}], + "usage": {"cost_in_usd_ticks": ticks}, + }, + ) + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "pixel-fixture", + "litellm_params": { + "model": "xai/grok-imagine-image-2.0", + "api_key": "test", + "api_base": "https://image.test/v1", + }, + "model_info": {"input_cost_per_pixel": 0.000001}, + } + ] + ) + logging_obj: Final = Logging( + model="grok-imagine-image-2.0", + messages=[], + stream=False, + call_type="aimage_generation", + start_time=datetime(2026, 1, 1), + litellm_call_id="pixel-override", + function_id="pixel-override", + ) + async with AsyncOpenAI( + api_key="test", base_url="https://image.test/v1", http_client=httpx.AsyncClient(transport=transport) + ) as client: + response: Final = await router.aimage_generation( + model="pixel-fixture", + prompt="circle", + quality="low", + n=1, + extra_body={"resolution": "2k"}, + client=client, + litellm_logging_obj=logging_obj, + ) + assert getattr(response.usage, "cost_in_usd_ticks") == ticks + assert logging_obj.response_cost_calculator(result=response) == pytest.approx(1024 * 1024 * 0.000001) + + +def test_shared_resolution_prices_are_used_when_ticks_are_absent() -> None: + pricing: Final = litellm.get_model_info("xai/grok-imagine-image-2.0") + assert CostCalculatorUtils.route_image_generation_cost_calculator( + model="grok-imagine-image-2.0", + custom_llm_provider="xai", + completion_response=_response(), + n=2, + optional_params={"resolution": "2k"}, + ) == pytest.approx(pricing["output_cost_per_image_2K"] * 2) + + +def test_missing_all_prices_keeps_missing_cost_error() -> None: + with pytest.raises(Exception, match="Model not found in cost map"): + CostCalculatorUtils.route_image_generation_cost_calculator( + model="unmapped-image-fixture", + custom_llm_provider="xai", + completion_response=_response(), + ) + + +def test_non_xai_usage_ticks_do_not_override_flat_prices() -> None: + assert CostCalculatorUtils.route_image_generation_cost_calculator( + model="unmapped-image-fixture", + custom_llm_provider="openai", + completion_response=_response(1234567890), + n=2, + model_info={"output_cost_per_image": 0.37}, + ) == pytest.approx(0.37 * 2) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True)) +async def test_generation_sdk_preserves_billed_ticks_through_normalization(is_async: bool) -> None: + ticks: Final = 1234567890 + transport: Final = httpx.MockTransport( + lambda request: httpx.Response( + 200, + json={ + "created": 123, + "data": [{"url": "https://image.test/output.png"}], + "usage": {"cost_in_usd_ticks": ticks}, + }, + ) + ) + if is_async: + async with AsyncOpenAI( + api_key="test", base_url="https://image.test/v1", http_client=httpx.AsyncClient(transport=transport) + ) as client: + response: Final = await litellm.aimage_generation( + model="xai/grok-imagine-image-2.0", + prompt="circle", + api_key="test", + api_base="https://image.test/v1", + client=client, + ) + assert getattr(response.usage, "cost_in_usd_ticks") == ticks + assert ( + litellm.completion_cost(completion_response=response, model="xai/grok-imagine-image-2.0", n=4) + == ticks / 10_000_000_000 + ) + else: + with OpenAI( + api_key="test", base_url="https://image.test/v1", http_client=httpx.Client(transport=transport) + ) as client: + sync_response: Final = litellm.image_generation( + model="xai/grok-imagine-image-2.0", + prompt="circle", + api_key="test", + api_base="https://image.test/v1", + client=client, + ) + assert getattr(sync_response.usage, "cost_in_usd_ticks") == ticks + assert ( + litellm.completion_cost(completion_response=sync_response, model="xai/grok-imagine-image-2.0", n=4) + == ticks / 10_000_000_000 + ) diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index dd532c0ccf5..78ce9728f3b 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -2401,6 +2401,37 @@ async def test_deny_by_default_ignores_tools_that_name_no_vector_store( ) +@pytest.mark.asyncio +@pytest.mark.parametrize("deny_by_default", [False, True], ids=["legacy", "deny-by-default"]) +@pytest.mark.parametrize( + ("key_vector_stores", "denied_by"), + [(["KBSTOREA"], _KEY_DENIED), (["KBSTOREB"], None)], + ids=["ungranted", "granted"], +) +async def test_file_search_inside_additional_tools_input_item_needs_a_vector_store_grant( + deny_by_default: bool, key_vector_stores: list[str], denied_by: ProxyErrorTypes | None +): + await _assert_rag_query_outcome( + _common_checks_for_rag_query( + {"vector_store_deny_by_default": deny_by_default}, + _virtual_key(object_permission_id="key-permission"), + SimpleNamespace(vector_stores=key_vector_stores), + request_body={ + "model": "gpt-4o-mini", + "input": [ + { + "type": "additional_tools", + "role": "developer", + "tools": [{"type": "file_search", "vector_store_ids": ["KBSTOREB"]}], + }, + {"type": "message", "role": "user", "content": "what is in this KB?"}, + ], + }, + ), + denied_by, + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( "request_body", diff --git a/tests/unit/proxy/auth/test_model_checks.py b/tests/unit/proxy/auth/test_model_checks.py index 250c83fa61a..cd5410ffd18 100644 --- a/tests/unit/proxy/auth/test_model_checks.py +++ b/tests/unit/proxy/auth/test_model_checks.py @@ -35,6 +35,15 @@ def test_get_team_models_for_all_models_and_team_only_models(): assert set(result) == set(combined_models) +def test_get_team_models_lists_models_in_the_order_the_team_set_them() -> None: + from litellm.proxy.auth.model_checks import get_team_models + + team_models: Final = ["zeta", "alpha", "mu", "beta", "omega", "kappa", "delta", "all-proxy-models"] + result: Final = get_team_models(team_models, ["proxy-b", "proxy-a"], {"grp": ["g1"]}, True) + + assert result == [*team_models, "proxy-b", "proxy-a", "grp", "g1"] + + def test_get_team_models_all_proxy_models_includes_access_groups(): """ When a team has 'all-proxy-models' and include_model_access_groups=True, diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 741669337cd..8440d3d2ab2 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -12752,6 +12752,19 @@ def _user_defined_pass_through_endpoint(): return endpoint +@pytest.mark.asyncio +async def test_router_settings_model_group_alias_leaves_a_deployed_name_alone_on_provider_pass_through(monkeypatch): + router = _alias_router() + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + data = {"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}]} + route = "/anthropic/v1/messages" + request = _alias_request(route, data, path_params={"endpoint": "v1/messages"}) + token = _alias_token(monkeypatch, "key", {"claude-sonnet": "claude-haiku"}, ["claude-haiku", "claude-sonnet"]) + await _enforce_alias_access(token, data, route, request, router) + assert data["model"] == "claude-sonnet" + assert (await request.json())["model"] == "claude-sonnet" + + @pytest.mark.asyncio @pytest.mark.parametrize("user_defined", [False, True]) async def test_router_settings_model_group_alias_leaves_pass_through_bodies_alone(monkeypatch, user_defined): diff --git a/tests/unit/proxy/client/test_credentials.py b/tests/unit/proxy/client/test_credentials.py index 9de990da487..d49223be710 100644 --- a/tests/unit/proxy/client/test_credentials.py +++ b/tests/unit/proxy/client/test_credentials.py @@ -316,3 +316,25 @@ def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server): client.list() assert time.monotonic() - started < 10 + + +@responses.activate +def test_update_display_name_returns_the_updated_credential(client): + updated = {"credential_name": "azure1", "credential_info": {"display_name": "Azure EU"}} + responses.add( + responses.PATCH, + f"{client._base_url}/credentials/azure1", + json=updated, + status=200, + match=[responses.matchers.json_params_matcher({"display_name": "Azure EU", "credential_info": {}})], + ) + + assert client.update_display_name("azure1", "Azure EU") == updated + + +@responses.activate +def test_update_display_name_unauthorized_error(client): + responses.add(responses.PATCH, f"{client._base_url}/credentials/azure1", status=401, json={"error": "Unauthorized"}) + + with pytest.raises(UnauthorizedError): + client.update_display_name("azure1", "Azure EU") diff --git a/tests/unit/proxy/client/test_keys.py b/tests/unit/proxy/client/test_keys.py index 3bff944144c..14a7891f367 100644 --- a/tests/unit/proxy/client/test_keys.py +++ b/tests/unit/proxy/client/test_keys.py @@ -524,3 +524,27 @@ def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server): client.list() assert time.monotonic() - started < 10 + + +@responses.activate +def test_update_sends_only_the_given_fields_and_returns_the_updated_key(client): + updated = {"key": "key-abc", "key_alias": "renamed", "models": ["gpt-4o"]} + responses.add( + responses.POST, + f"{client._base_url}/key/update", + json=updated, + status=200, + match=[ + responses.matchers.json_params_matcher({"key": "key-abc", "key_alias": "renamed", "models": ["gpt-4o"]}) + ], + ) + + assert client.update(key="key-abc", key_alias="renamed", models=["gpt-4o"]) == updated + + +@responses.activate +def test_update_failure_carries_the_server_response_text(client): + responses.add(responses.POST, f"{client._base_url}/key/update", status=400, body="key not found") + + with pytest.raises(Exception, match="Error updating key: key not found"): + client.update(key="key-abc") diff --git a/tests/unit/proxy/client/test_teams.py b/tests/unit/proxy/client/test_teams.py index b61091ca44b..7aff5c62ff5 100644 --- a/tests/unit/proxy/client/test_teams.py +++ b/tests/unit/proxy/client/test_teams.py @@ -2,7 +2,9 @@ import time import pytest import requests +import responses +from litellm.proxy.client.exceptions import UnauthorizedError from litellm.proxy.client.teams import TeamsManagementClient @@ -18,3 +20,42 @@ def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server): client.list() assert time.monotonic() - started < 10 + + +@responses.activate +def test_list_v2_sends_filters_and_returns_the_page(): + client = TeamsManagementClient(base_url="http://localhost:8000", api_key="sk-test") + page = {"teams": [{"team_id": "team-1"}], "total": 1, "page": 2, "page_size": 5, "total_pages": 1} + responses.add( + responses.GET, + "http://localhost:8000/v2/team/list", + json=page, + status=200, + match=[ + responses.matchers.query_param_matcher( + { + "page": "2", + "page_size": "5", + "sort_order": "desc", + "user_id": "user-1", + "team_alias": "search", + "sort_by": "created_at", + } + ) + ], + ) + + result = client.list_v2( + user_id="user-1", team_alias="search", page=2, page_size=5, sort_by="created_at", sort_order="desc" + ) + + assert result == page + + +@responses.activate +def test_list_v2_unauthorized_error(): + client = TeamsManagementClient(base_url="http://localhost:8000", api_key="sk-bad") + responses.add(responses.GET, "http://localhost:8000/v2/team/list", status=401, json={"error": "Unauthorized"}) + + with pytest.raises(UnauthorizedError): + client.list_v2() diff --git a/tests/unit/proxy/common_utils/test_codex_catalog_source.py b/tests/unit/proxy/common_utils/test_codex_catalog_source.py new file mode 100644 index 00000000000..e3009e19dd7 --- /dev/null +++ b/tests/unit/proxy/common_utils/test_codex_catalog_source.py @@ -0,0 +1,257 @@ +import asyncio +import json +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from typing import Final + +import httpx +import pytest + +from litellm.proxy.common_utils.codex_catalog_source import ( + CODEX_RELEASE_FILE_URL, + CodexCatalogSource, + bundled_codex_stock, +) +from litellm.proxy.common_utils.codex_model_catalog import bundled_codex_models + +_TTL: Final = 3600.0 + + +def _release_files(slug: str, prompt: str) -> Mapping[str, bytes]: + entry: Final = {**bundled_codex_models()["gpt-5.5"].model_dump(mode="json", exclude_unset=True), "slug": slug} + return {"models.json": json.dumps({"models": [entry]}).encode(), "prompt.md": prompt.encode()} + + +@dataclass +class _GitHub: + """Serves each tagged release's files, 404 for any other tag, and records every URL asked for.""" + + releases: Mapping[str, Mapping[str, bytes]] + fail: Callable[[httpx.Request], httpx.Response] | None = None + requested: list[str] = field(default_factory=list) + + async def handle(self, request: httpx.Request) -> httpx.Response: + self.requested.append(str(request.url)) + if self.fail is not None: + return self.fail(request) + files: Final = { + CODEX_RELEASE_FILE_URL.format(version=version, file=name): body + for version, release in self.releases.items() + for name, body in release.items() # comprehension-ok: flattens release -> file in one lookup table + } + body: Final = files.get(str(request.url)) + return httpx.Response(404, text="404: Not Found") if body is None else httpx.Response(200, content=body) + + +@dataclass +class _Clock: + now: float = 0.0 + + def __call__(self) -> float: + return self.now + + +def _source( + github: _GitHub, + clock: _Clock | None = None, + *, + failure_ttl: float = _TTL, + max_versions: int = 8, + fetches_per_minute: int = 100, +) -> CodexCatalogSource: + client: Final = httpx.AsyncClient(transport=httpx.MockTransport(github.handle)) + return CodexCatalogSource( + lambda: client, + offline=lambda: False, + clock=clock or _Clock(), + failure_ttl=failure_ttl, + max_versions=max_versions, + fetches_per_minute=fetches_per_minute, + ) + + +async def test_release_catalog_and_prompt_are_served_for_the_asking_codex_version() -> None: + github: Final = _GitHub({"0.170.0": _release_files("gpt-from-0.170.0", "prompt of 0.170.0")}) + + stock: Final = await _source(github).stock("0.170.0") + + assert (list(stock.models), stock.instructions) == (["gpt-from-0.170.0"], "prompt of 0.170.0") + assert stock.models["gpt-from-0.170.0"].display_name == bundled_codex_models()["gpt-5.5"].display_name + + +async def test_each_codex_version_gets_its_own_release_and_a_missing_tag_does_not_spill_over() -> None: + github: Final = _GitHub( + { + "0.166.0": _release_files("gpt-from-0.166.0", "prompt of 0.166.0"), + "0.171.0-alpha.2": _release_files("gpt-from-alpha", "prompt of alpha"), + } + ) + source: Final = _source(github) + + served: Final = {version: await source.stock(version) for version in ("0.166.0", "0.171.0-alpha.2", "0.0.0")} + + assert {version: (list(stock.models), stock.instructions) for version, stock in served.items()} == { + "0.166.0": (["gpt-from-0.166.0"], "prompt of 0.166.0"), + "0.171.0-alpha.2": (["gpt-from-alpha"], "prompt of alpha"), + "0.0.0": (list(bundled_codex_models()), bundled_codex_stock().instructions), + } + + +def _unreachable(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("egress blocked", request=request) + + +def _server_error(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, request=request) + + +@pytest.mark.parametrize("fail", [_unreachable, _server_error], ids=["unreachable", "server-error"]) +async def test_a_failed_fetch_serves_the_bundled_catalog(fail: Callable[[httpx.Request], httpx.Response]) -> None: + github: Final = _GitHub({"0.170.0": _release_files("gpt-from-0.170.0", "p")}, fail=fail) + + assert await _source(github).stock("0.170.0") is bundled_codex_stock() + + +async def test_a_release_that_never_answers_serves_the_bundled_catalog_within_the_timeout() -> None: + async def hang(_: httpx.Request) -> httpx.Response: + await asyncio.Event().wait() + raise AssertionError("unreachable") + + client: Final = httpx.AsyncClient(transport=httpx.MockTransport(hang)) + source: Final = CodexCatalogSource(lambda: client, offline=lambda: False, clock=_Clock(), timeout=0.01) + + assert await asyncio.wait_for(source.stock("0.170.0"), timeout=30) is bundled_codex_stock() + + +@pytest.mark.parametrize( + "files", + [ + {"models.json": b'{"models": [{"slug": "gpt-x"}]}', "prompt.md": b"p"}, + {"models.json": b"rate limited", "prompt.md": b"p"}, + {"models.json": _release_files("gpt-x", "p")["models.json"], "prompt.md": b"\xff\xfe"}, + {"models.json": _release_files("gpt-x", "p")["models.json"]}, + ], + ids=["entry-missing-fields", "not-json", "prompt-not-utf8", "prompt-missing"], +) +async def test_a_release_whose_files_do_not_validate_serves_the_bundled_catalog(files: Mapping[str, bytes]) -> None: + assert await _source(_GitHub({"0.170.0": files})).stock("0.170.0") is bundled_codex_stock() + + +async def test_a_failed_fetch_is_not_retried_until_the_failure_expires() -> None: + github: Final = _GitHub({}, fail=_unreachable) + clock: Final = _Clock() + source: Final = _source(github, clock, failure_ttl=_TTL) + + first: Final = await source.stock("0.170.0") + after_first: Final = len(github.requested) + clock.now = _TTL - 1 + within_ttl: Final = await source.stock("0.170.0") + after_within: Final = len(github.requested) + clock.now = _TTL + 1 + await source.stock("0.170.0") + + assert first is within_ttl is bundled_codex_stock() + assert (after_first > 0, after_within == after_first, len(github.requested) > after_within) == (True, True, True) + + +async def test_a_fetched_release_is_not_fetched_again() -> None: + github: Final = _GitHub({"0.170.0": _release_files("gpt-from-0.170.0", "p")}) + clock: Final = _Clock() + source: Final = _source(github, clock) + + first: Final = await source.stock("0.170.0") + fetches: Final = len(github.requested) + clock.now = _TTL * 100 + + assert await source.stock("0.170.0") is first + assert len(github.requested) == fetches + + +async def test_only_the_most_recent_versions_are_kept() -> None: + github: Final = _GitHub({version: _release_files(version, "p") for version in ("0.1.0", "0.2.0", "0.3.0")}) + source: Final = _source(github, max_versions=2) + for version in ("0.1.0", "0.2.0", "0.3.0"): + await source.stock(version) + fetches: Final = len(github.requested) + + await source.stock("0.3.0") + kept_was_cached: Final = len(github.requested) == fetches + await source.stock("0.1.0") + + assert (kept_was_cached, len(github.requested) > fetches) == (True, True) + + +async def test_versions_that_fail_never_push_out_a_fetched_one() -> None: + github: Final = _GitHub({"0.170.0": _release_files("gpt-from-0.170.0", "p")}) + source: Final = _source(github, max_versions=2) + fetched: Final = await source.stock("0.170.0") + for made_up in ("9.9.1", "9.9.2", "9.9.3"): + await source.stock(made_up) + fetches: Final = len(github.requested) + + assert await source.stock("0.170.0") is fetched + assert len(github.requested) == fetches + + +async def test_a_failed_version_keeps_its_cooldown_past_the_fetched_catalog_cap() -> None: + github: Final = _GitHub({}, fail=_unreachable) + source: Final = _source(github, max_versions=2) + for made_up in ("9.9.1", "9.9.2", "9.9.3", "9.9.4"): + await source.stock(made_up) + fetches: Final = len(github.requested) + + await source.stock("9.9.1") + + assert len(github.requested) == fetches, github.requested + + +async def test_a_worker_starts_at_most_its_fetch_budget_each_minute() -> None: + github: Final = _GitHub({}) + clock: Final = _Clock() + source: Final = _source(github, clock, fetches_per_minute=2) + for made_up in ("9.9.1", "9.9.2", "9.9.3"): + await source.stock(made_up) + within_minute: Final = len(github.requested) + clock.now = 61 + await source.stock("9.9.3") + + assert (within_minute, len(github.requested)) == (4, 6), github.requested + + +async def test_requests_for_a_version_being_fetched_share_that_fetch() -> None: + answer: Final = asyncio.Event() + files: Final = _release_files("gpt-from-0.170.0", "p") + requested: Final[list[str]] = [] + + async def slow_github(request: httpx.Request) -> httpx.Response: + requested.append(str(request.url)) + await answer.wait() + return httpx.Response(200, content=files[str(request.url).rsplit("/", 1)[1]]) + + client: Final = httpx.AsyncClient(transport=httpx.MockTransport(slow_github)) + source: Final = CodexCatalogSource(lambda: client, offline=lambda: False, clock=_Clock()) + waiting: Final = [asyncio.ensure_future(source.stock("0.170.0")) for _ in range(5)] + await asyncio.sleep(0) + answer.set() + served: Final = await asyncio.gather(*waiting) + + assert [list(stock.models) for stock in served] == [["gpt-from-0.170.0"]] * 5 + assert len(requested) == 2, requested + + +async def test_offline_mode_never_fetches(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_OFFLINE_MODE", "true") + github: Final = _GitHub({"0.170.0": _release_files("gpt-from-0.170.0", "p")}) + client: Final = httpx.AsyncClient(transport=httpx.MockTransport(github.handle)) + + stock: Final = await CodexCatalogSource(lambda: client, clock=_Clock()).stock("0.170.0") + + assert (stock, github.requested) == (bundled_codex_stock(), []) + + +@pytest.mark.parametrize("client_version", ["../../main", "0.170.0/../../x", "latest", "", "\u0660.\u0661.\u0662"]) +async def test_a_version_that_is_no_release_tag_is_never_fetched(client_version: str) -> None: + github: Final = _GitHub({}) + + assert await _source(github).stock(client_version) is bundled_codex_stock() + assert github.requested == [] diff --git a/tests/unit/proxy/common_utils/test_codex_model_catalog.py b/tests/unit/proxy/common_utils/test_codex_model_catalog.py index 760f79f4ec2..d4cc2fcb006 100644 --- a/tests/unit/proxy/common_utils/test_codex_model_catalog.py +++ b/tests/unit/proxy/common_utils/test_codex_model_catalog.py @@ -1,4 +1,5 @@ import json +import logging from collections.abc import Mapping from dataclasses import replace from types import MappingProxyType @@ -7,6 +8,7 @@ from typing import Final, Literal import pytest import litellm +from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.llm_cost_calc.utils import get_service_tier_cost_key from litellm.proxy.common_utils.codex_model_catalog import ( @@ -15,11 +17,14 @@ from litellm.proxy.common_utils.codex_model_catalog import ( CodexStockModel, CodexStockUpgrade, bundled_codex_models, + codex_catalog_etag, codex_catalog_rows, + codex_entry_names_default_reasoning_level, codex_model_list_body, codex_models_response_json, configured_service_tiers, deployment_capabilities, + merged_codex_catalog, ) from litellm.proxy.common_utils.codex_model_translation import CodexDeploymentCapabilities, CodexServiceTier @@ -628,6 +633,15 @@ def test_default_service_tier_survives_only_while_still_offered(configured, expe assert entry["default_service_tier"] == expected + +@pytest.mark.parametrize(("alias_defaults", "expected"), [("codex", "priority"), ("deployment", None)]) +def test_alias_with_offered_tiers_takes_codexs_default_tier_only_under_codex_defaults(alias_defaults, expected): + (entry,) = _models( + _row("team-sol", upstream_model="gpt-6-sol", service_tiers=["priority"]), alias_defaults=alias_defaults + ) + + assert (_tiers(entry), entry["default_service_tier"]) == ([("priority", "Fast", "2x speed, increased usage")], expected) + def test_non_chat_and_wildcard_rows_are_left_out_of_the_catalog(): body = _body( _row("gpt-6-astra"), @@ -705,6 +719,55 @@ def test_an_entry_too_large_for_the_bytes_left_is_passed_over_and_the_smaller_on assert len(body.json.encode()) <= limit +def test_merged_catalog_keeps_upstream_whole_and_adds_only_the_entries_that_fit_the_byte_limit(): + upstream = {"models": [{"slug": "upstream-model", "priority": 4, "base_instructions": "u" * 500}]} + upstream_only = len(json.dumps(upstream, separators=(",", ":"))) + rows = tuple(_row(f"model-{index}") for index in range(4)) + + def merged(*merged_rows, byte_limit): + return merged_codex_catalog( + json.dumps(upstream).encode(), merged_rows, stock=_STOCK, instructions=_PROMPT, byte_limit=byte_limit + ) + + unlimited = json.loads(merged(*rows, byte_limit=upstream_only * 100))["models"] + limit = len(merged(*rows[:2], byte_limit=upstream_only * 100).encode()) + + body = merged(*rows, byte_limit=limit) + + assert len(body.encode()) <= limit + assert json.loads(body)["models"] == unlimited[:3] + assert merged(*rows, byte_limit=upstream_only) is None + + +def test_merged_catalog_gives_an_alias_no_service_tier_or_upgrade_of_codexs_own(): + upstream = json.dumps({"models": [{"slug": "upstream-model", "priority": 0}]}).encode() + rows = (_row("team-sol", upstream_model="openai/gpt-6-sol"), _row("team-55", upstream_model="openai/gpt-5.5")) + + added = json.loads(merged_codex_catalog(upstream, rows, stock=_STOCK, instructions=_PROMPT))["models"][1:] + + assert [ + (entry["slug"], entry["service_tiers"], entry["default_service_tier"], entry["upgrade"]) for entry in added + ] == [ + ("team-sol", [], None, None), + ("team-55", [], None, None), + ] + + +def test_merged_catalog_warns_once_with_how_many_models_the_byte_limit_left_out(caplog, monkeypatch): + upstream = json.dumps({"models": [{"slug": "upstream-model", "priority": 0}]}).encode() + rows = tuple(_row(f"model-{index}") for index in range(4)) + limit = len(merged_codex_catalog(upstream, rows[:1], stock=_STOCK, instructions=_PROMPT, byte_limit=10**6).encode()) + monkeypatch.setattr(verbose_proxy_logger, "propagate", True) + + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + merged_codex_catalog(upstream, rows, stock=_STOCK, instructions=_PROMPT, byte_limit=limit) + merged_codex_catalog(upstream, rows[:1], stock=_STOCK, instructions=_PROMPT, byte_limit=limit) + + assert [record.getMessage() for record in caplog.records] == [ + f"Codex native model catalog left out 3 of 4 LiteLLM models to stay within Codex's {limit} byte limit" + ] + + def test_unlimited_body_keeps_every_entry(): rows = tuple(_row(f"model-{index}") for index in range(3)) @@ -947,6 +1010,38 @@ def test_catalog_rows_without_a_router_carry_only_the_listing(): assert codex_catalog_rows(listing, (("plain", "plain"),), None) == (CodexCatalogRow(id="plain", mode="chat"),) +@pytest.mark.parametrize("client_version", [None, "0.165.0"]) +def test_entry_names_a_default_reasoning_level_exactly_when_the_native_entry_carries_one(client_version): + router = litellm.Router( + model_list=[ + {"model_name": "open-weights", "litellm_params": {"model": "fireworks_ai/open-weights"}}, + {"model_name": "coding-model", "litellm_params": {"model": "openai/gpt-5.5"}}, + {"model_name": "pinned-model", "litellm_params": {"model": "openai/gpt-5.5", "reasoning_effort": "high"}}, + {"model_name": "split-pin", "litellm_params": {"model": "openai/gpt-5.5", "reasoning_effort": "high"}}, + {"model_name": "split-pin", "litellm_params": {"model": "openai/gpt-5.5", "reasoning_effort": "low"}}, + { + "model_name": "max-only", + "litellm_params": {"model": "hosted_vllm/qa-local-coder", "api_base": "https://vllm.example/v1"}, + "model_info": {"supports_reasoning": True, "reasoning_effort_levels": ["max"]}, + }, + ] + ) + names = ("open-weights", "coding-model", "pinned-model", "split-pin", "max-only") + listing = tuple( + {"id": name, "object": "model", "created": 0, "owned_by": "openai", "mode": "chat"} for name in names + ) + upstream = json.dumps({"models": [{"slug": "upstream-model", "priority": 0}]}).encode() + rows = codex_catalog_rows(listing, tuple((name, name) for name in names), router) + served = json.loads(merged_codex_catalog(upstream, rows, client_version=client_version))["models"][1:] + served_names_one = [entry.get("default_reasoning_level") is not None for entry in served] + + assert [ + codex_entry_names_default_reasoning_level(entry["slug"], router, None, client_version) for entry in served + ] == served_names_one + assert served_names_one == [False, False, True, False, client_version is not None] + assert served[2]["default_reasoning_level"] == "high" + + def test_codex_tool_search_is_offered_off_the_route_and_model_info_and_only_when_every_deployment_answers_it(): claude = "anthropic/claude-sonnet-5-5" converse = "bedrock/us.anthropic.claude-sonnet-5-5" @@ -995,6 +1090,34 @@ def test_codex_tool_search_is_offered_off_the_route_and_model_info_and_only_when assert models[0]["context_window"] == 1000000 +def test_native_merged_catalog_offers_codex_tool_search_as_the_proxy_listing_does(): + claude = "anthropic/claude-sonnet-5-5" + router = litellm.Router( + model_list=[ + {"model_name": "claude", "litellm_params": {"model": claude}}, + { + "model_name": "claude-off", + "litellm_params": {"model": claude}, + "model_info": {"supports_search_tool": False}, + }, + {"model_name": "coding-model", "litellm_params": {"model": "openai/gpt-5.5"}}, + ] + ) + names = ("claude", "claude-off", "coding-model") + listing = tuple({"id": name, "object": "model", "created": 0, "owned_by": "x", "mode": "chat"} for name in names) + entries = tuple((name, name) for name in names) + upstream = json.dumps({"models": [{"slug": "upstream-model", "priority": 0}]}).encode() + + listed = json.loads(codex_model_list_body(listing, entries, router))["models"] + merged = json.loads(merged_codex_catalog(upstream, codex_catalog_rows(listing, entries, router)))["models"][1:] + + assert [entry["slug"] for entry in merged] == list(names) + assert [entry.get("supports_search_tool") for entry in merged] == [ + entry.get("supports_search_tool") for entry in listed + ] + assert [entry["supports_search_tool"] for entry in merged[:2]] == [True, False] + + _SOL_SLUG: Final = next( slug for slug, model in bundled_codex_models().items() if model.default_service_tier is not None ) @@ -1282,3 +1405,48 @@ def test_model_info_that_does_not_validate_declares_no_capability_but_keeps_its_ ) assert capabilities == CodexDeploymentCapabilities(priced_service_tier_keys=frozenset({priced})) + + +def test_model_list_body_takes_the_asking_releases_catalog_and_prompt() -> None: + router: Final = litellm.Router( + model_list=[ + {"model_name": "gpt-release-only", "litellm_params": {"model": "openai/gpt-release-only"}}, + {"model_name": "in-house", "litellm_params": {"model": "openai/in-house-coder"}}, + ] + ) + release_entry: Final = bundled_codex_models()["gpt-5.5"].model_copy( + update={"slug": "gpt-release-only", "display_name": "Release Only"} + ) + names: Final = ("gpt-release-only", "in-house") + + body: Final = codex_model_list_body( + _listing(*names), + tuple((name, name) for name in names), + router, + client_version="0.170.0", + stock={"gpt-release-only": release_entry}, + instructions="release prompt", + ) + + served: Final = {entry["slug"]: entry for entry in json.loads(body)["models"]} + assert (served["gpt-release-only"]["display_name"], served["in-house"]["base_instructions"]) == ( + "Release Only", + "release prompt", + ), served + + +class _CollidingName(str): + def __hash__(self) -> int: + return 1 + + +def test_catalog_etag_ignores_the_order_a_process_iterates_priced_service_tiers_in() -> None: + def row_with_tiers(*tiers: str) -> CodexCatalogRow: + capabilities: Final = CodexDeploymentCapabilities.model_construct( + priced_service_tier_keys=frozenset(_CollidingName(tier) for tier in tiers) + ) + return replace(_row("gpt-6-astra"), deployment_capabilities=(capabilities,)) + + assert codex_catalog_etag("upstream", [row_with_tiers("flex", "priority")]) == codex_catalog_etag( + "upstream", [row_with_tiers("priority", "flex")] + ) diff --git a/tests/unit/proxy/common_utils/test_http_parsing_utils.py b/tests/unit/proxy/common_utils/test_http_parsing_utils.py index 0fc7bb15a32..e4e32eeeac9 100644 --- a/tests/unit/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/unit/proxy/common_utils/test_http_parsing_utils.py @@ -9,7 +9,7 @@ import orjson import pytest from fastapi import HTTPException from fastapi.testclient import TestClient -from starlette.datastructures import FormData +from starlette.datastructures import FormData, Headers, UploadFile from starlette.requests import Request @@ -24,6 +24,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( safe_get_request_query_params, safe_set_request_parsed_body, coerce_numeric_form_fields, + convert_upload_files_to_file_data, get_form_data, get_request_body, get_tags_from_request_body, @@ -1764,3 +1765,24 @@ async def test_read_request_body_closes_the_otel_server_span_when_the_client_dis with pytest.raises(HTTPException): await read_request_body(request) assert (span.end.called, request.state.parent_otel_span) == (True, None), "the auth-time span must not dangle" + + +def _upload(filename: str, content: bytes, content_type: str) -> UploadFile: + return UploadFile(file=io.BytesIO(content), filename=filename, headers=Headers({"content-type": content_type})) + + +@pytest.mark.asyncio +async def test_convert_upload_files_to_file_data_reads_every_upload_into_file_tuples(): + form_data: Final = { + "files": [_upload("a.txt", b"alpha", "text/plain"), _upload("b.json", b"{}", "application/json")], + "file": _upload("skill.zip", b"PK", "application/zip"), + "tags": ["red", "blue"], + "display_title": "My skill", + } + + assert await convert_upload_files_to_file_data(form_data) == { + "files": [("a.txt", b"alpha", "text/plain"), ("b.json", b"{}", "application/json")], + "file": [("skill.zip", b"PK", "application/zip")], + "tags": ["red", "blue"], + "display_title": "My skill", + } diff --git a/tests/unit/proxy/common_utils/test_openai_endpoint_utils.py b/tests/unit/proxy/common_utils/test_openai_endpoint_utils.py index 3e3e89f8117..3eb901144a3 100644 --- a/tests/unit/proxy/common_utils/test_openai_endpoint_utils.py +++ b/tests/unit/proxy/common_utils/test_openai_endpoint_utils.py @@ -77,9 +77,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( ), ], ) -def test_remove_sensitive_info_from_deployment( - model_config: dict, expected_config: dict -): +def test_remove_sensitive_info_from_deployment(model_config: dict, expected_config: dict): sanitized_config = remove_sensitive_info_from_deployment(model_config) assert sanitized_config == expected_config @@ -110,10 +108,7 @@ def test_remove_sensitive_info_from_deployment_with_excluded_keys(): sanitized_config = remove_sensitive_info_from_deployment( copy.deepcopy(base_config), excluded_keys={"litellm_credentials_name"} ) - assert ( - sanitized_config["litellm_params"]["litellm_credentials_name"] - == "my-credential-name" - ) + assert sanitized_config["litellm_params"]["litellm_credentials_name"] == "my-credential-name" # access_token should still be masked (not in excluded_keys) assert sanitized_config["litellm_params"]["access_token"] != "token-12345" @@ -121,3 +116,24 @@ def test_remove_sensitive_info_from_deployment_with_excluded_keys(): # api_key should still be removed (popped) regardless of excluded_keys assert "api_key" not in sanitized_config["litellm_params"] + + +def test_remove_sensitive_info_from_deployment_keeps_nested_nulls(): + router_config = { + "models": {"capable": "strong", "efficient": "fast", "judge": None}, + "options": {"confidence_threshold": None, "tool_semantics": None}, + "fallbacks": [None, "fast"], + } + deployment = { + "model_name": "route", + "litellm_params": { + "model": "auto_router/switchyard", + "api_key": "sk-sensitive-key-123", + "switchyard_router_config": copy.deepcopy(router_config), + }, + } + + sanitized = remove_sensitive_info_from_deployment(deployment) + + assert "api_key" not in sanitized["litellm_params"] + assert sanitized["litellm_params"]["switchyard_router_config"] == router_config diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py b/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py index c6a7fbd7cab..4ac3e769030 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py @@ -1,3 +1,4 @@ +import json from unittest.mock import AsyncMock, patch import httpx @@ -227,3 +228,57 @@ async def test_guardrails_ai_process_input(): # Should keep original content when no output fields are present assert result["messages"][0]["content"] == "Original message" + + +class _RecordingGuardrailsClient: + def __init__(self, status_code: int, body: dict[str, object]) -> None: + self.status_code = status_code + self.body = body + self.posts: list[dict[str, object]] = [] + + async def post(self, **kwargs: object) -> httpx.Response: + self.posts.append(kwargs) + return httpx.Response(self.status_code, json=self.body) + + +def _guardrails_ai() -> GuardrailsAI: + return GuardrailsAI(guardrail_name="test_guard", api_base="http://guardrails.test", guard_name="gibberish_guard") + + +@pytest.mark.asyncio +async def test_guardrails_ai_api_request_returns_the_passing_verdict(monkeypatch: pytest.MonkeyPatch) -> None: + verdict = {"callId": "call-1", "rawLlmOutput": "hello", "validatedOutput": "hello", "validationPassed": True} + client = _RecordingGuardrailsClient(200, verdict) + monkeypatch.setattr(litellm, "module_level_aclient", client) + + result = await _guardrails_ai().make_guardrails_ai_api_request(llm_output="hello", request_data={}) + + assert result == verdict + assert client.posts[0]["url"] == "http://guardrails.test/guards/gibberish_guard/validate" + assert json.loads(client.posts[0]["data"]) == {"llmOutput": "hello"} + + +@pytest.mark.asyncio +async def test_guardrails_ai_api_request_blocks_a_failed_validation(monkeypatch: pytest.MonkeyPatch) -> None: + verdict = {"callId": "call-2", "rawLlmOutput": "asdfgh", "validationPassed": False} + monkeypatch.setattr(litellm, "module_level_aclient", _RecordingGuardrailsClient(200, verdict)) + + with pytest.raises(HTTPException) as exc_info: + await _guardrails_ai().make_guardrails_ai_api_request(llm_output="asdfgh", request_data={}) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": "Violated guardrail policy", "guardrails_ai_response": verdict} + + +@pytest.mark.asyncio +async def test_guardrails_ai_pre_call_request_returns_the_first_output_text(monkeypatch: pytest.MonkeyPatch) -> None: + body = {"outputs": [{"name": "text", "shape": [1], "datatype": "BYTES", "data": ["cleaned text"]}]} + client = _RecordingGuardrailsClient(200, body) + monkeypatch.setattr(litellm, "module_level_aclient", client) + + result = await _guardrails_ai().make_guardrails_ai_api_request_pre_call_request( + text_input="raw text", request_data={} + ) + + assert result == "cleaned text" + assert json.loads(client.posts[0]["data"])["inputs"][0]["data"] == ["raw text"] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py index 23288a6136f..dfb72b42000 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -10,6 +10,7 @@ from collections.abc import Mapping from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException @@ -2461,3 +2462,38 @@ async def test_lakera_post_call_masks_pii_and_allows(): result_dict = result.model_dump() assert result_dict["choices"][0]["message"]["content"] != "Your email is test@example.com" assert "[MASKED" in result_dict["choices"][0]["message"]["content"] + + +class _RecordingLakeraClient: + def __init__(self, body: dict[str, object]) -> None: + self.body = body + self.posts: list[dict[str, object]] = [] + + async def post(self, **kwargs: object) -> httpx.Response: + self.posts.append(kwargs) + return httpx.Response(200, json=self.body) + + +@pytest.mark.asyncio +async def test_call_v2_guard_returns_the_lakera_verdict_and_logs_it_without_the_payload() -> None: + guardrail: Final = LakeraAIGuardrail(api_key="test_key", guardrail_name="lakera-guard") + verdict: Final = { + "payload": [{"detector_type": "pii/email", "start": 0, "end": 16, "message_id": 0}], + "flagged": True, + "breakdown": [{"detector_type": "pii/email", "detected": True, "message_id": 0}], + } + client: Final = _RecordingLakeraClient(verdict) + guardrail.async_handler = client + request_data: Final[dict[str, object]] = {"metadata": {}} + + response, masked_entity_count = await guardrail.call_v2_guard( + messages=[{"role": "user", "content": "test@example.com"}], + request_data=request_data, + event_type="pre_call", + ) + + assert (response, masked_entity_count) == (verdict, {}) + assert client.posts[0]["url"] == "https://api.lakera.ai/v2/guard" + logged: Final = request_data["metadata"]["standard_logging_guardrail_information"][0] + assert logged["guardrail_status"] == "success" + assert logged["guardrail_response"] == {"flagged": True, "breakdown": verdict["breakdown"]} diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py index 34221720205..ea6892379c2 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py @@ -2803,3 +2803,11 @@ class TestRegistration: guardrail_class_registry["microsoft_purview"] is MicrosoftPurviewDLPGuardrail ) + + +@pytest.mark.parametrize("odata_type", [None, 7], ids=["null", "number"]) +def test_should_block_rejects_a_policy_action_whose_odata_type_is_not_text(odata_type: object) -> None: + response: Final = {"policyActions": [{"@odata.type": odata_type, "restrictionAction": "block"}]} + + with pytest.raises(TypeError, match=f"argument of type '{type(odata_type).__name__}'"): + PurviewGuardrailBase._should_block(response) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py index f40465e8d7b..bac470b1365 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py @@ -2,13 +2,16 @@ Unit tests for ToolPolicyGuardrail. """ -from typing import Any +from collections.abc import Mapping, Sequence +from types import SimpleNamespace +from typing import Any, Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException +from litellm.proxy.db import tool_registry_writer from litellm.proxy.guardrails.guardrail_hooks.tool_policy.tool_policy_guardrail import ( ToolPolicyGuardrail, ) @@ -222,3 +225,104 @@ async def test_blocked_tool_loaded_by_a_replayed_tool_search_raises_next_to_allo ) assert exc_info.value.status_code == 400 assert exc_info.value.detail["blocked_tools"] == ["run_shell"] + + +async def _install_tool_policies(monkeypatch: pytest.MonkeyPatch, input_policies: Mapping[str, str]) -> None: + prisma_client = MagicMock() + prisma_client.db.litellm_tooltable.find_many = AsyncMock( + return_value=[ + SimpleNamespace(tool_name=name, input_policy=policy, output_policy="untrusted") + for name, policy in input_policies.items() + ] + ) + prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock(return_value=[]) + registry = tool_registry_writer.ToolPolicyRegistry() + await registry.sync_tool_policy_from_db(prisma_client) + monkeypatch.setattr(tool_registry_writer, "_tool_policy_registry", registry) + + +def _responses_request( + tools: Sequence[Mapping[str, object]], item_tools: Sequence[Mapping[str, object]] +) -> dict[str, object]: + return { + "user_api_key_request_route": "/v1/responses", + "tools": tools, + "input": [ + {"type": "additional_tools", "role": "developer", "tools": item_tools}, + {"role": "user", "content": "clean up the disk"}, + ], + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_data", + [ + _responses_request( + [{"type": "function", "name": "get_weather"}], + [{"type": "function", "name": "drop_database"}], + ), + _responses_request( + [{"type": "function", "name": "get_weather"}], + [ + { + "type": "namespace", + "name": "mcp__db", + "tools": [{"type": "function", "name": "drop_database"}], + } + ], + ), + _responses_request( + [ + {"type": "function", "name": "get_weather"}, + {"type": "mcp", "server_label": "drop_database"}, + ], + [], + ), + _responses_request( + [ + {"type": "function", "name": "get_weather"}, + {"type": "namespace", "name": "mcp__db", "tools": [{"type": "function", "name": "drop_database"}]}, + ], + [], + ), + ], + ids=["additional_tools_item", "namespace_in_item", "top_level_mcp", "top_level_namespace"], +) +async def test_blocked_tool_raises_next_to_a_top_level_function_tool( + guardrail: ToolPolicyGuardrail, request_data: dict[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + await _install_tool_policies(monkeypatch, {"get_weather": "untrusted", "drop_database": "blocked"}) + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_tool_request_inputs(["get_weather"]), + request_data=request_data, + input_type="request", + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["blocked_tools"] == ["drop_database"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("name", [["drop_database"], {"name": "drop_database"}], ids=["list", "dict"]) +async def test_function_tool_name_that_is_not_a_string_is_left_for_the_provider_to_refuse( + guardrail: ToolPolicyGuardrail, name: object, monkeypatch: pytest.MonkeyPatch +) -> None: + await _install_tool_policies(monkeypatch, {"get_weather": "untrusted"}) + inputs: Final = {"tools": [{"type": "function", "function": {"name": name}}]} + assert await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") == inputs + + +@pytest.mark.asyncio +async def test_blocked_tool_defined_at_the_top_level_and_in_an_item_is_named_once( + guardrail: ToolPolicyGuardrail, monkeypatch: pytest.MonkeyPatch +) -> None: + tool = {"type": "function", "name": "drop_database"} + await _install_tool_policies(monkeypatch, {"drop_database": "blocked"}) + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_tool_request_inputs(["drop_database"]), + request_data=_responses_request([tool], [tool]), + input_type="request", + ) + assert exc_info.value.detail["blocked_tools"] == ["drop_database"] diff --git a/tests/unit/proxy/guardrails/test_pillar_guardrails.py b/tests/unit/proxy/guardrails/test_pillar_guardrails.py index 86c1b1de8a1..1f167d34822 100644 --- a/tests/unit/proxy/guardrails/test_pillar_guardrails.py +++ b/tests/unit/proxy/guardrails/test_pillar_guardrails.py @@ -1571,3 +1571,35 @@ async def test_mcp_call_masking( if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) + + +class _RecordingPillarClient: + def __init__(self) -> None: + self.posted: list[dict[str, object]] = [] + + async def post(self, **kwargs: object) -> Response: + self.posted.append(kwargs) + return Response(200, json={"flagged": False}, request=Request("POST", "https://api.pillar.security")) + + +@pytest.mark.asyncio +async def test_messages_that_are_not_a_list_take_the_api_error_fallback_without_calling_pillar(env_setup): + guardrail = PillarGuardrail( + guardrail_name="pillar-fallback-block", + api_key="test-pillar-key", + api_base="https://api.pillar.security", + fallback_on_error="block", + ) + pillar_client = _RecordingPillarClient() + guardrail.async_handler = pillar_client + + with pytest.raises(HTTPException) as excinfo: + await guardrail.async_pre_call_hook( + data={"model": "gpt-4", "messages": 5}, + cache=DualCache(), + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + ) + + assert excinfo.value.status_code == 503 + assert pillar_client.posted == [] diff --git a/tests/unit/proxy/hooks/test_key_management_event_hooks.py b/tests/unit/proxy/hooks/test_key_management_event_hooks.py index 8952951351c..4ba6fb95a9d 100644 --- a/tests/unit/proxy/hooks/test_key_management_event_hooks.py +++ b/tests/unit/proxy/hooks/test_key_management_event_hooks.py @@ -779,3 +779,34 @@ async def test_deleting_many_keys_writes_every_audit_row_in_one_insert(): assert {(row["table_name"], row["action"], row["changed_by"]) for row in rows} == { (LitellmTableNames.KEY_TABLE_NAME, "deleted", "admin") } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_metadata", "expected"), + [ + ({"secret_manager_settings": {"prefix": "team-a/"}}, {"prefix": "team-a/"}), + ({"secret_manager_settings": {}}, None), + ({"owner": "platform"}, None), + (None, None), + ], + ids=["team_settings", "empty_settings", "no_settings_key", "no_metadata"], +) +async def test_secret_manager_optional_params_come_from_the_team_metadata( + monkeypatch: pytest.MonkeyPatch, + team_metadata: dict[str, object] | None, + expected: dict[str, object] | None, +) -> None: + from litellm.models.team import LiteLLM_TeamTable + from litellm.proxy import proxy_server + + async def fake_get_team_object( + team_id: str, prisma_client: object, user_api_key_cache: object + ) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id=team_id, metadata=team_metadata) + + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", fake_get_team_object) + + assert await KeyManagementEventHooks._get_secret_manager_optional_params("team-a") == expected diff --git a/tests/unit/proxy/hooks/test_max_iterations_limiter.py b/tests/unit/proxy/hooks/test_max_iterations_limiter.py index 5eb499042b5..0fc6d1ce543 100644 --- a/tests/unit/proxy/hooks/test_max_iterations_limiter.py +++ b/tests/unit/proxy/hooks/test_max_iterations_limiter.py @@ -150,3 +150,42 @@ async def test_max_iterations_no_agent_id_passes(): call_type="", ) assert result is None + + +class _SharedCounterRedis: + def __init__(self) -> None: + self.counts: dict[str, int] = {} + self.script_args: list[list[int]] = [] + + def async_register_script(self, script: str): + async def increment(keys: list[str], args: list[int]) -> int: + self.script_args.append(args) + self.counts[keys[0]] = self.counts.get(keys[0], 0) + 1 + return self.counts[keys[0]] + + return increment + + +@pytest.mark.asyncio +async def test_max_iterations_counts_one_session_across_proxy_instances_through_redis(): + shared_redis = _SharedCounterRedis() + first, second = ( + PROXY_MaxIterationsHandler(internal_usage_cache=InternalUsageCache(DualCache(redis_cache=shared_redis))) + for _ in range(2) + ) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test-key-1234", metadata={"max_iterations": 3}) + data = {"metadata": {"session_id": "session-shared"}} + + for handler in (first, second, first): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type="" + ) + + with pytest.raises(HTTPException) as exc_info: + await second.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type="" + ) + + assert exc_info.value.status_code == 429 + assert "Current count: 4" in str(exc_info.value.detail) + assert shared_redis.script_args == [[first.ttl]] * 4 diff --git a/tests/unit/proxy/management_endpoints/test_common_daily_activity.py b/tests/unit/proxy/management_endpoints/test_common_daily_activity.py index 5fe4cd2edcd..1cfe295b15d 100644 --- a/tests/unit/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/unit/proxy/management_endpoints/test_common_daily_activity.py @@ -1,11 +1,14 @@ from collections.abc import Mapping, Sequence +from dataclasses import asdict, fields from datetime import date, datetime -from types import SimpleNamespace +from types import FunctionType, SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock +import httpx import pytest from fastapi import HTTPException +from prisma.errors import DataError, PrismaError, TableNotFoundError import litellm.proxy.management_endpoints.common_daily_activity as common_daily_activity_module from litellm.constants import USAGE_TOP_API_KEYS_DEFAULT @@ -32,10 +35,301 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_E from litellm.proxy.utils import PrismaClient, hash_token from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendMetadata, + KeyMetadata, SpendAnalyticsPaginatedResponse, SpendMetrics, ) -from litellm.types.repositories.daily_activity import GroupingSetsRow, KeyMetadataRow +from litellm.types.repositories.daily_activity import ( + AggregatedRows, + DailyActivityScope, + EntityRollupRow, + GroupingSetsRow, + KeyMetadataRow, +) + + +def _raise_missing_error_meta() -> None: + raise AttributeError("'NoneType' object has no attribute 'get'") + + +def _query_engine_crash_on_closed_connection() -> AttributeError: + raise_inside_query_engine: Final = FunctionType( + _raise_missing_error_meta.__code__, {"__name__": "prisma.engine.utils"} + ) + with pytest.raises(AttributeError) as crash: + raise_inside_query_engine() + return crash.value + + +DATABASE_FAILURES: Final = pytest.mark.parametrize( + "failure", + [ + DataError( + {"user_facing_error": {"error_code": "P1001", "message": "Can't reach database server at `db:5432`"}} + ), + PrismaError("metadata unavailable"), + httpx.ReadTimeout("the query engine stopped answering"), + httpx.ConnectError("the query engine refused the connection"), + _query_engine_crash_on_closed_connection(), + ], + ids=[ + "unreachable-database", + "prisma-error", + "engine-read-timeout", + "engine-connect-error", + "engine-crash-on-closed-connection", + ], +) + +CODE_BUGS: Final = ( + pytest.param(TypeError("not a database failure"), id="type-error"), + pytest.param(AttributeError("not a database failure"), id="attribute-error"), +) +REJECTED_QUERIES: Final = ( + pytest.param( + DataError({"user_facing_error": {"error_code": "P2010", "message": "not a database failure"}}), + id="rejected-query", + ), + pytest.param( + TableNotFoundError( + { + "user_facing_error": { + "error_code": "P2021", + "message": "not a database failure", + "meta": {"table": "LiteLLM_VerificationToken"}, + } + } + ), + id="missing-table", + ), +) +NOT_DATABASE_OUTAGES: Final = pytest.mark.parametrize("bug", (*CODE_BUGS, *REJECTED_QUERIES)) + + +@pytest.mark.asyncio +@DATABASE_FAILURES +@pytest.mark.parametrize( + ("key_failure", "user_failure", "empty"), + [(False, False, False), (True, False, False), (False, True, False), (True, True, False), (False, False, True)], +) +async def test_entity_metadata_outages_preserve_daily_spend_and_other_metadata( + failure: Exception, key_failure: bool, user_failure: bool, empty: bool +) -> None: + prisma: Final = MagicMock() + key: Final = SimpleNamespace(token="key", key_alias="Alias", team_id=None, user_id=None, metadata=None) + user: Final = SimpleNamespace(user_id="user", user_email="user@example.com", user_alias="User") + main_rows: Final = tuple( + asdict(_grouping_row(level, api_key=api_key, spend=2.0)) + for level, api_key in ((127, None), (63, None), (31, "key")) + ) + entity_rows: Final = tuple( + {**asdict(_grouping_row(63, api_key=api_key, spend=2.0)), "entity_id": "user", "api_key_rolled": rolled} + for rolled, api_key in ((1, None), (0, "key")) + ) + prisma.db.query_raw = AsyncMock(side_effect=[main_rows, entity_rows]) + prisma.db.litellm_verificationtoken.find_many = AsyncMock( + side_effect=[[key], failure if key_failure else ([] if empty else [key])] + ) + prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_usertable.find_many = AsyncMock( + side_effect=failure if user_failure else None, + return_value=[] if empty else [user], + ) + + result: Final = await get_daily_activity_aggregated( + prisma_client=prisma, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, + start_date="2024-01-01", + end_date="2024-01-01", + model=None, + api_key=None, + include_entity_breakdown=True, + ) + + day: Final = result.results[0] + entity: Final = day.breakdown.entities["user"] + assert result.metadata.total_spend == day.metrics.spend == entity.metrics.spend == 2.0 + assert day.breakdown.api_keys["key"].metadata.key_alias == "Alias" + assert entity.metadata == ( + {} if user_failure or empty else {"user_email": "user@example.com", "user_alias": "User"} + ) + assert entity.api_key_breakdown["key"].metrics.spend == 2.0 + assert entity.api_key_breakdown["key"].metadata == ( + KeyMetadata() if key_failure else KeyMetadata(key_alias=None if empty else "Alias", key_exists=not empty) + ) + + +def _usage_scope() -> DailyActivityScope: + return daily_activity_scope( + "litellm_dailyuserspend", "user_id", None, None, None, "2024-01-01", "2024-01-01", None, None + ) + + +def _entity_row() -> EntityRollupRow: + grouping: Final = asdict(_grouping_row(63, api_key="key", spend=2.0)) + return EntityRollupRow( + **{ + **{field.name: grouping[field.name] for field in fields(EntityRollupRow) if field.name in grouping}, + "entity_id": "user", + "api_key_rolled": 0, + } + ) + + +def _usage_repository() -> MagicMock: + repository: Final = MagicMock() + repository.aggregated = AsyncMock( + return_value=AggregatedRows( + grouping_rows=tuple( + _grouping_row(level, api_key=api_key, spend=2.0) + for level, api_key in ((127, None), (63, None), (31, "key")) + ), + entity_rows=(_entity_row(),), + distinct_api_keys=1, + ) + ) + repository.key_metadata = AsyncMock(return_value={}) + repository.user_metadata = AsyncMock(return_value={}) + repository.guardrail_spend = AsyncMock(return_value=()) + return repository + + +@pytest.mark.asyncio +@DATABASE_FAILURES +async def test_usage_keeps_spend_when_the_first_key_lookup_hits_a_database_failure(failure: Exception) -> None: + repository: Final = _usage_repository() + repository.key_metadata = AsyncMock(side_effect=failure) + + result: Final = await _get_daily_activity_aggregated(repository, _usage_scope()) + + day: Final = result.results[0] + assert result.metadata.total_spend == day.metrics.spend == 2.0 + assert day.breakdown.api_keys["key"].metrics.spend == 2.0 + assert day.breakdown.api_keys["key"].metadata == KeyMetadata() + + +@pytest.mark.asyncio +@DATABASE_FAILURES +async def test_usage_keeps_spend_when_the_guardrail_lookup_hits_a_database_failure(failure: Exception) -> None: + repository: Final = _usage_repository() + repository.guardrail_spend = AsyncMock(side_effect=failure) + + result: Final = await _get_daily_activity_aggregated(repository, _usage_scope(), include_guardrail_breakdown=True) + + assert result.metadata.total_spend == result.results[0].metrics.spend == 2.0 + assert result.results[0].breakdown.guardrails == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", REJECTED_QUERIES) +async def test_usage_keeps_spend_when_the_guardrail_table_rejects_the_query(failure: Exception) -> None: + repository: Final = _usage_repository() + repository.guardrail_spend = AsyncMock(side_effect=failure) + + result: Final = await _get_daily_activity_aggregated(repository, _usage_scope(), include_guardrail_breakdown=True) + + assert result.metadata.total_spend == result.results[0].metrics.spend == 2.0 + assert result.results[0].breakdown.guardrails == {} + + +async def _assert_usage_reports_the_bug(repository: MagicMock) -> None: + with pytest.raises(HTTPException) as error: + await _get_daily_activity_aggregated( + repository, _usage_scope(), include_entity_breakdown=True, include_guardrail_breakdown=True + ) + + assert error.value.status_code == 500 + assert error.value.detail == {"error": "Failed to fetch analytics: not a database failure"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lookup", ["key_metadata", "user_metadata"]) +@NOT_DATABASE_OUTAGES +async def test_usage_still_reports_a_lookup_bug_that_is_not_a_database_failure(lookup: str, bug: Exception) -> None: + repository: Final = _usage_repository() + setattr(repository, lookup, AsyncMock(side_effect=bug)) + + await _assert_usage_reports_the_bug(repository) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bug", CODE_BUGS) +async def test_usage_still_reports_a_guardrail_lookup_bug_that_is_not_a_database_error(bug: Exception) -> None: + repository: Final = _usage_repository() + repository.guardrail_spend = AsyncMock(side_effect=bug) + + await _assert_usage_reports_the_bug(repository) + + +def _daily_user_spend_prisma(key_lookup: AsyncMock) -> MagicMock: + prisma: Final = MagicMock() + records: Final = [_daily_user_spend_record(user_id="user", api_key="key", spend=7.0)] + prisma.db.litellm_dailyuserspend.count = AsyncMock(return_value=len(records)) + prisma.db.litellm_dailyuserspend.find_many = AsyncMock(return_value=records) + prisma.db.litellm_verificationtoken.find_many = key_lookup + prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + return prisma + + +async def _daily_user_activity(prisma: MagicMock, user_emails: AsyncMock) -> SpendAnalyticsPaginatedResponse: + return await get_daily_activity( + prisma_client=prisma, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, + entity_metadata_field=None, + start_date="2024-01-01", + end_date="2024-01-01", + model=None, + api_key=None, + page=1, + page_size=1000, + resolve_entity_metadata=user_emails, + ) + + +@pytest.mark.asyncio +@DATABASE_FAILURES +@pytest.mark.parametrize(("key_failure", "email_failure"), [(True, False), (False, True)]) +async def test_daily_activity_keeps_spend_when_a_metadata_lookup_hits_a_database_failure( + failure: Exception, key_failure: bool, email_failure: bool +) -> None: + key: Final = SimpleNamespace(token="key", key_alias="Alias", team_id=None, user_id=None, metadata=None) + prisma: Final = _daily_user_spend_prisma( + AsyncMock(side_effect=failure if key_failure else None, return_value=[key]) + ) + user_emails: Final = AsyncMock( + side_effect=failure if email_failure else None, return_value={"user": {"user_email": "user@example.com"}} + ) + + result: Final = await _daily_user_activity(prisma, user_emails) + + day: Final = result.results[0] + assert result.metadata.total_spend == day.metrics.spend == 7.0 + assert day.breakdown.api_keys["key"].metadata == ( + KeyMetadata() if key_failure else KeyMetadata(key_alias="Alias", key_exists=True) + ) + assert day.breakdown.entities["user"].metadata == ({} if email_failure else {"user_email": "user@example.com"}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lookup", ["key metadata", "user emails"]) +@NOT_DATABASE_OUTAGES +async def test_daily_activity_still_reports_a_metadata_lookup_bug_that_is_not_a_database_failure( + lookup: str, bug: Exception +) -> None: + prisma: Final = _daily_user_spend_prisma( + AsyncMock(side_effect=bug if lookup == "key metadata" else None, return_value=[]) + ) + user_emails: Final = AsyncMock(side_effect=bug if lookup == "user emails" else None, return_value={}) + + with pytest.raises(HTTPException) as error: + await _daily_user_activity(prisma, user_emails) + + assert error.value.status_code == 500 + assert error.value.detail == {"error": "Failed to fetch analytics: not a database failure"} async def _run_aggregated_daily_activity( diff --git a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py index d658e59b76c..4ad17272fad 100644 --- a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py +++ b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py @@ -1,15 +1,18 @@ import csv import io from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +from contextlib import contextmanager from dataclasses import dataclass, fields from itertools import chain from types import MappingProxyType, SimpleNamespace from typing import Final from unittest.mock import AsyncMock +import httpx import pytest from fastapi import FastAPI, Request from fastapi.testclient import TestClient +from prisma.errors import DataError from litellm import constants from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY @@ -520,6 +523,14 @@ class _PrismaTable: return next((row for row in self._rows if _matches(row, where)), None) +class _FailingPrismaTable: + def __init__(self, failure: Exception) -> None: + self._failure: Final = failure + + async def find_many(self, **kwargs: object) -> tuple[object, ...]: + raise self._failure + + def _matches(row: object, where: Mapping[str, object]) -> bool: return all( getattr(row, field_name, None) in value["in"] @@ -529,7 +540,7 @@ def _matches(row: object, where: Mapping[str, object]) -> bool: ) -def _prisma_client() -> object: +def _prisma_client(**table_overrides: object) -> object: user: Final = LiteLLM_UserTable( user_id="user-a", user_email="user@example.test", @@ -552,7 +563,7 @@ def _prisma_client() -> object: team_alias="Team B", members_with_roles=[Member(user_id="user-b", role="user")], ) - db: Final = SimpleNamespace( + default_tables: Final = SimpleNamespace( litellm_usertable=_PrismaTable((user, other_user)), litellm_teamtable=_PrismaTable((team, other_team)), litellm_verificationtoken=_PrismaTable((SimpleNamespace(token="key-alpha", user_id="user-a"),)), @@ -581,13 +592,12 @@ def _prisma_client() -> object: ) ), ) + db: Final = SimpleNamespace(**{**vars(default_tables), **table_overrides}) return SimpleNamespace(db=db, writer_db=db) -@pytest.fixture -def daily_activity_client() -> Iterator[tuple[TestClient, _FakeRepository]]: - repository: Final = _FakeRepository(_seeded_activity()) - prisma_client: Final = _prisma_client() +@contextmanager +def _daily_activity_test_client(repository: _FakeRepository, prisma_client: object) -> Iterator[TestClient]: app: Final = FastAPI() app.include_router(router) @@ -607,6 +617,13 @@ def daily_activity_client() -> Iterator[tuple[TestClient, _FakeRepository]]: app.dependency_overrides[get_daily_activity_llm_router] = lambda: None app.dependency_overrides[user_api_key_auth] = resolve_auth with TestClient(app) as client: + yield client + + +@pytest.fixture +def daily_activity_client() -> Iterator[tuple[TestClient, _FakeRepository]]: + repository: Final = _FakeRepository(_seeded_activity()) + with _daily_activity_test_client(repository, _prisma_client()) as client: yield client, repository @@ -1015,6 +1032,178 @@ def test_user_page_route_ranks_users_by_spend_and_attaches_metadata( assert repository.user_metadata.call_args.args[0] == frozenset(("user-a", "user-b")) +_LOOKUP_DATABASE_FAILURES: Final = pytest.mark.parametrize( + "failure", + [ + DataError({"user_facing_error": {"error_code": "P1001", "message": "Can't reach database server"}}), + httpx.ConnectError("the query engine refused the connection"), + ], + ids=["unreachable-database", "engine-connect-error"], +) +_KEY_ROWS_ROUTES: Final = pytest.mark.parametrize( + ("suffix", "extra_params"), + [ + ("keys", {}), + ("search", {"search": "key"}), + ("model_top_keys", {"model_group": "rare-group"}), + ("cache_leakage_keys", {}), + ], + ids=["keys", "search", "model_top_keys", "cache_leakage_keys"], +) + + +@_LOOKUP_DATABASE_FAILURES +@_KEY_ROWS_ROUTES +def test_key_routes_keep_spend_when_the_key_lookup_hits_a_database_failure( + daily_activity_client: tuple[TestClient, _FakeRepository], + suffix: str, + extra_params: Mapping[str, str], + failure: Exception, +) -> None: + client, repository = daily_activity_client + repository.key_metadata.side_effect = failure + + response: Final = client.get( + f"/user/daily/activity/aggregated/{suffix}", + params={**_entity_params("user_id", "user-a"), **extra_params}, + ) + + assert response.status_code == 200, response.text + rows: Final = response.json()["api_keys"] + assert rows + assert all(row["metrics"]["spend"] > 0 for row in rows) + assert {row["metadata"]["key_alias"] for row in rows} == {None} + + +@_KEY_ROWS_ROUTES +def test_key_routes_still_report_a_lookup_bug_that_is_not_a_database_failure( + daily_activity_client: tuple[TestClient, _FakeRepository], + suffix: str, + extra_params: Mapping[str, str], +) -> None: + client, repository = daily_activity_client + repository.key_metadata.side_effect = AttributeError("not a database failure") + + response: Final = client.get( + f"/user/daily/activity/aggregated/{suffix}", + params={**_entity_params("user_id", "user-a"), **extra_params}, + ) + + assert response.status_code == 500, response.text + assert "not a database failure" in response.text + + +@_LOOKUP_DATABASE_FAILURES +def test_user_page_route_keeps_spend_when_the_user_lookup_hits_a_database_failure( + daily_activity_client: tuple[TestClient, _FakeRepository], + failure: Exception, +) -> None: + client, repository = daily_activity_client + repository.user_metadata.side_effect = failure + + response: Final = client.get("/user/daily/activity/aggregated/users", params=_DATE_PARAMS) + + assert response.status_code == 200, response.text + assert [(user["user_id"], user["spend"], user["user_email"]) for user in response.json()["users"]] == [ + ("user-b", 100.0, None), + ("user-a", 14.5, None), + ] + + +def test_user_page_route_still_reports_a_lookup_bug_that_is_not_a_database_failure( + daily_activity_client: tuple[TestClient, _FakeRepository], +) -> None: + client, repository = daily_activity_client + repository.user_metadata.side_effect = AttributeError("not a database failure") + + response: Final = client.get("/user/daily/activity/aggregated/users", params=_DATE_PARAMS) + + assert response.status_code == 500, response.text + assert "not a database failure" in response.text + + +_ENTITY_ALIAS_LOOKUPS: Final = pytest.mark.parametrize( + ("prefix", "query_name", "entity_id", "alias_table", "alias"), + [ + ("/team", "team_ids", "team-a", "litellm_teamtable", {"team_alias": "Team A"}), + ("/organization", "organization_ids", "org-a", "litellm_organizationtable", {"organization_alias": "Org A"}), + ("/customer", "end_user_ids", "customer-a", "litellm_endusertable", {"alias": "Customer A"}), + ("/agent", "agent_ids", "agent-a", "litellm_agentstable", {"agent_name": "Agent A"}), + ], + ids=["team", "organization", "customer", "agent"], +) + + +def _aggregated_entity_metadata( + prefix: str, query_name: str, entity_id: str, prisma_client: object +) -> tuple[httpx.Response, tuple[Mapping[str, object], ...]]: + with _daily_activity_test_client(_FakeRepository(_seeded_activity()), prisma_client) as client: + response: Final = client.get(f"{prefix}/daily/activity/aggregated", params=_entity_params(query_name, entity_id)) + if response.status_code != 200: + return response, () + entities: Final = chain.from_iterable( + result["breakdown"]["entities"].values() for result in response.json()["results"] + ) + return response, tuple(entity["metadata"] for entity in entities) + + +@_LOOKUP_DATABASE_FAILURES +@_ENTITY_ALIAS_LOOKUPS +def test_entity_routes_keep_spend_and_blank_the_alias_when_the_alias_lookup_hits_a_database_failure( + prefix: str, + query_name: str, + entity_id: str, + alias_table: str, + alias: Mapping[str, object], + failure: Exception, +) -> None: + healthy, healthy_metadata = _aggregated_entity_metadata(prefix, query_name, entity_id, _prisma_client()) + failed, failed_metadata = _aggregated_entity_metadata( + prefix, query_name, entity_id, _prisma_client(**{alias_table: _FailingPrismaTable(failure)}) + ) + + assert failed.status_code == 200, failed.text + assert failed.json()["metadata"] == healthy.json()["metadata"] + assert healthy_metadata and set(map(str, healthy_metadata)) == {str(alias)} + assert failed_metadata and set(map(str, failed_metadata)) == {str({})} + + +@_ENTITY_ALIAS_LOOKUPS +def test_entity_routes_still_report_an_alias_lookup_bug_that_is_not_a_database_failure( + prefix: str, + query_name: str, + entity_id: str, + alias_table: str, + alias: Mapping[str, object], +) -> None: + failed, _ = _aggregated_entity_metadata( + prefix, + query_name, + entity_id, + _prisma_client(**{alias_table: _FailingPrismaTable(AttributeError("not a database failure"))}), + ) + + assert failed.status_code == 500, failed.text + assert "not a database failure" in failed.text + + +@_LOOKUP_DATABASE_FAILURES +def test_team_route_fails_closed_for_a_team_member_when_the_team_lookup_hits_a_database_failure( + failure: Exception, +) -> None: + repository: Final = _FakeRepository(_seeded_activity()) + prisma_client: Final = _prisma_client(litellm_teamtable=_FailingPrismaTable(failure)) + with _daily_activity_test_client(repository, prisma_client) as client: + response: Final = client.get( + "/team/daily/activity/aggregated", + params=_entity_params("team_ids", "team-a"), + headers={"x-user-role": LitellmUserRoles.INTERNAL_USER.value, "x-user-id": "user-a"}, + ) + + assert response.status_code == 500, response.text + repository.aggregated.assert_not_called() + + def test_user_page_route_forwards_offset_and_limit( daily_activity_client: tuple[TestClient, _FakeRepository], ) -> None: diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 32c9535b4c8..2fdb77af8f9 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -1085,7 +1085,7 @@ class TestClearCache: Test that clear_cache successfully clears router model caches and reloads models. """ mock_router = MagicMock() - plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}} + plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}} mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__ mock_router.registry.return_value.all.return_value = ["openai/gpt-4o", "openai/gpt-4o-mini"] @@ -1125,7 +1125,7 @@ class TestClearCache: # entries are auto_router/* deployments (so their router-map entries should be # cleared for reload); the config-defined router is preserved. mock_router = MagicMock() - plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}} + plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}} mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__ mock_router.registry.return_value.all.return_value = [ { @@ -1209,7 +1209,7 @@ class TestClearCache: ) mock_router = MagicMock() - plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}} + plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}} mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__ mock_router.registry.return_value.all.return_value = [ { @@ -1270,7 +1270,7 @@ class TestClearCachePreservesConfigRouters: ) mock_router = MagicMock() - plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}} + plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}} mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__ mock_router.registry.return_value.all.return_value = [ { @@ -1315,7 +1315,7 @@ class TestClearCachePreservesConfigRouters: from litellm.proxy.management_endpoints.model_management_endpoints import clear_cache mock_router = MagicMock() - plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}} + plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}} mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__ mock_router.registry.return_value.all.return_value = [ { @@ -1346,15 +1346,15 @@ class TestClearCachePreservesConfigRouters: assert "shared-name" in plugins["complexity"] @pytest.mark.asyncio - async def test_db_quality_and_adaptive_routers_are_evicted(self): - """The auto_router/ prefix also covers quality_router/ and adaptive_router/. Their + async def test_db_quality_adaptive_and_switchyard_routers_are_evicted(self): + """The auto_router/ prefix also covers quality_router/, adaptive_router/ and switchyard. Their registry entries must be popped too, or reload's init raises 'already exists' (quality) or leaves a stale entry (adaptive). """ from litellm.proxy.management_endpoints.model_management_endpoints import clear_cache mock_router = MagicMock() - plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}} + plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}} mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__ mock_router.registry.return_value.all.return_value = [ { @@ -1367,11 +1367,17 @@ class TestClearCachePreservesConfigRouters: "model_info": {"id": "db-a", "db_model": True}, "litellm_params": {"model": "auto_router/adaptive_router/a1"}, }, + { + "model_name": "s1", + "model_info": {"id": "db-s", "db_model": True}, + "litellm_params": {"model": "auto_router/switchyard"}, + }, ] plugins["auto"] = {} plugins["complexity"] = {} plugins["quality"] = {"q1": MagicMock()} plugins["adaptive"] = {"a1": MagicMock()} + plugins["switchyard"] = {"s1": MagicMock()} mock_config = MagicMock() mock_config.add_deployment_locked = AsyncMock( @@ -1389,6 +1395,7 @@ class TestClearCachePreservesConfigRouters: assert "q1" not in plugins["quality"] assert "a1" not in plugins["adaptive"] + assert "s1" not in plugins["switchyard"] class TestDeleteModelClearsRouterRegistry: @@ -1503,7 +1510,7 @@ class TestDeleteModelClearsRouterRegistry: mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) mock_router = MagicMock() - plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}} + plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}} mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__ mock_router.apply.return_value = ApplyResult( version=1, @@ -11431,6 +11438,98 @@ async def test_add_new_model_reports_a_database_outage_during_the_insert_as_503( assert (exc_info.value.code, exc_info.value.type) == ("503", "no_db_connection") +@pytest.mark.parametrize("invalid", [ + {"algorithm": "random"}, + {"classifier": {}}, + {"options": {"capable_system_prompt": "ignored"}}, +]) +def test_switchyard_write_rejects_unloadable_configuration(invalid: dict[str, object]) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import _strategy_router_write_violation + from litellm.types.router import updateLiteLLMParams + + violation: Final = _strategy_router_write_violation( + incoming_params=updateLiteLLMParams( + model="auto_router/switchyard", + switchyard_router_config={"models": {"capable": "strong", "efficient": "fast"}, **invalid}, + ), + existing_params=None, + ) + assert violation is not None + assert "Invalid switchyard_router_config" in violation + + +SWITCHYARD_PARAMS: Final = { + "model": "auto_router/switchyard", + "switchyard_router_config": {"models": {"capable": "strong", "efficient": "fast"}}, +} + + +@pytest.mark.parametrize( + ("model_name", "incoming", "existing"), + [ + ("strong", SWITCHYARD_PARAMS, None), + ("fast", {"model": "auto_router/switchyard"}, SWITCHYARD_PARAMS), + ("fast", None, SWITCHYARD_PARAMS), + ( + "smart", + { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "smart", "COMPLEX": "strong"}}, + }, + None, + ), + ], +) +def test_strategy_router_that_lists_itself_is_rejected_before_saving( + model_name: str, incoming: dict[str, object] | None, existing: dict[str, object] | None +) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _raise_on_strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + with pytest.raises(ProxyException) as exc_info: + _raise_on_strategy_router_write_violation( + incoming_params=updateLiteLLMParams(**incoming) if incoming is not None else None, + existing_params=updateLiteLLMParams(**existing) if existing is not None else None, + model_name=model_name, + ) + assert exc_info.value.code == "400" + assert model_name in exc_info.value.message + + +def test_strategy_router_with_other_targets_saves() -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import _self_routing_violation + from litellm.types.router import updateLiteLLMParams + + assert _self_routing_violation("route", updateLiteLLMParams(**SWITCHYARD_PARAMS), None) is None + + +def test_self_routing_is_caught_when_the_stored_model_is_encrypted(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _raise_on_strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-for-encryption") + stored: Final = updateLiteLLMParams( + **{**SWITCHYARD_PARAMS, "model": encrypt_value_helper("auto_router/switchyard")} + ) + assert stored.model != "auto_router/switchyard" + with pytest.raises(ProxyException) as exc_info: + _raise_on_strategy_router_write_violation( + incoming_params=updateLiteLLMParams( + switchyard_router_config={"models": {"capable": "route", "efficient": "fast"}} + ), + existing_params=stored, + model_name="route", + ) + assert exc_info.value.code == "400" + + @pytest.mark.asyncio async def test_rename_rewrites_only_in_a_transaction_holding_every_locked_teams_advisory_lock(): """A team that starts listing the old name between the advisory snapshot and the row lock must not be diff --git a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py index 531bf98b60e..813875bce7c 100644 --- a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py @@ -16,6 +16,7 @@ from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.proxy_server import app from litellm.proxy.utils import PrismaClient from litellm.types.tag_management import TagNewRequest +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters client = TestClient(app) @@ -230,6 +231,138 @@ async def test_update_tag(): app.dependency_overrides.clear() +@pytest.mark.asyncio +@pytest.mark.parametrize("spend_fields", [{"spend": 0}, {}, {"spend": None}]) +async def test_update_tag_changes_spend_and_counter_only_when_provided( + spend_fields: Mapping[str, float | None], +) -> None: + from litellm.proxy.management_endpoints.tag_management_endpoints import update_tag + from litellm.types.tag_management import TagUpdateRequest + + updated_tag: Final = SimpleNamespace( + tag_name="reset-tag", + description=None, + models=[], + budget_id=None, + created_at=datetime(2024, 1, 1), + updated_at=datetime(2024, 1, 1), + created_by="admin", + ) + update_row: Final = AsyncMock(return_value=updated_tag) + spend: Final = FakeSpendCounters(redis=None) + mock_prisma: Final = SimpleNamespace( + db=SimpleNamespace( + litellm_tagtable=SimpleNamespace( + find_unique=AsyncMock(return_value=updated_tag), update=update_row + ), + litellm_proxymodeltable=SimpleNamespace(find_many=AsyncMock(return_value=[])), + ) + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + _tag_cache_doubles(), + patch("litellm.proxy.proxy_server.invalidate_spend_counter", new_callable=AsyncMock) as invalidate, + patch("litellm.proxy.proxy_server.spend_counter_cache") as counter_cache, + patch("litellm.proxy.proxy_server.spend_counter_store", spend.store()), + ): + counter_cache.redis_cache = None + response: Final = await update_tag( + tag=TagUpdateRequest.model_validate({"name": "reset-tag", **spend_fields}), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert response["message"] == "Tag reset-tag updated successfully" + update_row.assert_awaited_once_with( + where={"tag_name": "reset-tag"}, + data={ + "description": None, + "models": [], + "model_info": "{}", + **({"spend": spend_fields["spend"]} if spend_fields.get("spend") is not None else {}), + }, + ) + if spend_fields.get("spend") is not None: + assert spend.memory["spend:tag:reset-tag"] == 0 + else: + assert "spend:tag:reset-tag" not in spend.memory + counter_cache.in_memory_cache.set_cache.assert_not_called() + invalidate.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("spend", [float("nan"), float("inf"), float("-inf")]) +async def test_update_tag_rejects_non_finite_spend_before_writing(spend: float) -> None: + from litellm.proxy.management_endpoints.tag_management_endpoints import update_tag + from litellm.types.tag_management import TagUpdateRequest + + update_row: Final = AsyncMock() + mock_prisma: Final = SimpleNamespace( + db=SimpleNamespace(litellm_tagtable=SimpleNamespace(update=update_row)) + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.invalidate_spend_counter", new_callable=AsyncMock) as invalidate, + pytest.raises(HTTPException) as error, + ): + await update_tag( + tag=TagUpdateRequest(name="reset-tag", spend=spend), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert error.value.status_code == 400 + assert "spend must be a finite number" in str(error.value.detail) + update_row.assert_not_awaited() + invalidate.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("redis_state", ["available", "absent", "error"]) +async def test_update_tag_overwrites_stale_counter_repopulated_during_tag_cache_eviction(redis_state: str) -> None: + from litellm.proxy.management_endpoints.tag_management_endpoints import update_tag + from litellm.types.tag_management import TagUpdateRequest + + counter_key: Final = "spend:tag:overlap-reset-tag" + spend: Final = FakeSpendCounters(redis=None if redis_state == "absent" else {}, fail=redis_state == "error") + updated_tag: Final = SimpleNamespace( + tag_name="overlap-reset-tag", + description=None, + models=[], + budget_id=None, + created_at=datetime(2024, 1, 1), + updated_at=datetime(2024, 1, 1), + created_by="admin", + ) + mock_prisma: Final = SimpleNamespace( + db=SimpleNamespace( + litellm_tagtable=SimpleNamespace( + find_unique=AsyncMock(return_value=updated_tag), update=AsyncMock(return_value=updated_tag) + ), + litellm_proxymodeltable=SimpleNamespace(find_many=AsyncMock(return_value=[])), + ) + ) + + async def overlapping_completion(cache_keys: Sequence[str]) -> None: + spend.set_local(counter_key, 1.5, ttl_seconds=60) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch( + "litellm.proxy.proxy_server.spend_counter_store", spend.store(), + ), + patch( + "litellm.proxy.management_endpoints.tag_management_endpoints._evict_tag_cache_keys", + overlapping_completion, + ), + ): + await update_tag( + tag=TagUpdateRequest(name="overlap-reset-tag", spend=0), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert spend.memory[counter_key] == 0 + if redis_state == "available": + assert spend.redis is not None + assert spend.redis[counter_key] == 0 + + @pytest.mark.asyncio async def test_new_tag_persists_a_budget(): from datetime import datetime diff --git a/tests/unit/proxy/management_endpoints/test_user_agent_analytics_endpoints.py b/tests/unit/proxy/management_endpoints/test_user_agent_analytics_endpoints.py new file mode 100644 index 00000000000..a56c3095d03 --- /dev/null +++ b/tests/unit/proxy/management_endpoints/test_user_agent_analytics_endpoints.py @@ -0,0 +1,54 @@ +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock + +import pytest + +from litellm.proxy import proxy_server +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import ( + ActiveUsersAnalyticsResponse, + DistinctTagResponse, + DistinctTagsResponse, + TagActiveUsersResponse, + get_daily_active_users, + get_distinct_user_agent_tags, +) + + +def _prisma_returning(rows: list[dict[str, object]]) -> SimpleNamespace: + return SimpleNamespace(db=SimpleNamespace(query_raw=AsyncMock(return_value=rows))) + + +@pytest.mark.asyncio +async def test_distinct_user_agent_tags_are_read_from_the_raw_query_rows(monkeypatch: pytest.MonkeyPatch) -> None: + prisma: Final = _prisma_returning([{"tag": "User-Agent:curl", "usage_count": 7}, {"tag": "prod", "usage_count": 2}]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + response: Final = await get_distinct_user_agent_tags(user_api_key_dict=UserAPIKeyAuth()) + + assert response == DistinctTagsResponse( + results=[DistinctTagResponse(tag="User-Agent:curl"), DistinctTagResponse(tag="prod")] + ) + assert len(prisma.db.query_raw.await_args.args) == 1 + + +@pytest.mark.asyncio +async def test_daily_active_users_forwards_each_tag_filter_as_a_query_parameter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + prisma: Final = _prisma_returning([{"tag": "User-Agent:curl", "active_users": 3, "date": "2026-08-01"}]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + response: Final = await get_daily_active_users( + tag_filter=None, + tag_filters=["User-Agent:curl", "prod"], + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert response == ActiveUsersAnalyticsResponse( + results=[TagActiveUsersResponse(tag="User-Agent:curl", active_users=3, date="2026-08-01")] + ) + sql_query, *params = prisma.db.query_raw.await_args.args + assert params[2:] == ["User-Agent:curl", "prod"] + assert "dts.tag = $3 OR dts.tag = $4" in sql_query diff --git a/tests/unit/proxy/management_helpers/test_management_helpers_utils.py b/tests/unit/proxy/management_helpers/test_management_helpers_utils.py index 644f0656c7a..640e17aa58e 100644 --- a/tests/unit/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/unit/proxy/management_helpers/test_management_helpers_utils.py @@ -249,6 +249,35 @@ async def test_add_new_member_links_default_team_budget_id(): assert create_data["budget_id"] == test_default_budget_id +@pytest.mark.asyncio +async def test_handle_budget_for_entity_updates_with_existing_budget_id_in_request() -> None: + from litellm.proxy.management_helpers.utils import handle_budget_for_entity + from litellm.types.tag_management import TagUpdateRequest + + budget_update: Final = AsyncMock() + prisma_client: Final = cast( + PrismaClient, + SimpleNamespace(db=SimpleNamespace(litellm_budgettable=SimpleNamespace(update=budget_update))), + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + ): + budget_id: Final = await handle_budget_for_entity( + data=TagUpdateRequest(name="budget-tag", budget_id="budget-1", max_budget=20), + existing_budget_id="budget-1", + user_api_key_dict=UserAPIKeyAuth(user_id="admin"), + prisma_client=prisma_client, + litellm_proxy_admin_name="admin", + ) + + assert budget_id == "budget-1" + budget_update.assert_awaited_once_with( + where={"budget_id": "budget-1"}, + data={"budget_id": "budget-1", "max_budget": 20.0, "updated_by": "admin"}, + ) + + @pytest.mark.parametrize( ("request_fields", "cleared_budget_fields", "should_update", "expected_max_budget"), cast( diff --git a/tests/unit/proxy/management_helpers/test_object_permission_utils.py b/tests/unit/proxy/management_helpers/test_object_permission_utils.py index 5814d8d2a97..12fb390e995 100644 --- a/tests/unit/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/unit/proxy/management_helpers/test_object_permission_utils.py @@ -1054,6 +1054,26 @@ async def test_resolve_team_allowed_mcp_servers_dict_tool_permissions( assert result == {"server-a"} +@pytest.mark.asyncio +@pytest.mark.parametrize("stored_tool_permissions", ["[]", "null", '"server-2"'], ids=["array", "null", "string"]) +@patch( + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_resolve_team_allowed_mcp_servers_rejects_stored_tool_permissions_that_are_not_an_object( + mock_access_groups, + stored_tool_permissions: str, +): + mock_perm = MagicMock(spec=LiteLLM_ObjectPermissionTable) + mock_perm.mcp_servers = ["server-1"] + mock_perm.mcp_access_groups = [] + mock_perm.mcp_tool_permissions = stored_tool_permissions + + with pytest.raises(AttributeError, match="object has no attribute 'keys'"): + await _resolve_team_allowed_mcp_servers(mock_perm) + + # ---- Tests for the all-proxy-mcpservers sentinel (team scoped to every server) ---- diff --git a/tests/unit/proxy/mcp_server/test_semantic_tool_filter.py b/tests/unit/proxy/mcp_server/test_semantic_tool_filter.py index c04323d7f7c..65e59fab0da 100644 --- a/tests/unit/proxy/mcp_server/test_semantic_tool_filter.py +++ b/tests/unit/proxy/mcp_server/test_semantic_tool_filter.py @@ -2238,3 +2238,62 @@ async def test_semantic_filter_hook_narrows_only_references_the_gateway_serves() ) assert narrowed == [{**gateway_reference, "allowed_tools": ["srv-tool_1"]}, external_tool] + + +@pytest.mark.asyncio +async def test_expand_mcp_tools_returns_every_expanded_tool_as_a_dict(): + from typing import Optional + + from pydantic import BaseModel, ConfigDict + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.proxy.mcp_server.semantic_tool_filter import SemanticMCPToolFilter + from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler + + class _FunctionTool(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str + name: str + description: Optional[str] = None + + class _LegacyFunctionTool: + def __init__(self, name: str) -> None: + self.name = name + + def dict(self, *, exclude_none: bool) -> dict[str, object]: + fields = {"type": "function", "name": self.name, "description": None} + return {key: value for key, value in fields.items() if value is not None or not exclude_none} + + plain_tool = {"type": "function", "name": "plain_tool", "parameters": {"type": "object"}} + unknown_tool = "opaque-tool" + pipeline_output = [ + _FunctionTool(type="function", name="v2_tool"), + _LegacyFunctionTool("legacy_tool"), + plain_tool, + unknown_tool, + ] + auth = UserAPIKeyAuth(api_key="hashed-key") + gateway_reference = {"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"} + + hook = SemanticToolFilterHook( + SemanticMCPToolFilter(embedding_model="text-embedding-3-small", litellm_router_instance=Mock()) + ) + with patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_to_openai_format", + new=AsyncMock(return_value=(pipeline_output, {})), + ) as process_tools: + expanded = await hook._expand_mcp_tools([gateway_reference], auth) + + assert expanded == [ + {"type": "function", "name": "v2_tool"}, + {"type": "function", "name": "legacy_tool"}, + plain_tool, + unknown_tool, + ] + assert process_tools.await_args.kwargs == { + "user_api_key_auth": auth, + "mcp_tools_with_litellm_proxy": [gateway_reference], + } diff --git a/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py index 354c65776dc..1a4751b4598 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py @@ -1,17 +1,39 @@ import asyncio +import base64 import json -from collections.abc import Iterator +import time +from collections.abc import Callable, Iterator, Sequence from datetime import datetime from typing import Final import httpx +import jwt import pytest import respx +from cryptography.hazmat.primitives.asymmetric import rsa from fastapi.testclient import TestClient +from jwt.algorithms import RSAAlgorithm import litellm +from litellm import Router +from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.chatgpt.common_utils import OPENAI_AUTH_CLAIM_KEY from litellm.proxy import proxy_server +from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_TeamTableCachedObj, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.handle_jwt import JWTHandler +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.codex_model_catalog import bundled_codex_models +from litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints import ( + _ACCOUNT_SLUGS_FETCHES_PER_CALLER, + codex_shared_cache, +) +from litellm.proxy.utils import hash_token +from litellm.responses.litellm_completion_transformation.reasoning_items import ( + encode_thinking_blocks, + mint_reasoning_item_id, +) +from litellm.responses.utils import ResponsesAPIRequestUtils from tests._master_key import MASTER_KEY _BASE: Final = "https://chatgpt.com/backend-api/codex" @@ -22,6 +44,7 @@ _HEADERS: Final = {"authorization": f"Bearer {_TOKEN}", "x-litellm-api-key": MAS @pytest.fixture def client(monkeypatch: pytest.MonkeyPatch, httpx_transport: None) -> Iterator[TestClient]: monkeypatch.setattr(proxy_server, "master_key", MASTER_KEY) + monkeypatch.setattr(proxy_server.proxy_logging_obj.internal_usage_cache, "dual_cache", DualCache()) yield TestClient(proxy_server.app) @@ -183,3 +206,1860 @@ async def test_usage_logged_without_credentials( serialized: Final = json.dumps(logged, default=str) for secret in (_TOKEN, MASTER_KEY, "test-hidden-token", "test-hidden-cookie"): assert secret not in serialized + + +_PROVIDER_BASE: Final = "https://open-weights.test/v1" +_CHATGPT_JWT: Final = ( + ".".join( + base64.urlsafe_b64encode(json.dumps(part).encode()).decode().rstrip("=") + for part in ({"alg": "none"}, {OPENAI_AUTH_CLAIM_KEY: {"chatgpt_account_id": "account-a"}}) + ) + + "." +) +_PROVIDER_KEY: Final = "test-provider-key" +_CODEX_USER_AGENT: Final = "codex_cli_rs/0.165.0 (Ubuntu 24.4.0; x86_64) xterm-256color" +_ACCOUNT_HEADERS: Final = { + **_HEADERS, + "authorization": f"Bearer {_CHATGPT_JWT}", + "chatgpt-account-id": "account-a", + "originator": "codex_cli_rs", + "user-agent": _CODEX_USER_AGENT, +} +_BUNDLED_SLUG: Final = next(iter(bundled_codex_models())) +_PRICED_CHATGPT_SLUG: Final = "priced-subscription-model" +_UPSTREAM_CATALOG: Final = { + "models": [ + {"slug": "subscription-model", "priority": 0, "supported_reasoning_levels": ["high"]}, + {"slug": "shared-name", "priority": 7, "visibility": "list"}, + {"slug": "subscription-hidden", "priority": 3, "visibility": "hide"}, + ] +} + + +def _open_weights(model_name: str, mode: str | None = None, upstream: str | None = None) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": f"fireworks_ai/{upstream or model_name}", + "api_base": _PROVIDER_BASE, + "api_key": _PROVIDER_KEY, + }, + "model_info": {"mode": mode} if mode else {}, + } + + +@pytest.fixture +def mixed_router(monkeypatch: pytest.MonkeyPatch) -> Router: + router: Final = Router( + model_list=[ + { + "model_name": "chatgpt/gpt-*", + "litellm_params": {"model": "chatgpt/gpt-*"}, + "model_info": {"access_groups": ["chatgpt-native"]}, + }, + { + "model_name": "chatgpt/codex-*", + "litellm_params": {"model": "chatgpt/codex-*"}, + "model_info": {"access_groups": ["chatgpt-native"]}, + }, + {"model_name": "chatgpt-subscription", "litellm_params": {"model": "chatgpt/subscription-model"}}, + _open_weights("open-weights"), + _open_weights("shared-name"), + _open_weights(_BUNDLED_SLUG), + _open_weights(_PRICED_CHATGPT_SLUG), + _open_weights("open-embedder", mode="embedding"), + _open_weights("*"), + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + unpriced_bundled: Final = { + name: price for name, price in litellm.model_cost.items() if name != f"chatgpt/{_BUNDLED_SLUG}" + } + monkeypatch.setattr( + litellm, + "model_cost", + {**unpriced_bundled, f"chatgpt/{_PRICED_CHATGPT_SLUG}": {"litellm_provider": "chatgpt", "mode": "responses"}}, + ) + return router + + +def test_native_catalog_lists_the_keys_other_chat_models_after_chatgpts( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router +) -> None: + respx_mock.get(f"{_BASE}/models").respond(json=_UPSTREAM_CATALOG, headers={"etag": '"upstream-v1"'}) + response: Final = client.get("/backend-api/codex/models?client_version=1.0.0", headers=_HEADERS) + assert response.status_code == 200 + models: Final = response.json()["models"] + assert models[: len(_UPSTREAM_CATALOG["models"])] == _UPSTREAM_CATALOG["models"] + added: Final = models[len(_UPSTREAM_CATALOG["models"]) :] + assert [model["slug"] for model in added] == ["open-weights"] + assert added[0]["visibility"] == "list" + assert added[0]["priority"] > max(model["priority"] for model in _UPSTREAM_CATALOG["models"]) + + +@pytest.mark.parametrize( + ("client_version", "expected_efforts"), + [ + pytest.param("0.165.0", ["low", "max"], id="a codex that parses custom efforts gets them all"), + pytest.param("0.137.0", ["low"], id="an older codex gets only the efforts it parses"), + ], +) +def test_native_catalog_offers_the_declared_reasoning_efforts_the_requesting_codex_parses( + client: TestClient, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + client_version: str, + expected_efforts: list[str], +) -> None: + declared: Final = { + **_open_weights("open-weights"), + "model_info": {"supports_reasoning": True, "reasoning_effort_levels": ["low", "max"]}, + } + monkeypatch.setattr(proxy_server, "llm_router", Router(model_list=[declared])) + respx_mock.get(f"{_BASE}/models").respond(json=_UPSTREAM_CATALOG, headers={"etag": '"upstream-v1"'}) + response: Final = client.get(f"/backend-api/codex/models?client_version={client_version}", headers=_HEADERS) + assert response.status_code == 200 + (added,) = response.json()["models"][len(_UPSTREAM_CATALOG["models"]) :] + assert [level["effort"] for level in added["supported_reasoning_levels"]] == expected_efforts + + +@pytest.mark.parametrize( + "upstream", + [ + pytest.param(httpx.Response(503, json={"detail": "busy"}, headers={"etag": '"e"'}), id="non-200"), + pytest.param(httpx.Response(200, content=b"login", headers={"etag": '"e"'}), id="not json"), + pytest.param(httpx.Response(200, json={"models": "none"}, headers={"etag": '"e"'}), id="no model list"), + ], +) +def test_native_catalog_relays_an_unmergeable_upstream_answer_unchanged( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router, upstream: httpx.Response +) -> None: + respx_mock.get(f"{_BASE}/models").mock(return_value=upstream) + response: Final = client.get("/backend-api/codex/models?client_version=1.0.0", headers=_HEADERS) + assert response.status_code == upstream.status_code + assert response.content == upstream.content + assert response.headers["etag"] == '"e"' + + +def _provider_completed(text: str) -> dict[str, object]: + return { + "id": "resp_open_weights", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "open-weights", + "output": [ + { + "type": "message", + "id": "msg_open_weights", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ], + "usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4}, + } + + +def _provider_stream(text: str) -> bytes: + completed: Final = _provider_completed(text) + events: Final = ( + { + "type": "response.created", + "sequence_number": 0, + "response": {**completed, "status": "in_progress", "output": []}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_open_weights", + "output_index": 0, + "content_index": 0, + "delta": text, + }, + {"type": "response.completed", "sequence_number": 2, "response": completed}, + ) + return "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events).encode() + + +def _event_types(stream: bytes) -> list[str]: + payloads: Final = [line.removeprefix(b"data: ") for line in stream.splitlines() if line.startswith(b"data: {")] + return [json.loads(payload)["type"] for payload in payloads] + + +class PreCallRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.requests: Final[list[tuple[str | None, str]]] = [] + self.raw_header_names: Final[list[frozenset[str]]] = [] + + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: object, data: dict[str, object], call_type: str + ) -> dict[str, object]: + self.requests.append((user_api_key_dict.request_route, repr(data))) + secret_fields: Final = data.get("secret_fields") + raw_headers: Final = secret_fields.get("raw_headers") if isinstance(secret_fields, dict) else None + self.raw_header_names.append(frozenset(raw_headers) if isinstance(raw_headers, dict) else frozenset()) + return data + + +def _provider_reply(request: httpx.Request) -> httpx.Response: + if json.loads(request.content).get("stream") is not True: + return httpx.Response(200, json=_provider_completed("Hi")) + return httpx.Response(200, content=_provider_stream("Hi"), headers={"content-type": "text/event-stream"}) + + +def _open_weights_provider(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post(f"{_PROVIDER_BASE}/responses").mock(side_effect=_provider_reply) + + +_UPSTREAM_MODELS_ETAG: Final = '"upstream-v2"' + + +def _chatgpt_responses(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post(f"{_BASE}/responses").respond( + content=b"event: response.completed\ndata: {}\n\n", + headers={"content-type": "text/event-stream", "x-models-etag": _UPSTREAM_MODELS_ETAG}, + ) + + +def _account_catalog(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.get(f"{_BASE}/models").respond(json=_UPSTREAM_CATALOG) + + +def test_litellm_model_on_native_route_streams_like_v1_without_chatgpt_credentials( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router, monkeypatch: pytest.MonkeyPatch +) -> None: + recorder: Final = PreCallRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + _account_catalog(respx_mock) + chatgpt: Final = respx_mock.post(f"{_BASE}/responses").respond(json={}) + provider: Final = respx_mock.post(f"{_PROVIDER_BASE}/responses").mock( + side_effect=lambda _: httpx.Response( + 200, content=_provider_stream("Hi"), headers={"content-type": "text/event-stream"} + ) + ) + body: Final = {"model": "open-weights", "input": "Say hi", "stream": True, "store": False} + native: Final = client.post("/backend-api/codex/responses", json=body, headers=_ACCOUNT_HEADERS) + v1: Final = client.post("/v1/responses", json=body, headers={"authorization": f"Bearer {MASTER_KEY}"}) + assert native.status_code == v1.status_code == 200 + assert native.headers["content-type"].startswith("text/event-stream") + assert _event_types(native.content) == _event_types(v1.content) + assert "response.completed" in _event_types(native.content) + assert not chatgpt.called + provider_headers: Final = provider.calls[0].request.headers + assert provider_headers["authorization"] == f"Bearer {_PROVIDER_KEY}" + assert "chatgpt-account-id" not in provider_headers + assert all(_CHATGPT_JWT not in value for value in provider_headers.values()) + (native_route, native_call), (v1_route, _) = recorder.requests + assert native_route == v1_route == "/v1/responses" + assert _CHATGPT_JWT not in native_call + assert "account-a" not in native_call + + +def test_litellm_model_on_native_route_keeps_chatgpt_credentials_out_of_hook_raw_headers( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router, monkeypatch: pytest.MonkeyPatch +) -> None: + recorder: Final = PreCallRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + _account_catalog(respx_mock) + provider: Final = _open_weights_provider(respx_mock) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "open-weights", "input": "Say hi", "stream": True, "store": False}, + headers=_ACCOUNT_HEADERS, + ) + assert response.status_code == 200 + assert provider.called + (raw_header_names,) = recorder.raw_header_names + assert "x-litellm-api-key" in raw_header_names + assert not raw_header_names & {"authorization", "chatgpt-account-id"}, sorted(raw_header_names) + + +@pytest.mark.parametrize( + ("model", "sent", "drop_params", "expected"), + [ + pytest.param("open-weights", "none", None, "none", id="none without drop_params"), + pytest.param("open-weights", "none", "global", None, id="none under litellm_settings drop_params"), + pytest.param("open-weights", "none", "deployment", None, id="none under the deployment's drop_params"), + pytest.param("open-weights", "none", "request", None, id="none under the request's drop_params"), + pytest.param("open-weights", "high", "global", "high", id="effort set for a model without a default effort"), + pytest.param("coding-model", "none", "global", None, id="none on an alias whose native entry names no effort"), + pytest.param("api-model", "none", "global", "none", id="none on a model that supports it"), + pytest.param("glm-model", "none", "global", None, id="none on a model that does not support it"), + pytest.param("declared-model", "none", "global", "none", id="none on a deployment whose model_info allows it"), + ], +) +def test_litellm_model_on_native_route_gets_the_reasoning_effort_codex_m_would_send( + client: TestClient, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + model: str, + sent: str, + drop_params: str | None, + expected: str | None, +) -> None: + monkeypatch.setattr(litellm, "drop_params", drop_params == "global") + open_weights: Final = _open_weights("open-weights") + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + **open_weights, + "litellm_params": {**open_weights["litellm_params"], "drop_params": drop_params == "deployment"}, + }, + { + **_open_weights("coding-model", upstream="gpt-5.5"), + "model_info": {"supports_none_reasoning_effort": False}, + }, + _open_weights("glm-model", upstream="glm-5p3"), + {**_open_weights("declared-model"), "model_info": {"supports_none_reasoning_effort": True}}, + { + "model_name": "api-model", + "litellm_params": {"model": "openai/gpt-5.4", "api_base": _PROVIDER_BASE, "api_key": _PROVIDER_KEY}, + }, + ] + ), + ) + _account_catalog(respx_mock) + provider: Final = respx_mock.post(f"{_PROVIDER_BASE}/responses").mock( + side_effect=lambda _: httpx.Response( + 200, content=_provider_stream("Hi"), headers={"content-type": "text/event-stream"} + ) + ) + body: Final = { + "model": model, + "input": "Say hi", + "stream": True, + "store": False, + **({"drop_params": True} if drop_params == "request" else {}), + } + expected_reasoning: Final = {"summary": "auto"} if expected is None else {"effort": expected, "summary": "auto"} + native: Final = client.post( + "/backend-api/codex/responses", + json={**body, "reasoning": {"effort": sent, "summary": "auto"}}, + headers=_ACCOUNT_HEADERS, + ) + v1: Final = client.post( + "/v1/responses", + json={**body, "reasoning": expected_reasoning}, + headers={"authorization": f"Bearer {MASTER_KEY}"}, + ) + assert native.status_code == v1.status_code == 200 + native_request, v1_request = (json.loads(call.request.content) for call in provider.calls) + assert native_request == v1_request + assert native_request["reasoning"].get("effort") == expected + + +@pytest.mark.parametrize( + "model", ["gpt-unlisted", "chatgpt-subscription", _BUNDLED_SLUG, _PRICED_CHATGPT_SLUG, "newly-released"] +) +def test_chatgpt_model_keeps_passthrough_without_other_providers_reasoning( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router, model: str +) -> None: + provider: Final = respx_mock.post(f"{_PROVIDER_BASE}/responses").respond(json={}) + chatgpt: Final = respx_mock.post(f"{_BASE}/responses").respond( + content=b"event: response.completed\ndata: {}\n\n", + headers={"content-type": "text/event-stream", "x-models-etag": '"upstream-v2"'}, + ) + user_turn: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Hello"}]} + chatgpt_reasoning: Final = { + "type": "reasoning", + "id": "rs_upstream", + "summary": [], + "encrypted_content": "opaque-upstream-blob", + } + assistant_turn: Final = {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Hi"}]} + other_providers_reasoning: Final = [ + {"type": "reasoning", "id": mint_reasoning_item_id(), "summary": []}, + { + "type": "reasoning", + "summary": [], + "encrypted_content": encode_thinking_blocks([{"type": "thinking", "thinking": "t", "signature": "s"}]), + }, + {"type": "reasoning", "summary": [], "content": [{"type": "reasoning_text", "text": "t"}]}, + { + "type": "reasoning", + "summary": [], + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("blob", "deployment"), + }, + ] + body: Final = { + "model": model, + "input": [user_turn, *other_providers_reasoning, chatgpt_reasoning, assistant_turn], + "stream": True, + } + response: Final = client.post("/backend-api/codex/responses", json=body, headers=_HEADERS) + assert response.status_code == 200 + assert not provider.called + assert json.loads(chatgpt.calls.last.request.content) == { + **body, + "input": [user_turn, chatgpt_reasoning, assistant_turn], + } + + +@pytest.fixture +def restricted_key(monkeypatch: pytest.MonkeyPatch) -> Callable[..., str]: + monkeypatch.setattr(proxy_server, "prisma_client", "test-database") + + def mint( + models: list[str], + allowed_routes: tuple[str, ...] = (), + user_role: LitellmUserRoles | None = None, + allowed_tools: tuple[str, ...] = (), + ) -> str: + key: Final = f"sk-test-{'-'.join((*models, *allowed_routes, str(user_role), *allowed_tools))}" + proxy_server.user_api_key_cache.set_cache( + key=hash_token(key), + value=UserAPIKeyAuth( + token=hash_token(key), + models=models, + allowed_routes=list(allowed_routes), + user_role=user_role, + metadata={"allowed_tools": list(allowed_tools)} if allowed_tools else {}, + ), + ) + return key + + return mint + + +@pytest.mark.parametrize( + ("models", "added"), + [ + pytest.param(["chatgpt-native", "open-weights"], ["open-weights"], id="key with a LiteLLM model"), + pytest.param(["chatgpt-native"], [], id="key with ChatGPT models only"), + ], +) +def test_native_catalog_adds_only_the_litellm_models_the_key_can_call( + client: TestClient, + respx_mock: respx.MockRouter, + mixed_router: Router, + restricted_key: Callable[..., str], + models: list[str], + added: list[str], +) -> None: + respx_mock.get(f"{_BASE}/models").respond(json=_UPSTREAM_CATALOG, headers={"etag": '"upstream-v1"'}) + response: Final = client.get( + "/backend-api/codex/models?client_version=1.0.0", + headers={**_HEADERS, "x-litellm-api-key": restricted_key(models)}, + ) + assert response.status_code == 200 + assert [model["slug"] for model in response.json()["models"]] == [ + *(model["slug"] for model in _UPSTREAM_CATALOG["models"]), + *added, + ] + assert (response.headers["etag"] == '"upstream-v1"') is not added + + +@pytest.mark.parametrize( + ("models", "model", "allowed"), + [ + pytest.param(["chatgpt-native", "open-weights"], "gpt-unlisted", True, id="chatgpt access group"), + pytest.param(["chatgpt-subscription", "open-weights"], "gpt-unlisted", False, id="single subscription model"), + pytest.param(["chatgpt-native"], "open-weights", False, id="litellm model outside the key"), + ], +) +def test_native_route_serves_only_the_models_the_key_allows( + client: TestClient, + respx_mock: respx.MockRouter, + mixed_router: Router, + restricted_key: Callable[..., str], + models: list[str], + model: str, + allowed: bool, +) -> None: + provider: Final = respx_mock.post(f"{_PROVIDER_BASE}/responses").respond(json={}) + chatgpt: Final = respx_mock.post(f"{_BASE}/responses").respond( + content=b"event: response.completed\ndata: {}\n\n", headers={"content-type": "text/event-stream"} + ) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": model, "input": "Hello", "stream": True}, + headers={**_HEADERS, "x-litellm-api-key": restricted_key(models)}, + ) + assert (response.status_code == 200) is allowed + assert chatgpt.called is allowed + assert not provider.called + assert [json.loads(call.request.content)["model"] for call in chatgpt.calls] == ([model] if allowed else []) + + +def _key_with_alias_to_budgeted_target(monkeypatch: pytest.MonkeyPatch, target_budget: float) -> str: + monkeypatch.setattr(proxy_server, "prisma_client", "test-database") + monkeypatch.setattr( + proxy_server, "llm_router", Router(model_list=[_open_weights("group-a"), _open_weights("group-b")]) + ) + key: Final = f"sk-test-alias-budget-{target_budget}" + proxy_server.user_api_key_cache.set_cache( + key=hash_token(key), + value=UserAPIKeyAuth( + token=hash_token(key), + models=["group-a", "group-b"], + router_settings={"model_group_alias": {"group-a": "group-b"}}, + model_max_budget={"group-b": {"max_budget": target_budget, "budget_duration": "1d"}}, + ), + ) + return key + + +def test_native_route_refuses_a_saved_alias_when_its_target_is_over_budget( + client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + key: Final = _key_with_alias_to_budgeted_target(monkeypatch, target_budget=0.0) + _account_catalog(respx_mock) + provider: Final = _open_weights_provider(respx_mock) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "group-a", "input": "Hello"}, + headers={**_ACCOUNT_HEADERS, "x-litellm-api-key": key}, + ) + assert response.json()["error"]["type"] == "budget_exceeded", response.text + assert "model=group-b" in response.json()["error"]["message"] + assert not provider.called + + +def test_native_route_serves_a_saved_alias_target_within_budget_under_the_requested_name( + client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + key: Final = _key_with_alias_to_budgeted_target(monkeypatch, target_budget=5.0) + _account_catalog(respx_mock) + provider: Final = _open_weights_provider(respx_mock) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "group-a", "input": "Hello"}, + headers={**_ACCOUNT_HEADERS, "x-litellm-api-key": key}, + ) + assert response.status_code == 200, response.text + assert response.json()["model"] == "group-a" + assert json.loads(provider.calls.last.request.content)["model"].endswith("/group-b") + + +def test_native_route_keeps_a_chatgpt_turn_on_the_subscription_despite_a_saved_alias( + client: TestClient, + respx_mock: respx.MockRouter, + mixed_router: Router, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(proxy_server, "prisma_client", "test-database") + key: Final = "sk-test-alias-on-a-chatgpt-name" + proxy_server.user_api_key_cache.set_cache( + key=hash_token(key), + value=UserAPIKeyAuth( + token=hash_token(key), + models=[_BUNDLED_SLUG, "open-weights"], + router_settings={"model_group_alias": {_BUNDLED_SLUG: "open-weights"}}, + ), + ) + provider: Final = _open_weights_provider(respx_mock) + chatgpt: Final = _chatgpt_responses(respx_mock) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": _BUNDLED_SLUG, "input": "Hello", "stream": True}, + headers={**_HEADERS, "x-litellm-api-key": key}, + ) + assert response.status_code == 200, response.text + assert not provider.called + assert [json.loads(call.request.content)["model"] for call in chatgpt.calls] == [_BUNDLED_SLUG] + + +def test_native_route_keeps_a_name_the_account_lists_on_the_subscription_despite_a_saved_alias( + client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(proxy_server, "prisma_client", "test-database") + monkeypatch.setattr( + proxy_server, "llm_router", Router(model_list=[_open_weights("gpt-reserve"), _open_weights("group-b")]) + ) + key: Final = "sk-test-alias-on-a-name-the-account-lists" + proxy_server.user_api_key_cache.set_cache( + key=hash_token(key), + value=UserAPIKeyAuth( + token=hash_token(key), + models=["gpt-reserve", "group-b"], + router_settings={"model_group_alias": {"gpt-reserve": "group-b"}}, + model_max_budget={"group-b": {"max_budget": 0.0, "budget_duration": "1d"}}, + ), + ) + catalog: Final = respx_mock.get(f"{_BASE}/models").respond(json={"models": [{"slug": "gpt-reserve"}]}) + provider: Final = _open_weights_provider(respx_mock) + chatgpt: Final = _chatgpt_responses(respx_mock) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "gpt-reserve", "input": "Hello", "stream": True}, + headers={**_ACCOUNT_HEADERS, "x-litellm-api-key": key}, + ) + assert response.status_code == 200, response.text + assert not provider.called + assert [json.loads(call.request.content)["model"] for call in chatgpt.calls] == ["gpt-reserve"] + assert catalog.call_count == 1 + + +_BEARER_KEY: Final = "sk-test-alias-key-sent-as-the-bearer" + + +@pytest.mark.parametrize( + "gateway_header", + [{}, {"x-litellm-api-key": _BEARER_KEY}], + ids=["key only in the bearer", "same key in both headers"], +) +def test_native_route_never_sends_a_litellm_key_sent_as_the_bearer_to_chatgpt( + client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, gateway_header: dict[str, str] +) -> None: + monkeypatch.setattr(proxy_server, "prisma_client", "test-database") + monkeypatch.setattr( + proxy_server, "llm_router", Router(model_list=[_open_weights("gpt-reserve"), _open_weights("group-b")]) + ) + proxy_server.user_api_key_cache.set_cache( + key=hash_token(_BEARER_KEY), + value=UserAPIKeyAuth( + token=hash_token(_BEARER_KEY), + models=["gpt-reserve", "group-b"], + router_settings={"model_group_alias": {"gpt-reserve": "group-b"}}, + ), + ) + catalog: Final = respx_mock.get(f"{_BASE}/models").respond(json={"models": []}) + provider: Final = _open_weights_provider(respx_mock) + chatgpt: Final = _chatgpt_responses(respx_mock) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "gpt-reserve", "input": "Hello", "stream": True}, + headers={ + "authorization": f"Bearer {_BEARER_KEY}", + "chatgpt-account-id": "account-a", + "originator": "codex_cli_rs", + "user-agent": _CODEX_USER_AGENT, + **gateway_header, + }, + ) + assert response.status_code == 401, response.text + assert not catalog.called + assert not provider.called + assert not chatgpt.called + + +def test_v1_request_for_a_gpt_name_still_reaches_the_catch_all_with_the_access_group_recipe( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router +) -> None: + provider: Final = _open_weights_provider(respx_mock) + chatgpt: Final = _chatgpt_responses(respx_mock) + response: Final = client.post( + "/v1/responses", + json={"model": "gpt-unlisted", "input": "Hello", "stream": True}, + headers={"authorization": f"Bearer {MASTER_KEY}"}, + ) + assert response.status_code == 200, response.text + assert provider.called + assert not chatgpt.called + + +_NATIVE_ROUTES: Final = ("/backend-api/codex/models", "/backend-api/codex/responses") + + +@pytest.mark.parametrize( + ("allowed_routes", "added"), + [ + pytest.param(_NATIVE_ROUTES, [], id="key limited to the native routes"), + pytest.param((*_NATIVE_ROUTES, "/v1/responses"), ["open-weights"], id="key that may call /v1/responses"), + ], +) +def test_native_catalog_adds_litellm_models_only_for_a_key_that_may_call_v1_responses( + client: TestClient, + respx_mock: respx.MockRouter, + mixed_router: Router, + restricted_key: Callable[..., str], + allowed_routes: tuple[str, ...], + added: list[str], +) -> None: + respx_mock.get(f"{_BASE}/models").respond(json=_UPSTREAM_CATALOG) + response: Final = client.get( + "/backend-api/codex/models?client_version=1.0.0", + headers={**_HEADERS, "x-litellm-api-key": restricted_key(["chatgpt-native", "open-weights"], allowed_routes)}, + ) + assert response.status_code == 200, response.text + assert [model["slug"] for model in response.json()["models"]] == [ + *(model["slug"] for model in _UPSTREAM_CATALOG["models"]), + *added, + ] + + +@pytest.mark.parametrize( + ("allowed_routes", "through_litellm"), + [ + pytest.param(_NATIVE_ROUTES, False, id="key limited to the native routes"), + pytest.param((*_NATIVE_ROUTES, "/v1/responses"), True, id="key that may call /v1/responses"), + ], +) +def test_native_responses_reach_litellm_models_only_for_a_key_that_may_call_v1_responses( + client: TestClient, + respx_mock: respx.MockRouter, + mixed_router: Router, + restricted_key: Callable[..., str], + allowed_routes: tuple[str, ...], + through_litellm: bool, +) -> None: + _account_catalog(respx_mock) + provider: Final = _open_weights_provider(respx_mock) + chatgpt: Final = _chatgpt_responses(respx_mock) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "open-weights", "input": "Hello", "stream": True}, + headers={ + **_ACCOUNT_HEADERS, + "x-litellm-api-key": restricted_key(["chatgpt-native", "open-weights"], allowed_routes), + }, + ) + assert response.status_code == 200, response.text + assert provider.called is through_litellm + assert chatgpt.called is not through_litellm + + + +@pytest.mark.parametrize( + ("user_role", "through_litellm"), + [ + pytest.param(LitellmUserRoles.INTERNAL_USER, False, id="non-admin key"), + pytest.param(LitellmUserRoles.PROXY_ADMIN, True, id="proxy admin key"), + ], +) +def test_admin_only_v1_responses_keeps_litellm_models_to_proxy_admins_on_native_routes( + client: TestClient, + respx_mock: respx.MockRouter, + mixed_router: Router, + restricted_key: Callable[..., str], + monkeypatch: pytest.MonkeyPatch, + user_role: LitellmUserRoles, + through_litellm: bool, +) -> None: + monkeypatch.setattr(proxy_server, "general_settings", {"admin_only_routes": ["/v1/responses"]}) + monkeypatch.setattr(proxy_server, "premium_user", True) + _account_catalog(respx_mock) + provider: Final = _open_weights_provider(respx_mock) + chatgpt: Final = _chatgpt_responses(respx_mock) + key: Final = restricted_key(["chatgpt-native", "open-weights"], user_role=user_role) + catalog: Final = client.get( + "/backend-api/codex/models?client_version=1.0.0", headers={**_ACCOUNT_HEADERS, "x-litellm-api-key": key} + ) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "open-weights", "input": "Hello", "stream": True}, + headers={**_ACCOUNT_HEADERS, "x-litellm-api-key": key}, + ) + assert (catalog.status_code, response.status_code) == (200, 200), response.text + assert ("open-weights" in [model["slug"] for model in catalog.json()["models"]]) is through_litellm + assert provider.called is through_litellm + assert chatgpt.called is not through_litellm + + +@pytest.mark.parametrize( + ("allowed_routes", "through_litellm"), + [ + pytest.param(["/backend-api/codex/models", "/backend-api/codex/responses"], False, id="native routes only"), + pytest.param( + ["/backend-api/codex/models", "/backend-api/codex/responses", "/v1/responses"], + True, + id="v1 responses allowed too", + ), + ], +) +def test_proxy_allowed_routes_keep_litellm_models_off_native_routes_unless_v1_responses_is_allowed( + client: TestClient, + respx_mock: respx.MockRouter, + mixed_router: Router, + restricted_key: Callable[..., str], + monkeypatch: pytest.MonkeyPatch, + allowed_routes: list[str], + through_litellm: bool, +) -> None: + monkeypatch.setattr(proxy_server, "general_settings", {"allowed_routes": allowed_routes}) + monkeypatch.setattr(proxy_server, "premium_user", True) + _account_catalog(respx_mock) + provider: Final = _open_weights_provider(respx_mock) + chatgpt: Final = _chatgpt_responses(respx_mock) + key: Final = restricted_key(["chatgpt-native", "open-weights"]) + catalog: Final = client.get( + "/backend-api/codex/models?client_version=1.0.0", headers={**_ACCOUNT_HEADERS, "x-litellm-api-key": key} + ) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "open-weights", "input": "Hello", "stream": True}, + headers={**_ACCOUNT_HEADERS, "x-litellm-api-key": key}, + ) + assert (catalog.status_code, response.status_code) == (200, 200), response.text + assert ("open-weights" in [model["slug"] for model in catalog.json()["models"]]) is through_litellm + assert provider.called is through_litellm + assert chatgpt.called is not through_litellm + + +_JWT_ISSUER: Final = "https://idp.codex-test.example" +_JWT_JWKS_URL: Final = f"{_JWT_ISSUER}/keys" +_JWT_AUDIENCE: Final = "litellm-proxy" +_JWT_KID: Final = "codex-test-key" +_JWT_TEAM: Final = "codex-team" + + +def _team_jwt(monkeypatch: pytest.MonkeyPatch, team_allowed_routes: list[str]) -> str: + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + public_jwk: Final = { + **json.loads(RSAAlgorithm.to_jwk(private_key.public_key())), + "kid": _JWT_KID, + "alg": "RS256", + "use": "sig", + } + keys_cache: Final = DualCache() + keys_cache.set_cache(key=f"litellm_jwt_auth_keys_{_JWT_JWKS_URL}", value=[public_jwk]) + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=None, + user_api_key_cache=keys_cache, + litellm_jwtauth=LiteLLM_JWTAuth( + issuers=[{"issuer": _JWT_ISSUER, "jwks_url": _JWT_JWKS_URL, "audience": _JWT_AUDIENCE}], + team_ids_jwt_field="groups", + team_allowed_routes=team_allowed_routes, + ), + ) + monkeypatch.setattr(proxy_server, "jwt_handler", handler) + monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": True}) + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setattr(proxy_server, "prisma_client", "test-database") + proxy_server.user_api_key_cache.set_cache( + key=f"team_id:{_JWT_TEAM}", + value=LiteLLM_TeamTableCachedObj(team_id=_JWT_TEAM, models=["chatgpt-native", "open-weights"]), + ) + issued_at: Final = int(time.time()) + return jwt.encode( + { + "sub": "codex-user", + "iss": _JWT_ISSUER, + "aud": _JWT_AUDIENCE, + "iat": issued_at, + "exp": issued_at + 300, + "groups": [_JWT_TEAM], + }, + private_key, + algorithm="RS256", + headers={"kid": _JWT_KID}, + ) + + +@pytest.mark.parametrize( + ("team_allowed_routes", "through_litellm"), + [ + pytest.param(["/backend-api/codex/*"], False, id="team allowed only the native routes"), + pytest.param(["/backend-api/codex/*", "openai_routes"], True, id="team allowed openai routes too"), + ], +) +def test_jwt_team_routes_keep_litellm_models_off_native_routes_unless_v1_responses_is_allowed( + client: TestClient, + respx_mock: respx.MockRouter, + mixed_router: Router, + monkeypatch: pytest.MonkeyPatch, + team_allowed_routes: list[str], + through_litellm: bool, +) -> None: + token: Final = _team_jwt(monkeypatch, team_allowed_routes) + _account_catalog(respx_mock) + provider: Final = _open_weights_provider(respx_mock) + chatgpt: Final = _chatgpt_responses(respx_mock) + catalog: Final = client.get( + "/backend-api/codex/models?client_version=1.0.0", headers={**_ACCOUNT_HEADERS, "x-litellm-api-key": token} + ) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "open-weights", "input": "Hello", "stream": True}, + headers={**_ACCOUNT_HEADERS, "x-litellm-api-key": token}, + ) + assert (catalog.status_code, response.status_code) == (200, 200), response.text + assert ("open-weights" in [model["slug"] for model in catalog.json()["models"]]) is through_litellm + assert provider.called is through_litellm + assert chatgpt.called is not through_litellm + + +_USER_TURN: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Hello"}]} +_ASSISTANT_TURN: Final = {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Hi"}]} +_FIREWORKS_REASONING: Final = { + "type": "reasoning", + "id": f"rs_{'0f' * 16}", + "summary": [{"type": "summary_text", "text": "Thinking it over"}], +} + + +@pytest.mark.parametrize( + ("item", "store", "forwarded"), + [ + pytest.param(_FIREWORKS_REASONING, False, False, id="another provider's summary-only item"), + pytest.param({**_FIREWORKS_REASONING, "encrypted_content": None}, None, False, id="null blob, store unset"), + pytest.param( + {"type": "reasoning", "id": "rs_upstream", "summary": [], "encrypted_content": "gAAAAB-test-blob"}, + False, + True, + id="ChatGPT's own item", + ), + pytest.param({**_FIREWORKS_REASONING, "id": None}, False, True, id="summary-only item without an id"), + pytest.param(_FIREWORKS_REASONING, True, True, id="stored conversation"), + ], +) +def test_chatgpt_passthrough_drops_reasoning_ids_chatgpt_never_persisted( + client: TestClient, + respx_mock: respx.MockRouter, + item: dict[str, object], + store: bool | None, + forwarded: bool, +) -> None: + chatgpt: Final = _chatgpt_responses(respx_mock) + body: Final = { + "model": "subscription-model", + "input": [_USER_TURN, item, _ASSISTANT_TURN], + "stream": True, + **({} if store is None else {"store": store}), + } + response: Final = client.post("/backend-api/codex/responses", json=body, headers=_HEADERS) + assert response.status_code == 200 + assert json.loads(chatgpt.calls.last.request.content)["input"] == ( + [_USER_TURN, item, _ASSISTANT_TURN] if forwarded else [_USER_TURN, _ASSISTANT_TURN] + ) + + +def test_exact_deployment_named_like_an_account_slug_stays_on_the_subscription_after_the_catalog_read( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router +) -> None: + catalog: Final = respx_mock.get(f"{_BASE}/models").respond(json=_UPSTREAM_CATALOG) + provider: Final = _open_weights_provider(respx_mock) + chatgpt: Final = _chatgpt_responses(respx_mock) + assert client.get("/backend-api/codex/models?client_version=0.165.0", headers=_ACCOUNT_HEADERS).status_code == 200 + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "shared-name", "input": "Hello", "stream": True}, + headers=_ACCOUNT_HEADERS, + ) + assert response.status_code == 200 + assert chatgpt.called + assert not provider.called + assert catalog.call_count == 1 + + +_LISTED: Final = httpx.Response(200, json=_UPSTREAM_CATALOG) + + +@pytest.mark.parametrize( + ("model", "user_agent", "catalog_answer", "subscription"), + [ + pytest.param("shared-name", _CODEX_USER_AGENT, _LISTED, True, id="listed"), + pytest.param("open-weights", _CODEX_USER_AGENT, _LISTED, False, id="unlisted"), + pytest.param("open-weights", _CODEX_USER_AGENT, httpx.Response(503), True, id="catalog unavailable"), + pytest.param("open-weights", _CODEX_USER_AGENT, httpx.ReadTimeout("slow"), True, id="catalog timeout"), + pytest.param("open-weights", "codex_cli_rs", _LISTED, True, id="no version"), + pytest.param("open-weights", "curl/8.5.0", _LISTED, True, id="user agent of another client"), + ], +) +def test_native_responses_ask_chatgpt_which_models_the_account_lists_on_a_cache_miss( + client: TestClient, + respx_mock: respx.MockRouter, + mixed_router: Router, + model: str, + user_agent: str, + catalog_answer: httpx.Response | httpx.TimeoutException, + subscription: bool, +) -> None: + catalog: Final = respx_mock.get(f"{_BASE}/models").mock(side_effect=[catalog_answer]) + provider: Final = _open_weights_provider(respx_mock) + chatgpt: Final = _chatgpt_responses(respx_mock) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": model, "input": "Hello", "stream": True}, + headers={**_ACCOUNT_HEADERS, "user-agent": user_agent}, + ) + assert response.status_code == 200 + assert chatgpt.called is subscription + assert provider.called is not subscription + assert catalog.called is (user_agent == _CODEX_USER_AGENT) + + +def test_a_failed_account_catalog_read_keeps_the_next_turns_on_the_subscription_without_asking_again( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router +) -> None: + catalog: Final = respx_mock.get(f"{_BASE}/models").respond(503) + provider: Final = _open_weights_provider(respx_mock) + chatgpt: Final = _chatgpt_responses(respx_mock) + responses: Final = tuple( + client.post( + "/backend-api/codex/responses", + json={"model": "open-weights", "input": "Hello", "stream": True}, + headers=_ACCOUNT_HEADERS, + ) + for _ in range(2) + ) + assert [response.status_code for response in responses] == [200, 200] + assert chatgpt.call_count == 2 + assert not provider.called + assert catalog.call_count == 1 + + +def test_a_failed_account_catalog_read_with_another_login_leaves_the_accounts_own_list_reachable( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router +) -> None: + own_login: Final = f"Bearer {_CHATGPT_JWT}" + catalog: Final = respx_mock.get(f"{_BASE}/models").mock( + side_effect=lambda request: ( + httpx.Response(200, json=_UPSTREAM_CATALOG) + if request.headers["authorization"] == own_login + else httpx.Response(401) + ) + ) + provider: Final = _open_weights_provider(respx_mock) + chatgpt: Final = _chatgpt_responses(respx_mock) + turn: Final = {"model": "open-weights", "input": "Hello", "stream": True} + other_login: Final = client.post( + "/backend-api/codex/responses", + json=turn, + headers={**_ACCOUNT_HEADERS, "authorization": "Bearer test-someone-elses-login"}, + ) + own: Final = client.post("/backend-api/codex/responses", json=turn, headers=_ACCOUNT_HEADERS) + assert other_login.status_code == own.status_code == 200 + assert chatgpt.call_count == 1 + assert provider.call_count == 1 + assert [call.request.headers["authorization"] for call in catalog.calls] == [ + "Bearer test-someone-elses-login", + own_login, + ] + + +def test_rotating_chatgpt_logins_on_one_key_reads_the_account_catalog_a_bounded_number_of_times_per_key( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router, restricted_key: Callable[..., str] +) -> None: + catalog: Final = respx_mock.get(f"{_BASE}/models").respond(401) + provider: Final = _open_weights_provider(respx_mock) + chatgpt: Final = _chatgpt_responses(respx_mock) + turn: Final = {"model": "open-weights", "input": "Hello", "stream": True} + rotated: Final = tuple( + client.post( + "/backend-api/codex/responses", + json=turn, + headers={**_ACCOUNT_HEADERS, "authorization": f"Bearer test-rotated-login-{attempt}"}, + ) + for attempt in range(_ACCOUNT_SLUGS_FETCHES_PER_CALLER + 2) + ) + teammate: Final = client.post( + "/backend-api/codex/responses", + json=turn, + headers={ + **_ACCOUNT_HEADERS, + "authorization": "Bearer test-teammate-login", + "x-litellm-api-key": restricted_key(["open-weights"]), + }, + ) + assert {response.status_code for response in (*rotated, teammate)} == {200} + assert chatgpt.call_count == len(rotated) + 1 + assert not provider.called + assert [call.request.headers["authorization"] for call in catalog.calls][-1] == "Bearer test-teammate-login" + assert catalog.call_count == _ACCOUNT_SLUGS_FETCHES_PER_CALLER + 1 + + +def test_account_catalog_lookup_sends_the_callers_chatgpt_login_and_codex_version( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router +) -> None: + catalog: Final = respx_mock.get(f"{_BASE}/models").respond(json=_UPSTREAM_CATALOG) + _chatgpt_responses(respx_mock) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "shared-name", "input": "Hello", "stream": True}, + headers=_ACCOUNT_HEADERS, + ) + assert response.status_code == 200 + fetched: Final = catalog.calls.last.request + assert fetched.url.params["client_version"] == "0.165.0" + assert {name: fetched.headers[name] for name in ("authorization", "chatgpt-account-id", "originator")} == { + "authorization": f"Bearer {_CHATGPT_JWT}", + "chatgpt-account-id": "account-a", + "originator": "codex_cli_rs", + } + assert fetched.headers["user-agent"] == _CODEX_USER_AGENT + assert "x-litellm-api-key" not in fetched.headers + + +def _models_etags(client: TestClient, headers: dict[str, str]) -> tuple[str | None, str | None]: + catalog: Final = client.get("/backend-api/codex/models?client_version=0.165.0", headers=headers) + responses: Final = client.post( + "/backend-api/codex/responses", + json={"model": "gpt-unlisted", "input": "Hello", "stream": True}, + headers=headers, + ) + assert catalog.status_code == responses.status_code == 200, (catalog.text, responses.text) + return catalog.headers.get("etag"), responses.headers.get("x-models-etag") + + +@pytest.mark.parametrize( + ("models", "merged"), + [ + pytest.param(["chatgpt-native", "open-weights"], True, id="catalog with LiteLLM models"), + pytest.param(["chatgpt-native"], False, id="catalog of ChatGPT models only"), + ], +) +def test_native_catalog_and_responses_carry_the_same_models_etag( + client: TestClient, + respx_mock: respx.MockRouter, + mixed_router: Router, + restricted_key: Callable[..., str], + models: list[str], + merged: bool, +) -> None: + respx_mock.get(f"{_BASE}/models").respond(json=_UPSTREAM_CATALOG, headers={"etag": _UPSTREAM_MODELS_ETAG}) + _chatgpt_responses(respx_mock) + catalog_etag, responses_etag = _models_etags(client, {**_HEADERS, "x-litellm-api-key": restricted_key(models)}) + assert catalog_etag is not None + assert catalog_etag == responses_etag + assert (catalog_etag != _UPSTREAM_MODELS_ETAG) is merged + + +def test_native_models_etag_is_the_same_on_replicas_with_empty_caches( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem(proxy_server.app.dependency_overrides, codex_shared_cache, lambda: DualCache()) + respx_mock.get(f"{_BASE}/models").respond(json=_UPSTREAM_CATALOG, headers={"etag": _UPSTREAM_MODELS_ETAG}) + _chatgpt_responses(respx_mock) + first_replica: Final = _models_etags(client, _HEADERS) + second_replica: Final = _models_etags(client, _HEADERS) + assert first_replica[0] not in (None, _UPSTREAM_MODELS_ETAG) + assert len({*first_replica, *second_replica}) == 1 + + +def test_native_models_etag_changes_with_the_keys_litellm_models_and_with_upstreams_etag( + client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + open_weights: Final = _open_weights("open-weights") + wider_window: Final = {**open_weights, "model_info": {"max_input_tokens": 64000}} + + def catalog_etag(deployment: dict[str, object], upstream_etag: str) -> str | None: + monkeypatch.setattr(proxy_server, "llm_router", Router(model_list=[deployment])) + respx_mock.get(f"{_BASE}/models").respond(json=_UPSTREAM_CATALOG, headers={"etag": upstream_etag}) + return client.get("/backend-api/codex/models?client_version=0.165.0", headers=_HEADERS).headers.get("etag") + + baseline: Final = catalog_etag(open_weights, '"upstream-v1"') + assert baseline == catalog_etag(open_weights, '"upstream-v1"') + assert ( + len({baseline, catalog_etag(wider_window, '"upstream-v1"'), catalog_etag(open_weights, '"upstream-v3"')}) == 3 + ) + assert baseline not in (None, '"upstream-v1"') + + +class FailingModelListing(CustomLogger): + async def async_filter_listed_models( + self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] + ) -> Sequence[str]: + raise RuntimeError("model listing backend unavailable") + + +def test_native_catalog_relays_chatgpts_answer_when_listing_the_keys_models_fails( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "callbacks", [FailingModelListing()]) + respx_mock.get(f"{_BASE}/models").respond(json=_UPSTREAM_CATALOG, headers={"etag": '"upstream-v1"'}) + response: Final = client.get("/backend-api/codex/models?client_version=0.165.0", headers=_HEADERS) + assert response.status_code == 200 + assert response.json() == _UPSTREAM_CATALOG + assert response.headers["etag"] == '"upstream-v1"' + + +def test_native_responses_keep_chatgpts_models_etag_when_listing_the_keys_models_fails( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "callbacks", [FailingModelListing()]) + chatgpt: Final = _chatgpt_responses(respx_mock) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "gpt-unlisted", "input": "Hello", "stream": True}, + headers=_HEADERS, + ) + assert response.status_code == 200 + assert chatgpt.called + assert response.headers["x-models-etag"] == _UPSTREAM_MODELS_ETAG + + +_OPENAI_BASE: Final = "https://openai-api.test/v1" +_OPENAI_BLOB: Final = "gAAAAB-openai-reasoning-blob" +_OPENAI_REASONING: Final = {"type": "reasoning", "id": "rs_openai", "summary": [], "encrypted_content": _OPENAI_BLOB} + + +def _openai_reply(request: httpx.Request) -> httpx.Response: + completed: Final = { + "id": "resp_openai", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.6", + "output": [ + _OPENAI_REASONING, + { + "type": "message", + "id": "msg_openai", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Hi", "annotations": []}], + }, + ], + "usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4}, + } + if json.loads(request.content).get("stream") is not True: + return httpx.Response(200, json=completed) + events: Final = ( + { + "type": "response.created", + "sequence_number": 0, + "response": {**completed, "status": "in_progress", "output": []}, + }, + {"type": "response.output_item.done", "sequence_number": 1, "output_index": 0, "item": _OPENAI_REASONING}, + {"type": "response.completed", "sequence_number": 2, "response": completed}, + ) + return httpx.Response( + 200, + content="".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events).encode(), + headers={"content-type": "text/event-stream"}, + ) + + +def _reasoning_item_of(response: httpx.Response, stream: bool) -> dict[str, object]: + if not stream: + return response.json()["output"][0] + payloads: Final = ( + json.loads(line.removeprefix(b"data: ")) + for line in response.content.splitlines() + if line.startswith(b"data: {") + ) + return next(event["item"] for event in payloads if event["type"] == "response.output_item.done") + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_reasoning_from_the_native_route_is_dropped_for_chatgpt_and_unwrapped_for_litellm( + client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stream: bool +) -> None: + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "api-model", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_base": _OPENAI_BASE, + "api_key": "test-openai-key", + }, + } + ] + ), + ) + _account_catalog(respx_mock) + provider: Final = respx_mock.post(f"{_OPENAI_BASE}/responses").mock(side_effect=_openai_reply) + chatgpt: Final = _chatgpt_responses(respx_mock) + turn: Final = {"model": "api-model", "input": [_USER_TURN], "stream": stream, "store": False} + first: Final = client.post("/backend-api/codex/responses", json=turn, headers=_ACCOUNT_HEADERS) + assert first.status_code == 200, first.text + reasoning: Final = _reasoning_item_of(first, stream) + model_id, blob = ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id( + str(reasoning["encrypted_content"]) + ) + assert (model_id is not None, blob) == (True, _OPENAI_BLOB) + history: Final = [_USER_TURN, reasoning, _ASSISTANT_TURN, _USER_TURN] + to_chatgpt: Final = client.post( + "/backend-api/codex/responses", + json={**turn, "model": "subscription-model", "input": history, "stream": True}, + headers=_ACCOUNT_HEADERS, + ) + back_to_litellm: Final = client.post( + "/backend-api/codex/responses", json={**turn, "input": history}, headers=_ACCOUNT_HEADERS + ) + v1: Final = client.post( + "/v1/responses", json={**turn, "stream": False}, headers={"authorization": f"Bearer {MASTER_KEY}"} + ) + assert to_chatgpt.status_code == back_to_litellm.status_code == v1.status_code == 200 + assert json.loads(chatgpt.calls.last.request.content)["input"] == [_USER_TURN, _ASSISTANT_TURN, _USER_TURN] + replayed: Final = json.loads(provider.calls[1].request.content)["input"][1] + assert (replayed["id"], replayed["encrypted_content"]) == ("rs_openai", _OPENAI_BLOB) + assert v1.json()["output"][0]["encrypted_content"] == _OPENAI_BLOB + + +_XAI_BASE: Final = "https://xai-api.test/v1" +_LUNA_ID: Final = "oai-luna-deployment" +_LUNA_BLOB: Final = "gAAAAB-luna-reasoning-blob" +_LUNA_REASONING: Final = { + "type": "reasoning", + "summary": [], + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id(_LUNA_BLOB, _LUNA_ID), +} +_CHATGPT_REASONING: Final = {"type": "reasoning", "summary": [], "encrypted_content": "gAAAAB-chatgpt-reasoning-blob"} + + +@pytest.fixture +def xai_fireworks_and_openai_router(monkeypatch: pytest.MonkeyPatch) -> Router: + router: Final = Router( + model_list=[ + { + "model_name": "grok-4.7", + "litellm_params": {"model": "xai/grok-4.7", "api_base": _XAI_BASE, "api_key": "test-xai-key"}, + }, + _open_weights("open-weights"), + { + "model_name": "oai-luna", + "litellm_params": { + "model": "openai/gpt-6-luna", + "api_base": _OPENAI_BASE, + "api_key": "test-openai-key", + }, + "model_info": {"id": _LUNA_ID}, + }, + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + return router + + +@pytest.mark.parametrize( + ("model", "replayed", "expected"), + [ + pytest.param("open-weights", _CHATGPT_REASONING, None, id="chatgpt reasoning to fireworks is dropped"), + pytest.param("open-weights", _LUNA_REASONING, None, id="another deployment's reasoning to fireworks is dropped"), + pytest.param("oai-luna", _LUNA_REASONING, _LUNA_BLOB, id="the minting deployment's reasoning is kept"), + ], +) +def test_litellm_model_on_native_route_gets_only_the_reasoning_it_can_decrypt( + client: TestClient, + respx_mock: respx.MockRouter, + xai_fireworks_and_openai_router: Router, + model: str, + replayed: dict[str, object], + expected: str | None, +) -> None: + _account_catalog(respx_mock) + fireworks: Final = _open_weights_provider(respx_mock) + luna: Final = respx_mock.post(f"{_OPENAI_BASE}/responses").mock(side_effect=_provider_reply) + response: Final = client.post( + "/backend-api/codex/responses", + json={ + "model": model, + "input": [_USER_TURN, replayed, _ASSISTANT_TURN, _USER_TURN], + "stream": False, + "store": False, + }, + headers=_ACCOUNT_HEADERS, + ) + assert response.status_code == 200, response.text + provider: Final = fireworks if model == "open-weights" else luna + sent: Final = json.loads(provider.calls.last.request.content)["input"] + assert [item.get("encrypted_content") for item in sent if item.get("type") == "reasoning"] == ( + [] if expected is None else [expected] + ) + + +_WHOLE_FIREWORKS_REASONING: Final = {**_FIREWORKS_REASONING, "status": "completed"} +_SUMMARY_ONLY_REASONING: Final = {key: value for key, value in _FIREWORKS_REASONING.items() if key != "id"} +_OTHER_DEPLOYMENTS_BLOB: Final = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + "gAAAAB-other-deployment-blob", "other-deployment" +) + + +@pytest.mark.parametrize( + ("model", "replayed", "store", "expected"), + [ + pytest.param( + "oai-luna", _WHOLE_FIREWORKS_REASONING, False, [_SUMMARY_ONLY_REASONING], id="whole fireworks item to openai" + ), + pytest.param( + "oai-luna", + {**_WHOLE_FIREWORKS_REASONING, "encrypted_content": _OTHER_DEPLOYMENTS_BLOB}, + False, + [_SUMMARY_ONLY_REASONING], + id="whole item with another deployment's blob to openai", + ), + pytest.param( + "oai-luna", + {**_SUMMARY_ONLY_REASONING, "content": [], "encrypted_content": _OTHER_DEPLOYMENTS_BLOB}, + False, + [{**_SUMMARY_ONLY_REASONING, "content": []}], + id="codex's replay of another deployment's item to openai", + ), + pytest.param( + "oai-luna", + {**_SUMMARY_ONLY_REASONING, "content": None, "encrypted_content": None}, + False, + [_SUMMARY_ONLY_REASONING], + id="codex's replay of a blobless item to openai", + ), + pytest.param( + "oai-luna", + {"type": "reasoning", "id": "rs_unreadable", "summary": []}, + False, + [], + id="an unresolvable id with nothing readable is dropped", + ), + pytest.param( + "open-weights", + {**_SUMMARY_ONLY_REASONING, "id": "rs_openai"}, + False, + [_SUMMARY_ONLY_REASONING], + id="openai summary-only item to fireworks", + ), + pytest.param( + "oai-luna", + {**_LUNA_REASONING, "id": "rs_luna", "content": []}, + False, + [{**_LUNA_REASONING, "id": "rs_luna", "content": [], "encrypted_content": _LUNA_BLOB}], + id="the minting deployment gets its item whole", + ), + pytest.param( + "oai-luna", _FIREWORKS_REASONING, True, [_FIREWORKS_REASONING], id="a stored conversation keeps the id" + ), + ], +) +def test_reasoning_the_target_cannot_resolve_reaches_it_without_its_id( + client: TestClient, + respx_mock: respx.MockRouter, + xai_fireworks_and_openai_router: Router, + model: str, + replayed: dict[str, object], + store: bool, + expected: list[dict[str, object]], +) -> None: + _account_catalog(respx_mock) + fireworks: Final = _open_weights_provider(respx_mock) + luna: Final = respx_mock.post(f"{_OPENAI_BASE}/responses").mock(side_effect=_provider_reply) + response: Final = client.post( + "/backend-api/codex/responses", + json={ + "model": model, + "input": [_USER_TURN, replayed, _ASSISTANT_TURN, _USER_TURN], + "stream": False, + "store": store, + }, + headers=_ACCOUNT_HEADERS, + ) + assert response.status_code == 200, response.text + sent: Final = json.loads((fireworks if model == "open-weights" else luna).calls.last.request.content)["input"] + assert [item for item in sent if item.get("type") == "reasoning"] == expected + + +def test_reasoning_the_chat_bridge_target_cannot_resolve_reaches_it_as_its_summary( + client: TestClient, respx_mock: respx.MockRouter, xai_fireworks_and_openai_router: Router +) -> None: + _account_catalog(respx_mock) + chat_endpoint: Final = respx_mock.post(f"{_XAI_BASE}/chat/completions").mock(side_effect=_xai_search_call_reply) + response: Final = client.post( + "/backend-api/codex/responses", + json={ + "model": "grok-4.7", + "input": [_USER_TURN, _WHOLE_FIREWORKS_REASONING, _ASSISTANT_TURN, _USER_TURN], + "stream": False, + "store": False, + }, + headers=_ACCOUNT_HEADERS, + ) + assert response.status_code == 200, response.text + messages: Final = json.loads(chat_endpoint.calls.last.request.content)["messages"] + assert [message.get("reasoning_content") for message in messages if message["role"] == "assistant"] == [ + "Thinking it over" + ] + + +def _namespace(name: str, *members: str) -> dict[str, object]: + return { + "type": "namespace", + "name": name, + "description": "", + "tools": [{"type": "function", "name": member, "parameters": {"type": "object"}} for member in members], + } + + +_EXEC_TOOL: Final = {"type": "function", "name": "exec_command", "parameters": {"type": "object"}} +_ADDITIONAL_TOOLS: Final = { + "type": "additional_tools", + "role": "developer", + "tools": [_namespace("collaboration", "followup_task"), _namespace("mcp__docs__", "search")], +} + + +@pytest.mark.parametrize( + ("model", "expected_input", "expected_tool_names"), + [ + pytest.param( + "open-weights", + [_USER_TURN], + ["exec_command", "collaboration", "mcp__docs__"], + id="fireworks gets the item's tools as top-level tools", + ), + pytest.param("oai-luna", [_ADDITIONAL_TOOLS, _USER_TURN], ["exec_command"], id="the openai api keeps the item"), + ], +) +def test_codex_additional_tools_item_reaches_each_provider_in_a_shape_it_accepts( + client: TestClient, + respx_mock: respx.MockRouter, + xai_fireworks_and_openai_router: Router, + model: str, + expected_input: list[dict[str, object]], + expected_tool_names: list[str], +) -> None: + _account_catalog(respx_mock) + fireworks: Final = _open_weights_provider(respx_mock) + luna: Final = respx_mock.post(f"{_OPENAI_BASE}/responses").mock(side_effect=_provider_reply) + response: Final = client.post( + "/backend-api/codex/responses", + json={ + "model": model, + "input": [_ADDITIONAL_TOOLS, _USER_TURN], + "tools": [_EXEC_TOOL], + "stream": False, + "store": False, + }, + headers=_ACCOUNT_HEADERS, + ) + assert response.status_code == 200, response.text + sent: Final = json.loads((fireworks if model == "open-weights" else luna).calls.last.request.content) + assert sent["input"] == expected_input + assert [tool["name"] for tool in sent["tools"]] == expected_tool_names + + +_FORBIDDEN_TOOL: Final = {"type": "function", "name": "delete_user", "parameters": {"type": "object"}} + + +@pytest.mark.parametrize( + ("tools", "input_items", "allowed"), + [ + pytest.param([_EXEC_TOOL, _FORBIDDEN_TOOL], [_USER_TURN], False, id="a forbidden top-level tool"), + pytest.param( + [_EXEC_TOOL], + [{"type": "additional_tools", "role": "developer", "tools": [_FORBIDDEN_TOOL]}, _USER_TURN], + False, + id="a forbidden tool in an additional_tools item", + ), + pytest.param([_EXEC_TOOL], [_USER_TURN], True, id="only allowed tools"), + ], +) +def test_native_litellm_turn_enforces_the_key_tool_allowlist( + client: TestClient, + respx_mock: respx.MockRouter, + xai_fireworks_and_openai_router: Router, + restricted_key: Callable[..., str], + tools: list[dict[str, object]], + input_items: list[dict[str, object]], + allowed: bool, +) -> None: + _account_catalog(respx_mock) + fireworks: Final = _open_weights_provider(respx_mock) + key: Final = restricted_key( + ["open-weights"], (*_NATIVE_ROUTES, "/v1/responses"), allowed_tools=(str(_EXEC_TOOL["name"]),) + ) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "open-weights", "input": input_items, "tools": tools, "stream": False, "store": False}, + headers={**_ACCOUNT_HEADERS, "x-litellm-api-key": key}, + ) + assert response.status_code == (200 if allowed else 403), response.text + assert fireworks.called is allowed + + +_ANTHROPIC_BASE: Final = "https://anthropic-api.test" + + +def test_codex_namespace_tool_called_through_the_chat_bridge_comes_back_with_its_namespace( + client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "claude-sonnet-5-5", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "api_base": _ANTHROPIC_BASE, + "api_key": "test-anthropic-key", + }, + } + ] + ), + ) + _account_catalog(respx_mock) + anthropic: Final = respx_mock.post(f"{_ANTHROPIC_BASE}/v1/messages").respond( + json={ + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [ + {"type": "tool_use", "id": "toolu_test", "name": "mcp__docs____search", "input": {"query": "proxy"}} + ], + "stop_reason": "tool_use", + "usage": {"input_tokens": 3, "output_tokens": 1}, + } + ) + response: Final = client.post( + "/backend-api/codex/responses", + json={ + "model": "claude-sonnet-5-5", + "input": [_ADDITIONAL_TOOLS, _USER_TURN], + "tools": [_EXEC_TOOL], + "stream": False, + "store": False, + }, + headers=_ACCOUNT_HEADERS, + ) + assert response.status_code == 200, response.text + assert anthropic.called + (call,) = (item for item in response.json()["output"] if item["type"] == "function_call") + assert (call["namespace"], call["name"]) == ("mcp__docs__", "search") + + +_CLOCK_NAMESPACE: Final = _namespace("clock", "sleep") +_CLOCK_TOOLS: Final = {"type": "additional_tools", "role": "developer", "tools": [_CLOCK_NAMESPACE]} +def test_codex_namespace_tool_reaches_fireworks_as_a_namespace( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router +) -> None: + _account_catalog(respx_mock) + provider: Final = _open_weights_provider(respx_mock) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "open-weights", "input": [_CLOCK_TOOLS, _USER_TURN], "stream": False, "store": False}, + headers=_ACCOUNT_HEADERS, + ) + assert response.status_code == 200, response.text + assert json.loads(provider.calls.last.request.content)["tools"] == [_CLOCK_NAMESPACE] + + + +@pytest.mark.parametrize( + ("drop_params", "expected_tools"), + [ + pytest.param({}, [_CLOCK_NAMESPACE], id="codex's drop_params default merges the copies"), + pytest.param({"drop_params": False}, [_CLOCK_NAMESPACE, _CLOCK_NAMESPACE], id="the caller's off sends each"), + ], +) +def test_codex_tool_list_repeated_on_every_turn_reaches_fireworks_merged_only_under_drop_params( + client: TestClient, + respx_mock: respx.MockRouter, + mixed_router: Router, + drop_params: dict[str, bool], + expected_tools: list[dict[str, object]], +) -> None: + _account_catalog(respx_mock) + provider: Final = _open_weights_provider(respx_mock) + response: Final = client.post( + "/backend-api/codex/responses", + json={ + "model": "open-weights", + "input": [_CLOCK_TOOLS, _USER_TURN, _CLOCK_TOOLS], + "stream": False, + "store": False, + **drop_params, + }, + headers=_ACCOUNT_HEADERS, + ) + assert response.status_code == 200, response.text + assert json.loads(provider.calls.last.request.content)["tools"] == expected_tools + +_CUSTOM_EXEC_TOOL: Final = {"type": "custom", "name": "exec", "description": "Run a command", "format": {"type": "text"}} +_XAI_SEARCH_CALL: Final = { + "id": "call_search", + "type": "function", + "function": {"name": "mcp__docs____search", "arguments": '{"query":"proxy"}'}, +} +_XAI_USAGE: Final = {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4} + + +def _xai_chat_chunk(choice: dict[str, object], **fields: object) -> str: + chunk: Final = { + "id": "chatcmpl-xai", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "grok-4.7", + "choices": [{"index": 0, **choice}], + **fields, + } + return f"data: {json.dumps(chunk)}\n\n" + + +def _xai_search_call_reply(request: httpx.Request) -> httpx.Response: + if json.loads(request.content).get("stream") is not True: + return httpx.Response( + 200, + json={ + "id": "chatcmpl-xai", + "object": "chat.completion", + "created": 1700000000, + "model": "grok-4.7", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": None, "tool_calls": [_XAI_SEARCH_CALL]}, + "finish_reason": "tool_calls", + } + ], + "usage": _XAI_USAGE, + }, + ) + stream: Final = ( + _xai_chat_chunk( + {"delta": {"role": "assistant", "tool_calls": [{"index": 0, **_XAI_SEARCH_CALL}]}, "finish_reason": None} + ) + + _xai_chat_chunk({"delta": {}, "finish_reason": "tool_calls"}, usage=_XAI_USAGE) + + "data: [DONE]\n\n" + ) + return httpx.Response(200, content=stream.encode(), headers={"content-type": "text/event-stream"}) + + +def _data_lines(stream: bytes) -> list[bytes]: + return [line.removeprefix(b"data: ") for line in stream.splitlines() if line.startswith(b"data: {")] + + +def _function_calls_of(response: httpx.Response, stream: bool) -> list[dict[str, object]]: + items: Final = ( + response.json()["output"] + if not stream + else [ + event["item"] + for event in map(json.loads, _data_lines(response.content)) + if event["type"] == "response.output_item.done" + ] + ) + return [item for item in items if item["type"] == "function_call"] + + +@pytest.mark.parametrize("stream", [False, True]) +def test_codex_turn_for_xai_goes_through_chat_completions_and_its_namespaced_call_comes_back_whole( + client: TestClient, respx_mock: respx.MockRouter, xai_fireworks_and_openai_router: Router, stream: bool +) -> None: + _account_catalog(respx_mock) + responses_endpoint: Final = respx_mock.post(f"{_XAI_BASE}/responses").respond(json=_provider_completed("Hi")) + chat_endpoint: Final = respx_mock.post(f"{_XAI_BASE}/chat/completions").mock(side_effect=_xai_search_call_reply) + response: Final = client.post( + "/backend-api/codex/responses", + json={ + "model": "grok-4.7", + "input": [_ADDITIONAL_TOOLS, _USER_TURN], + "tools": [_EXEC_TOOL, _CUSTOM_EXEC_TOOL], + "stream": stream, + "store": False, + }, + headers=_ACCOUNT_HEADERS, + ) + assert response.status_code == 200, response.text + assert not responses_endpoint.called + sent_tools: Final = json.loads(chat_endpoint.calls.last.request.content)["tools"] + assert [(tool["type"], tool["function"]["name"]) for tool in sent_tools] == [ + ("function", "exec_command"), + ("function", "exec"), + ("function", "collaboration__followup_task"), + ("function", "mcp__docs____search"), + ] + (call,) = _function_calls_of(response, stream) + assert (call["name"], call["namespace"], call["call_id"]) == ("search", "mcp__docs__", "call_search") + + +def test_encrypted_reasoning_stays_only_when_every_deployment_the_router_may_pick_can_decrypt_it( + client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + second_base: Final = "https://second-openai-org.test/v1" + + def replayed_blobs(optional_pre_call_checks: list[str]) -> list[object]: + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "oai-pair", + "litellm_params": {"model": "openai/gpt-6-luna", "api_base": base, "api_key": "test-key"}, + "model_info": {"id": model_id}, + } + for base, model_id in ((_OPENAI_BASE, _LUNA_ID), (second_base, "second-org-deployment")) + ], + optional_pre_call_checks=optional_pre_call_checks, + ), + ) + providers: Final = tuple( + respx_mock.post(f"{base}/responses").mock(side_effect=_provider_reply) + for base in (_OPENAI_BASE, second_base) + ) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "oai-pair", "input": [_USER_TURN, _LUNA_REASONING, _USER_TURN], "store": False}, + headers=_ACCOUNT_HEADERS, + ) + assert response.status_code == 200, response.text + (sent,) = (json.loads(call.request.content)["input"] for provider in providers for call in provider.calls) + return [item.get("encrypted_content") for item in sent if item.get("type") == "reasoning"] + + _account_catalog(respx_mock) + assert replayed_blobs([]) == [] + respx_mock.reset() + _account_catalog(respx_mock) + assert replayed_blobs(["encrypted_content_affinity"]) == [_LUNA_BLOB] + + +def test_encrypted_reasoning_stays_on_a_global_alias_of_a_group_pinned_to_its_origin( + client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + second_base: Final = "https://second-openai-org.test/v1" + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "oai-pair", + "litellm_params": {"model": "openai/gpt-6-luna", "api_base": base, "api_key": "test-key"}, + "model_info": {"id": model_id}, + } + for base, model_id in ((_OPENAI_BASE, _LUNA_ID), (second_base, "second-org-deployment")) + ], + model_group_alias={"pair-alias": "oai-pair"}, + model_group_affinity_config={"oai-pair": ["encrypted_content_affinity"]}, + ), + ) + _account_catalog(respx_mock) + providers: Final = tuple( + respx_mock.post(f"{base}/responses").mock(side_effect=_provider_reply) for base in (_OPENAI_BASE, second_base) + ) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "pair-alias", "input": [_USER_TURN, _LUNA_REASONING, _USER_TURN], "store": False}, + headers=_ACCOUNT_HEADERS, + ) + assert response.status_code == 200, response.text + (sent,) = (json.loads(call.request.content)["input"] for provider in providers for call in provider.calls) + assert [item.get("encrypted_content") for item in sent if item.get("type") == "reasoning"] == [_LUNA_BLOB] + + +class ModelListingCounter(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.listings: Final[list[tuple[str, ...]]] = [] + + async def async_filter_listed_models( + self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] + ) -> Sequence[str]: + self.listings.append(tuple(model_names)) + return model_names + + +def test_jwt_caller_builds_the_models_etag_once_across_chatgpt_turns( + client: TestClient, respx_mock: respx.MockRouter, mixed_router: Router, monkeypatch: pytest.MonkeyPatch +) -> None: + counter: Final = ModelListingCounter() + monkeypatch.setattr(litellm, "callbacks", [counter]) + monkeypatch.setitem( + proxy_server.app.dependency_overrides, + user_api_key_auth, + lambda: UserAPIKeyAuth(api_key=None, user_id="jwt-user", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + _chatgpt_responses(respx_mock) + turns: Final = tuple( + client.post( + "/backend-api/codex/responses", + json={"model": "gpt-unlisted", "input": "Hello", "stream": True}, + headers={**_HEADERS, "x-litellm-api-key": "test-jwt-credential"}, + ) + for _ in range(2) + ) + assert [turn.status_code for turn in turns] == [200, 200] + assert len(counter.listings) == 1 + first_etag, second_etag = (turn.headers["x-models-etag"] for turn in turns) + assert first_etag == second_etag != _UPSTREAM_MODELS_ETAG diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index cd6c1ff11b5..45abb880f92 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -61,6 +61,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( WebsocketFrameModelRefused, WebsocketFrameModelVerdict, _build_passthrough_failure_request_payload, + _parse_request_data_by_content_type, _registered_pass_through_routes, _truncate_upstream_error_body, _update_metadata_with_tags_in_header, @@ -13242,3 +13243,106 @@ async def test_passthrough_header_and_callback_share_finalized_cost( finally: cleanup() await client.aclose() + + +def _stub_adapter_proxy(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + proxy_logging: Final = MagicMock() + proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + proxy_logging.post_call_failure_hook = AsyncMock() + + async def fake_add_litellm_data_to_request(**kwargs: object) -> object: + return kwargs["data"] + + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging) + monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + return proxy_logging + + +@pytest.mark.asyncio +async def test_chat_completion_pass_through_endpoint_reads_a_body_only_json_can_parse( + monkeypatch: pytest.MonkeyPatch, +): + proxy_logging: Final = _stub_adapter_proxy(monkeypatch) + request: Final = MagicMock(spec=Request) + request.body = AsyncMock( + return_value=b'{"model": "unknown-model", "stream": false, "user": null, "messages": [{"role": "user", "content": "hi"}]}' + ) + + with pytest.raises(ProxyException): + await chat_completion_pass_through_endpoint( + fastapi_response=Response(), + request=request, + adapter_id="anthropic", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + forwarded: Final = proxy_logging.pre_call_hook.call_args.kwargs["data"] + assert (forwarded["model"], forwarded["stream"], forwarded["user"], forwarded["adapter_id"]) == ( + "unknown-model", + False, + None, + "anthropic", + ) + assert forwarded["messages"] == [{"role": "user", "content": "hi"}] + + +@pytest.mark.asyncio +async def test_chat_completion_pass_through_endpoint_rejects_a_body_that_is_not_an_object( + monkeypatch: pytest.MonkeyPatch, +): + proxy_logging: Final = _stub_adapter_proxy(monkeypatch) + request: Final = MagicMock(spec=Request) + request.body = AsyncMock(return_value=b'["not", "an", "object"]') + + with pytest.raises(ProxyException) as raised: + await chat_completion_pass_through_endpoint( + fastapi_response=Response(), + request=request, + adapter_id="anthropic", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert raised.value.message + proxy_logging.pre_call_hook.assert_not_awaited() + + +def _request_with_body(content_type: str, body: bytes) -> Request: + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + return Request( + { + "type": "http", + "method": "POST", + "path": "/pass-through", + "query_string": b"", + "headers": [(b"content-type", content_type.encode())], + }, + receive, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("body", "expected"), + [ + (b'{"model": "m", "stream": true}', (None, {"model": "m", "stream": True}, None, True)), + ( + b'{"query_params": {"a": "1"}, "custom_body": {"b": 2}, "stream": false}', + ({"a": "1"}, {"b": 2}, None, False), + ), + (b"{}", (None, None, None, None)), + (b"[1, 2]", (None, None, None, None)), + (b'--x\r\nContent-Disposition: form-data; name="f"\r\n\r\nv\r\n--x--\r\n', (None, None, None, None)), + ], + ids=["json_body_is_the_custom_body", "explicit_envelope", "empty_object", "json_array", "real_multipart"], +) +async def test_parse_request_data_by_content_type_reads_json_sent_as_multipart( + body: bytes, expected: tuple[object, object, None, object] +): + request: Final = _request_with_body("multipart/form-data; boundary=x", body) + + assert await _parse_request_data_by_content_type(request) == expected diff --git a/tests/unit/proxy/proxy_server/test_routes_model_info.py b/tests/unit/proxy/proxy_server/test_routes_model_info.py index 6ad5ef1a4eb..79f3e697a4f 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -490,6 +490,51 @@ def test_model_info_reports_null_cost_for_unpriced_deployment_and_zero_for_decla assert input_cost > 0 and output_cost > 0 +@pytest.mark.parametrize( + ("deployment_model_info", "native"), + [ + ({}, True), + ({"supports_bedrock_native_structured_output": False}, False), + ({"supports_native_structured_output": True, "supports_bedrock_native_structured_output": False}, True), + ], + ids=["native-by-default", "bedrock-opt-out", "explicit-generic-flag"], +) +def test_model_info_reports_a_deployment_bedrock_native_output_opt_out( + local_model_cost_map: None, deployment_model_info: dict[str, object], native: bool +) -> None: + info: Final = proxy_server._enrich_model_info_with_litellm_data( + { + "model_name": "sonnet", + "litellm_params": {"model": "bedrock/us.anthropic.claude-sonnet-5"}, + "model_info": {"id": "sonnet-deployment", **deployment_model_info}, + } + )["model_info"] + + assert info["supports_native_structured_output"] is native + + +@pytest.mark.parametrize( + ("model", "by_default"), + [("bedrock/us.anthropic.claude-opus-4-7", True), ("bedrock/us.anthropic.claude-sonnet-5", False)], + ids=["set-on-a-model-the-cost-map-leaves-on-the-json-tool", "unset-on-a-model-the-cost-map-makes-native"], +) +def test_model_info_ignores_a_deployment_by_default_key_the_router_reads_only_from_the_cost_map( + local_model_cost_map: None, model: str, by_default: bool +) -> None: + def reported_native(deployment_model_info: dict[str, object]) -> object: + return proxy_server._enrich_model_info_with_litellm_data( + { + "model_name": "claude", + "litellm_params": {"model": model}, + "model_info": {"id": "claude-deployment", **deployment_model_info}, + } + )["model_info"]["supports_native_structured_output"] + + assert reported_native({"supports_bedrock_native_structured_output_by_default": by_default}) is reported_native( + {} + ) + + def test_model_info_id_lookup_reports_the_same_cost_as_the_list( client: TestClient, auth_as: Callable[[], AbstractContextManager[object]], diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index d1afcf80ec9..27ba8968200 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -987,6 +987,57 @@ def test_public_model_hub_reports_the_newest_check_of_a_group_not_the_last_row_r app.dependency_overrides.clear() +def test_public_model_hub_lists_each_deployments_prices_without_exposing_the_deployment() -> None: + llm_router: Final = litellm.Router( + model_list=[ + { + "model_name": "hub-public", + "litellm_params": { + "model": "openai/hub-secret-backend-a", + "api_base": "https://secret-a.internal", + "api_key": "sk-secret-a", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + }, + "model_info": {"id": "hub-secret-id-a"}, + }, + { + "model_name": "hub-public", + "litellm_params": { + "model": "openai/hub-secret-backend-b", + "api_base": "https://secret-b.internal", + "api_key": "sk-secret-b", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 4e-06, + }, + "model_info": {"id": "hub-secret-id-b"}, + }, + ] + ) + app: Final = FastAPI() + app.include_router(router) + + with ( + patch("litellm.public_model_groups", ["hub-public"]), + patch("litellm.proxy.proxy_server.llm_router", llm_router), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + response: Final = TestClient(app).get("/public/model_hub") + + assert response.status_code == 200, response.text + groups: Final = TypeAdapter(list[ModelGroupInfoProxy]).validate_json(response.text) + assert [dict(p.costs) for p in groups[0].deployment_pricing] == [ + {"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}, + {"input_cost_per_token": 3e-06, "output_cost_per_token": 4e-06}, + ] + leaked: Final = [ + secret + for secret in ("secret-a.internal", "secret-b.internal", "sk-secret", "hub-secret-backend", "hub-secret-id") + if secret in response.text + ] + assert leaked == [], response.text + + # --------------------------------------------------------------------------- # /public/agent_hub # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/router_entry_point_stubs.py b/tests/unit/proxy/router_entry_point_stubs.py index 1d30cb99b44..1445fe517c6 100644 --- a/tests/unit/proxy/router_entry_point_stubs.py +++ b/tests/unit/proxy/router_entry_point_stubs.py @@ -125,6 +125,8 @@ def route_entry_points_to_stubs(router: _RouterDouble) -> _RouterDouble: return router.get_configured_display_name(name) if field == "service_tiers": return router.get_configured_service_tiers(name, **_given(team_id=team_id)) + if field == "pinned_reasoning_efforts": + return router.get_pinned_reasoning_efforts(name, **_given(team_id=team_id)) if field == "max_input_tokens": return router.get_configured_max_input_tokens(name, **_given(team_id=team_id)) if field == "deployment_model_info": diff --git a/tests/unit/proxy/test_caching_routes.py b/tests/unit/proxy/test_caching_routes.py index 8a59793495a..e139593976e 100644 --- a/tests/unit/proxy/test_caching_routes.py +++ b/tests/unit/proxy/test_caching_routes.py @@ -397,3 +397,20 @@ def test_extract_cache_params_is_empty_when_a_backend_attribute_name_is_not_a_st _use_redis_backed_cache(monkeypatch, {attribute_name: "redis.internal"}) assert _extract_cache_params() == {}, attribute_name + + +@pytest.fixture +def mock_redis_delete(mocker): + mock_cache = mocker.MagicMock() + mock_cache.type = "redis" + mock_cache.delete_cache_keys = mocker.AsyncMock() + mocker.patch.object(litellm, "cache", mock_cache) + return mock_cache + + +def test_cache_delete_removes_the_requested_keys(mock_redis_delete): + response = client.post("/cache/delete", headers={"Authorization": "Bearer sk-9876"}, json={"keys": ["k1", "k2"]}) + + assert response.status_code == 200 + assert response.json() == {"status": "success"} + mock_redis_delete.delete_cache_keys.assert_awaited_once_with(keys=["k1", "k2"]) diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index 3b000ba741f..f74ff9b1fca 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -427,8 +427,7 @@ def test_enrich_http_exception_with_guardrail_context_dict_detail(): assert exc.detail["guardrail_mode"] == "post_call" -def test_enrich_http_exception_string_detail_noop(): - """L2: string-detail HTTPException is not mutated (can't add fields to a str).""" +def test_enrich_http_exception_string_detail_preserves_body_and_identifies_guardrail(): from litellm.proxy.guardrails.exception_utils import enrich_http_exception_with_guardrail_context class StubCallback: @@ -438,6 +437,7 @@ def test_enrich_http_exception_string_detail_noop(): exc = HTTPException(status_code=400, detail="Content blocked") enrich_http_exception_with_guardrail_context(exc, StubCallback()) assert exc.detail == "Content blocked" + assert getattr(exc, "guardrail_name") == "x" def test_enrich_http_exception_setdefault_does_not_overwrite(): diff --git a/tests/unit/proxy/test_tools_allowlist_enforcement.py b/tests/unit/proxy/test_tools_allowlist_enforcement.py index e306f7be4f8..efaf25c480c 100644 --- a/tests/unit/proxy/test_tools_allowlist_enforcement.py +++ b/tests/unit/proxy/test_tools_allowlist_enforcement.py @@ -242,6 +242,42 @@ class TestCheckToolsAllowlist: assert exc_info.value.type == ProxyErrorTypes.tool_access_denied assert "restricted_tool" in str(exc_info.value.message) + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("allowed_tools", "denied"), + [ + pytest.param(["get_weather"], True, id="forbidden"), + pytest.param(["get_weather", "delete_user"], False, id="allowed"), + ], + ) + async def test_tool_defined_in_an_additional_tools_input_item_is_checked_on_responses_route( + self, allowed_tools: list[str], denied: bool + ) -> None: + body = { + "model": "fireworks_ai/glm-5p3", + "input": [ + { + "type": "additional_tools", + "role": "developer", + "tools": [{"type": "function", "name": "delete_user", "parameters": {"type": "object"}}], + }, + {"type": "message", "role": "user", "content": "Delete user 42"}, + ], + } + check = check_tools_allowlist( + request_body=body, + valid_token=_token(metadata={"allowed_tools": allowed_tools}), + team_object=None, + route="/v1/responses", + ) + if not denied: + await check + return + with pytest.raises(ProxyException) as exc_info: + await check + assert exc_info.value.type == ProxyErrorTypes.tool_access_denied + assert "delete_user" in str(exc_info.value.message) + @pytest.mark.parametrize( "loaded_tool", [ diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py index 05c4f9d8a67..3f3795dd33c 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -659,3 +659,14 @@ async def test_authoritative_combined_key_view_uses_writer_through_rotation( assert response.token == "current-token" reader.query_first.assert_not_awaited() assert writer.query_first.await_count == (2 if rotated else 1) + + +@pytest.mark.asyncio +async def test_get_data_all_users_pages_the_raw_user_query(prisma_client: PrismaClient) -> None: + rows: Final = [{"user_id": "u-1", "spend": 3.0, "key_aliases": ["alias-a"]}] + prisma_client.db.query_raw = AsyncMock(return_value=rows) + + result: Final = await prisma_client.get_data(table_name="user", query_type="find_all", limit=10, offset=20) + + assert result == rows + assert prisma_client.db.query_raw.await_args.args[1:] == (10, 20) diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 5fd40696224..50b5684f5e3 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -699,6 +699,7 @@ def _moderation_guardrail() -> MagicMock: cb.should_run_guardrail = MagicMock(return_value=True) cb.async_moderation_hook = AsyncMock(return_value=None) cb.async_post_call_success_hook = AsyncMock(return_value=None) + cb.post_call_already_checked = MagicMock(return_value=False) cb.run_in_parallel = False return cb diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 73d3414abb0..fa4a1656bc3 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -3450,3 +3450,131 @@ def test_proxy_vector_store_writes_send_only_the_callers_metadata_when_the_raw_b assert response.status_code == 200, response.text assert len(sent_bodies) == 1 assert sent_bodies[0].get("metadata") == CALLERS_METADATA + + +class _RecordingManagedVectorStores: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + async def acreate_vector_store(self, **kwargs: object) -> dict[str, object]: + self.calls.append(kwargs) + return {"id": "vs_managed", "object": "vector_store"} + + +def _vector_store_create_request(body: dict[str, object]) -> Request: + encoded = json.dumps(body).encode() + + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": encoded, "more_body": False} + + return Request( + { + "type": "http", + "method": "POST", + "path": "/v1/vector_stores", + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + }, + receive, + ) + + +def _install_managed_vector_stores( + monkeypatch: pytest.MonkeyPatch, managed: _RecordingManagedVectorStores +) -> MagicMock: + from litellm.proxy import proxy_server + + llm_router = MagicMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.get_proxy_hook.side_effect = lambda name: managed if name == "managed_vector_stores" else None + monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging_obj) + monkeypatch.setattr(proxy_server, "llm_router", llm_router) + return llm_router + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("target_model_names", "expected_targets"), + [("embed-a, embed-b", ["embed-a", "embed-b"]), (["embed-a"], ["embed-a"])], + ids=["comma_separated_string", "list"], +) +async def test_vector_store_create_with_target_model_names_creates_a_managed_store( + monkeypatch: pytest.MonkeyPatch, target_model_names: object, expected_targets: list[str] +) -> None: + from fastapi import Response + + from litellm.proxy.vector_store_endpoints.endpoints import vector_store_create + + managed = _RecordingManagedVectorStores() + llm_router = _install_managed_vector_stores(monkeypatch, managed) + + response = await vector_store_create( + request=_vector_store_create_request({"name": "docs", "target_model_names": target_model_names}), + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert response == {"id": "vs_managed", "object": "vector_store"} + (call,) = managed.calls + assert call["target_model_names_list"] == expected_targets + assert call["llm_router"] is llm_router + assert call["create_request"]["name"] == "docs" + assert "target_model_names" not in call["create_request"] + + +@pytest.mark.asyncio +async def test_vector_store_create_rejects_target_model_names_that_are_not_a_string_or_list( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import Response + + from litellm.proxy.vector_store_endpoints.endpoints import vector_store_create + + managed = _RecordingManagedVectorStores() + _install_managed_vector_stores(monkeypatch, managed) + + with pytest.raises(HTTPException) as raised: + await vector_store_create( + request=_vector_store_create_request({"name": "docs", "target_model_names": 5}), + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert raised.value.status_code == 400 + assert managed.calls == [] + + +@pytest.mark.asyncio +async def test_update_vector_store_persists_submitted_metadata_as_json() -> None: + from litellm.proxy.vector_store_endpoints.management_endpoints import update_vector_store + from litellm.types.vector_stores import VectorStoreUpdateRequest + + stored_row = MagicMock() + stored_row.model_dump = MagicMock(return_value={"vector_store_id": "vs_owned", "team_id": "team-A"}) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=stored_row) + mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock(return_value=stored_row) + + with ( + patch( + "litellm.proxy.vector_store_endpoints.management_endpoints.check_feature_access_for_user", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.vector_store_endpoints.management_endpoints._check_vector_store_access", + new_callable=AsyncMock, + return_value=True, + ), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.vector_store_registry", None), + ): + await update_vector_store( + data=VectorStoreUpdateRequest( + vector_store_id="vs_owned", + vector_store_metadata={"owner": "search-team", "tier": 2}, + ), + user_api_key_dict=UserAPIKeyAuth(user_id="owner", team_id="team-A"), + ) + + written = mock_prisma_client.db.litellm_managedvectorstorestable.update.await_args.kwargs["data"] + assert json.loads(written["vector_store_metadata"]) == {"owner": "search-team", "tier": 2} diff --git a/tests/unit/proxy/video_endpoints/test_endpoints.py b/tests/unit/proxy/video_endpoints/test_endpoints.py index ecc7bccfa06..967f8089c70 100644 --- a/tests/unit/proxy/video_endpoints/test_endpoints.py +++ b/tests/unit/proxy/video_endpoints/test_endpoints.py @@ -260,6 +260,23 @@ async def test_generation__route_type_data_and_no_provider_default(harness): harness.batch_to_bytesio.assert_not_called() +@pytest.mark.asyncio +async def test_generation_preserves_first_class_video_fields(harness: Harness) -> None: + body: Final = { + "model": "queue-video/model", + "prompt": "a cat", + "duration": 8, + "aspect_ratio": "16:9", + "frame_images": [{"type": "image_url", "image_url": {"url": "https://example.com/first.png"}, "frame_type": "first_frame"}], + "input_references": [{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}}], + "generate_audio": False, + "seed": 0, + } + response: Final = await call_generation(harness, body=body) + assert response is SENTINEL + assert harness.processor_data() == body + + @pytest.mark.asyncio async def test_generation__input_reference_attached(harness): body = {"model": "sora-2", "prompt": "a sunset"} diff --git a/tests/unit/responses/litellm_completion_transformation/test_reasoning_items.py b/tests/unit/responses/litellm_completion_transformation/test_reasoning_items.py index 093d1744418..580ee672500 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_reasoning_items.py +++ b/tests/unit/responses/litellm_completion_transformation/test_reasoning_items.py @@ -1,12 +1,18 @@ import json +import pytest + from litellm.responses.litellm_completion_transformation.reasoning_items import ( decode_thinking_blocks, encode_thinking_blocks, is_litellm_minted_reasoning_item, is_minted_reasoning_item_id, + is_other_provider_reasoning_item, mint_reasoning_item_id, ) +from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams A_PROVIDER_OWNED_REASONING_ITEM_ID = "rs_08d3a89dbb92277a006abf04f4266087d0b4eedacd7848f306" A_PROVIDER_OWNED_ENCRYPTED_BLOB = "gAAAAABo-opaque-provider-blob" @@ -58,5 +64,82 @@ def test_a_reasoning_item_is_litellm_minted_by_its_id_or_by_its_encoded_thinking "encrypted_content": A_PROVIDER_OWNED_ENCRYPTED_BLOB, } ) - assert not is_litellm_minted_reasoning_item({"type": "message", "id": mint_reasoning_item_id(), "role": "assistant"}) + assert not is_litellm_minted_reasoning_item( + {"type": "message", "id": mint_reasoning_item_id(), "role": "assistant"} + ) assert not is_litellm_minted_reasoning_item("a bare string input") + + +SUMMARY = [{"type": "summary_text", "text": "Paris first."}] +PROVIDER_OWNED_ITEM = { + "type": "reasoning", + "id": A_PROVIDER_OWNED_REASONING_ITEM_ID, + "summary": SUMMARY, + "encrypted_content": A_PROVIDER_OWNED_ENCRYPTED_BLOB, +} + + +@pytest.mark.parametrize( + "item", + [ + pytest.param({"type": "reasoning", "id": mint_reasoning_item_id(), "summary": SUMMARY}, id="minted id"), + pytest.param( + {"type": "reasoning", "summary": SUMMARY, "encrypted_content": encode_thinking_blocks([SIGNED_BLOCK])}, + id="encoded thinking blocks", + ), + pytest.param( + {"type": "reasoning", "summary": [], "content": [{"type": "reasoning_text", "text": "Paris first."}]}, + id="plain-text reasoning content", + ), + pytest.param( + { + **PROVIDER_OWNED_ITEM, + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + A_PROVIDER_OWNED_ENCRYPTED_BLOB, "another-deployment" + ), + }, + id="blob wrapped with a litellm deployment id", + ), + ], +) +def test_reasoning_items_another_provider_produced_are_recognized(item): + assert is_other_provider_reasoning_item(item) + + +@pytest.mark.parametrize( + "item", + [ + pytest.param(PROVIDER_OWNED_ITEM, id="provider-owned blob"), + pytest.param({"type": "reasoning", "summary": SUMMARY}, id="summary only, no id"), + pytest.param( + {"type": "reasoning", "summary": SUMMARY, "content": None, "encrypted_content": None}, id="nulled replay" + ), + pytest.param({"type": "reasoning", "summary": SUMMARY, "content": []}, id="empty content"), + pytest.param( + {"type": "message", "id": mint_reasoning_item_id(), "role": "user", "content": "hi"}, id="not reasoning" + ), + pytest.param("a stray string", id="not an object"), + ], +) +def test_reasoning_items_the_provider_can_replay_are_kept(item): + assert not is_other_provider_reasoning_item(item) + + +def test_signed_thinking_survives_the_model_id_wrap_on_its_way_back_to_the_chat_completions_bridge(): + wrapped = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + encode_thinking_blocks([SIGNED_BLOCK]), "anthropic-deployment" + ) + messages = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=[ + {"role": "user", "content": "Capital of France?"}, + {"type": "reasoning", "id": mint_reasoning_item_id(), "summary": [], "encrypted_content": wrapped}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Paris."}]}, + {"role": "user", "content": "And of Italy?"}, + ], + responses_api_request=ResponsesAPIOptionalRequestParams(), + )["messages"] + + assert [message.get("thinking_blocks") for message in messages if message["role"] == "assistant"] == [ + [SIGNED_BLOCK] + ] diff --git a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py index a16661c0174..44ba9007d92 100644 --- a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -18,7 +18,9 @@ from openai.types.responses.tool_param import Mcp import litellm from litellm.caching.caching import DualCache +from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.mcp_server import operations as mcp_operations from litellm.proxy.common_utils.user_api_key_cache import object_permission_cache_key @@ -505,6 +507,139 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey assert post_call_failure_hook.await_args.kwargs.get("route") == "/responses/mcp/call_tool" +@pytest.mark.asyncio +@pytest.mark.parametrize("behavior", [None, "tool_result", "raise"]) +@pytest.mark.parametrize("error_kind", ["pii", "guardrail", "http", "string_http", "unattributed_http"]) +@pytest.mark.parametrize("mode", ["pre_mcp_call", "during_mcp_call"]) +async def test_mcp_block_behavior_preserves_failure_logging_and_propagates_only_opted_in_blocks( + monkeypatch: pytest.MonkeyPatch, + behavior: Literal["tool_result", "raise"] | None, + error_kind: str, + mode: str, +) -> None: + call_tool: Final = _setup_mcp_call_environment(monkeypatch) + failure_hook: Final = _setup_proxy_logging(monkeypatch) + callback: Final = CustomGuardrail(guardrail_name="blocked", event_hook=mode, default_on=True) + callback.mcp_block_behavior = behavior + callback.use_native_lifecycle_hooks = True + unrelated: Final = CustomGuardrail(guardrail_name="blocked", event_hook=mode, default_on=True) + unrelated.mcp_block_behavior = "tool_result" if behavior == "raise" else "raise" + monkeypatch.setattr(litellm, "callbacks", [unrelated, callback]) + error: Final = ( + BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="blocked") + if error_kind == "pii" + else GuardrailRaisedException(guardrail_name="blocked", message="denied") + if error_kind == "guardrail" + else HTTPException(status_code=400, detail="denied") + if error_kind == "string_http" + else HTTPException( + status_code=400, + detail={"error": "denied", **({"guardrail_name": "blocked"} if error_kind == "http" else {})}, + ) + ) + logger: Final = ProxyLogging(user_api_key_cache=DualCache()) + settings: Final = types.SimpleNamespace( + guardrail_list=[{"guardrail_name": "blocked", "callback": item} for item in (unrelated, callback)], + is_set=lambda name: name == "guardrail_list", + guardrail=lambda name: {"callback": callback}, + ) + monkeypatch.setattr( + sys.modules["litellm.proxy.proxy_server"], + "llm_router", + types.SimpleNamespace(settings=lambda: settings), + raising=False, + ) + auth: Final = UserAPIKeyAuth(api_key="test-key") + + async def block(**kwargs: object) -> None: + raise error + + callback.async_pre_call_hook = block + callback.async_moderation_hook = block + + async def call_with_guardrail(**kwargs: object) -> None: + if error_kind == "unattributed_http": + raise error + data: Final = {"mcp_arguments": {"query": "denied"}} + if mode == "pre_mcp_call": + await logger._process_guardrail_callback( + unrelated, data, auth, "call_mcp_tool", GuardrailEventHooks.pre_mcp_call + ) + else: + await logger._run_during_call_guardrail(callback, data, auth, auth, "call_mcp_tool") + + call_tool.side_effect = call_with_guardrail + tool_calls: Final = [{"id": "blocked-call", "function": {"name": "search", "arguments": "{}"}}] + + if behavior == "raise" and error_kind != "unattributed_http": + with pytest.raises(type(error)) as caught: + await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={"search": "local"}, tool_calls=tool_calls, user_api_key_auth=auth + ) + assert caught.value is error + else: + results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={"search": "local"}, tool_calls=tool_calls, user_api_key_auth=auth + ) + assert len(results) == 1 + assert results[0]["tool_call_id"] == "blocked-call" + expected_prefix: Final = "Tool call failed" if "http" in error_kind else "Tool call blocked" + assert results[0]["result"].startswith(expected_prefix) + failure_hook.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("behavior", [None, "tool_result", "raise"]) +async def test_initialized_content_filter_blocks_mcp_according_to_its_config( + monkeypatch: pytest.MonkeyPatch, behavior: Literal["tool_result", "raise"] | None +) -> None: + for name in ( + "callbacks", "input_callback", "success_callback", "failure_callback", + "_async_success_callback", "_async_failure_callback", + ): + monkeypatch.setattr(litellm, name, []) + handler: Final = InMemoryGuardrailHandler() + handler.initialize_guardrail( + guardrail={ + "guardrail_name": "configured-block", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_mcp_call", + "default_on": True, + "blocked_words": [{"keyword": "forbiddenmarker", "action": "BLOCK"}], + **({"mcp_block_behavior": behavior} if behavior is not None else {}), + }, + } + ) + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="local", name="local", transport=MCPTransport.http, url="http://upstream") + manager.registry = {server.server_id: server} + manager.tool_name_to_mcp_server_name_mapping = {"local-search": "local"} + client: Final = AsyncMock() + client.call_tool.return_value = CallToolResult(content=[TextContent(type="text", text="ok")]) + manager.create_mcp_client = AsyncMock(return_value=client) + logger: Final = ProxyLogging(user_api_key_cache=DualCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", logger) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", manager) + tool_calls: Final = [{"id": "blocked", "function": {"name": "search", "arguments": '{"query":"forbiddenmarker"}'}}] + kwargs: Final = { + "tool_server_map": {"search": "local"}, + "tool_calls": tool_calls, + "user_api_key_auth": UserAPIKeyAuth(api_key="test-key"), + "served_tools": [MCPTool(name="search", inputSchema={"type": "object"})], + } + if behavior == "raise": + with pytest.raises(HTTPException, match="forbiddenmarker"): + await LiteLLM_Proxy_MCP_Handler.execute_tool_calls(**kwargs) + else: + results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls(**kwargs) + assert len(results) == 1 + assert results[0]["tool_call_id"] == "blocked" + assert "Content blocked: keyword 'forbiddenmarker' detected" in results[0]["result"] + client.call_tool.assert_not_awaited() + + @pytest.mark.asyncio async def test_execute_tool_calls_logs_each_tool_call_under_its_own_id_on_the_request_trace( monkeypatch, diff --git a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py index f745b5bfadf..f113d00a471 100644 --- a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py +++ b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py @@ -6,9 +6,12 @@ from typing import Final, cast from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import HTTPException from mcp.types import CallToolResult, TextContent import litellm # noqa: F401 - ensures litellm.responses.main is registered in sys.modules +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails.exception_utils import enrich_http_exception_with_guardrail_context from litellm.responses.mcp.mcp_streaming_iterator import ( MCPEnhancedStreamingIterator, ) @@ -293,6 +296,27 @@ async def test_continuation_id_is_final_round_not_interim_tool_call(monkeypatch) assert completed[-1].response.id != "resp-interim" +@pytest.mark.asyncio +async def test_opted_in_guardrail_block_emits_terminal_error_without_follow_up(monkeypatch): + call_tool: Final = _mock_mcp_environment(monkeypatch) + callback: Final = CustomGuardrail(guardrail_name="blocked") + callback.mcp_block_behavior = "raise" + error: Final = HTTPException(status_code=400, detail={"error": "guardrail denied", "guardrail_name": "blocked"}) + enrich_http_exception_with_guardrail_context(error, callback) + call_tool.side_effect = error + follow_up: Final = AsyncMock(return_value=_text_only_stream("unguarded follow-up")) + monkeypatch.setattr(responses_main_module, "aresponses", follow_up) + iterator: Final = _make_iterator([_completed_chunk([_function_call("call_1", "read_wiki_contents")])]) + + chunks: Final = [chunk async for chunk in iterator] + + assert chunks[-1].type == ResponsesAPIStreamEvents.ERROR + assert chunks[-1].error.code == "400" + assert "guardrail denied" in chunks[-1].error.message + assert all(chunk.type != ResponsesAPIStreamEvents.RESPONSE_COMPLETED for chunk in chunks) + follow_up.assert_not_awaited() + + @pytest.mark.asyncio async def test_follow_up_call_failure_emits_terminal_error_event(monkeypatch): """ diff --git a/tests/unit/responses/test_additional_tools.py b/tests/unit/responses/test_additional_tools.py index 290670a7c0a..d2c8cbacf78 100644 --- a/tests/unit/responses/test_additional_tools.py +++ b/tests/unit/responses/test_additional_tools.py @@ -1,8 +1,9 @@ from typing import Final +import pytest from pydantic import JsonValue -from litellm.responses.additional_tools import hoist_additional_tools +from litellm.responses.additional_tools import hoist_additional_tools, with_additional_tools_hoisted _EXEC_TOOL = { "type": "custom", @@ -134,3 +135,10 @@ def test_tool_list_repeated_on_every_turn_is_sent_as_is_without_drop_params() -> assert hoisted.tools == (_TOP_LEVEL_TOOL, _function("exec", "old"), _function("exec", "new")) assert hoisted.input == [_USER_MESSAGE, _USER_MESSAGE] + + +@pytest.mark.parametrize("tools", [7, "exec", {"type": "function"}], ids=["int", "str", "dict"]) +def test_request_whose_top_level_tools_is_not_a_list_stays_as_sent_for_the_provider_to_refuse(tools: object) -> None: + request: Final = {"input": [_tool_list(_function("exec")), _USER_MESSAGE], "tools": tools} + + assert with_additional_tools_hoisted(request, drop_params=True) == request diff --git a/tests/unit/responses/test_responses_utils.py b/tests/unit/responses/test_responses_utils.py index 0863abd1f1e..49ff7918849 100644 --- a/tests/unit/responses/test_responses_utils.py +++ b/tests/unit/responses/test_responses_utils.py @@ -415,8 +415,8 @@ class TestResponseAPILoggingUtils: assert result.prompt_tokens_details.cache_creation_tokens == 10059 assert result.prompt_tokens_details.cached_tokens == 0 - @pytest.mark.parametrize("as_model", [False, True]) - def test_transform_response_api_usage_prices_12h_cache_writes(self, as_model: bool) -> None: + @pytest.mark.parametrize("usage_shape", ["dict", "validated_model", "unvalidated_model"]) + def test_transform_response_api_usage_prices_12h_cache_writes(self, usage_shape: str) -> None: usage_data: Final = { "input_tokens": 5_000, "output_tokens": 0, @@ -426,7 +426,11 @@ class TestResponseAPILoggingUtils: "cache_write_tokens_details": {"ttl_30m_tokens": 3_000, "ttl_12h_tokens": 2_000}, }, } - usage_input: Final = ResponseAPIUsage(**usage_data) if as_model else usage_data + usage_input: Final = { + "dict": usage_data, + "validated_model": ResponseAPIUsage(**usage_data), + "unvalidated_model": ResponseAPIUsage.model_construct(**usage_data), + }[usage_shape] usage: Final = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage_input) model_info: Final[ModelInfo] = { "input_cost_per_token": 1e-6, diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index 14ecb185f06..0acf5d9ce62 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -37,6 +37,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.responses.streaming_iterator import ( _ERROR_CODE_HTTP_STATUS, BaseResponsesAPIStreamingIterator, @@ -75,6 +76,83 @@ def _make_error_chunk(error_type: str, code: str, message: str = "err") -> Error return ErrorEvent(type=ResponsesAPIStreamEvents.ERROR, sequence_number=0, error=error_obj) +@pytest.mark.parametrize( + "error_type,error_code", + ( + ("resource_unavailable", "flex_unavailable"), + ("resource_unavailable", None), + (None, "flex_unavailable"), + ), +) +def test_streamed_flex_unavailable_maps_to_429(error_type: str | None, error_code: str | None) -> None: + assert _status_code_for_error_fields(error_type, error_code) == 429 + + +def test_streamed_unrelated_server_error_still_maps_to_500() -> None: + assert _status_code_for_error_fields("server_error", "server_error") == 500 + + +FLEX_UNAVAILABLE_ERROR: Final = { + "type": "resource_unavailable", + "code": "flex_unavailable", + "message": "Flex processing is temporarily unavailable. Please try again later or use standard processing.", +} + + +def _flex_outage_sse(kind: str) -> str: + pending: Final = { + "id": "resp_flex", + "object": "response", + "created_at": 1700000000, + "model": "gpt-5", + "status": "in_progress", + "output": [], + } + failure: Final = ( + {"type": "error", "sequence_number": 1, "error": FLEX_UNAVAILABLE_ERROR} + if kind == "error" + else { + "type": "response.failed", + "sequence_number": 1, + "response": {**pending, "status": "failed", "error": FLEX_UNAVAILABLE_ERROR}, + } + ) + events: Final = ({"type": "response.created", "sequence_number": 0, "response": pending}, failure) + return "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ("error", "response.failed")) +async def test_streamed_flex_outage_raises_rate_limit_like_the_non_streamed_path( + kind: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + async_client: Final = AsyncHTTPHandler( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 200, headers={"content-type": "text/event-stream"}, text=_flex_outage_sse(kind) + ) + ) + ) + try: + stream: Final = await litellm.aresponses( + model="openai/gpt-5", + input="hi", + service_tier="flex", + stream=True, + api_key="sk-test", + client=async_client, + ) + with pytest.raises(MidStreamFallbackError) as raised: + _ = [event async for event in stream] + finally: + await async_client.close() + assert raised.value.status_code == 429, raised.value + assert isinstance(raised.value.original_exception, litellm.RateLimitError), raised.value.original_exception + assert "Flex processing is temporarily unavailable" in str(raised.value.original_exception) + + @pytest.mark.parametrize("provider", ("openai", "azure", "bedrock")) @pytest.mark.parametrize("header", ("x-amzn-RequestId", "x-request-id", None)) @pytest.mark.parametrize( diff --git a/tests/unit/router_strategy/adaptive_router/test_hooks.py b/tests/unit/router_strategy/adaptive_router/test_hooks.py index 0ed22658144..d199f2e1922 100644 --- a/tests/unit/router_strategy/adaptive_router/test_hooks.py +++ b/tests/unit/router_strategy/adaptive_router/test_hooks.py @@ -1,5 +1,6 @@ """Unit tests for the AdaptiveRouterPostCallHook.""" +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest @@ -15,6 +16,7 @@ from litellm.router_strategy.adaptive_router.hooks import ( _resolve_session_key, ) from litellm.router_strategy.adaptive_router.signals import Turn +from litellm.types.utils import ModelResponse def _make_hook() -> AdaptiveRouterPostCallHook: @@ -236,6 +238,45 @@ async def test_hook_passes_tool_calls_through(): assert turn.tool_calls == [tc] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "expected_tool_calls"), + [ + ( + ModelResponse( + choices=[ + { + "message": { + "role": "assistant", + "content": "calling tool", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "search", "arguments": '{"q":"x"}'}, + } + ], + } + } + ] + ), + [{"id": "call_1", "type": "function", "function": {"name": "search", "arguments": '{"q":"x"}'}}], + ), + ( + _resp_with_content("calling tool", tool_calls=[SimpleNamespace(name="search")]), + [{"name": "search", "arguments": ""}], + ), + ], + ids=["model_response_tool_call", "tool_call_without_model_dump"], +) +async def test_hook_records_tool_call_objects_as_dicts(response: object, expected_tool_calls: list[dict[str, object]]): + hook = _make_hook() + await hook.async_log_success_event(_kwargs(chosen="fast"), response, 0.0, 1.0) + turn: Turn = hook.adaptive_router.record_turn.await_args.kwargs["turn"] + assert turn.assistant_content == "calling tool" + assert turn.tool_calls == expected_tool_calls + + # ---- _recent_tool_results ------------------------------------------------ diff --git a/tests/unit/router_strategy/adaptive_router/test_signals.py b/tests/unit/router_strategy/adaptive_router/test_signals.py index bd58b4c668d..f4b1298cd46 100644 --- a/tests/unit/router_strategy/adaptive_router/test_signals.py +++ b/tests/unit/router_strategy/adaptive_router/test_signals.py @@ -4,7 +4,7 @@ from typing import List, Tuple import pytest -from litellm.router_strategy.adaptive_router.config import TOOL_CALL_HISTORY_MAX +from litellm.router_strategy.adaptive_router.config import LOOP_REPEAT_THRESHOLD, TOOL_CALL_HISTORY_MAX from litellm.router_strategy.adaptive_router.signals import ( SessionState, SignalDelta, @@ -184,3 +184,33 @@ def test_apply_turn_is_o1_does_not_grow_history_unbounded(): Turn(tool_calls=[{"name": f"tool_{i}", "arguments": {}}]), ) assert len(state.tool_call_history) <= TOOL_CALL_HISTORY_MAX + + +def _function_tool_call(name: str, arguments: str) -> dict[str, object]: + return {"id": "call", "type": "function", "function": {"name": name, "arguments": arguments}} + + +@pytest.mark.parametrize( + ("names", "arguments", "expected_loops"), + [ + (("read_file",) * LOOP_REPEAT_THRESHOLD, ('{"path": "/x"}',) * LOOP_REPEAT_THRESHOLD, 1), + ( + tuple(f"tool_{i}" for i in range(LOOP_REPEAT_THRESHOLD)), + ('{"path": "/x"}',) * LOOP_REPEAT_THRESHOLD, + 0, + ), + ( + ("read_file",) * LOOP_REPEAT_THRESHOLD, + tuple(f'{{"path": "/{i}"}}' for i in range(LOOP_REPEAT_THRESHOLD)), + 0, + ), + ], + ids=["same_call", "different_names", "different_arguments"], +) +def test_loop_detection_reads_name_and_arguments_of_function_tool_calls( + names: Tuple[str, ...], arguments: Tuple[str, ...], expected_loops: int +): + state, _ = _replay( + [{"tool_calls": [_function_tool_call(name, args)]} for name, args in zip(names, arguments, strict=True)] + ) + assert state.loop_count == expected_loops diff --git a/tests/unit/router_strategy/switchyard_router/__init__.py b/tests/unit/router_strategy/switchyard_router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_strategy/switchyard_router/test_config.py b/tests/unit/router_strategy/switchyard_router/test_config.py new file mode 100644 index 00000000000..c24023ddb79 --- /dev/null +++ b/tests/unit/router_strategy/switchyard_router/test_config.py @@ -0,0 +1,100 @@ +from typing import Final + +import pytest +from pydantic import ValidationError + +from litellm.router_strategy.switchyard_router.config import SwitchyardRouterConfig, validate_switchyard_config_write + + +@pytest.mark.parametrize("classifier", [None, {}]) +def test_default_confidence_depends_on_classifier(classifier: dict[str, object] | None) -> None: + config: Final = SwitchyardRouterConfig.model_validate( + { + "models": { + "capable": "strong", + "efficient": "fast", + **({"judge": "judge"} if classifier is not None else {}), + }, + "classifier": classifier, + } + ) + assert config.confidence_threshold == (0.7 if classifier is not None else 0.5) + + +@pytest.mark.parametrize("classifier", [None, {}]) +def test_explicit_confidence_is_preserved(classifier: dict[str, object] | None) -> None: + config: Final = SwitchyardRouterConfig.model_validate( + { + "models": { + "capable": "strong", + "efficient": "fast", + **({"judge": "judge"} if classifier is not None else {}), + }, + "classifier": classifier, + "options": {"confidence_threshold": 0.63, "picker": "capable_first"}, + } + ) + assert config.confidence_threshold == 0.63 + + +@pytest.mark.parametrize( + "option", ["escalation_note", "deescalation_note", "capable_system_prompt", "efficient_system_prompt", "unknown"] +) +def test_request_rewrite_and_unknown_options_are_rejected(option: str) -> None: + with pytest.raises(ValidationError, match=option): + SwitchyardRouterConfig.model_validate( + {"models": {"capable": "strong", "efficient": "fast"}, "options": {option: "ignored"}} + ) + + +@pytest.mark.parametrize( + "bad", + [ + {"algorithm": "random"}, + {"classifier": {}}, + {"options": {"confidence_threshold": 2}}, + {"models": {"capable": "", "efficient": "fast"}}, + {"models": {"capable": "same", "efficient": "same"}}, + {"models": {"capable": "strong", "efficient": "fast", "judge": "j"}, "classifier": {"session_affinity": True}}, + {"models": {"capable": "strong", "efficient": "fast", "judge": "judge"}}, + {"models": {"capable": "strong", "efficient": "fast", "judge": "j"}, "classifier": {"kind": "other"}}, + { + "models": {"capable": "strong", "efficient": "fast", "judge": "j"}, + "classifier": {"kind": "decision", "cutoff": 1.5}, + }, + { + "models": {"capable": "strong", "efficient": "fast", "judge": "j"}, + "classifier": {"kind": "decision", "base_threshold": 0.5}, + }, + { + "models": {"capable": "strong", "efficient": "fast", "judge": "j"}, + "classifier": {"kind": "llm", "cutoff": 0.5}, + }, + ], +) +def test_invalid_config_cannot_be_saved(bad: dict[str, object]) -> None: + violation: Final = validate_switchyard_config_write({"models": {"capable": "strong", "efficient": "fast"}, **bad}) + assert violation is not None + assert "Invalid switchyard_router_config" in violation + + +def test_missing_package_rejects_save_with_install_instructions() -> None: + import sys + from unittest.mock import patch + + with patch.dict(sys.modules, {"switchyard_rust.libsy": None}): + violation: Final = validate_switchyard_config_write({"models": {"capable": "strong", "efficient": "fast"}}) + assert violation is not None + assert "litellm[switchyard]" in violation + + +@pytest.mark.parametrize( + "classifier", [{"kind": "decision"}, {"kind": "decision", "cutoff": 0.4, "evidence": {"a": 1}}] +) +def test_decision_classifier_config_can_be_saved(classifier: dict[str, object]) -> None: + assert ( + validate_switchyard_config_write( + {"models": {"capable": "strong", "efficient": "fast", "judge": "jev"}, "classifier": classifier} + ) + is None + ) diff --git a/tests/unit/router_strategy/switchyard_router/test_switchyard_router.py b/tests/unit/router_strategy/switchyard_router/test_switchyard_router.py new file mode 100644 index 00000000000..8a26b4610db --- /dev/null +++ b/tests/unit/router_strategy/switchyard_router/test_switchyard_router.py @@ -0,0 +1,438 @@ +import asyncio +import json +from itertools import chain +from collections.abc import AsyncIterator, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_api import ChangeSet +from litellm.router_strategy.switchyard_router.config import SwitchyardRouterConfig, build_algorithm +from litellm.router_strategy.switchyard_router.switchyard_router import SwitchyardRouter +from litellm.types.decisions import OpenAIDecisionResponse +from litellm.types.router import Deployment +from litellm.types.utils import ModelResponse + +libsy = pytest.importorskip("switchyard_rust.libsy") + +HELLO: Final = [{"role": "user", "content": "Hello"}] + + +async def unexpected_judge(request: Mapping[str, object]) -> ModelResponse: + raise AssertionError("This route must not call an LLM judge") + + +async def unexpected_decision(request: Mapping[str, object]) -> OpenAIDecisionResponse: + raise AssertionError("This route must not call a decision model") + + +def config( + *, classifier: Mapping[str, object] | None = None, picker: str = "efficient_first", threshold: float | None = None +) -> SwitchyardRouterConfig: + return SwitchyardRouterConfig.model_validate( + { + "models": { + "capable": "strong", + "efficient": "fast", + **({"judge": "judge"} if classifier is not None else {}), + }, + "options": {"picker": picker, **({"confidence_threshold": threshold} if threshold is not None else {})}, + **({"classifier": classifier} if classifier is not None else {}), + } + ) + + +def router( + route_config: SwitchyardRouterConfig, + *, + completion=unexpected_judge, + decide=unexpected_decision, + algorithm: object = None, +) -> SwitchyardRouter: + return SwitchyardRouter("route", route_config, algorithm or build_algorithm(route_config), completion, decide) + + +async def route( + route_config: SwitchyardRouterConfig, + messages: object, + *, + call_type: str | None = "acompletion", + request_kwargs: Mapping[str, object] | None = None, + **router_kwargs: object, +) -> tuple[str, Mapping[str, object]]: + result: Final = await router(route_config, **router_kwargs).async_pre_routing_hook( # pyright: ignore[reportArgumentType] # test doubles + "route", request_kwargs or {}, messages, call_type=call_type + ) + assert result is not None + assert result.routing_decision is not None + return result.model, result.routing_decision + + +def anthropic_turn(call_id: str, tool: str, result: str, marked: bool) -> tuple[dict[str, object], ...]: + error: Final = {"is_error": True} if marked else {} + return ( + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "Next step", "signature": "sig"}, + {"type": "tool_use", "id": call_id, "name": tool, "input": {"file_path": "a.py"}}, + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": call_id, "content": result, **error}]}, + ) + + +def responses_turn(call_id: str, tool: str, result: str, marked: bool) -> tuple[dict[str, object], ...]: + return ( + {"type": "function_call", "call_id": call_id, "name": tool, "arguments": '{"file_path":"a.py"}'}, + {"type": "function_call_output", "call_id": call_id, "output": result}, + ) + + +def chat_turn(call_id: str, tool: str, result: str, marked: bool) -> tuple[dict[str, object], ...]: + call: Final = {"id": call_id, "type": "function", "function": {"name": tool, "arguments": '{"file_path":"a.py"}'}} + return ( + {"role": "assistant", "content": None, "tool_calls": [call]}, + {"role": "tool", "tool_call_id": call_id, "content": result}, + ) + + +TURN_FORMATS: Final = {"anthropic_messages": anthropic_turn, "aresponses": responses_turn, "acompletion": chat_turn} + + +def transcript(endpoint: str, tool: str, results: tuple[str, ...], *, marked: bool = False) -> list[dict[str, object]]: + turn: Final = TURN_FORMATS[endpoint] + turns: Final = chain.from_iterable( + turn(f"call{index}", tool, result, marked) for index, result in enumerate(results) + ) + return [{"role": "user", "content": "Fix the helper"}, *turns] + + +def llm_router( + *, route_config: Mapping[str, object], judge_fails: bool = False, num_retries: int = 0 +) -> litellm.Router: + def deployment(name: str, reply: object) -> Mapping[str, object]: + return { + "model_name": name, + "litellm_params": {"model": f"openai/{name}", "api_key": "k", "mock_response": reply}, + } + + return litellm.Router( + num_retries=num_retries, + model_list=[ + deployment("strong", "strong"), + deployment("fast", "fast"), + *((deployment("judge", Exception("judge down")),) if judge_fails else ()), + switchyard_deployment(route_config), + ], + ) + + +def switchyard_deployment(route_config: Mapping[str, object]) -> dict[str, object]: + return { + "model_name": "route", + "litellm_params": {"model": "auto_router/switchyard", "switchyard_router_config": route_config}, + "model_info": {"id": "route-id"}, + } + + +@pytest.mark.asyncio +async def test_critical_tool_error_selects_capable_without_rewriting_messages() -> None: + messages: Final = transcript("acompletion", "bash", ("CUDA out of memory",)) + result: Final = await router(config()).async_pre_routing_hook("route", {}, messages, call_type="acompletion") + assert result is not None + assert result.model == "strong" + assert result.messages == messages + + +async def routed_reply(proxy: litellm.Router, endpoint: str, edits: tuple[str, ...]) -> object: + match endpoint: + case "anthropic_messages": + message: Final = await proxy.aanthropic_messages( + model="route", messages=transcript(endpoint, "Edit", edits), max_tokens=16 + ) + return message["content"][0]["text"] # pyright: ignore[reportIndexIssue] # non-streaming response + case "aresponses": + response: Final = await proxy.aresponses(model="route", input=transcript(endpoint, "Edit", edits)) + return response.output[0].content[0].text # pyright: ignore[reportAttributeAccessIssue] # non-streaming response + case _: + chat: Final = await proxy.acompletion(model="route", messages=transcript(endpoint, "Edit", edits)) + return chat.choices[0].message.content # pyright: ignore[reportAttributeAccessIssue] # non-streaming response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint", ["acompletion", "anthropic_messages", "aresponses"]) +async def test_router_routes_agent_traffic_from_every_endpoint(endpoint: str) -> None: + proxy: Final = llm_router(route_config={"models": {"capable": "strong", "efficient": "fast"}}) + assert await routed_reply(proxy, endpoint, ("Edit applied",) * 3) == "fast" + + +@pytest.mark.asyncio +async def test_router_follows_an_edited_switchyard_router_and_forgets_a_deleted_one() -> None: + proxy: Final = llm_router(route_config={"models": {"capable": "strong", "efficient": "fast"}}) + assert await routed_reply(proxy, "acompletion", ("Edit applied",) * 3) == "fast" + + swapped: Final = switchyard_deployment({"models": {"capable": "fast", "efficient": "strong"}}) + _ = proxy.apply(ChangeSet(upsert=[Deployment(**swapped)])) + assert await routed_reply(proxy, "acompletion", ("Edit applied",) * 3) == "strong" + + _ = proxy.apply(ChangeSet(delete=["route-id"])) + with pytest.raises(litellm.BadRequestError): + await proxy.acompletion(model="route", messages=HELLO) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("marked", "expected"), [(True, "strong"), (False, "fast")]) +async def test_anthropic_tool_failure_marker_reaches_switchyard(marked: bool, expected: str) -> None: + failures: Final = transcript("anthropic_messages", "Edit", ("String to replace not found.",) * 4, marked=marked) + model, decision = await route(config(), failures, call_type="anthropic_messages") + assert model == expected + assert "switchyard_fallback_reason" not in decision + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_kwargs", "expected"), + [ + ({"litellm_session_id": "s1"}, ["strong", "strong", "strong", "fast"]), + ( + {"proxy_server_request": {"headers": {"X-Session-Id": "s1", "session-id": "s1"}}}, + ["strong", "fast", "fast", "fast"], + ), + ({}, ["strong", "fast", "fast", "fast"]), + ], +) +async def test_capable_hold_follows_the_session(request_kwargs: Mapping[str, object], expected: list[str]) -> None: + switchyard: Final = router(config()) + turns: Final = (("CUDA out of memory",), *((("All tests passed",) * count) for count in (2, 3, 4))) + picks: Final = [await pick(switchyard, request_kwargs, results) for results in turns] + assert picks == expected + + +async def pick(switchyard: SwitchyardRouter, request_kwargs: Mapping[str, object], results: tuple[str, ...]) -> str: + result: Final = await switchyard.async_pre_routing_hook( + "route", request_kwargs, transcript("acompletion", "bash", results), call_type="acompletion" + ) + assert result is not None + return result.model + + +@pytest.mark.asyncio +async def test_a_session_hold_never_crosses_api_keys() -> None: + switchyard: Final = router(config()) + owner: Final = {"litellm_session_id": "shared", "metadata": {"user_api_key_hash": "owner"}} + other: Final = {"litellm_session_id": "shared", "metadata": {"user_api_key_hash": "other"}} + assert await pick(switchyard, owner, ("CUDA out of memory",)) == "strong" + assert await pick(switchyard, other, ("All tests passed",) * 2) == "fast" + assert await pick(switchyard, owner, ("All tests passed",) * 2) == "strong" + + +@pytest.mark.asyncio +async def test_a_client_metadata_key_hash_cannot_join_another_keys_session() -> None: + switchyard: Final = router(config()) + owner: Final = {"litellm_session_id": "shared", "litellm_metadata": {"user_api_key_hash": "owner"}} + spoofer: Final = { + "litellm_session_id": "shared", + "metadata": {"user_api_key_hash": "owner"}, + "litellm_metadata": {"user_api_key_hash": "spoofer"}, + } + assert await pick(switchyard, owner, ("CUDA out of memory",)) == "strong" + assert await pick(switchyard, spoofer, ("All tests passed",) * 2) == "fast" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("p_solve", "expected"), [(0.0, "strong"), (1.0, "fast")]) +async def test_llm_judge_verdict_selects_the_target(p_solve: float, expected: str) -> None: + async def judge(request: Mapping[str, object]) -> ModelResponse: + assert "autorouter_classifier" in str(request["metadata"]) + verdict: Final = {"crux": "c", "primary_rule": "none", "capability_boundary": "unmatched", "p_solve": p_solve} + return ModelResponse(choices=[{"message": {"role": "assistant", "content": json.dumps(verdict)}}]) + + model, _ = await route(config(classifier={}), HELLO, completion=judge) + assert model == expected + + +def decision_answer(advantage: float) -> OpenAIDecisionResponse: + return OpenAIDecisionResponse.model_validate( + { + "model": "jev", + "answers": [ + { + "type": "choice", + "name": "route", + "choice": "advantage", + "probabilities": [ + {"value": "advantage", "probability": advantage}, + {"value": "no_advantage", "probability": 1 - advantage}, + ], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 0, "total_tokens": 1}, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("parent_bucket", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize("classifier", [{}, {"kind": "decision"}]) +async def test_judge_calls_carry_the_callers_tags_and_origin_where_their_endpoint_reads_them( + parent_bucket: str, classifier: Mapping[str, object] +) -> None: + seen: Final[list[Mapping[str, object]]] = [] # mutable-ok: collects the judge request + + async def judge(request: Mapping[str, object]) -> ModelResponse: + seen.append(request) + verdict: Final = {"crux": "c", "primary_rule": "none", "capability_boundary": "unmatched", "p_solve": 1.0} + return ModelResponse(choices=[{"message": {"role": "assistant", "content": json.dumps(verdict)}}]) + + async def decide(request: Mapping[str, object]) -> OpenAIDecisionResponse: + seen.append(request) + return decision_answer(0.0) + + caller: Final = {parent_bucket: {"tags": ["eu-only"], "inherited_tags": ["team-eu"]}} + await route( + config(classifier=classifier, threshold=1.0), HELLO, completion=judge, decide=decide, request_kwargs=caller + ) + (request,) = seen + read_by_endpoint: Final = "litellm_metadata" if classifier else "metadata" + assert set(request) & {"metadata", "litellm_metadata"} == {read_by_endpoint} + metadata: Final = request[read_by_endpoint] + assert metadata["tags"] == ["eu-only"] # pyright: ignore[reportIndexIssue] # the judge request's metadata dict + assert metadata["inherited_tags"] == ["team-eu"] # pyright: ignore[reportIndexIssue] # same + assert metadata["internal_call_origin"] == "autorouter_classifier" # pyright: ignore[reportIndexIssue] # same + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("advantage", "expected"), [(0.61, "strong"), (0.6, "fast")]) +async def test_decision_model_advantage_must_exceed_the_cutoff(advantage: float, expected: str) -> None: + async def decide(request: Mapping[str, object]) -> OpenAIDecisionResponse: + (question,) = request["questions"] # pyright: ignore[reportGeneralTypeIssues] # the decisions request shape + assert {choice["value"] for choice in question["choices"]} == {"advantage", "no_advantage"} + return decision_answer(advantage) + + model, decision = await route( + config(classifier={"kind": "decision", "cutoff": 0.6}, threshold=1.0), HELLO, decide=decide + ) + assert model == expected + assert decision["switchyard_evidence"]["score"] == advantage + + +async def judge_down(request: Mapping[str, object]) -> ModelResponse: + raise ConnectionError("judge down") + + +async def decision_refused(request: Mapping[str, object]) -> OpenAIDecisionResponse: + return OpenAIDecisionResponse.model_validate( + { + "model": "jev", + "answers": [{"type": "refusal", "name": "route"}], + "usage": {"input_tokens": 1, "output_tokens": 0, "total_tokens": 1}, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("classifier", "router_kwargs", "picker", "expected"), + [ + ({}, {"completion": judge_down}, "efficient_first", "fast"), + ({}, {"completion": judge_down}, "capable_first", "strong"), + ({"kind": "decision"}, {"decide": decision_refused}, "efficient_first", "fast"), + ], +) +async def test_classifier_failure_falls_open_to_the_picker_default( + classifier: Mapping[str, object], router_kwargs: Mapping[str, object], picker: str, expected: str +) -> None: + model, decision = await route(config(classifier=classifier, picker=picker, threshold=1.0), HELLO, **router_kwargs) + assert model == expected + assert "switchyard_fallback_reason" not in decision + + +class EndsWithoutDecision: + async def _steps(self) -> AsyncIterator[object]: + for step in (): + yield step + + def run_stream(self, *args: object, **kwargs: object) -> AsyncIterator[object]: + return self._steps() + + +class RaisesLibsyError: + async def _steps(self) -> AsyncIterator[object]: + raise libsy.LibsyError("algorithm failed") + yield + + def run_stream(self, *args: object, **kwargs: object) -> AsyncIterator[object]: + return self._steps() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route_config", "messages", "call_type", "algorithm", "reason"), + [ + ( + config(), + [{"role": "user", "content": [{"type": "input_audio"}]}], + "acompletion", + None, + "unsupported_content", + ), + (config(), HELLO, None, None, "unsupported_request_type"), + (config(), HELLO, "acompletion", EndsWithoutDecision(), "no_decision"), + (config(), HELLO, "acompletion", RaisesLibsyError(), "libsy_error"), + ( + config(classifier={}, threshold=1.0), + HELLO, + "acompletion", + build_algorithm(config(classifier={"kind": "decision"}, threshold=1.0)), + "unexpected_decision_call", + ), + ( + config(classifier={"kind": "decision"}, threshold=1.0), + HELLO, + "acompletion", + build_algorithm(config(classifier={}, threshold=1.0)), + "unexpected_judge_call", + ), + ], +) +async def test_anything_switchyard_cannot_route_goes_to_capable( + route_config: SwitchyardRouterConfig, messages: object, call_type: str | None, algorithm: object, reason: str +) -> None: + model, decision = await route(route_config, messages, call_type=call_type, algorithm=algorithm) + assert model == "strong" + assert decision["switchyard_fallback_reason"] == reason + + +@pytest.mark.asyncio +async def test_failing_judge_is_called_once_even_when_the_router_retries(monkeypatch: pytest.MonkeyPatch) -> None: + attempts: Final[list[str]] = [] + + class CountAttempts(CustomLogger): + def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None: + attempts.append(model) + + monkeypatch.setattr(litellm, "callbacks", [CountAttempts()]) + proxy: Final = llm_router( + route_config={ + "models": {"capable": "strong", "efficient": "fast", "judge": "judge"}, + "classifier": {}, + "options": {"confidence_threshold": 1.0}, + }, + judge_fails=True, + num_retries=3, + ) + await proxy.acompletion(model="route", messages=HELLO) + assert attempts.count("judge") == 1 + + +@pytest.mark.asyncio +async def test_judge_cancellation_propagates() -> None: + async def cancelled(request: Mapping[str, object]) -> ModelResponse: + raise asyncio.CancelledError() + + with pytest.raises(asyncio.CancelledError): + await route(config(classifier={}), HELLO, completion=cancelled) diff --git a/tests/unit/router_strategy/test_router_tag_routing.py b/tests/unit/router_strategy/test_router_tag_routing.py index 0d3876c64a3..e916e5b235d 100644 --- a/tests/unit/router_strategy/test_router_tag_routing.py +++ b/tests/unit/router_strategy/test_router_tag_routing.py @@ -3023,7 +3023,7 @@ class _RewriteToTierStrategy: self.rewrite_to = rewrite_to async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False + self, model, request_kwargs, messages=None, input=None, specific_deployment=False, call_type=None ): from litellm.types.router import PreRoutingHookResponse diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 5169642bc43..76e9a70614f 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -2,6 +2,9 @@ import asyncio import copy import os import functools +import json +import logging +import threading import uuid from typing import Final, cast @@ -10,7 +13,11 @@ import pytest import litellm from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache -from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, PROMPT_CACHE_LOOKBACK_POSITIONS +from litellm.constants import ( + DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, + PROMPT_CACHE_LOOKBACK_POSITIONS, + REDACTED_BY_LITELLM, +) from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup @@ -19,6 +26,7 @@ from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import _honors_one_hour_cache_ttl, _get_min_token_count_for_deployments, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.prompt_caching_cache import PROMPT_CACHE_PIN_TTL_SECONDS, PromptCachingCache from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import StandardLoggingPayload @@ -138,10 +146,13 @@ def test_bedrock_application_profile_uses_its_base_model_for_one_hour_ttl() -> N }, ) - assert _honors_one_hour_cache_ttl( - standard_logging_object, - {"model": None, "base_model": "bedrock/us.anthropic.claude-sonnet-5-5"}, - ) is True + assert ( + _honors_one_hour_cache_ttl( + standard_logging_object, + {"model": None, "base_model": "bedrock/us.anthropic.claude-sonnet-5-5"}, + ) + is True + ) def test_bedrock_application_profile_cache_hit_reads_base_model_from_metadata_model_info() -> None: @@ -244,6 +255,16 @@ async def test_async_filter_deployments_narrows_prompt_above_model_minimum(): assert filtered == [deployments[1]] +class _ReadTrackingMessages(list[AllMessageValues]): + def __init__(self, messages: list[AllMessageValues]) -> None: + super().__init__(messages) + self.read = False + + def __iter__(self): + self.read = True + return super().__iter__() + + class _PinLookupCounter(DualCache): def __init__(self) -> None: super().__init__() @@ -275,17 +296,19 @@ async def test_async_filter_deployments_skips_prefix_hash_for_a_single_deploymen prompt that hash alone is ~30 ms of GIL-holding work per request. """ cache = _PinLookupCounter() - check = PromptCachingDeploymentCheck(cache=cache) + check = PromptCachingDeploymentCheck(cache=cache, group_has_alternatives=lambda group: True) deployments = _deployments("anthropic/claude-opus-4-6") messages = _messages(word_count=5000) await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-1", messages=messages, tools=None) + tracked_messages = _ReadTrackingMessages(messages) filtered = await check.async_filter_deployments( - model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=messages + model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=tracked_messages ) assert filtered == deployments assert cache.pin_lookups == 0 + assert tracked_messages.read is False two = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") assert await check.async_filter_deployments( @@ -313,10 +336,20 @@ async def test_async_filter_deployments_does_not_pin_when_target_order_is_set(): assert filtered == deployments +_RETRIES_ON_DEP_1: Final = tuple({"deployment_id": "dep-1"} for _ in range(4)) + + @pytest.mark.asyncio @pytest.mark.parametrize( "failed_in_request", - [{"_excluded_deployment_ids": ["dep-2"]}, {"_retry_skipped_deployment_ids": ("dep-2",)}], + [ + {"_excluded_deployment_ids": ["dep-2"]}, + {"_retry_skipped_deployment_ids": ("dep-2",)}, + {"metadata": {"previous_models": [{"deployment_id": "dep-2"}], "request_retry_count": 1}}, + {"litellm_metadata": {"previous_models": ({"deployment_id": "dep-2"},), "request_retry_count": 1}}, + {"metadata": {"previous_models": _RETRIES_ON_DEP_1, "request_retry_count": 5}}, + ], + ids=["excluded", "retry_skipped", "retried", "retried_litellm_metadata", "retries_rolled_over"], ) async def test_async_filter_deployments_does_not_pin_to_a_deployment_that_failed_this_request( failed_in_request: dict[str, object], @@ -334,6 +367,34 @@ async def test_async_filter_deployments_does_not_pin_to_a_deployment_that_failed ) == _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") +@pytest.mark.asyncio +@pytest.mark.parametrize( + "retry_metadata", + [ + {"previous_models": _RETRIES_ON_DEP_1, "request_retry_count": 4}, + {"previous_models": "dep-2"}, + {"previous_models": [None, "dep-2", {"deployment_id": None}], "request_retry_count": 3}, + ], + ids=["every_retry_on_record", "malformed_breadcrumbs", "malformed_records"], +) +async def test_async_filter_deployments_keeps_the_pin_while_every_retry_is_on_record( + retry_metadata: dict[str, object], +): + cache: Final = DualCache() + check: Final = PromptCachingDeploymentCheck(cache=cache) + messages: Final = _messages(word_count=5000) + await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) + + filtered: Final = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=_deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6"), + messages=messages, + request_kwargs={"metadata": retry_metadata}, + ) + + assert [deployment["model_info"]["id"] for deployment in filtered] == ["dep-2"] + + @pytest.mark.asyncio async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is_lower(): """ @@ -1535,3 +1596,920 @@ async def test_pin_is_written_with_message_logging_off_while_loggers_still_see_r await _eventually(lambda: len(probe.logged_messages) == 5) redacted = [{"role": "user", "content": "redacted-by-litellm"}] assert probe.logged_messages == [(redacted, redacted)] * 5, probe.logged_messages + + +RESPONSES_BRIDGED_MODEL: Final = "anthropic/claude-sonnet-4-6" +RESPONSES_NATIVE_MODEL: Final = "openai/gpt-5.4-mini" +_RESPONSES_TOOL: Final = { + "type": "function", + "name": "get_time", + "parameters": {"type": "object", "properties": {}}, +} + + +def _marked_input_text(text: str) -> dict[str, object]: + return {"type": "input_text", "text": text, "cache_control": {"type": "ephemeral"}} + + +def _marked_responses_input() -> list[dict[str, object]]: + return [ + {"role": "system", "content": [_marked_input_text(LONG_PROMPT)]}, + {"role": "user", "content": "hello"}, + ] + + +def _responses_router(model: str) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": model, "api_key": "sk-fake"}, + "model_info": {"id": model_id}, + } + for model_id in ("dep-1", "dep-2") + ], + optional_pre_call_checks=["prompt_caching"], + ) + + +async def _serve_response(router: litellm.Router, request_input: object) -> str: + response = await router.aresponses( + model=MODEL_GROUP_ALIAS, + input=copy.deepcopy(request_input), + instructions="Be terse.", + tools=[_RESPONSES_TOOL], + mock_response="ok", + ) + return response._hidden_params["model_id"] + + +async def _deployments_serving_repeats(router: litellm.Router, request_input: object) -> set[str]: + await _serve_response(router, request_input) + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + return {await _serve_response(router, request_input) for _ in range(8)} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", [RESPONSES_BRIDGED_MODEL, RESPONSES_NATIVE_MODEL], ids=["bridged", "native"]) +async def test_repeated_responses_request_pins_to_one_deployment(model: str): + assert len(await _deployments_serving_repeats(_responses_router(model), _marked_responses_input())) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", [RESPONSES_BRIDGED_MODEL, RESPONSES_NATIVE_MODEL], ids=["bridged", "native"]) +async def test_responses_thread_stays_on_the_deployment_holding_its_prefix(model: str): + router = _responses_router(model) + first_turn = _marked_responses_input() + pinned_deployment = await _serve_response(router, first_turn) + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + + later_turns = ( + [*first_turn, {"role": "assistant", "content": "ok"}, {"role": "user", "content": f"question {turn}"}] + for turn in range(8) + ) + + assert {await _serve_response(router, turn) for turn in later_turns} == {pinned_deployment} + + +@pytest.mark.asyncio +async def test_responses_request_pins_after_prompt_caching_is_turned_off_and_on_again(): + router = _responses_router(RESPONSES_NATIVE_MODEL) + await _serve_response(router, _marked_responses_input()) + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + router.update_settings(optional_pre_call_checks=[]) + router.update_settings(optional_pre_call_checks=["prompt_caching"]) + router.prompt_caching_pin_cache.in_memory_cache.flush_cache() + + assert len(await _deployments_serving_repeats(router, _marked_responses_input())) == 1 + + +def _team_tagged_router(model: str) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": model, "api_key": "sk-fake", "tags": [team]}, + "model_info": {"id": model_id}, + } + for model_id, team in (("dep-1", "team-a"), ("dep-2", "team-b")) + ], + optional_pre_call_checks=["prompt_caching"], + enable_tag_filtering=True, + ) + + +async def _serve_team_response(router: litellm.Router, team: str) -> str: + response = await router.aresponses( + model=MODEL_GROUP_ALIAS, + input=_marked_responses_input(), + instructions="Be terse.", + tools=[_RESPONSES_TOOL], + litellm_metadata={"tags": [team]}, + mock_response="ok", + ) + return response._hidden_params["model_id"] + + +async def _serve_team_completion(router: litellm.Router, team: str) -> str: + response = await router.acompletion( + model=MODEL_GROUP_ALIAS, + messages=_messages(word_count=5000), + metadata={"tags": [team]}, + mock_response="ok", + ) + return response._hidden_params["model_id"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model, serve", + [(RESPONSES_NATIVE_MODEL, _serve_team_response), (RESPONSES_BRIDGED_MODEL, _serve_team_completion)], + ids=["responses", "chat"], +) +async def test_prompt_cache_pin_never_overrides_the_request_tags(model: str, serve): + router = _team_tagged_router(model) + assert await serve(router, "team-a") == "dep-1" + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + + assert await serve(router, "team-b") == "dep-2" + + +@pytest.mark.asyncio +async def test_retry_moves_off_the_pinned_deployment_that_failed(): + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": { + "model": RESPONSES_BRIDGED_MODEL, + "api_key": "sk-fake", + "mock_response": mock_response, + "weight": weight, + }, + "model_info": {"id": model_id}, + } + for model_id, mock_response, weight in ( + ("dep-1", "litellm.InternalServerError", 1), + ("dep-2", "ok", 1000), + ) + ], + optional_pre_call_checks=["prompt_caching"], + num_retries=3, + retry_after=0, + disable_cooldowns=True, + ) + messages: Final = _messages(word_count=5000) + await PromptCachingCache(cache=router.prompt_caching_pin_cache).async_add_model_id( + model_id="dep-1", messages=messages, tools=None + ) + + response: Final = await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages) + + assert response._hidden_params["model_id"] == "dep-2" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reused_as", ["metadata", "litellm_metadata"]) +async def test_pin_holds_for_a_metadata_dict_reused_after_a_request_failed_on_the_pinned_deployment(reused_as: str): + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": RESPONSES_BRIDGED_MODEL, "api_key": "sk-fake", "weight": weight}, + "model_info": {"id": model_id}, + } + for model_id, weight in (("dep-1", 1000), ("dep-2", 1)) + ], + optional_pre_call_checks=["prompt_caching"], + num_retries=6, + retry_after=0, + allowed_fails=100, + ) + router.cooldown_cache.add_deployment_to_cooldown( + model_id="dep-1", + original_exception=litellm.RateLimitError( + message="upstream 429", llm_provider="anthropic", model=RESPONSES_BRIDGED_MODEL + ), + exception_status=429, + cooldown_time=60, + ) + reused_metadata: Final[dict[str, object]] = {} + with pytest.raises(litellm.InternalServerError): + await router.acompletion( + model=MODEL_GROUP_ALIAS, + messages=[{"role": "user", "content": "hello"}], + metadata=reused_metadata, + mock_response="litellm.InternalServerError", + ) + router.cooldown_cache.cooldown_store.evict_local("dep-1") + messages: Final = _messages(word_count=5000) + await PromptCachingCache(cache=router.prompt_caching_pin_cache).async_add_model_id( + model_id="dep-2", messages=messages, tools=None + ) + + response: Final = await router.acompletion( + model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok", **{reused_as: reused_metadata} + ) + + assert response._hidden_params["model_id"] == "dep-2" + + +@pytest.mark.asyncio +async def test_responses_mark_on_the_input_item_pins_to_one_deployment(): + request_input = [ + {"role": "system", "content": LONG_PROMPT, "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "hello"}, + ] + + assert len(await _deployments_serving_repeats(_responses_router(RESPONSES_BRIDGED_MODEL), request_input)) == 1 + + +async def _read_responses_pin( + check: PromptCachingDeploymentCheck, + request: dict[str, object], + models: tuple[str, ...] = (AUTO_CACHING_MODEL, AUTO_CACHING_MODEL), +) -> tuple[str, ...]: + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=_deployments(*models), + messages=None, + request_kwargs=request, + ) + return tuple(str(deployment["model_info"]["id"]) for deployment in filtered) + + +async def _log_responses_success( + check: PromptCachingDeploymentCheck, + logged_messages: object, + metadata: dict[str, object], + model_id: str = "dep-1", + model_map_value: dict[str, float | None] | None = None, + call_type: str = "aresponses", +) -> None: + event = _success_kwargs( + call_type=call_type, + model=AUTO_CACHING_MODEL, + model_id=model_id, + request_messages=logged_messages, + request_kwargs={}, + ) + await check.async_log_success_event( + kwargs={ + **event, + "standard_logging_object": { + **event["standard_logging_object"], + "model_map_information": {"model_map_key": AUTO_CACHING_MODEL, "model_map_value": model_map_value}, + }, + "litellm_params": {"litellm_metadata": metadata}, + }, + response_obj=None, + start_time=None, + end_time=None, + ) + + +async def _pinned_deployment_ids( + first_request: dict[str, object], logged_messages: object, next_request: dict[str, object] +) -> tuple[str, ...]: + check = PromptCachingDeploymentCheck(cache=DualCache()) + metadata: dict[str, object] = {} + await _read_responses_pin(check, {**first_request, "litellm_metadata": metadata}) + await _log_responses_success(check, logged_messages, metadata) + return await _read_responses_pin(check, {**next_request, "litellm_metadata": {}}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "logged_messages", + [ + _marked_responses_input(), + [ + {"role": "system", "content": "Be terse."}, + {"role": "system", "content": LONG_PROMPT, "cache_control": {"type": "ephemeral"}}, + {"role": "assistant", "tool_calls": [{"id": "call_1", "type": "function"}]}, + {"role": "tool", "tool_call_id": "call_1", "content": "12:00"}, + ], + [{"role": "user", "content": REDACTED_BY_LITELLM}], + ], + ids=["client_input", "bridged_chat_messages", "redacted"], +) +async def test_responses_pin_does_not_depend_on_what_the_success_event_logged( + logged_messages: list[dict[str, object]], +): + request = {"input": _marked_responses_input(), "instructions": "Be terse.", "tools": [_RESPONSES_TOOL]} + + pinned = await _pinned_deployment_ids(request, logged_messages, request) + + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "forged_pin", + [["forged-redis-key", 3600], {"cache_key": "forged-redis-key", "marked_ttl_seconds": 3600}], + ids=["array", "object"], +) +async def test_responses_pin_sent_in_the_request_metadata_is_never_written(forged_pin: object): + cache = DualCache() + + await _log_responses_success( + PromptCachingDeploymentCheck(cache=cache), + _marked_responses_input(), + {"responses_prompt_cache_pin": forged_pin}, + ) + + assert not cache.in_memory_cache.cache_dict + + +@pytest.mark.asyncio +async def test_interleaved_responses_requests_each_pin_their_own_prompt(): + check = PromptCachingDeploymentCheck(cache=DualCache()) + terse = {"input": _marked_responses_input(), "instructions": "Be terse."} + verbose = {"input": _marked_responses_input(), "instructions": "Be verbose."} + terse_metadata: dict[str, object] = {} + verbose_metadata: dict[str, object] = {} + await _read_responses_pin(check, {**terse, "litellm_metadata": terse_metadata}) + await _read_responses_pin(check, {**verbose, "litellm_metadata": verbose_metadata}) + + await _log_responses_success(check, _marked_responses_input(), terse_metadata, model_id="dep-1") + await _log_responses_success(check, _marked_responses_input(), verbose_metadata, model_id="dep-2") + + assert await _read_responses_pin(check, {**terse, "litellm_metadata": {}}) == ("dep-1",) + assert await _read_responses_pin(check, {**verbose, "litellm_metadata": {}}) == ("dep-2",) + + +@pytest.mark.asyncio +async def test_retried_responses_request_leaves_no_pin_in_flight_once_it_succeeds(): + pins_in_flight = InMemoryCache() + check = PromptCachingDeploymentCheck(cache=DualCache(), responses_pins_in_flight=pins_in_flight) + metadata: dict[str, object] = {} + request = {"input": _marked_responses_input(), "litellm_metadata": metadata} + await _read_responses_pin(check, request) + await _read_responses_pin(check, request) + + await _log_responses_success(check, _marked_responses_input(), metadata) + + assert not pins_in_flight.cache_dict + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["acompact_responses", "acreate_interaction"]) +async def test_request_sharing_the_responses_input_shape_leaves_no_pin_behind(call_type: str): + cache = DualCache() + pins_in_flight = InMemoryCache() + check = PromptCachingDeploymentCheck(cache=cache, responses_pins_in_flight=pins_in_flight) + metadata: dict[str, object] = {} + await _read_responses_pin(check, {"input": _marked_responses_input(), "litellm_metadata": metadata}) + assert pins_in_flight.cache_dict + + await _log_responses_success(check, _marked_responses_input(), metadata, call_type=call_type) + + assert not pins_in_flight.cache_dict + assert not cache.in_memory_cache.cache_dict + + +@pytest.mark.asyncio +async def test_responses_request_served_by_the_only_healthy_deployment_pins_it(): + cache = _PinLookupCounter() + check = PromptCachingDeploymentCheck(cache=cache, group_has_alternatives=lambda group: True) + request = {"input": _marked_responses_input(), "instructions": "Be terse."} + only_healthy = [_deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL)[1]] + metadata: dict[str, object] = {} + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=only_healthy, + messages=None, + request_kwargs={**request, "litellm_metadata": metadata}, + ) + assert filtered == only_healthy + assert cache.pin_lookups == 0 + + await _log_responses_success(check, _marked_responses_input(), metadata, model_id="dep-2") + + assert await _read_responses_pin(check, {**request, "litellm_metadata": {}}) == ("dep-2",) + + +@pytest.mark.asyncio +async def test_responses_request_to_a_single_deployment_group_skips_the_prefix_hash(): + cache = _PinLookupCounter() + pins_in_flight = InMemoryCache() + check = PromptCachingDeploymentCheck( + cache=cache, responses_pins_in_flight=pins_in_flight, group_has_alternatives=lambda group: False + ) + deployments = _deployments(AUTO_CACHING_MODEL) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=None, + request_kwargs={"input": _marked_responses_input(), "litellm_metadata": {}}, + ) + + assert filtered == deployments + assert cache.pin_lookups == 0 + assert not pins_in_flight.cache_dict + + +@pytest.mark.asyncio +async def test_responses_request_to_a_single_deployment_router_group_writes_no_pin(monkeypatch): + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": RESPONSES_NATIVE_MODEL, "api_key": "sk-fake"}, + "model_info": {"id": "dep-1"}, + } + ], + optional_pre_call_checks=["prompt_caching"], + ) + logged_after_the_check: list[str] = [] + + class _LoggedAfterTheCheck(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + logged_after_the_check.append(kwargs["litellm_call_id"]) + + monkeypatch.setattr(litellm, "callbacks", [*litellm.callbacks, _LoggedAfterTheCheck()]) + + assert await _serve_response(router, _marked_responses_input()) == "dep-1" + + assert await _eventually(lambda: logged_after_the_check) + assert not router.prompt_caching_pin_cache.in_memory_cache.cache_dict + + +def _distinct_marked_responses_input(index: int) -> list[dict[str, object]]: + return [ + {"role": "system", "content": [_marked_input_text(f"prompt {index} {LONG_PROMPT}")]}, + {"role": "user", "content": "hello"}, + ] + + +@pytest.mark.asyncio +async def test_responses_prompts_served_while_a_peer_cools_down_stay_pinned_once_it_recovers(): + router = _responses_router(RESPONSES_NATIVE_MODEL) + pins = router.prompt_caching_pin_cache.in_memory_cache.cache_dict + router.cooldown_cache.add_deployment_to_cooldown( + model_id="dep-1", + original_exception=litellm.RateLimitError( + message="upstream 429", llm_provider="openai", model=RESPONSES_NATIVE_MODEL + ), + exception_status=429, + cooldown_time=60, + ) + prompts = [_distinct_marked_responses_input(index) for index in range(8)] + for prompt in prompts: + pins_before = len(pins) + assert await _serve_response(router, prompt) == "dep-2" + assert await _eventually(lambda before=pins_before: len(pins) > before) + router.cooldown_cache.cooldown_store.evict_local("dep-1") + unpinned = { + ( + await router.aresponses(model=MODEL_GROUP_ALIAS, input=f"unpinned {index}", mock_response="ok") + )._hidden_params["model_id"] + for index in range(16) + } + assert unpinned == {"dep-1", "dep-2"} + + assert {await _serve_response(router, prompt) for prompt in prompts} == {"dep-2"} + + +@pytest.mark.asyncio +async def test_responses_with_other_instructions_are_a_different_prefix(): + pinned = await _pinned_deployment_ids( + {"input": _marked_responses_input(), "instructions": "Be terse."}, + _marked_responses_input(), + {"input": _marked_responses_input(), "instructions": "Be verbose."}, + ) + + assert len(pinned) == 2 + + +@pytest.mark.asyncio +async def test_responses_with_another_tool_set_are_a_different_prefix(): + other_tool = {**_RESPONSES_TOOL, "name": "get_weather"} + + pinned = await _pinned_deployment_ids( + {"input": _marked_responses_input(), "tools": [_RESPONSES_TOOL]}, + _marked_responses_input(), + {"input": _marked_responses_input(), "tools": [other_tool]}, + ) + + assert len(pinned) == 2 + + +@pytest.mark.asyncio +async def test_responses_custom_tools_with_another_grammar_are_a_different_prefix(): + custom_tool: Final = { + "type": "custom", + "name": "run_sql", + "description": "Run a query", + "format": {"type": "grammar", "syntax": "lark", "definition": "start: SELECT"}, + } + other_grammar: Final = {**custom_tool, "format": {"type": "grammar", "syntax": "lark", "definition": "start: UPDATE"}} + + pinned = await _pinned_deployment_ids( + {"input": _marked_responses_input(), "tools": [custom_tool]}, + _marked_responses_input(), + {"input": _marked_responses_input(), "tools": [other_grammar]}, + ) + + assert len(pinned) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "parameters", + [{"type": "object", "properties": {}}, {"properties": {"city": {"type": "string"}}}], + ids=["typed_parameters", "untyped_parameters"], +) +async def test_chat_and_responses_requests_sending_the_same_tool_share_one_prefix(parameters: dict[str, object]): + described_tool: Final = {**_RESPONSES_TOOL, "description": "Get the time", "parameters": parameters} + chat_tool: Final = { + "type": "function", + "function": {key: value for key, value in described_tool.items() if key != "type"}, + } + router: Final = _responses_router(RESPONSES_BRIDGED_MODEL) + await router.acompletion( + model=MODEL_GROUP_ALIAS, + messages=[ + {"role": "system", "content": [{**_marked_input_text(LONG_PROMPT), "type": "text"}]}, + {"role": "user", "content": "hello"}, + ], + tools=[chat_tool], + mock_response="ok", + ) + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + chat_pin_keys: Final = set(router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + router.prompt_caching_pin_cache.in_memory_cache.flush_cache() + + await router.aresponses( + model=MODEL_GROUP_ALIAS, input=_marked_responses_input(), tools=[described_tool], mock_response="ok" + ) + + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + assert set(router.prompt_caching_pin_cache.in_memory_cache.cache_dict) == chat_pin_keys + + +@pytest.mark.asyncio +async def test_responses_marked_tool_set_long_enough_to_cache_pins_on_its_own(): + marked_tool = {**_RESPONSES_TOOL, "description": LONG_PROMPT, "cache_control": {"type": "ephemeral"}} + request = {"input": "hello", "tools": [marked_tool]} + + pinned = await _pinned_deployment_ids(request, [{"role": "user", "content": "hello"}], request) + + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +async def test_responses_continuation_turn_is_left_to_the_deployment_holding_its_previous_response(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + new_thread = {"input": "hello", "instructions": LONG_PROMPT} + + pinned = await _pinned_deployment_ids( + new_thread, [{"role": "user", "content": "hello"}], {**new_thread, "previous_response_id": "resp_1"} + ) + + assert len(pinned) == 2 + + +@pytest.mark.asyncio +async def test_responses_answer_from_a_fallback_group_keeps_the_primary_groups_pin(): + check = PromptCachingDeploymentCheck(cache=DualCache()) + request = {"input": _marked_responses_input()} + assert await _serve_turn(check, request["input"]) == 2 + fallback_hop_metadata: dict[str, object] = {} + await _read_responses_pin(check, {**request, "litellm_metadata": fallback_hop_metadata}) + + await _log_responses_success(check, request["input"], fallback_hop_metadata, model_id="fallback-dep") + + pinned = await _read_responses_pin(check, {**request, "litellm_metadata": {}}) + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "answering_deployment,expected_pin", + [("dep-1", ("dep-1", "dep-2")), ("dep-2", ("dep-2",))], + ids=["below_its_minimum", "above_its_minimum"], +) +async def test_responses_pin_goes_only_to_a_deployment_whose_model_caches_the_prefix( + answering_deployment: str, expected_pin: tuple[str, ...] +): + mixed_group = ("anthropic/claude-opus-4-5", "anthropic/claude-sonnet-4-5") + check = PromptCachingDeploymentCheck(cache=DualCache()) + request = {"input": [{"role": "user", "content": [_marked_input_text("word " * 1400)]}]} + metadata: dict[str, object] = {} + await _read_responses_pin(check, {**request, "litellm_metadata": metadata}, mixed_group) + + await _log_responses_success(check, request["input"], metadata, model_id=answering_deployment) + + pinned = await _read_responses_pin(check, {**request, "litellm_metadata": {}}, mixed_group) + assert pinned == expected_pin + + +@pytest.mark.asyncio +async def test_embedding_request_never_takes_a_responses_pin(): + check = PromptCachingDeploymentCheck(cache=DualCache()) + pinned_input = _marked_responses_input() + assert await _serve_turn(check, pinned_input) == 2 + assert len(await _read_responses_pin(check, {"input": pinned_input, "litellm_metadata": {}})) == 1 + + filtered = await _read_responses_pin(check, {"input": pinned_input, "metadata": {}}) + + assert len(filtered) == 2 + + +class _InputReadOnlyOnceTheLoopTurns(list[dict[str, object]]): + def __init__(self, items: list[dict[str, object]], loop_turned: threading.Event) -> None: + super().__init__(items) + self.loop_turned: Final = loop_turned + self.loop_turned_before_first_read: bool | None = None + + def __iter__(self): + if self.loop_turned_before_first_read is None: + self.loop_turned_before_first_read = self.loop_turned.wait(timeout=5) + return super().__iter__() + + +@pytest.mark.asyncio +async def test_responses_thread_is_converted_off_the_event_loop(): + loop_turned = threading.Event() + thread = _InputReadOnlyOnceTheLoopTurns(_marked_responses_input(), loop_turned) + asyncio.get_running_loop().call_soon(loop_turned.set) + + await _read_responses_pin( + PromptCachingDeploymentCheck(cache=DualCache()), {"input": thread, "litellm_metadata": {}} + ) + + assert thread.loop_turned_before_first_read is True + + +@pytest.mark.asyncio +async def test_responses_long_instructions_count_toward_the_caching_minimum(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + request = {"input": "hello", "instructions": LONG_PROMPT} + + pinned = await _pinned_deployment_ids(request, [{"role": "user", "content": "hello"}], request) + + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_map_value,pinned_past_five_minutes", + [ + pytest.param({"cache_creation_input_token_cost_above_1hr": 6.6e-06}, True, id="one-hour-cache-writes"), + pytest.param({"cache_creation_input_token_cost_above_1hr": None}, False, id="no-one-hour-cache-writes"), + ], +) +async def test_responses_one_hour_pin_follows_the_deployments_one_hour_cache_support( + model_map_value: dict[str, float | None], pinned_past_five_minutes: bool +): + clock = _ManualClock() + check = PromptCachingDeploymentCheck(cache=DualCache(in_memory_cache=InMemoryCache(clock=clock))) + marked = {**_marked_input_text(LONG_PROMPT), "cache_control": {"type": "ephemeral", "ttl": "1h"}} + request = {"input": [{"role": "user", "content": [marked]}]} + metadata: dict[str, object] = {} + await _read_responses_pin(check, {**request, "litellm_metadata": metadata}) + await _log_responses_success(check, request["input"], metadata, model_map_value=model_map_value) + + clock.advance(PROMPT_CACHE_PIN_TTL_SECONDS + 1) + + assert (len(await _read_responses_pin(check, {**request, "litellm_metadata": {}})) == 1) is pinned_past_five_minutes + + +@pytest.mark.asyncio +async def test_responses_turn_without_input_leaves_routing_to_its_previous_response(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + request = {"previous_response_id": "resp_1", "instructions": LONG_PROMPT} + + pinned = await _pinned_deployment_ids(request, None, request) + + assert len(pinned) == 2 + + +async def _serve_turn(check: PromptCachingDeploymentCheck, request_input: object) -> int: + metadata: dict[str, object] = {} + pinned = await _read_responses_pin(check, {"input": request_input, "litellm_metadata": metadata}) + await _log_responses_success(check, request_input, metadata) + return len(pinned) + + +@pytest.mark.asyncio +async def test_responses_thread_longer_than_the_lookback_window_keeps_its_deployment(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + check = PromptCachingDeploymentCheck(cache=DualCache()) + turns = [ + {"role": "system", "content": LONG_PROMPT}, + *( + {"role": "assistant" if index % 2 else "user", "content": f"turn {index}"} + for index in range(2 * PROMPT_CACHE_LOOKBACK_POSITIONS) + ), + ] + + served = [await _serve_turn(check, turns[:end]) for end in range(2, len(turns) + 1, 2)] + + assert served[1:] == [1] * (len(served) - 1) + + +@pytest.mark.asyncio +async def test_responses_failover_writes_the_pin_for_the_deployment_that_answered(): + check = PromptCachingDeploymentCheck(cache=DualCache()) + request = {"input": _marked_responses_input()} + metadata: dict[str, object] = {} + await _read_responses_pin(check, {**request, "litellm_metadata": metadata}) + await _read_responses_pin(check, {**request, "litellm_metadata": metadata}) + await _log_responses_success(check, _marked_responses_input(), metadata, model_id="dep-2") + + pinned = await _read_responses_pin(check, {**request, "litellm_metadata": {}}) + + assert pinned == ("dep-2",) + + +@pytest.mark.asyncio +async def test_responses_pin_covers_only_what_follows_the_newest_compaction(): + compaction = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "summary"}), + } + prefix = {"role": "user", "content": [_marked_input_text(LONG_PROMPT)]} + + first_input = [{"role": "user", "content": "dropped turn"}, compaction, prefix] + + pinned = await _pinned_deployment_ids( + {"input": first_input}, + first_input, + {"input": [{"role": "user", "content": "another dropped turn"}, compaction, prefix]}, + ) + + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +async def test_responses_image_in_the_prefix_pins_to_one_deployment(): + request_input = [ + { + "role": "user", + "content": [{"type": "input_image", "image_url": ONE_PIXEL_PNG}, _marked_input_text(LONG_PROMPT)], + }, + {"role": "user", "content": "hello"}, + ] + + assert len(await _deployments_serving_repeats(_responses_router(RESPONSES_BRIDGED_MODEL), request_input)) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_input", + [[{"role": "system", "content": LONG_PROMPT}, {"role": "user", "content": "hello"}], LONG_PROMPT], + ids=["item_list", "plain_string"], +) +async def test_responses_auto_cached_prefix_pins_to_one_deployment(monkeypatch, request_input: object): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + + assert len(await _deployments_serving_repeats(_responses_router(RESPONSES_BRIDGED_MODEL), request_input)) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_input", + [ + [LONG_PROMPT, LONG_PROMPT], + [{"role": "user", "content": 7}], + 7, + [{"type": "web_search_call"}], + [{"type": "web_search_call", "id": "ws_1", "status": "completed", "action": {"type": "search"}}], + [{"type": ["message"], "role": "user", "content": LONG_PROMPT}], + [{"role": "user", "content": [{"type": "input_image"}]}], + ], + ids=[ + "embedding_strings", + "non_text_item_content", + "integer", + "web_search_call_without_id", + "search_action_without_query", + "unhashable_item_type", + "image_without_a_source", + ], +) +async def test_input_that_is_not_a_responses_prompt_never_narrows_or_raises(request_input: object): + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + metadata: dict[str, object] = {} + + filtered = await _read_responses_pin(check, {"input": request_input, "litellm_metadata": metadata}) + await _log_responses_success(check, request_input, metadata) + + assert filtered == ("dep-1", "dep-2") + assert not cache.in_memory_cache.cache_dict + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "next_tools,expected_pin", + [([{"type": "image_generation"}], ("dep-1",)), ([{"type": "web_search"}], ("dep-1", "dep-2"))], + ids=["same_native_tool", "other_native_tool"], +) +async def test_responses_native_tools_are_part_of_the_prefix( + next_tools: list[dict[str, object]], expected_pin: tuple[str, ...] +): + request_input = _marked_responses_input() + + pinned = await _pinned_deployment_ids( + {"input": request_input, "tools": [{"type": "image_generation"}]}, + request_input, + {"input": request_input, "tools": next_tools}, + ) + + assert pinned == expected_pin + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native_tool_type", ["shell", "image_generation", "computer_use", "local_shell"]) +async def test_responses_native_tool_routed_to_a_native_model_logs_no_dropped_tool_warning( + caplog: pytest.LogCaptureFixture, native_tool_type: str +): + router = _responses_router(RESPONSES_NATIVE_MODEL) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await router.aresponses( + model=MODEL_GROUP_ALIAS, + input=_marked_responses_input(), + tools=[{"type": native_tool_type}], + mock_response="ok", + ) + + assert not [record for record in caplog.records if "Dropping Responses API tool" in record.getMessage()] + + +class _MetadataSerializingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.serialized_metadata: list[str] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.serialized_metadata.append(json.dumps(kwargs["litellm_params"]["metadata"])) + + +@pytest.mark.asyncio +async def test_logger_serializing_the_request_metadata_never_sees_the_prompt_cache_key(monkeypatch): + serializing_logger = _MetadataSerializingLogger() + monkeypatch.setattr(litellm, "callbacks", [serializing_logger]) + router = _responses_router(RESPONSES_NATIVE_MODEL) + + assert len(await _deployments_serving_repeats(router, _marked_responses_input())) == 1 + + assert await _eventually(lambda: len(serializing_logger.serialized_metadata) == 9) + logged_metadata = "".join(serializing_logger.serialized_metadata) + pin_keys = tuple(router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + assert pin_keys + assert not [pin_key for pin_key in pin_keys if pin_key.split(":")[1] in logged_metadata] + + +@pytest.mark.asyncio +async def test_responses_call_hands_back_the_callers_metadata_without_the_pin_token(): + router: Final = _responses_router(RESPONSES_BRIDGED_MODEL) + caller_metadata: Final[dict[str, object]] = {"tags": ["team-a"]} + + await router.aresponses( + model=MODEL_GROUP_ALIAS, + input=_marked_responses_input(), + litellm_metadata=caller_metadata, + mock_response="ok", + ) + + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + assert "responses_prompt_cache_pin" not in caller_metadata + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failed_in_request,expected_deployments", + [({}, ("dep-2",)), ({"_excluded_deployment_ids": ["dep-2"]}, ("dep-1", "dep-2"))], + ids=["previous_deployment_healthy", "previous_deployment_failed_this_request"], +) +async def test_responses_continuation_turn_goes_to_the_deployment_that_answered_its_previous_response( + failed_in_request: dict[str, object], expected_deployments: tuple[str, ...] +): + previous_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id("anthropic", "dep-2", "resp_x") + request: dict[str, object] = { + "input": "hello", + "previous_response_id": previous_response_id, + "litellm_metadata": {}, + } + + pinned = await _read_responses_pin( + PromptCachingDeploymentCheck(cache=DualCache()), {**request, **failed_in_request} + ) + + assert pinned == expected_deployments + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bad_tool", ["get_time", None, 42], ids=["string", "null", "number"]) +async def test_responses_request_with_a_malformed_tool_is_left_to_the_provider(bad_tool: object): + request: dict[str, object] = {"input": _marked_responses_input(), "tools": [bad_tool], "litellm_metadata": {}} + + pinned = await _read_responses_pin(PromptCachingDeploymentCheck(cache=DualCache()), request) + + assert pinned == ("dep-1", "dep-2") diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 31fef208a8c..d285d6fa619 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -81,6 +81,7 @@ _IN_MEMORY_ONLY_CALLERS: Final = frozenset( "litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py", "litellm/responses/litellm_completion_transformation/transformation.py", "litellm/router_utils/client_initalization_utils.py", + "litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py", "litellm/router_utils/router_callbacks/track_deployment_metrics.py", "litellm/secret_managers/cyberark_secret_manager.py", "litellm/secret_managers/google_secret_manager.py", diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 034b76507cf..2a5c3c81f18 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -4218,6 +4218,41 @@ def test_anthropic_text_disable_url_suffix_env_var( assert str(disabled_route.calls.last.request.url) == "https://api.example.com/custom/complete" +def test_huggingface_text_completion_sends_each_token_prompt_with_the_callers_params( + respx_mock: respx.MockRouter, +) -> None: + from litellm.rust_bridge.tokenizer import get_encoding + + encoding: Final = get_encoding("p50k_base") + respx_mock.get(url__startswith="https://huggingface.co/").respond(404) + endpoint: Final = respx_mock.post("https://hf.example.com/v1/chat/completions").respond( + json={ + "id": "chatcmpl-hf", + "object": "chat.completion", + "created": 1, + "model": "tgi", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"}], + } + ) + + response: Final = litellm.text_completion( + model="huggingface/tgi", + prompt=[list(encoding.encode("Hello")), list(encoding.encode("world"))], + api_key="hf-test-key", + api_base="https://hf.example.com", + max_tokens=7, + ) + + assert [choice.text for choice in response.choices] == ["OK", "OK"] + sent: Final = sorted( + (json.loads(call.request.content) for call in endpoint.calls), key=lambda body: body["messages"][0]["content"] + ) + assert [(body["messages"], body["max_tokens"]) for body in sent] == [ + ([{"role": "user", "content": "Hello"}], 7), + ([{"role": "user", "content": "world"}], 7), + ] + + def test_image_edit_merges_headers_and_extra_headers(): from litellm.images.main import base_llm_http_handler diff --git a/tests/unit/test_router/core/scripted_host.py b/tests/unit/test_router/core/scripted_host.py index 33164bd4408..066faad4db8 100644 --- a/tests/unit/test_router/core/scripted_host.py +++ b/tests/unit/test_router/core/scripted_host.py @@ -166,6 +166,7 @@ def dynamic_config(**changes: object) -> RouterDynamicConfig: # kwargs-ok: Rout context_window_fallbacks=None, content_policy_fallbacks=None, drop_params=None, + disable_token_counter=False, provider_list=litellm.provider_list, expose_router_debug_in_errors=False, redact_user_api_key_info=None, diff --git a/tests/unit/test_router/core/test_construction_precall_characterization_precall.py b/tests/unit/test_router/core/test_construction_precall_characterization_precall.py index 9cdb3f55d17..c9e04117f4d 100644 --- a/tests/unit/test_router/core/test_construction_precall_characterization_precall.py +++ b/tests/unit/test_router/core/test_construction_precall_characterization_precall.py @@ -27,6 +27,7 @@ PICKS: Final = 60 # three-way shuffles miss a deployment in 20 picks about once SHORT: Final = [{"role": "user", "content": "hi"}] LONG_TEXT: Final = " ".join(["word"] * 100) LONG: Final = [{"role": "user", "content": LONG_TEXT}] +OPENAI_EMBEDDER: Final = "text-embedding-3-small" EMBEDDING_IMAGE: Final = f"data:image/png;base64,{base64.b64encode(bytes(range(256)) * 1024).decode()}" IMAGE: Final = [ { @@ -86,6 +87,14 @@ def _small_and_big( ) +def _small_and_big_openai_embedders(make_core: MakeCore, small: int) -> RouterCore: + counted_as: Final = {"model_info": {"base_model": OPENAI_EMBEDDER}} + return make_core( + model_list=[windowed("g", "small", small, **counted_as), windowed("g", "big", 10000, **counted_as)], + enable_pre_call_checks=True, + ) + + def _json_mode_group(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: constructor arguments return make_core( model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL), deployment("g", "new", model=JSON_MODE_MODEL)], @@ -279,6 +288,43 @@ async def test_an_embedding_image_stays_out_of_the_context_window_on_a_backend_t assert await _picked(core, times=2, input=EMBEDDING_IMAGE) == {"image"} +@pytest.mark.asyncio +async def test_an_embedding_input_over_the_smaller_window_in_bytes_but_not_in_tokens_fits_it( + make_core: MakeCore, +) -> None: + short_words: Final = "word " * 8 + assert litellm.token_counter(model=OPENAI_EMBEDDER, text=short_words) <= 10 < len(short_words.encode()) + assert await _picked(_small_and_big_openai_embedders(make_core, small=10), input=["hi", short_words]) == { + "small", + "big", + } + + +@pytest.mark.asyncio +async def test_an_embedding_input_that_fits_the_window_in_characters_but_not_in_tokens_is_counted( + make_core: MakeCore, +) -> None: + dense: Final = "\ua66e" * 40 + assert len(dense) <= 50 < litellm.token_counter(model=OPENAI_EMBEDDER, text=dense) + assert await _picked(_small_and_big_openai_embedders(make_core, small=50), times=10, input=["hi", dense]) == {"big"} + + +@pytest.mark.asyncio +async def test_an_embedding_batch_is_never_refused_on_the_default_tokenizers_count(make_core: MakeCore) -> None: + dense: Final = "\ua66e" * 40 + assert litellm.token_counter(text=dense) > 50 + assert await _picked(_small_and_big(make_core, small=50), input=["hi", dense]) == {"small", "big"} + + +@pytest.mark.asyncio +async def test_an_embedding_input_with_a_lone_surrogate_is_counted(make_core: MakeCore) -> None: + unpaired: Final = "\ud800" + "word " * 20 + assert litellm.token_counter(model=OPENAI_EMBEDDER, text=unpaired) > 10 + assert await _picked(_small_and_big_openai_embedders(make_core, small=10), times=10, input=["hi", unpaired]) == { + "big" + } + + @pytest.mark.asyncio async def test_an_embedding_image_counts_as_text_on_a_text_only_backend(make_core: MakeCore) -> None: core: Final = make_core( diff --git a/tests/unit/test_router/core/test_registry_characterization_model_info.py b/tests/unit/test_router/core/test_registry_characterization_model_info.py index d79aebf75b9..578183d7f02 100644 --- a/tests/unit/test_router/core/test_registry_characterization_model_info.py +++ b/tests/unit/test_router/core/test_registry_characterization_model_info.py @@ -273,6 +273,27 @@ def test_max_input_tokens_leave_out_blocked_deployments_and_other_teams(make_cor assert core.registry().group("mt", field="max_input_tokens", team_id="TA") == (2, 3) +def test_pinned_reasoning_efforts_are_each_routable_deployments_own_value_in_order(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + deployment("pe", "p1", reasoning_effort="high"), + deployment("pe", "p2"), + deployment("pe", "p3", reasoning_effort="low"), + deployment("pe", "p4", reasoning_effort="minimal", model_info={"blocked": True}), + deployment("pe", "p5", reasoning_effort="medium", model_info={"team_id": "TA"}), + ] + ) + + assert core.registry().group("pe", field="pinned_reasoning_efforts") == ("high", None, "low") + assert core.registry().group("pe", field="pinned_reasoning_efforts", team_id="TA") == ( + "high", + None, + "low", + "medium", + ) + assert core.registry().group("not-configured", field="pinned_reasoning_efforts") == () + + def test_max_input_tokens_skip_a_value_a_saved_model_info_copied_from_the_cost_map(make_core: MakeCore) -> None: catalog_limit = GetModelCostMap.loaded_model_cost_map()["gpt-4o"]["max_input_tokens"] core = make_core( diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 271615b18cf..ca61fcd0905 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -1889,6 +1889,281 @@ def test_model_group_info_cost_none_for_unpriced_deployment_but_zero_when_declar assert priced.output_cost_per_token is not None and priced.output_cost_per_token > 0 +_HUB_PRICES: Final = MappingProxyType( + { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 4e-06, + "cache_read_input_token_cost": 1e-07, + "cache_creation_input_token_cost": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2e-06, + "output_cost_per_token_above_200k_tokens": 8e-06, + "input_cost_per_token_priority": 3e-06, + "input_cost_per_token_flex": 5e-07, + "input_cost_per_token_batches": 5e-07, + } +) + + +def _cost_map_model(key: str, entry: Mapping[str, object]) -> str: + litellm.register_model( + {key: {"litellm_provider": "openai", "mode": "chat", **entry}}, + persist_across_reloads=False, + add_to_provider_model_sets=False, + ) + return f"openai/{key}" + + +def test_model_group_info_lists_every_cost_map_price_of_a_deployment(local_model_cost_map: None) -> None: + model: Final = _cost_map_model("hub-pricing-model", {**_HUB_PRICES, "max_input_tokens": 400000}) + router: Final = litellm.Router(model_list=[{"model_name": "hub", "litellm_params": {"model": model}}]) + + group: Final = router.registry().group("hub") + + assert group is not None + assert [(p.provider, dict(p.costs), p.tiered_pricing) for p in group.deployment_pricing] == [ + ("openai", dict(_HUB_PRICES), ()) + ] + + +def test_model_group_info_lists_each_deployments_price_once_next_to_the_unchanged_max( + local_model_cost_map: None, +) -> None: + model: Final = _cost_map_model("hub-pricing-model", _HUB_PRICES) + override: Final = 3 * _HUB_PRICES["input_cost_per_token"] + router: Final = litellm.Router( + model_list=[ + {"model_name": "hub", "litellm_params": {"model": model, "api_base": "https://a.example"}}, + {"model_name": "hub", "litellm_params": {"model": model, "api_base": "https://b.example"}}, + {"model_name": "hub", "litellm_params": {"model": model, "input_cost_per_token": override}}, + ] + ) + + group: Final = router.registry().group("hub") + + assert group is not None + assert [p.costs["input_cost_per_token"] for p in group.deployment_pricing] == [ + _HUB_PRICES["input_cost_per_token"], + override, + ] + assert group.input_cost_per_token == override + assert group.output_cost_per_token == _HUB_PRICES["output_cost_per_token"] + + +def _usage_billed_at(cost_key: str, prompt_tokens: int = 1000) -> tuple[Usage, str | None]: + match cost_key: + case "input_cost_per_token": + return Usage(prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens), None + case "output_cost_per_token": + return Usage(prompt_tokens=0, completion_tokens=1000, total_tokens=1000), None + case "cache_read_input_token_cost": + return Usage( + prompt_tokens=1000, + completion_tokens=0, + total_tokens=1000, + prompt_tokens_details={"cached_tokens": 1000}, + ), None + case "cache_creation_input_token_cost": + return Usage( + prompt_tokens=1000, completion_tokens=0, total_tokens=1000, cache_creation_input_tokens=1000 + ), None + case "input_cost_per_token_priority": + return Usage(prompt_tokens=1000, completion_tokens=0, total_tokens=1000), "priority" + case "input_cost_per_token_flex": + return Usage(prompt_tokens=1000, completion_tokens=0, total_tokens=1000), "flex" + case _: + raise AssertionError(f"no billing scenario for {cost_key}") + + +def _billed_rate(model: str, deployment_id: str, usage: Usage, service_tier: str | None) -> float: + cost: Final = litellm.completion_cost( + completion_response=litellm.ModelResponse(usage=usage), + model=model, + custom_llm_provider="openai", + custom_pricing=True, + router_model_id=deployment_id, + service_tier=service_tier, + ) + return cost / max(usage.prompt_tokens, usage.completion_tokens) + + +def test_model_group_info_shows_an_overridden_deployment_only_the_prices_billing_charges( + local_model_cost_map: None, +) -> None: + model: Final = _cost_map_model("hub-pricing-model", _HUB_PRICES) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "hub", + "litellm_params": { + "model": model, + "input_cost_per_token": 3 * _HUB_PRICES["input_cost_per_token"], + "output_cost_per_token": 3 * _HUB_PRICES["output_cost_per_token"], + }, + "model_info": {"id": "hub-override"}, + } + ] + ) + + group: Final = router.registry().group("hub") + + assert group is not None + costs: Final = group.deployment_pricing[0].costs + assert {key: _billed_rate(model, "hub-override", *_usage_billed_at(key)) for key in costs} == pytest.approx( + dict(costs) + ) + long_prompt: Final = _usage_billed_at("input_cost_per_token", prompt_tokens=250_000) + assert _billed_rate(model, "hub-override", *long_prompt) == pytest.approx(costs["input_cost_per_token"]) + shadowed: Final = { + "input_cost_per_token_above_200k_tokens", + "output_cost_per_token_above_200k_tokens", + "input_cost_per_token_batches", + } + assert shadowed.isdisjoint(costs), dict(costs) + + +@pytest.mark.parametrize("vertex_location", ["us-east5", "global"]) +def test_model_group_info_shows_a_vertex_deployment_the_token_prices_billing_charges_at_its_location( + local_model_cost_map: None, vertex_location: str +) -> None: + model: Final = "vertex_ai/hub-vertex-model" + litellm.register_model( + { + model: { + "litellm_provider": "vertex_ai-language-models", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 4e-06, + "cache_read_input_token_cost": 1e-07, + "regional_endpoint_uplift_multiplier": 1.1, + } + }, + persist_across_reloads=False, + add_to_provider_model_sets=False, + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "hub", + "litellm_params": {"model": model, "vertex_project": "p", "vertex_location": vertex_location}, + "model_info": {"id": "hub-vertex"}, + } + ] + ) + + group: Final = router.registry().group("hub") + + assert group is not None + costs: Final = group.deployment_pricing[0].costs + + def billed(cost_key: str) -> float: + usage, _ = _usage_billed_at(cost_key) + cost: Final = litellm.completion_cost( + completion_response=litellm.ModelResponse(usage=usage), + model=model, + custom_llm_provider="vertex_ai", + router_model_id="hub-vertex", + vertex_location=vertex_location, + ) + return cost / max(usage.prompt_tokens, usage.completion_tokens) + + shown: Final = ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost") + assert {key: costs[key] for key in shown} == pytest.approx({key: billed(key) for key in shown}) + + +def test_model_group_info_lists_tiered_pricing_ranges(local_model_cost_map: None) -> None: + tiers: Final = ( + {"range": [0, 256000], "input_cost_per_token": 5e-08, "output_cost_per_token": 4e-07}, + {"range": [256000, 1000000], "input_cost_per_token": 2.5e-07, "output_cost_per_token": 2e-06}, + ) + model: Final = _cost_map_model("hub-tiered-model", {"tiered_pricing": list(tiers)}) + router: Final = litellm.Router( + model_list=[ + {"model_name": "tiered", "litellm_params": {"model": model}}, + { + "model_name": "tiered", + "litellm_params": {"model": model, "input_cost_per_token": 1e-06, "output_cost_per_token": 4e-06}, + }, + ] + ) + + group: Final = router.registry().group("tiered") + + assert group is not None + assert [ + [(tier.range, dict(tier.costs)) for tier in pricing.tiered_pricing] for pricing in group.deployment_pricing + ] == [ + [ + ((0.0, 256000.0), {"input_cost_per_token": 5e-08, "output_cost_per_token": 4e-07}), + ((256000.0, 1000000.0), {"input_cost_per_token": 2.5e-07, "output_cost_per_token": 2e-06}), + ], + [], + ] + + +def test_model_group_info_reads_prices_written_as_text_the_way_billing_does(local_model_cost_map: None) -> None: + tier: Final = {"range": [0, 256000], "input_cost_per_token": "5e-08", "output_cost_per_token": "4e-07"} + model: Final = _cost_map_model( + "hub-text-priced-model", {"tiered_pricing": [tier], "cache_read_input_token_cost": "N/A"} + ) + router: Final = litellm.Router( + model_list=[ + {"model_name": "tiered", "litellm_params": {"model": model}}, + { + "model_name": "flat", + "litellm_params": {"model": model, "input_cost_per_token": "3e-7", "output_cost_per_token": "1.2e-6"}, + "model_info": {"id": "hub-text-override"}, + }, + ] + ) + + tiered: Final = router.registry().group("tiered") + flat: Final = router.registry().group("flat") + + assert tiered is not None + assert flat is not None + assert [dict(t.costs) for t in tiered.deployment_pricing[0].tiered_pricing] == [ + {"input_cost_per_token": 5e-08, "output_cost_per_token": 4e-07} + ] + assert "cache_read_input_token_cost" not in tiered.deployment_pricing[0].costs + flat_costs: Final = flat.deployment_pricing[0].costs + shown: Final = ("input_cost_per_token", "output_cost_per_token") + assert {key: flat_costs[key] for key in shown} == pytest.approx( + {key: _billed_rate(model, "hub-text-override", *_usage_billed_at(key)) for key in shown} + ) + + +def test_model_group_info_lists_custom_prices_of_an_unmapped_deployment_and_skips_an_unpriced_one() -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "self-hosted", + "litellm_params": {"model": "openai/hub-unpriced", "api_base": "http://localhost:8000/v1"}, + "model_info": {"cost_tracking_exempt": True}, + }, + { + "model_name": "self-hosted", + "litellm_params": { + "model": "openai/hub-priced", + "api_base": "http://localhost:8001/v1", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + }, + "model_info": {"cache_read_input_token_cost": 1e-07}, + }, + ] + ) + + group: Final = router.registry().group("self-hosted") + + assert group is not None + assert [(p.provider, dict(p.costs)) for p in group.deployment_pricing] == [ + ( + "openai", + {"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06, "cache_read_input_token_cost": 1e-07}, + ) + ] + + def test_model_group_info_description_from_model_info(): """ model_info.description set in the config should surface on ModelGroupInfo, @@ -4121,6 +4396,48 @@ def test_context_window_pre_call_check_counts_an_embedding_image_as_text_on_a_te ) +@pytest.mark.parametrize( + ("backend", "fits"), + [ + pytest.param("vertex_ai/gemini-embedding-2", True, id="image-backend-without-an-openai-tokenizer"), + pytest.param("text-embedding-3-small", False, id="text-only-backend-counts-the-image"), + ], +) +def test_context_window_pre_call_check_on_a_caption_next_to_an_embedding_image(backend: str, fits: bool, caplog): + caption: Final = "word " * 40 + router: Final = _embedding_context_router(backend, window=20) + deployments: Final = router.registry().deployments("embedder") + + with caplog.at_level(logging.ERROR, logger="LiteLLM Router"): + if fits: + kept: Final = router._pre_call_checks( + model="embedder", healthy_deployments=deployments, input=[caption, _EMBEDDING_IMAGE_INPUT] + ) + assert len(kept) == 1 + else: + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks( + model="embedder", healthy_deployments=deployments, input=[caption, _EMBEDDING_IMAGE_INPUT] + ) + + assert "failed to count tokens" not in caplog.text + + +@pytest.mark.asyncio +async def test_an_async_pre_call_check_reads_embedding_media_off_the_event_loop(): + from tests.unit.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags + + router: Final = _embedding_context_router("vertex_ai/gemini-embedding-2", window=20) + caption_and_images: Final = ["a red square", *[_EMBEDDING_IMAGE_INPUT] * 400] + + deployment, took, lags = await timed_with_loop_lags( + lambda: router.select("embedder", {}, input=caption_and_images) + ) + + assert deployment["litellm_params"]["model"] == "vertex_ai/gemini-embedding-2" + assert_loop_stayed_free(took, lags) + + async def test_a_wildcard_route_checks_the_context_window_with_the_requested_openai_models_tokenizer(): limit: Final = ( litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT) @@ -4272,6 +4589,173 @@ async def test_aresponses_enforces_context_window_pre_call_check(): ) +def _embedding_router(max_input_tokens: int) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "sk-test", + "mock_response": [0.1, 0.2], + }, + "model_info": {"max_input_tokens": max_input_tokens}, + } + ], + enable_pre_call_checks=True, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "embedding_input", + [ + pytest.param(["one two three four five six"] * 3, id="texts"), + pytest.param([list(range(6))] * 3, id="token-id-lists"), + pytest.param(list(range(10)), id="token-ids-at-the-limit"), + ], +) +async def test_aembedding_accepts_a_batch_whose_inputs_each_fit_the_context_window(embedding_input, caplog): + router: Final = _embedding_router(max_input_tokens=10) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Router"): + response: Final = await router.aembedding(model="embed", input=embedding_input) + + assert response.data + assert "failed to count tokens" not in caplog.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "embedding_input", + [ + pytest.param(["one", " ".join(["word"] * 20)], id="texts"), + pytest.param(list(range(11)), id="token-ids"), + pytest.param([[1], list(range(11))], id="token-id-lists"), + ], +) +async def test_aembedding_rejects_a_batch_with_one_input_over_the_context_window(embedding_input): + router: Final = _embedding_router(max_input_tokens=10) + + with pytest.raises(litellm.ContextWindowExceededError): + await router.aembedding(model="embed", input=embedding_input) + + +def test_embedding_rejects_a_batch_with_one_input_over_the_context_window(): + router: Final = _embedding_router(max_input_tokens=10) + + with pytest.raises(litellm.ContextWindowExceededError): + router.embedding(model="embed", input=["one", " ".join(["word"] * 20)]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "embedding_input", + [ + pytest.param(["one", " ".join(["word"] * 20)], id="texts"), + pytest.param(list(range(11)), id="token-ids"), + pytest.param([[1], list(range(11))], id="token-id-lists"), + ], +) +async def test_aembedding_skips_the_context_window_check_when_token_counting_is_disabled(embedding_input, monkeypatch): + monkeypatch.setattr(litellm, "disable_token_counter", True) + router: Final = _embedding_router(max_input_tokens=10) + + response: Final = await router.aembedding(model="embed", input=embedding_input) + + assert response.data + + +async def _until_every_count_slot_is_taken() -> None: + from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + from litellm.litellm_core_utils.token_counter import _count_offload_limiter_for_this_loop + + while _count_offload_limiter_for_this_loop().borrowed_tokens < TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_a_chat_pre_call_count_does_not_wait_behind_long_embedding_batches(): + from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + + window: Final = {"max_input_tokens": 10000} + router: Final = litellm.Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "sk-test", + "mock_response": [0.1], + }, + "model_info": window, + }, + { + "model_name": "chat", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "mock_response": "hi"}, + "model_info": window, + }, + ], + enable_pre_call_checks=True, + ) + batch: Final = ["word " * 2100] * 76 + batches: Final = tuple( + asyncio.create_task(router.aembedding(model="embed", input=batch)) + for _ in range(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS) + ) + await asyncio.wait_for(_until_every_count_slot_is_taken(), timeout=30) + chat: Final = asyncio.create_task(router.acompletion(model="chat", messages=[{"role": "user", "content": "hi"}])) + + first_done, _ = await asyncio.wait((chat, *batches), return_when=asyncio.FIRST_COMPLETED) + await asyncio.gather(chat, *batches) + + assert first_done == {chat} + + +def test_a_routing_plugin_counts_an_embedding_text_batch_by_its_longest_text(): + long_text: Final = "word " * 4000 + context: Final = RouterContext(_embedding_router(max_input_tokens=10)) + + assert context.count_pre_call_check_tokens(None, ["one", long_text]) == litellm.token_counter(text=long_text) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_model_cost_map", "isolate_router_model_cost_state") +@pytest.mark.parametrize( + "prompt_kwargs", + [ + pytest.param({"instructions": "you are a helpful assistant. " * 20}, id="instructions"), + pytest.param( + { + "tools": [ + { + "type": "function", + "name": "lookup", + "description": "look up a record by its identifier. " * 20, + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + } + ] + }, + id="tools", + ), + ], +) +async def test_aresponses_with_empty_input_still_counts_its_other_prompt_content(prompt_kwargs): + router: Final = litellm.Router( + model_list=[ + { + "model_name": "small-ctx", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, + "model_info": {"max_input_tokens": 50}, + } + ], + enable_pre_call_checks=True, + ) + + with pytest.raises(litellm.ContextWindowExceededError): + await router.aresponses(model="small-ctx", input=[], **prompt_kwargs) + + def test_get_deployment_model_info_base_model_flow(): """Test that get_deployment_model_info correctly handles the base model flow""" from unittest.mock import patch @@ -8522,7 +9006,7 @@ class TestConsumedRequestTagsStamp: self.rewrite_to = rewrite_to async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False + self, model, request_kwargs, messages=None, input=None, specific_deployment=False, call_type=None ): from litellm.types.router import PreRoutingHookResponse @@ -8613,7 +9097,7 @@ class TestClaudeCodeSubagentSessionRouterBinding: self.routed_model = routed_model async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False + self, model, request_kwargs, messages=None, input=None, specific_deployment=False, call_type=None ): from litellm.types.router import PreRoutingHookResponse @@ -9667,7 +10151,7 @@ class TestModelGroupAliasReachesPreRoutingStrategies: class _RewriteStrategy: async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False + self, model, request_kwargs, messages=None, input=None, specific_deployment=False, call_type=None ): from litellm.types.router import PreRoutingHookResponse @@ -9784,7 +10268,7 @@ class TestTeamPublicNameReachesPreRoutingStrategies: self.rewrite_to = rewrite_to async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False + self, model, request_kwargs, messages=None, input=None, specific_deployment=False, call_type=None ): from litellm.types.router import PreRoutingHookResponse @@ -10021,7 +10505,7 @@ class TestAutoRouterCompressionDecoupling: self.received_messages: list[dict] | None = None async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False + self, model, request_kwargs, messages=None, input=None, specific_deployment=False, call_type=None ): from litellm.types.router import PreRoutingHookResponse diff --git a/tests/unit/test_router/test_router_entry_points.py b/tests/unit/test_router/test_router_entry_points.py index c3d5d962123..06287ba7e0b 100644 --- a/tests/unit/test_router/test_router_entry_points.py +++ b/tests/unit/test_router/test_router_entry_points.py @@ -514,6 +514,81 @@ def test_attribute_errors_about_other_objects_pass_through_unchanged(router): assert _public_attribute_error(router, "_core_x", "x", other) is other +_MINTING_DEPLOYMENT = "dep-minting-org" +_OTHER_ORG_DEPLOYMENT = "dep-other-org" +_BOTH_ORGS = (_MINTING_DEPLOYMENT, _OTHER_ORG_DEPLOYMENT) +_GROUP_AFFINITY = ["encrypted_content_affinity"] + + +@pytest.mark.parametrize( + "affinity, target_ids, reasoning_stays", + [ + pytest.param({}, _BOTH_ORGS, False, id="no affinity, one target cannot decrypt"), + pytest.param({}, (_MINTING_DEPLOYMENT,), True, id="no affinity, the only target minted it"), + pytest.param( + {"optional_pre_call_checks": _GROUP_AFFINITY}, _BOTH_ORGS, True, id="global affinity, origin among targets" + ), + pytest.param( + {"optional_pre_call_checks": _GROUP_AFFINITY}, + (_OTHER_ORG_DEPLOYMENT,), + False, + id="global affinity, origin not among targets", + ), + pytest.param( + {"optional_pre_call_checks": [], "model_group_affinity_config": {"pair": _GROUP_AFFINITY}}, + _BOTH_ORGS, + True, + id="affinity on for this group", + ), + pytest.param( + {"optional_pre_call_checks": [], "model_group_affinity_config": {"other-group": _GROUP_AFFINITY}}, + _BOTH_ORGS, + False, + id="affinity on for another group only", + ), + ], +) +def test_decryptable_reasoning_keeps_encrypted_reasoning_only_when_the_router_sends_it_to_a_deployment_that_decrypts_it( + affinity, target_ids, reasoning_stays +): + from litellm.responses.utils import ResponsesAPIRequestUtils + + r = Router( + model_list=[ + { + "model_name": "pair", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_base": api_base, "api_key": "sk-pair"}, + "model_info": {"id": model_id}, + } + for api_base, model_id in ( + ("https://minting-org.test/v1", _MINTING_DEPLOYMENT), + ("https://other-org.test/v1", _OTHER_ORG_DEPLOYMENT), + ) + ], + **affinity, + ) + summary = [{"type": "summary_text", "text": "multiply 17 by 23"}] + user_turn = {"type": "message", "role": "user", "content": "What is 17*23?"} + reasoning = { + "type": "reasoning", + "id": "rs_minted", + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + "gAAAAA-blob", _MINTING_DEPLOYMENT + ), + "summary": summary, + } + request_input = [user_turn, reasoning] + targets = tuple(r.registry().deployment(model_id) for model_id in target_ids) + try: + kept = r.registry().decryptable_reasoning("pair", request_input, targets) + finally: + r.retire() + + readable_only = {"type": "reasoning", "summary": summary} + assert kept == (user_turn, reasoning if reasoning_stays else readable_only) + assert request_input == [user_turn, reasoning], "the caller's input must not change" + + # Router-family helpers whose Router parameter became a RoutingContext keep the old spellings. diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index c0c8a0655fe..cd71c3f9bee 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -542,7 +542,13 @@ def test_should_not_pollute_shared_key_with_custom_nonzero_pricing(): @pytest.mark.parametrize( "flag", - ["supports_regex_lookaround", "supports_assistant_prefill", "supports_sampling_params", "supports_continuation"], + [ + "supports_regex_lookaround", + "supports_assistant_prefill", + "supports_sampling_params", + "supports_continuation", + "supports_bedrock_native_structured_output_by_default", + ], ) def test_capability_flag_stays_on_the_deployment_that_set_it(flag: str) -> None: """A deployment's capability override must not land on the shared diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index f8ec0e56e4a..5eaabdd2c59 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1354,6 +1354,7 @@ _COST_MAP_ROW_SCHEMA: Final = { "supports_image_size": {"type": "boolean"}, "supports_native_structured_output": {"type": "boolean"}, "supports_bedrock_native_structured_output": {"type": "boolean"}, + "supports_bedrock_native_structured_output_by_default": {"type": "boolean"}, "use_openai_responses_path": {"type": "boolean"}, "off_peak_pricing": { "type": "object", @@ -11611,6 +11612,35 @@ async def test_success_enqueue_is_tracked_without_starting_callbacks_synchronous assert logging_obj.model_call_details["has_logged_async_success"] is True +def test_get_response_string_joins_the_text_parts_of_a_completed_responses_event() -> None: + from litellm.types.llms.openai import ResponseCompletedEvent + from litellm.utils import get_response_string + + response: Final = ResponsesAPIResponse( + id="resp_1", + created_at=1700000000, + model="gpt-5.6", + object="response", + status="completed", + output=[ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [ + {"type": "output_text", "text": "Hello ", "annotations": []}, + {"type": "refusal", "refusal": "not this part"}, + {"type": "output_text", "text": "world", "annotations": []}, + ], + }, + {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "lookup", "arguments": "{}"}, + ], + ) + + assert get_response_string(ResponseCompletedEvent(type="response.completed", response=response)) == "Hello world" + + class _RecordingPreCallDeploymentHook(CustomLogger): def __init__(self) -> None: super().__init__() diff --git a/tests/unit/test_utils_get_model_info.py b/tests/unit/test_utils_get_model_info.py index d08742be4af..7548001ce4c 100644 --- a/tests/unit/test_utils_get_model_info.py +++ b/tests/unit/test_utils_get_model_info.py @@ -13,7 +13,7 @@ from litellm import get_model_info from litellm.llms.bedrock.common_utils import BedrockModelInfo from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.utils import ModelInfoBase -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.utils import _invalidate_model_cost_lowercase_map, supports_native_structured_output import os from unittest.mock import MagicMock, patch @@ -388,6 +388,29 @@ def test_get_model_info_bedrock_cross_region_capability_parity(): assert checked > 0, "no cross-region bedrock profiles found - the filter is inert" +def test_bedrock_native_output_default_reaches_only_versions_that_read_its_key(local_model_cost_map): + flagged: Final = tuple( + key for key, row in litellm.model_cost.items() if row.get("supports_bedrock_native_structured_output_by_default") + ) + + assert flagged + assert [key for key in flagged if litellm.model_cost[key].get("supports_native_structured_output")] == [] + assert [key for key in flagged if not get_model_info(key)["supports_native_structured_output"]] == [] + + +@pytest.mark.parametrize( + ("registered", "native"), + [({"supports_bedrock_native_structured_output": False}, False), ({"max_output_tokens": 1024}, True)], + ids=["bedrock-opt-out", "unrelated-override"], +) +def test_a_registered_bedrock_opt_out_beats_the_native_output_default(local_model_cost_map, registered, native): + model: Final = "us.anthropic.claude-sonnet-5" + litellm.register_model({model: registered}, persist_across_reloads=False) + + assert get_model_info(f"bedrock/{model}")["supports_native_structured_output"] is native + assert supports_native_structured_output(f"bedrock/{model}") is native + + def _is_positive_cost(value: object) -> bool: return isinstance(value, (int, float)) and value > 0 diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index d71791b6f11..aae5dfe6472 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -168,6 +168,7 @@ OPTION_NAMES: Final = ( "tag_regex", "max_file_size_mb", "silent_model", + "switchyard_router_config", "auto_router_config_path", "auto_router_config", "auto_router_default_model", diff --git a/tests/unit/videos/test_main.py b/tests/unit/videos/test_main.py index 76e9b062476..d43141bdea3 100644 --- a/tests/unit/videos/test_main.py +++ b/tests/unit/videos/test_main.py @@ -35,6 +35,8 @@ from dataclasses import dataclass from typing import Any, Dict, Final from unittest.mock import MagicMock, patch +from litellm.types.videos.main import VideoCreateOptionalRequestParams + import pytest @@ -439,18 +441,39 @@ def test_video_generation__custom_provider_map_base_video_config_is_dispatched(m from litellm.llms.openai.videos.transformation import OpenAIVideoConfig class QueueVideoConfig(OpenAIVideoConfig): - pass + def map_openai_params( + self, video_create_optional_params: VideoCreateOptionalRequestParams, model: str, drop_params: bool + ) -> dict[str, object]: + return dict(video_create_optional_params) video_config: Final = QueueVideoConfig() monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "queue-video", "custom_handler": video_config}]) monkeypatch.setattr(litellm, "_custom_providers", ["queue-video"]) handler: Final = MagicMock(spec=BaseLLMHTTPHandler) with patch.object(videos_main, "base_llm_http_handler", handler): - videos_main.video_generation(model="queue-video/seedance-1", prompt="a cat", api_base="http://backend.test") + videos_main.video_generation( + model="queue-video/seedance-1", + prompt="a cat", + api_base="http://backend.test", + duration=8, + aspect_ratio="16:9", + frame_images=[{"type": "image_url", "image_url": {"url": "https://example.com/first.png"}, "frame_type": "first_frame"}], + input_references=[{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}}], + generate_audio=False, + seed=0, + ) kw: Final = handler.video_generation_handler.call_args.kwargs assert kw["video_generation_provider_config"] is video_config assert kw["custom_llm_provider"] == "queue-video" assert kw["model"] == "seedance-1" + assert kw["video_generation_optional_request_params"] == { + "duration": 8, + "aspect_ratio": "16:9", + "frame_images": [{"type": "image_url", "image_url": {"url": "https://example.com/first.png"}, "frame_type": "first_frame"}], + "input_references": [{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}}], + "generate_audio": False, + "seed": 0, + } def test_video_generation__custom_provider_without_a_video_config_is_a_400(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/unit/videos/test_utils.py b/tests/unit/videos/test_utils.py index 728644cdda5..ffcbde7381d 100644 --- a/tests/unit/videos/test_utils.py +++ b/tests/unit/videos/test_utils.py @@ -10,15 +10,29 @@ production. Every test asserts the exact resulting dict, never "ran without erro """ from unittest.mock import MagicMock +from typing import Final + +import pytest import litellm +from litellm.exceptions import UnsupportedParamsError +from litellm.llms.volcengine.videos.transformation import VolcEngineVideoConfig +from litellm.llms.vertex_ai.videos.transformation import VertexAIVideoConfig from litellm.videos.utils import VideoGenerationRequestUtils get_requested = ( VideoGenerationRequestUtils.get_requested_video_generation_optional_param ) + + +@pytest.mark.parametrize("drop_params", [False, True]) +def test_extra_body_video_fields_remain_raw(drop_params: bool) -> None: + extra: Final = {"seed": 0, "duration": 6, "aspect_ratio": "1:1", "generate_audio": False, + "frame_images": [], "input_references": []} + params: Final = get_requested({"extra_body": extra}) + assert get_optional("veo-3.1-generate-001", VertexAIVideoConfig(), params, drop_params=drop_params) == extra get_optional = VideoGenerationRequestUtils.get_optional_params_video_generation @@ -179,3 +193,23 @@ def test_optional__non_dict_extra_body_ignored(): ) assert result == {"seconds": "8"} + + +@pytest.mark.parametrize("extra_body", [ + {"content": [{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}, "role": "reference_image"}]}, + {"input_reference": "https://example.com/ref.png"}, +]) +def test_extra_body_overrides_mapped_images_after_dropping_conflicts(extra_body: dict[str, object]) -> None: + params: Final = get_requested({ + "kwargs": { + "frame_images": [{"type": "image_url", "image_url": {"url": "https://example.com/first.png"}, "frame_type": "first_frame"}], + "extra_body": extra_body, + }, + }) + config: Final = VolcEngineVideoConfig() + with pytest.raises(UnsupportedParamsError, match="cannot be combined"): + get_optional("model", config, params) + assert get_optional("model", config, params, drop_params=True) == { + "content": [{"type": "image_url", "image_url": {"url": "https://example.com/first.png"}, "role": "first_frame"}], + **extra_body, + } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index f9304bab0d2..d13bc3d3181 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -109,6 +109,7 @@ export interface AutoRouterDeployment extends AutoRouterCandidateDeployment { model?: string | null; base_model?: string | null; custom_llm_provider?: string | null; + switchyard_router_config?: unknown; complexity_router_config?: unknown; complexity_router_default_model?: string | null; auto_router_config?: unknown; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts index 2cdc13606bb..0fe320a318a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts @@ -87,6 +87,14 @@ const configManaged = (label: string, config: Record): Presenta /** How each strategy renders itself, given its own config object. */ const PRESENTERS: Record) => Presentation> = { + switchyard: (config) => ({ + typeLabel: "Switchyard", + targets: dedupe( + [asRecord(config.models).capable, asRecord(config.models).efficient].filter( + (model): model is string => typeof model === "string", + ), + ), + }), complexity: (config) => ({ typeLabel: complexityTypeLabel(config), targets: dedupe(Object.values(asRecord(config.tiers)).flatMap(normalizeTierModels)), @@ -131,7 +139,8 @@ export const toAutoRouterRow = ( canDelete: canDelete && mayActOnRow, editBlockedReason, createdAt: info.created_at ?? undefined, - defaultModel: (params[strategy.defaultModelKey] as string | null | undefined) ?? null, + defaultModel: + strategy.kind === "switchyard" ? null : (params[strategy.defaultModelKey] as string | null | undefined) ?? null, deployment, ...PRESENTERS[strategy.kind](asRecord(params[strategy.configKey])), }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/UsageStatStrip.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/UsageStatStrip.integration.test.tsx index 2e2c42cb218..c6b41cb9e31 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/UsageStatStrip.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/UsageStatStrip.integration.test.tsx @@ -1,6 +1,7 @@ import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it } from "vitest"; +import i18n from "@/i18n/i18n"; import type { OverviewTotals } from "./overviewData"; import UsageStatStrip from "./UsageStatStrip"; @@ -19,6 +20,15 @@ const totals: OverviewTotals = { successRate: 80, }; +const asRendered = (text: string) => text.replace(/\s+/g, " "); +const SECONDS = { + style: "unit", + unit: "second", + unitDisplay: "narrow", + minimumFractionDigits: 1, + maximumFractionDigits: 1, +} as const; + const defaultProps = { results: [], totals, @@ -58,3 +68,58 @@ describe("UsageStatStrip token breakdown", () => { expect(screen.queryByText("failed")).not.toBeInTheDocument(); }); }); + +describe("UsageStatStrip in another UI language", () => { + afterEach(async () => { + await i18n.changeLanguage("en"); + i18n.removeResourceBundle("de", "translation"); + }); + + it("should format every value in the selected language and keep money in US dollars", async () => { + i18n.addResourceBundle("de", "translation", {}); + await i18n.changeLanguage("de"); + const german = new Intl.NumberFormat("de"); + const dollars = new Intl.NumberFormat("de", { style: "currency", currency: "USD" }); + const euros = new Intl.NumberFormat("de", { style: "currency", currency: "EUR" }); + const dollarsToFourPlaces = new Intl.NumberFormat("de", { + style: "currency", + currency: "USD", + minimumFractionDigits: 4, + }); + const oneDecimal = new Intl.NumberFormat("de", { minimumFractionDigits: 1 }); + const milliseconds = new Intl.NumberFormat("de", { style: "unit", unit: "millisecond", unitDisplay: "narrow" }); + + render(); + + expect(screen.getByText(asRendered(dollars.format(42.5)))).toBeInTheDocument(); + expect(screen.getByText(asRendered(`of ${dollars.format(1500)} budget`))).toBeInTheDocument(); + expect(screen.getByText(asRendered(dollarsToFourPlaces.format(4.25)))).toBeInTheDocument(); + expect(screen.getByText(asRendered(milliseconds.format(250)))).toBeInTheDocument(); + expect(screen.getByText(`${oneDecimal.format(80)}% success rate`)).toBeInTheDocument(); + expect(screen.getByTitle(`${german.format(9123456)} tokens`)).toBeInTheDocument(); + expect(screen.queryByText(asRendered(euros.format(42.5)))).not.toBeInTheDocument(); + }); + + it("should format large counts, compact tokens, seconds and the token breakdown in the selected language", async () => { + const user = userEvent.setup(); + i18n.addResourceBundle("de", "translation", {}); + await i18n.changeLanguage("de"); + const german = new Intl.NumberFormat("de"); + const compact = new Intl.NumberFormat("de", { notation: "compact", maximumFractionDigits: 1 }); + const large = { ...totals, requests: 8123456, successful: 7654321, failed: 469135, avgLatencyMs: 1234 }; + + render(); + + expect(screen.getByText(german.format(8123456))).toBeInTheDocument(); + expect(screen.getByText(german.format(7654321))).toBeInTheDocument(); + expect(screen.getByText(german.format(469135))).toBeInTheDocument(); + expect(screen.getByText(asRendered(compact.format(9123456)))).toBeInTheDocument(); + expect(screen.getByText(asRendered(new Intl.NumberFormat("de", SECONDS).format(1.234)))).toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "View token breakdown" })); + const breakdown = await screen.findByRole("dialog", { name: "What counts toward total tokens?" }); + [8123456, 5123456, 2000000, 1000000, 9123456].forEach((count) => + expect(breakdown).toHaveTextContent(german.format(count)), + ); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/UsageStatStrip.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/UsageStatStrip.tsx index 15c33d8653c..48eb1e96691 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/UsageStatStrip.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/UsageStatStrip.tsx @@ -1,14 +1,19 @@ "use client"; import React, { useMemo } from "react"; +import type { TFunction } from "i18next"; import { Trans, useTranslation } from "react-i18next"; import type { DailyData } from "@/components/UsagePage/types"; import { cn } from "@/lib/cva.config"; import { Popover, PopoverContent, PopoverTitle, PopoverTrigger } from "@/components/ui/popover"; -import { dailyTotals, formatCompact, formatLatency, formatUsd, type OverviewTotals } from "./overviewData"; +import { dailyTotals, type OverviewTotals } from "./overviewData"; import { PANEL_INSET_X, Sparkline, Stat } from "./Primitives"; const BRAND = "#2b3fd6"; +const ONE_DECIMAL = { minimumFractionDigits: 1, maximumFractionDigits: 1 } as const; +const FOUR_DECIMALS = { minimumFractionDigits: 4, maximumFractionDigits: 4 } as const; +const MILLISECONDS = { style: "unit", unit: "millisecond", unitDisplay: "narrow", maximumFractionDigits: 0 } as const; +const SECONDS = { ...ONE_DECIMAL, style: "unit", unit: "second", unitDisplay: "narrow" } as const; interface UsageStatStripProps { results: readonly DailyData[]; @@ -33,15 +38,19 @@ export default function UsageStatStrip({ }> }> {!requestCountsPending && ( @@ -50,7 +59,9 @@ export default function UsageStatStrip({ {totals.successful.toLocaleString()}, + count: ( + {t("format.number", { val: totals.successful })} + ), }} /> @@ -60,7 +71,7 @@ export default function UsageStatStrip({ components={{ count: ( 0 ? "text-destructive tabular-nums" : "tabular-nums"}> - {totals.failed.toLocaleString()} + {t("format.number", { val: totals.failed })} ), }} @@ -72,11 +83,15 @@ export default function UsageStatStrip({ }> {!loading && ( @@ -89,23 +104,23 @@ export default function UsageStatStrip({
{t("usage.statStrip.tokenPopover.inputTokens")}
-
{totals.inputTokens.toLocaleString()}
+
{t("format.number", { val: totals.inputTokens })}
{t("usage.statStrip.tokenPopover.cacheReads")}
-
{totals.cacheReadTokens.toLocaleString()}
+
{t("format.number", { val: totals.cacheReadTokens })}
{t("usage.statStrip.tokenPopover.cacheWrites")}
-
{totals.cacheWriteTokens.toLocaleString()}
+
{t("format.number", { val: totals.cacheWriteTokens })}
{t("usage.statStrip.tokenPopover.outputTokens")}
-
{totals.outputTokens.toLocaleString()}
+
{t("format.number", { val: totals.outputTokens })}
{t("usage.statStrip.tokenPopover.totalTokens")}
-
{totals.tokens.toLocaleString()}
+
{t("format.number", { val: totals.tokens })}

{t("usage.statStrip.tokenPopover.pricingNote")}

@@ -118,17 +133,19 @@ export default function UsageStatStrip({
@@ -137,6 +154,12 @@ export default function UsageStatStrip({ ); } +function latency(ms: number, t: TFunction): string { + return ms < 1000 + ? t("format.number", { val: ms, ...MILLISECONDS }) + : t("format.number", { val: ms / 1000, ...SECONDS }); +} + function StatCell({ children, trend }: { children: React.ReactNode; trend?: React.ReactNode }) { return (
{ export const formatCompact = (value: number): string => Intl.NumberFormat("en-US", { notation: "compact", maximumFractionDigits: 1 }).format(value); -export const formatLatency = (ms: number): string => (ms < 1000 ? `${Math.round(ms)}ms` : `${(ms / 1000).toFixed(1)}s`); - export interface OverviewTotals { spend: number; requests: number; diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx index 123c7431732..1b30b060ab0 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx @@ -2,6 +2,7 @@ import * as networking from "@/components/networking"; import userEvent from "@testing-library/user-event"; import { act } from "@testing-library/react"; import type { MCPServerData } from "./MCPHubTableColumns"; +import type { DeploymentPricing } from "@/components/molecules/models/deploymentPricing"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, screen, testQueryClient, waitFor, within } from "../../../tests/test-utils"; import ModelHubTable from "./ModelHubTable"; @@ -374,7 +375,7 @@ describe("ModelHubTable", () => { describe("model details usage example", () => { const openDetails = async ( - model: { model_group: string; providers: string[]; mode: string }, + model: { model_group: string; providers: string[]; mode: string; deployment_pricing?: DeploymentPricing[] }, canOpenPlayground = true, ) => { vi.mocked(networking.modelHubCall).mockResolvedValue({ data: [model] }); @@ -425,6 +426,34 @@ describe("ModelHubTable", () => { expect(screen.queryByRole("link", { name: "Try it in the Playground" })).not.toBeInTheDocument(); }); + it("shows each deployment's own prices, by service tier, in the model details", async () => { + const gpt5OnOpenAIAndAzure = { + model_group: "gpt-5", + providers: ["openai", "azure"], + mode: "chat", + deployment_pricing: [ + { + provider: "openai", + costs: { input_cost_per_token: 2.5e-6, input_cost_per_token_priority: 5e-6 }, + tiered_pricing: [], + }, + { provider: "azure", costs: { input_cost_per_token: 2.75e-6 }, tiered_pricing: [] }, + ], + }; + const dialog = await openDetails(gpt5OnOpenAIAndAzure); + + const openai = within(dialog).getByRole("region", { name: "OpenAI pricing" }); + expect(within(openai).getByRole("row", { name: /Input \/ 1M tokens/ })).toHaveTextContent("$2.50$5.00"); + expect(within(dialog).getByRole("region", { name: "Azure pricing" })).toHaveTextContent("$2.75"); + }); + + it("leaves the pricing breakdown out when the proxy sends no deployment prices", async () => { + const dialog = await openDetails({ model_group: "gpt-5", providers: ["openai"], mode: "chat" }); + + expect(within(dialog).getByText("Input Cost per 1M Tokens:")).toBeInTheDocument(); + expect(within(dialog).queryByText("Pricing by Deployment:")).not.toBeInTheDocument(); + }); + it("keeps the chat completions example for a chat model", async () => { const dialog = await openDetails({ model_group: "claude-opus-4-8", providers: ["anthropic"], mode: "chat" }); diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index 4d44d0bb809..60259b1f7f0 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -57,6 +57,7 @@ import { isDecisionMode, } from "@/lib/decisionModels"; import DecisionModelGarden from "@/components/AIHub/DecisionModelGarden"; +import { DeploymentPricingBreakdown } from "@/components/molecules/models/DeploymentPricingBreakdown"; import { ModelCustomMetadata } from "./ModelCustomMetadata"; interface ModelHubTableProps { @@ -819,6 +820,11 @@ const ModelHubTable: React.FC = ({

+ {t("aiHub.modelHubTable.modelDialog.pricingByDeployment")}

} + /> {/* Capabilities */} diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx index 2e4af18455f..7f2f3a49d44 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx @@ -22,6 +22,7 @@ import { copyToClipboard } from "@/utils/dataUtils"; import { ModelCustomTags, type ModelCustomMetadataValue } from "./ModelCustomMetadata"; import { formatModelModeLabel, ProviderLogo } from "./ModelHubProviders"; import { getProviderLogoAndName } from "@/components/provider_info_helpers"; +import type { DeploymentPricing } from "@/components/molecules/models/deploymentPricing"; export interface ModelHubData { model_group: string; @@ -41,6 +42,7 @@ export interface ModelHubData { description?: string; custom_tags?: string[] | null; custom_metadata?: Record | null; + deployment_pricing?: DeploymentPricing[] | null; [key: string]: any; } diff --git a/ui/litellm-dashboard/src/components/PublicModelHubTableColumns.tsx b/ui/litellm-dashboard/src/components/PublicModelHubTableColumns.tsx index 08d6b973b87..eb89af22f17 100644 --- a/ui/litellm-dashboard/src/components/PublicModelHubTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/PublicModelHubTableColumns.tsx @@ -11,6 +11,7 @@ import i18n from "@/i18n/i18n"; import { displayMode } from "@/lib/decisionModels"; import type { ModelCustomMetadataValue } from "@/components/AIHub/ModelCustomMetadata"; import { formatModelModeLabel, ProviderLogo } from "@/components/AIHub/ModelHubProviders"; +import type { DeploymentPricing } from "@/components/molecules/models/deploymentPricing"; export interface ModelGroupInfo { model_group: string; @@ -32,6 +33,7 @@ export interface ModelGroupInfo { description?: string; custom_tags?: string[] | null; custom_metadata?: Record | null; + deployment_pricing?: DeploymentPricing[] | null; [key: string]: any; } diff --git a/ui/litellm-dashboard/src/components/add_model/AutoRouterClassifierTabs.tsx b/ui/litellm-dashboard/src/components/add_model/AutoRouterClassifierTabs.tsx index bd8b4eb9854..0dd80854312 100644 --- a/ui/litellm-dashboard/src/components/add_model/AutoRouterClassifierTabs.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AutoRouterClassifierTabs.tsx @@ -135,6 +135,7 @@ interface AutoRouterClassifierTabsProps { children: React.ReactNode; judgeModel?: React.ReactNode; decisionModel?: React.ReactNode; + switchyard?: { selected: boolean; onSelect: () => void; fields: React.ReactNode }; } const AutoRouterClassifierTabs: React.FC = ({ @@ -143,6 +144,7 @@ const AutoRouterClassifierTabs: React.FC = ({ children, judgeModel, decisionModel, + switchyard, }) => { const { t } = useTranslation(); const id = useId(); @@ -159,14 +161,16 @@ const AutoRouterClassifierTabs: React.FC = ({ jev: "jev", custom: "custom", }; - const family = familyByType[classifierType]; + const family = switchyard?.selected ? "switchyard" : familyByType[classifierType]; const hasCustomTiers = Boolean(value.custom_tier_set); const changeType = (next: ClassifierType) => { if (next !== classifierType) onChange(transitionClassifierType(value, next)); }; const changeFamily = (next: unknown) => { + if (next === "switchyard") return switchyard?.onSelect(); if (next === family) return; + if (next === familyByType[classifierType]) return onChange(value); if (next === "heuristics") changeType("heuristic"); if (next === "llm") changeType("llm"); if (next === "jev") changeType("jev"); @@ -207,7 +211,11 @@ const AutoRouterClassifierTabs: React.FC = ({ - + {[ { value: "heuristics", @@ -220,6 +228,9 @@ const AutoRouterClassifierTabs: React.FC = ({ label: t("addModel.autoRouterClassifierTabs.families.decisionModel"), description: t("addModel.autoRouterClassifierTabs.families.decisionModelDescription"), }, + ...(switchyard + ? [{ value: "switchyard", label: "Switchyard", description: "Route by the conversation’s tool history." }] + : []), ].map((option) => ( + {family === "switchyard" && switchyard?.fields} {family === "jev" && (decisionModel ?? (
@@ -355,7 +367,7 @@ const AutoRouterClassifierTabs: React.FC = ({

)} - {hasCustomTiers && ( + {family !== "switchyard" && hasCustomTiers && (

{t("addModel.autoRouterClassifierTabs.restoreStandardTiers")}

)} {availability.data?.error && !availability.isChecking && ( @@ -372,7 +384,7 @@ const AutoRouterClassifierTabs: React.FC = ({

)} - {children} + {family !== "switchyard" && children} ); }; diff --git a/ui/litellm-dashboard/src/components/add_model/BlockedReasonTooltip.tsx b/ui/litellm-dashboard/src/components/add_model/BlockedReasonTooltip.tsx new file mode 100644 index 00000000000..59588cf478f --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/BlockedReasonTooltip.tsx @@ -0,0 +1,17 @@ +import React from "react"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; + +const BlockedReasonTooltip: React.FC<{ reason: string | null; children: React.ReactElement }> = ({ + reason, + children, +}) => + reason === null ? ( + children + ) : ( + + + {reason} + + ); + +export default BlockedReasonTooltip; diff --git a/ui/litellm-dashboard/src/components/add_model/SwitchyardConfigFields.tsx b/ui/litellm-dashboard/src/components/add_model/SwitchyardConfigFields.tsx new file mode 100644 index 00000000000..2043e392928 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/SwitchyardConfigFields.tsx @@ -0,0 +1,342 @@ +"use client"; + +import { useId, useState } from "react"; +import { Input } from "@/components/ui/input"; +import { Textarea } from "@/components/ui/textarea"; +import { Slider } from "@/components/ui/slider"; +import { Switch } from "@/components/ui/switch"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import ModelChoiceCombobox, { type ModelChoice } from "./ModelChoiceCombobox"; +import { + stageConfidence, + toggleClassifier, + setClassifierKind, + parseToolSemantics, + type SwitchyardConfig, +} from "./switchyard_config"; + +interface Props { + value: SwitchyardConfig; + onChange: (value: SwitchyardConfig) => void; + models: ModelChoice[]; + decisionModels: ModelChoice[]; + onToolSemanticsError: (error: string | null) => void; +} + +const isStructured = (instructions: unknown): instructions is Record => + typeof instructions === "object" && instructions !== null; +const instructionsText = (instructions: string | Record | null | undefined): string => + isStructured(instructions) ? JSON.stringify(instructions, null, 2) : instructions ?? ""; + +export default function SwitchyardConfigFields({ + value, + onChange, + models, + decisionModels, + onToolSemanticsError, +}: Props) { + const id = useId(); + const [toolSemantics, setToolSemantics] = useState( + value.options.tool_semantics ? JSON.stringify(value.options.tool_semantics, null, 2) : "", + ); + const classifier = value.classifier; + const llm = classifier?.kind === "llm" ? classifier : undefined; + const decision = classifier?.kind === "decision" ? classifier : undefined; + const setOption = (key: "recent_window", text: string) => + onChange({ ...value, options: { ...value.options, [key]: text === "" ? undefined : Number(text) } }); + const setClassifierNumber = ( + key: "base_threshold" | "threshold_step" | "max_output_tokens" | "recent_turn_window", + text: string, + ) => { + if (llm) onChange({ ...value, classifier: { ...llm, [key]: text === "" ? undefined : Number(text) } }); + }; + const changeToolSemantics = (text: string) => { + setToolSemantics(text); + try { + const parsed = parseToolSemantics(text); + onChange({ ...value, options: { ...value.options, tool_semantics: parsed } }); + onToolSemanticsError(null); + } catch { + onToolSemanticsError("Tool semantics must be JSON mapping observe, mutate, plan, or new to arrays of tool names"); + } + }; + return ( +
+
+
+ + Heuristic + +

+ Stage Router +

+
+
+ + + {value.options.picker === "capable_first" && ( +

+ Switchyard has not benchmarked this bias, so it has no calibrated confidence thresholds. +

+ )} +
+ {(["capable", "efficient"] as const).map((tier) => ( +
+ + onChange({ ...value, models: { ...value.models, [tier]: model ?? "" } })} + choices={models} + placeholder="Select a model" + ariaInvalid={undefined} + ariaDescribedBy={undefined} + /> +
+ ))} +
+
+
+
+ +

+ Ask a judge model when the tool history does not give a confident routing decision. +

+
+ onChange(toggleClassifier(value, enabled))} + /> +
+ {classifier && ( +
+
+ + +
+
+ + onChange({ ...value, models: { ...value.models, judge: judge ?? "" } })} + choices={decision ? decisionModels : models} + placeholder={decision ? "Select a decision model" : "Select a judge model"} + ariaInvalid={undefined} + ariaDescribedBy={undefined} + /> +
+
+ )} +
+
+
+ + Confidence threshold + + {stageConfidence(value).toFixed(2)} +
+ + onChange({ + ...value, + options: { ...value.options, confidence_threshold: Array.isArray(next) ? next[0] : next }, + }) + } + /> +

+ Below this confidence, use {classifier ? "the judge model" : "the bias’s default model"}. Critical errors can + still select the capable model. +

+
+
+ Advanced settings +
+

Stage

+
+ + setOption("recent_window", event.target.value)} + /> +
+
+ +