From 0410abea8be391cb31796909800b6fe8b3ae2239 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 9 Oct 2026 16:29:44 -0700 Subject: [PATCH 001/104] feat(rust): type the AWS connection params and where Python reads them (#45610) AwsParams in auth-types holds the aws_* fields of CredentialLiteLLMParams with one ParamSpec per field: its wire name and the environment names it falls back to. fields() and secret_names() derive from the specs, the AWS helpers take the struct instead of a map, aws_auth_config and the region resolution read through the specs, and a missing region is AwsParams::REGION.missing("AWS"), whose message is rendered from the spec. Bedrock Converse, Bedrock transcription, Textract OCR and the Bedrock Messages config build the struct at their one untyped boundary. --- litellm-rust/Cargo.lock | 2 + litellm-rust/crates/auth-aws/Cargo.toml | 3 +- litellm-rust/crates/auth-aws/src/aws.rs | 164 ++++++++++------- litellm-rust/crates/auth-aws/src/constants.rs | 14 -- litellm-rust/crates/auth-types/Cargo.toml | 1 + litellm-rust/crates/auth-types/src/error.rs | 5 + litellm-rust/crates/auth-types/src/lib.rs | 2 + .../crates/auth-types/src/params/aws.rs | 169 ++++++++++++++++++ .../crates/auth-types/src/params/mod.rs | 46 +++++ .../crates/auth-types/tests/aws_params.rs | 26 +++ .../crates/auth-types/tests/params.rs | 72 ++++++++ .../inference-ocr/tests/ocr/aws_textract.rs | 21 +++ .../ocr/analyze_transformation.rs | 2 +- .../llms/src/aws_textract/ocr/common_utils.rs | 13 +- .../src/aws_textract/ocr/transformation.rs | 2 +- .../crates/llms/src/base_llm/ocr/error.rs | 1 + .../src/bedrock/audio_transcription/mod.rs | 21 ++- .../bedrock/chat/converse_transformation.rs | 21 ++- .../anthropic_claude3_transformation.rs | 5 +- .../tests/bedrock_converse_transformation.rs | 46 ++--- 20 files changed, 500 insertions(+), 136 deletions(-) create mode 100644 litellm-rust/crates/auth-types/src/params/aws.rs create mode 100644 litellm-rust/crates/auth-types/src/params/mod.rs create mode 100644 litellm-rust/crates/auth-types/tests/aws_params.rs create mode 100644 litellm-rust/crates/auth-types/tests/params.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 10c09d3b77f..c1ee2c2e2f3 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3378,6 +3378,7 @@ dependencies = [ "moka", "reqwest 0.12.28", "rstest", + "serde", "serde_json", "sha2 0.10.9", "thiserror 2.0.19", @@ -3421,6 +3422,7 @@ version = "0.1.0" dependencies = [ "rstest", "serde", + "serde_json", "subtle", "thiserror 2.0.19", "tokio", diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index 2a9a9e4768c..bf757c6744e 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -10,7 +10,8 @@ litellm-auth-types.workspace = true litellm-http.workspace = true moka = { workspace = true, features = ["sync"] } -serde_json.workspace = true +serde.workspace = true +serde_json = { workspace = true, features = ["preserve_order"] } sha2.workspace = true thiserror.workspace = true diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index 69eb4265159..8fdbba21711 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -3,7 +3,6 @@ use std::time::Duration; use std::time::{SystemTime, UNIX_EPOCH}; use moka::sync::Cache; -use serde_json::{Map, Value}; use sha2::{Digest, Sha256}; use aws_credential_types::Credentials; @@ -14,10 +13,12 @@ use aws_sigv4::http_request::{ use aws_sigv4::sign::v4; use aws_smithy_runtime_api::client::identity::Identity; +use litellm_auth_types::AwsParams; + use super::Error; use super::constants::{ - AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION, AWS_REGION_NAME, - AWS_ROLE_ARN, AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN, + AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION_NAME, AWS_ROLE_ARN, + AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN, AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE, DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX, SIGV4_COMPUTED_HEADER_NAMES, }; @@ -543,50 +544,39 @@ fn is_bedrock_region(value: &str) -> bool { /// region, then the environment. Each service decides what a missing one means. pub fn resolve_aws_region( model_region: Option<&str>, - optional_params: &Map, + params: &AwsParams, env_lookup: &dyn Fn(&str) -> Option, ) -> Option { - optional_params - .get("aws_region_name") - .and_then(Value::as_str) - .or(model_region) - .map(str::to_string) - .or_else(|| env_lookup(AWS_REGION_NAME)) - .or_else(|| env_lookup(AWS_REGION)) + params + .resolve(&AwsParams::REGION, &|_| None) + .or_else(|| model_region.map(str::to_string)) + .or_else(|| AwsParams::REGION.resolve(&|_| None, env_lookup)) } pub fn resolve_bedrock_region( model_region: Option<&str>, - optional_params: &Map, + params: &AwsParams, env_lookup: &dyn Fn(&str) -> Option, ) -> String { - resolve_aws_region(model_region, optional_params, env_lookup) + resolve_aws_region(model_region, params, env_lookup) .unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string()) } pub fn aws_auth_config( - optional_params: &Map, + params: &AwsParams, env_lookup: &dyn Fn(&str) -> Option, ) -> AwsAuthConfig { - let value = |key: &str| { - optional_params - .get(key) - .and_then(Value::as_str) - .map(str::to_string) - }; - let env = |key: &str| env_lookup(key); AwsAuthConfig { - access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")), - secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")), - session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")), - region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)), - session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")), - profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")), - role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")), - web_identity_token: value("aws_web_identity_token") - .or_else(|| env("AWS_WEB_IDENTITY_TOKEN")), - sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")), - external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")), + access_key_id: params.resolve(&AwsParams::ACCESS_KEY_ID, env_lookup), + secret_access_key: params.resolve(&AwsParams::SECRET_ACCESS_KEY, env_lookup), + session_token: params.resolve(&AwsParams::SESSION_TOKEN, env_lookup), + region_name: params.resolve(&AwsParams::REGION, env_lookup), + session_name: params.resolve(&AwsParams::SESSION_NAME, env_lookup), + profile_name: params.resolve(&AwsParams::PROFILE_NAME, env_lookup), + role_name: params.resolve(&AwsParams::ROLE_NAME, env_lookup), + web_identity_token: params.resolve(&AwsParams::WEB_IDENTITY_TOKEN, env_lookup), + sts_endpoint: params.resolve(&AwsParams::STS_ENDPOINT, env_lookup), + external_id: params.resolve(&AwsParams::EXTERNAL_ID, env_lookup), } } @@ -599,13 +589,10 @@ pub enum AwsCredentialSource { } impl AwsCredentialSource { - pub fn from_params( - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Self { - match host_supplied_credentials(optional_params) { + pub fn from_params(params: &AwsParams, env_lookup: &dyn Fn(&str) -> Option) -> Self { + match host_supplied_credentials(params) { Some(credentials) => Self::HostSupplied(credentials), - None => Self::Chain(aws_auth_config(optional_params, env_lookup)), + None => Self::Chain(aws_auth_config(params, env_lookup)), } } @@ -629,20 +616,20 @@ impl AwsCredentialSource { /// state, where an unrelated `AWS_ROLE_NAME` or `AWS_PROFILE_NAME` in the /// environment outranks explicit keys in [`classify_auth`] and the two sides /// would sign as different principals. -pub fn host_supplied_credentials(optional_params: &Map) -> Option { - let value = |key: &str| { - optional_params - .get(key) - .and_then(Value::as_str) +pub fn host_supplied_credentials(params: &AwsParams) -> Option { + let value = |value: &Option| { + value + .as_deref() .map(str::trim) .filter(|value| !value.is_empty()) + .map(str::to_string) }; - let access_key_id = value("aws_access_key_id")?; - let secret_access_key = value("aws_secret_access_key")?; + let access_key_id = value(¶ms.aws_access_key_id)?; + let secret_access_key = value(¶ms.aws_secret_access_key)?; Some(Credentials::new( access_key_id, secret_access_key, - value("aws_session_token").map(str::to_string), + value(¶ms.aws_session_token), None, "litellm-host-supplied", )) @@ -651,7 +638,7 @@ pub fn host_supplied_credentials(optional_params: &Map) -> Option #[cfg(test)] mod tests { use super::*; - use crate::constants::BEDROCK_SERVICE; + use crate::constants::{AWS_REGION, BEDROCK_SERVICE}; fn no_env(_: &str) -> Option { None @@ -667,37 +654,80 @@ mod tests { recorded.lock().unwrap().insert(name.to_string()); None }; - resolve_aws_region(None, &Map::new(), &env); - aws_auth_config(&Map::new(), &env); + resolve_aws_region(None, &AwsParams::default(), &env); + aws_auth_config(&AwsParams::default(), &env); assert!( seen.lock() .unwrap() .iter() - .all(|name| crate::constants::SECRET_NAMES.contains(&name.as_str())) + .all(|name| AwsParams::secret_names().contains(&name.as_str())) ); } - #[test] - fn a_region_comes_from_the_call_then_the_model_then_the_environment() { - let params = Map::from_iter([("aws_region_name".to_string(), Value::from("eu-west-1"))]); - let region_name = |key: &str| (key == AWS_REGION_NAME).then(|| "ap-south-1".to_string()); - let region = |key: &str| (key == AWS_REGION).then(|| "sa-east-1".to_string()); - - let resolved = [ - resolve_aws_region(Some("us-east-2"), ¶ms, ®ion_name), - resolve_aws_region(Some("us-east-2"), &Map::new(), ®ion_name), - resolve_aws_region(None, &Map::new(), ®ion_name), - resolve_aws_region(None, &Map::new(), ®ion), - resolve_aws_region(None, &Map::new(), &no_env), - ]; + #[rstest::rstest] + #[case::param_wins(Some("from-params"), &[(AWS_SECRET_ACCESS_KEY, "from-env")], Some("from-params"))] + #[case::blank_param_falls_to_environment(Some(" "), &[(AWS_SECRET_ACCESS_KEY, "from-env")], Some("from-env"))] + #[case::absent_param_falls_to_environment(None, &[(AWS_SECRET_ACCESS_KEY, "from-env")], Some("from-env"))] + #[case::nothing(None, &[], None)] + fn auth_config_reads_each_spec_from_params_then_the_environment( + #[case] aws_secret_access_key: Option<&str>, + #[case] environment: &[(&str, &str)], + #[case] expected: Option<&str>, + ) { + let params = AwsParams { + aws_secret_access_key: aws_secret_access_key.map(str::to_string), + ..AwsParams::default() + }; + let env = |key: &str| { + environment + .iter() + .find(|(name, _)| *name == key) + .map(|(_, value)| value.to_string()) + }; assert_eq!( - resolved.map(|region| region.unwrap_or_else(|| "none".into())), - ["eu-west-1", "us-east-2", "ap-south-1", "sa-east-1", "none"] + aws_auth_config(¶ms, &env).secret_access_key.as_deref(), + expected ); assert_eq!( - resolve_bedrock_region(None, &Map::new(), &no_env), - DEFAULT_BEDROCK_REGION + aws_auth_config(&AwsParams::default(), &|key: &str| (key == AWS_REGION) + .then(|| "sa-east-1".to_string())) + .region_name + .as_deref(), + Some("sa-east-1") + ); + } + + #[rstest::rstest] + #[case::call_params(Some("us-east-2"), Some("eu-west-1"), &[(AWS_REGION_NAME, "ap-south-1")], Some("eu-west-1"))] + #[case::model_region(Some("us-east-2"), None, &[(AWS_REGION_NAME, "ap-south-1")], Some("us-east-2"))] + #[case::aws_region_name(None, None, &[(AWS_REGION_NAME, "ap-south-1")], Some("ap-south-1"))] + #[case::aws_region(None, None, &[(AWS_REGION, "sa-east-1")], Some("sa-east-1"))] + #[case::nothing(None, None, &[], None)] + fn a_region_comes_from_the_call_then_the_model_then_the_environment( + #[case] model_region: Option<&str>, + #[case] aws_region_name: Option<&str>, + #[case] environment: &[(&str, &str)], + #[case] expected: Option<&str>, + ) { + let params = AwsParams { + aws_region_name: aws_region_name.map(str::to_string), + ..AwsParams::default() + }; + let env = |key: &str| { + environment + .iter() + .find(|(name, _)| *name == key) + .map(|(_, value)| value.to_string()) + }; + + assert_eq!( + resolve_aws_region(model_region, ¶ms, &env).as_deref(), + expected + ); + assert_eq!( + resolve_bedrock_region(model_region, ¶ms, &env), + expected.unwrap_or(DEFAULT_BEDROCK_REGION) ); } diff --git a/litellm-rust/crates/auth-aws/src/constants.rs b/litellm-rust/crates/auth-aws/src/constants.rs index 26df4f2a350..61a379f8d66 100644 --- a/litellm-rust/crates/auth-aws/src/constants.rs +++ b/litellm-rust/crates/auth-aws/src/constants.rs @@ -14,20 +14,6 @@ pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT"; pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID"; pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK"; -pub const SECRET_NAMES: &[&str] = &[ - AWS_ACCESS_KEY_ID, - AWS_SECRET_ACCESS_KEY, - AWS_SESSION_TOKEN, - AWS_REGION_NAME, - AWS_REGION, - AWS_SESSION_NAME, - AWS_PROFILE_NAME, - AWS_ROLE_NAME, - AWS_WEB_IDENTITY_TOKEN, - AWS_STS_ENDPOINT, - AWS_EXTERNAL_ID, -]; - /// Headers SigV4 covers, beyond the `x-amz-` / `x-amzn-` prefixes. Mirrors /// Python's `_filter_headers_for_aws_signature` allowlist. pub const AWS_SIGNED_HEADER_NAMES: &[&str] = &[ diff --git a/litellm-rust/crates/auth-types/Cargo.toml b/litellm-rust/crates/auth-types/Cargo.toml index 8bdd52a8e84..b54c04f0229 100644 --- a/litellm-rust/crates/auth-types/Cargo.toml +++ b/litellm-rust/crates/auth-types/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] serde.workspace = true +serde_json.workspace = true subtle.workspace = true thiserror.workspace = true veil.workspace = true diff --git a/litellm-rust/crates/auth-types/src/error.rs b/litellm-rust/crates/auth-types/src/error.rs index 36e016ccc2a..ff17dcf3810 100644 --- a/litellm-rust/crates/auth-types/src/error.rs +++ b/litellm-rust/crates/auth-types/src/error.rs @@ -19,6 +19,11 @@ pub enum Error { provider: &'static str, environment_variable: &'static str, }, + #[error("Missing {provider} {} - {}", .spec.setting, .spec.guidance())] + MissingParam { + provider: &'static str, + spec: &'static crate::ParamSpec, + }, #[error("Missing {provider} API Base - {guidance}")] MissingApiBase { provider: &'static str, diff --git a/litellm-rust/crates/auth-types/src/lib.rs b/litellm-rust/crates/auth-types/src/lib.rs index c242a49af64..38b13988448 100644 --- a/litellm-rust/crates/auth-types/src/lib.rs +++ b/litellm-rust/crates/auth-types/src/lib.rs @@ -3,6 +3,7 @@ mod credential; mod error; pub mod http; +mod params; mod policy; mod secret; mod token; @@ -52,6 +53,7 @@ pub use credential::{ }; pub use error::{Error, ErrorDetail, ErrorSource}; pub use http::CredentialPlacement; +pub use params::{AwsParams, ParamSpec}; pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; pub use secret::SecretValue; pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; diff --git a/litellm-rust/crates/auth-types/src/params/aws.rs b/litellm-rust/crates/auth-types/src/params/aws.rs new file mode 100644 index 00000000000..f3653e969d6 --- /dev/null +++ b/litellm-rust/crates/auth-types/src/params/aws.rs @@ -0,0 +1,169 @@ +use std::sync::LazyLock; + +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +use super::ParamSpec; + +const fn spec( + setting: &'static str, + wire: &'static [&'static str], + env: &'static [&'static str], +) -> ParamSpec { + ParamSpec { + setting, + wire, + module_global: None, + env, + } +} + +/// The `aws_*` fields of Python's `GenericLiteLLMParams`. +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct AwsParams { + #[serde(default)] + pub aws_access_key_id: Option, + #[serde(default)] + pub aws_secret_access_key: Option, + #[serde(default)] + pub aws_session_token: Option, + #[serde(default)] + pub aws_region_name: Option, + #[serde(default)] + pub aws_session_name: Option, + #[serde(default)] + pub aws_profile_name: Option, + #[serde(default)] + pub aws_role_name: Option, + #[serde(default)] + pub aws_web_identity_token: Option, + #[serde(default)] + pub aws_sts_endpoint: Option, + #[serde(default)] + pub aws_external_id: Option, + #[serde(default)] + pub aws_bedrock_runtime_endpoint: Option, +} + +impl AwsParams { + pub const ACCESS_KEY_ID: ParamSpec = spec( + "access key id", + &["aws_access_key_id"], + &["AWS_ACCESS_KEY_ID"], + ); + pub const SECRET_ACCESS_KEY: ParamSpec = spec( + "secret access key", + &["aws_secret_access_key"], + &["AWS_SECRET_ACCESS_KEY"], + ); + pub const SESSION_TOKEN: ParamSpec = spec( + "session token", + &["aws_session_token"], + &["AWS_SESSION_TOKEN"], + ); + pub const REGION: ParamSpec = spec( + "region", + &["aws_region_name"], + &["AWS_REGION_NAME", "AWS_REGION"], + ); + pub const SESSION_NAME: ParamSpec = + spec("session name", &["aws_session_name"], &["AWS_SESSION_NAME"]); + pub const PROFILE_NAME: ParamSpec = + spec("profile name", &["aws_profile_name"], &["AWS_PROFILE_NAME"]); + pub const ROLE_NAME: ParamSpec = spec("role name", &["aws_role_name"], &["AWS_ROLE_NAME"]); + pub const WEB_IDENTITY_TOKEN: ParamSpec = spec( + "web identity token", + &["aws_web_identity_token"], + &["AWS_WEB_IDENTITY_TOKEN"], + ); + pub const STS_ENDPOINT: ParamSpec = + spec("STS endpoint", &["aws_sts_endpoint"], &["AWS_STS_ENDPOINT"]); + pub const EXTERNAL_ID: ParamSpec = + spec("external id", &["aws_external_id"], &["AWS_EXTERNAL_ID"]); + pub const BEDROCK_RUNTIME_ENDPOINT: ParamSpec = spec( + "Bedrock runtime endpoint", + &["aws_bedrock_runtime_endpoint"], + &["AWS_BEDROCK_RUNTIME_ENDPOINT"], + ); + + /// Every `aws_*` param Python reads, in declaration order, with the environment names + /// each falls back to. + pub const SPECS: [ParamSpec; 11] = [ + Self::ACCESS_KEY_ID, + Self::SECRET_ACCESS_KEY, + Self::SESSION_TOKEN, + Self::REGION, + Self::SESSION_NAME, + Self::PROFILE_NAME, + Self::ROLE_NAME, + Self::WEB_IDENTITY_TOKEN, + Self::STS_ENDPOINT, + Self::EXTERNAL_ID, + Self::BEDROCK_RUNTIME_ENDPOINT, + ]; + + /// The wire names a host projects out of a caller's kwargs, derived from [`Self::SPECS`]. + pub fn fields() -> impl Iterator { + Self::SPECS + .iter() + .flat_map(|spec| spec.wire.iter().copied()) + } + + /// The environment names the specs read, for hosts that resolve secrets up front. + pub fn secret_names() -> &'static [&'static str] { + static NAMES: LazyLock> = LazyLock::new(|| { + AwsParams::SPECS + .iter() + .flat_map(|spec| spec.env.iter().copied()) + .fold(Vec::new(), |mut names, name| { + if !names.contains(&name) { + names.push(name); + } + names + }) + }); + &NAMES + } + + /// The value under a wire name, read through serde so the names can never drift from + /// the struct. + pub fn get(&self, wire: &str) -> Option { + serde_json::to_value(self) + .ok()? + .get(wire)? + .as_str() + .map(str::to_string) + } + + /// The spec's value from these params, then the environment. + pub fn resolve( + &self, + spec: &ParamSpec, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Option { + spec.resolve(&|name| self.get(name), env_lookup) + } + + /// Reads the string-valued `aws_*` keys of an untyped params map, ignoring anything else. + pub fn from_optional_params(optional_params: &Map) -> Self { + let value = |key: &str| { + optional_params + .get(key) + .and_then(Value::as_str) + .map(str::to_string) + }; + Self { + aws_access_key_id: value("aws_access_key_id"), + aws_secret_access_key: value("aws_secret_access_key"), + aws_session_token: value("aws_session_token"), + aws_region_name: value("aws_region_name"), + aws_session_name: value("aws_session_name"), + aws_profile_name: value("aws_profile_name"), + aws_role_name: value("aws_role_name"), + aws_web_identity_token: value("aws_web_identity_token"), + aws_sts_endpoint: value("aws_sts_endpoint"), + aws_external_id: value("aws_external_id"), + aws_bedrock_runtime_endpoint: value("aws_bedrock_runtime_endpoint"), + } + } +} diff --git a/litellm-rust/crates/auth-types/src/params/mod.rs b/litellm-rust/crates/auth-types/src/params/mod.rs new file mode 100644 index 00000000000..093cdb42d64 --- /dev/null +++ b/litellm-rust/crates/auth-types/src/params/mod.rs @@ -0,0 +1,46 @@ +mod aws; + +pub use aws::AwsParams; + +/// One connection param of Python's `CredentialLiteLLMParams` and every place Python reads +/// it from, in order: the wire names on a call or a deployment, the `litellm.` module +/// global a host folds in, then the environment names. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ParamSpec { + pub setting: &'static str, + pub wire: &'static [&'static str], + pub module_global: Option<&'static str>, + pub env: &'static [&'static str], +} + +impl ParamSpec { + /// The first non-blank value: each wire name through `params`, then each environment + /// name through `env`. + pub fn resolve( + &self, + params: &dyn Fn(&str) -> Option, + env: &dyn Fn(&str) -> Option, + ) -> Option { + self.wire + .iter() + .map(|name| params(name)) + .chain(self.env.iter().map(|name| env(name))) + .find_map(|value| value.filter(|value| !value.trim().is_empty())) + } + + pub fn missing(&'static self, provider: &'static str) -> crate::Error { + crate::Error::MissingParam { + provider, + spec: self, + } + } + + pub(crate) fn guidance(&self) -> String { + let wire = self.wire.join(" or "); + let env = self.env.join(" or "); + match self.module_global { + Some(global) => format!("pass {wire}, set litellm.{global}, or set {env}"), + None => format!("pass {wire} or set {env}"), + } + } +} diff --git a/litellm-rust/crates/auth-types/tests/aws_params.rs b/litellm-rust/crates/auth-types/tests/aws_params.rs new file mode 100644 index 00000000000..ed9a4f78ec2 --- /dev/null +++ b/litellm-rust/crates/auth-types/tests/aws_params.rs @@ -0,0 +1,26 @@ +use litellm_auth_types::AwsParams; +use rstest::rstest; +use serde_json::{Map, Value}; + +#[rstest] +fn fields_name_every_param_once_and_in_declaration_order() { + let filled: Map = AwsParams::fields() + .map(|name| (name.to_string(), Value::from(format!("value-of-{name}")))) + .collect(); + let typed = AwsParams::from_optional_params(&filled); + let serialized = serde_json::to_value(&typed).unwrap(); + assert_eq!( + serialized + .as_object() + .unwrap() + .keys() + .map(String::as_str) + .collect::>(), + AwsParams::fields().collect::>() + ); + assert_eq!(serialized, Value::Object(filled.clone())); + assert_eq!( + serde_json::from_value::(Value::Object(filled)).unwrap(), + typed + ); +} diff --git a/litellm-rust/crates/auth-types/tests/params.rs b/litellm-rust/crates/auth-types/tests/params.rs new file mode 100644 index 00000000000..9d488b170e1 --- /dev/null +++ b/litellm-rust/crates/auth-types/tests/params.rs @@ -0,0 +1,72 @@ +use litellm_auth_types::{AwsParams, ParamSpec}; +use rstest::rstest; + +const LOCATION: ParamSpec = ParamSpec { + setting: "location", + wire: &["vertex_location", "vertex_ai_location"], + module_global: Some("vertex_location"), + env: &["VERTEXAI_LOCATION", "VERTEX_LOCATION"], +}; + +fn lookup(entries: &'static [(&'static str, &'static str)]) -> impl Fn(&str) -> Option { + move |name| { + entries + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + } +} + +#[rstest] +#[case::current_spelling_first( + &[("vertex_location", "eu"), ("vertex_ai_location", "us")], + &[("VERTEXAI_LOCATION", "env")], + Some("eu") +)] +#[case::legacy_spelling_before_the_environment( + &[("vertex_ai_location", "us")], + &[("VERTEXAI_LOCATION", "env")], + Some("us") +)] +#[case::blank_values_are_skipped( + &[("vertex_location", " "), ("vertex_ai_location", "")], + &[("VERTEXAI_LOCATION", " "), ("VERTEX_LOCATION", "fallback")], + Some("fallback") +)] +#[case::environment_names_in_order( + &[], + &[("VERTEX_LOCATION", "second"), ("VERTEXAI_LOCATION", "first")], + Some("first") +)] +#[case::nothing_set(&[], &[], None)] +fn resolution_walks_wire_names_then_environment_names( + #[case] params: &'static [(&'static str, &'static str)], + #[case] env: &'static [(&'static str, &'static str)], + #[case] expected: Option<&str>, +) { + assert_eq!( + LOCATION.resolve(&lookup(params), &lookup(env)).as_deref(), + expected + ); +} + +#[rstest] +#[case::without_a_module_global( + &AwsParams::REGION, + "Missing AWS region - pass aws_region_name or set AWS_REGION_NAME or AWS_REGION" +)] +#[case::with_a_module_global( + &LOCATION, + "Missing AWS location - pass vertex_location or vertex_ai_location, set litellm.vertex_location, or set VERTEXAI_LOCATION or VERTEX_LOCATION" +)] +fn a_missing_param_names_every_place_the_spec_reads( + #[case] spec: &'static ParamSpec, + #[case] expected: &str, +) { + let message = spec.missing("AWS").to_string(); + + assert_eq!(message, expected); + assert!(spec.wire.iter().all(|name| message.contains(name))); + assert!(spec.env.iter().all(|name| message.contains(name))); + assert_eq!(message.contains("litellm."), spec.module_global.is_some()); +} diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs b/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs index 790e16a95ec..5336d67895d 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs @@ -144,6 +144,27 @@ async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() { assert_eq!(response.pages[0].markdown, "# Quarterly Report"); } +#[rstest] +#[tokio::test] +async fn a_missing_region_is_the_callers_error_and_names_every_way_to_set_it() { + let upstream = upstream([textract_response()]).await; + let request = ocr_request_with_document( + DETECT, + &format!("{}/", upstream.uri()), + json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}), + json!({"aws_access_key_id": ACCESS_KEY_ID, "aws_secret_access_key": SECRET_ACCESS_KEY}), + ); + + let error = perform_with(LocalOcrHost::new(request)).await.unwrap_err(); + + assert!(error.is_request(), "{error:?}"); + assert_eq!( + error.to_string(), + "Missing AWS region - pass aws_region_name or set AWS_REGION_NAME or AWS_REGION" + ); + assert!(upstream.received_requests().await.unwrap().is_empty()); +} + #[rstest] #[case::detect(DETECT)] #[case::analyze(ANALYZE)] diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs index 5506c86c17f..e813d4ac8cc 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs @@ -41,7 +41,7 @@ impl BaseOcrConfig for TextractAnalyzeDocumentConfig { type Environment = TextractEnvironment; fn secret_names(&self) -> Vec<&'static str> { - litellm_auth_aws::constants::SECRET_NAMES.to_vec() + litellm_auth::AwsParams::secret_names().to_vec() } fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs index 4139ee860cf..bc199f40938 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs @@ -1,4 +1,5 @@ use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_auth::AwsParams; use litellm_auth_aws::{AwsCredentialSource, SigV4Signer, resolve_aws_region}; use litellm_http::outbound::RequestSigner; use serde::{Deserialize, Serialize}; @@ -237,18 +238,14 @@ pub(super) async fn environment( operation: TextractOperation, ) -> Result { let env_lookup = |name: &str| request.connection.secret(name); - let region = - resolve_aws_region(None, &request.optional_params, &env_lookup).ok_or_else(|| { - Error::InvalidRequest( - "Missing AWS region - pass aws_region_name or set AWS_REGION_NAME or AWS_REGION" - .into(), - ) - })?; + let params = AwsParams::from_optional_params(&request.optional_params); + let region = resolve_aws_region(None, ¶ms, &env_lookup) + .ok_or_else(|| Error::Auth(AwsParams::REGION.missing("AWS")))?; let signer = SigV4Signer::resolve( auth, region.clone(), TEXTRACT_SERVICE, - AwsCredentialSource::from_params(&request.optional_params, &env_lookup), + AwsCredentialSource::from_params(¶ms, &env_lookup), &env_lookup, ) .await diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs index 6bef577e6f8..4b27a191734 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs @@ -30,7 +30,7 @@ impl BaseOcrConfig for TextractDetectTextConfig { type Environment = TextractEnvironment; fn secret_names(&self) -> Vec<&'static str> { - litellm_auth_aws::constants::SECRET_NAMES.to_vec() + litellm_auth::AwsParams::secret_names().to_vec() } fn get_health_check_document(&self) -> OcrDocument { diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index 5d7be8b10df..392dd56412c 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -166,6 +166,7 @@ impl Error { | Self::Params(_) | Self::Headers(_) | Self::Http(_) + | Self::Auth(litellm_auth::Error::MissingParam { .. }) ) } diff --git a/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs b/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs index c0b2c7aa3b7..307e7a09498 100644 --- a/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs +++ b/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs @@ -1,3 +1,4 @@ +use litellm_auth::AwsParams; use litellm_auth_aws::{ AwsCredentialSource, bedrock_model_id_and_region, constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}, @@ -85,7 +86,7 @@ fn optional_string<'a>(params: &'a Map, key: &str) -> Option<&'a impl BaseAudioTranscriptionConfig for BedrockAudioTranscriptionConfig { fn secret_names(&self) -> Vec<&'static str> { - litellm_auth_aws::constants::SECRET_NAMES.to_vec() + litellm_auth::AwsParams::secret_names().to_vec() } fn get_supported_openai_params(&self) -> &'static [&'static str] { @@ -154,7 +155,11 @@ impl BaseAudioTranscriptionConfig for BedrockAudioTranscriptionConfig { env_lookup: &dyn Fn(&str) -> Option, ) -> Result { let (model_id, model_region) = bedrock_model_id_and_region(model); - let region = resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup); + let region = resolve_bedrock_region( + model_region.as_deref(), + &AwsParams::from_optional_params(optional_params), + env_lookup, + ); let endpoint = optional_params .get("aws_bedrock_runtime_endpoint") .and_then(Value::as_str) @@ -181,19 +186,13 @@ impl BaseAudioTranscriptionConfig for BedrockAudioTranscriptionConfig { env_lookup: &dyn Fn(&str) -> Option, ) -> Result { let (_, model_region) = bedrock_model_id_and_region(model); + let params = AwsParams::from_optional_params(optional_params); Ok(ValidatedEnvironment { headers, auth: AuthScheme::AwsSigV4 { - region: resolve_bedrock_region( - model_region.as_deref(), - optional_params, - env_lookup, - ), + region: resolve_bedrock_region(model_region.as_deref(), ¶ms, env_lookup), service: BEDROCK_SERVICE, - credentials: Box::new(AwsCredentialSource::from_params( - optional_params, - env_lookup, - )), + credentials: Box::new(AwsCredentialSource::from_params(¶ms, env_lookup)), }, }) } diff --git a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs index bcd1077b830..73e02d4b44e 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -1,3 +1,4 @@ +use litellm_auth::AwsParams; use litellm_auth::{CredentialPlacement, SecretValue}; use litellm_auth_aws::{ AwsCredentialSource, bedrock_model_id_and_region, @@ -168,7 +169,7 @@ pub const BEDROCK_CHAT_COMPLETIONS_CONFIG: AmazonConverseConfig = AmazonConverse impl BaseConfig for AmazonConverseConfig { fn secret_names(&self) -> Vec<&'static str> { - litellm_auth_aws::constants::SECRET_NAMES + litellm_auth::AwsParams::secret_names() .iter() .copied() .chain([AWS_BEARER_TOKEN_BEDROCK]) @@ -187,7 +188,11 @@ impl BaseConfig for AmazonConverseConfig { env_lookup: &dyn Fn(&str) -> Option, ) -> Result { let (model_id, model_region) = bedrock_model_id_and_region(model); - let region = resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup); + let region = resolve_bedrock_region( + model_region.as_deref(), + &AwsParams::from_optional_params(optional_params), + env_lookup, + ); let endpoint = optional_params .get(AWS_BEDROCK_RUNTIME_ENDPOINT) .and_then(Value::as_str) @@ -316,19 +321,13 @@ impl BaseConfig for AmazonConverseConfig { }); } let (_, model_region) = bedrock_model_id_and_region(model); + let params = AwsParams::from_optional_params(optional_params); Ok(ValidatedEnvironment { headers, auth: AuthScheme::AwsSigV4 { - region: resolve_bedrock_region( - model_region.as_deref(), - optional_params, - env_lookup, - ), + region: resolve_bedrock_region(model_region.as_deref(), ¶ms, env_lookup), service: BEDROCK_SERVICE, - credentials: Box::new(AwsCredentialSource::from_params( - optional_params, - env_lookup, - )), + credentials: Box::new(AwsCredentialSource::from_params(¶ms, env_lookup)), }, }) } diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs index beae62ed065..c5bf7d8fc95 100644 --- a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs @@ -5,6 +5,7 @@ use crate::{ base_llm::messages::context::MessagesTransformContext, }; use futures_util::StreamExt; +use litellm_auth::AwsParams; use litellm_auth::{CredentialPlacement, SecretValue}; use litellm_auth_aws::{ AwsCredentialSource, bedrock_model_id_and_region, @@ -78,7 +79,7 @@ fn invoke_url( ) -> String { let (model_id, model_region) = bedrock_model_id_and_region(model.strip_prefix(INVOKE_MODEL_PREFIX).unwrap_or(model)); - let region = resolve_bedrock_region(model_region.as_deref(), &Map::new(), env_lookup); + let region = resolve_bedrock_region(model_region.as_deref(), &AwsParams::default(), env_lookup); let endpoint = api_base .map(str::trim) .filter(|value| !value.is_empty()) @@ -149,7 +150,7 @@ impl BaseMessagesConfig for AmazonAnthropicClaudeMessagesConfig { } let (_, model_region) = bedrock_model_id_and_region(model.strip_prefix(INVOKE_MODEL_PREFIX).unwrap_or(model)); - let params = Map::new(); + let params = AwsParams::default(); Ok(ValidatedEnvironment { headers, auth: AuthScheme::AwsSigV4 { diff --git a/litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs b/litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs index aa6920d58f6..dc3590c384f 100644 --- a/litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs +++ b/litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs @@ -582,28 +582,34 @@ fn leaves_a_complete_converse_url_untouched() { ); } -#[test] -fn host_supplied_credentials_outrank_ambient_profile_and_role_state() { +#[rstest] +#[case::full_static_pair( + json!({"aws_access_key_id": "AKIAHOST", "aws_secret_access_key": "hostsecret", "aws_session_token": "hosttoken"}), + Some(("AKIAHOST", "hostsecret", Some("hosttoken"))) +)] +#[case::pair_without_session_token( + json!({"aws_access_key_id": "AKIAHOST", "aws_secret_access_key": "hostsecret"}), + Some(("AKIAHOST", "hostsecret", None)) +)] +#[case::key_id_alone(json!({"aws_access_key_id": "AKIA"}), None)] +#[case::blank_key_id(json!({"aws_access_key_id": " ", "aws_secret_access_key": "s"}), None)] +#[case::nothing(json!({}), None)] +fn host_supplied_credentials_need_a_full_static_pair( + #[case] optional_params: Value, + #[case] expected: Option<(&str, &str, Option<&str>)>, +) { + use litellm_auth::AwsParams; use litellm_auth_aws::host_supplied_credentials; - let supplied = params(json!({ - "aws_access_key_id": "AKIAHOST", - "aws_secret_access_key": "hostsecret", - "aws_session_token": "hosttoken" - })); - let credentials = host_supplied_credentials(&supplied).expect("host credentials"); - assert_eq!(credentials.access_key_id(), "AKIAHOST"); - assert_eq!(credentials.secret_access_key(), "hostsecret"); - assert_eq!(credentials.session_token(), Some("hosttoken")); + let credentials = + host_supplied_credentials(&AwsParams::from_optional_params(¶ms(optional_params))); - // Without a full static pair there is nothing to honor, so the core falls - // back to deriving credentials itself. - assert!(host_supplied_credentials(¶ms(json!({"aws_access_key_id": "AKIA"}))).is_none()); - assert!( - host_supplied_credentials(¶ms( - json!({"aws_access_key_id": " ", "aws_secret_access_key": "s"}) - )) - .is_none() + assert_eq!( + credentials.as_ref().map(|credentials| ( + credentials.access_key_id(), + credentials.secret_access_key(), + credentials.session_token(), + )), + expected ); - assert!(host_supplied_credentials(&Map::new()).is_none()); } From 04a95f60c15137fdf134298199f48df7066048eb Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 9 Oct 2026 16:29:44 -0700 Subject: [PATCH 002/104] feat(rust): type the Vertex AI connection params and where Python reads them (#45641) VertexParams in auth-types holds the vertex_* fields of CredentialLiteLLMParams in both spellings, with one ParamSpec per setting: the wire names, the litellm module global VertexBase.get_vertex_ai_* consults, and the environment names. A credential deserializes from text or the JSON object a config spells, an empty object counts as absent, and both credential fields are redacted in Debug. auth-gcp is split so lib.rs is the entrypoint (config, auth, constants), its lookups resolve through the specs, secret_names() derives from them, and get_vertex_ai_project_from_credentials reads a service account's project_id for hosts that need it before a token exchange. The OCR host folds the module globals in by spec at projection, so OcrSettings no longer carries vertex_project and vertex_location. --- litellm-rust/Cargo.lock | 3 + litellm-rust/crates/auth-aws/src/constants.rs | 26 +- litellm-rust/crates/auth-gcp/Cargo.toml | 3 + litellm-rust/crates/auth-gcp/src/auth.rs | 274 +++++++ litellm-rust/crates/auth-gcp/src/config.rs | 311 ++++++++ litellm-rust/crates/auth-gcp/src/constants.rs | 36 + litellm-rust/crates/auth-gcp/src/lib.rs | 716 +----------------- .../crates/auth-gcp/tests/vertex_auth.rs | 120 +++ .../crates/auth-gcp/tests/vertex_config.rs | 95 +++ litellm-rust/crates/auth-types/src/lib.rs | 2 +- .../crates/auth-types/src/params/mod.rs | 2 + .../crates/auth-types/src/params/vertex.rs | 118 +++ .../crates/auth-types/tests/vertex_params.rs | 126 +++ .../inference-ocr/tests/ocr/vertex_ai.rs | 25 - .../crates/inference-ocr/tests/resources.rs | 7 +- .../crates/llms/src/base_llm/ocr/settings.rs | 4 - .../llms/src/vertex_ai/ocr/common_utils.rs | 7 +- .../llms/src/vertex_ai/ocr/transformation.rs | 2 +- .../python-bridge/src/python_settings.rs | 2 +- .../python-bridge/src/routes/ocr/host.rs | 13 +- .../python-bridge/src/routes/ocr/mod.rs | 40 +- .../python-bridge/src/routes/ocr/project.rs | 165 ++++ tests/test_litellm_rust/ocr/test_requests.py | 28 +- 23 files changed, 1337 insertions(+), 788 deletions(-) create mode 100644 litellm-rust/crates/auth-gcp/src/auth.rs create mode 100644 litellm-rust/crates/auth-gcp/src/config.rs create mode 100644 litellm-rust/crates/auth-gcp/src/constants.rs create mode 100644 litellm-rust/crates/auth-gcp/tests/vertex_auth.rs create mode 100644 litellm-rust/crates/auth-gcp/tests/vertex_config.rs create mode 100644 litellm-rust/crates/auth-types/src/params/vertex.rs create mode 100644 litellm-rust/crates/auth-types/tests/vertex_params.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c1ee2c2e2f3..4cbde047a71 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3411,9 +3411,12 @@ dependencies = [ "litellm-auth-types", "moka", "rstest", + "serde", "serde_json", "sha2 0.10.9", + "tempfile", "tokio", + "veil", ] [[package]] diff --git a/litellm-rust/crates/auth-aws/src/constants.rs b/litellm-rust/crates/auth-aws/src/constants.rs index 61a379f8d66..789fde156bd 100644 --- a/litellm-rust/crates/auth-aws/src/constants.rs +++ b/litellm-rust/crates/auth-aws/src/constants.rs @@ -1,18 +1,20 @@ -pub const AWS_ACCESS_KEY_ID: &str = "AWS_ACCESS_KEY_ID"; -pub const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY"; -pub const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN"; -pub const AWS_REGION_NAME: &str = "AWS_REGION_NAME"; -pub const AWS_REGION: &str = "AWS_REGION"; +use litellm_auth_types::AwsParams; + +pub const AWS_ACCESS_KEY_ID: &str = AwsParams::ACCESS_KEY_ID.env[0]; +pub const AWS_SECRET_ACCESS_KEY: &str = AwsParams::SECRET_ACCESS_KEY.env[0]; +pub const AWS_SESSION_TOKEN: &str = AwsParams::SESSION_TOKEN.env[0]; +pub const AWS_REGION_NAME: &str = AwsParams::REGION.env[0]; +pub const AWS_REGION: &str = AwsParams::REGION.env[1]; pub const AWS_DEFAULT_REGION: &str = "AWS_DEFAULT_REGION"; -pub const AWS_BEDROCK_RUNTIME_ENDPOINT: &str = "AWS_BEDROCK_RUNTIME_ENDPOINT"; -pub const AWS_SESSION_NAME: &str = "AWS_SESSION_NAME"; -pub const AWS_PROFILE_NAME: &str = "AWS_PROFILE_NAME"; -pub const AWS_ROLE_NAME: &str = "AWS_ROLE_NAME"; -pub const AWS_WEB_IDENTITY_TOKEN: &str = "AWS_WEB_IDENTITY_TOKEN"; +pub const AWS_BEDROCK_RUNTIME_ENDPOINT: &str = AwsParams::BEDROCK_RUNTIME_ENDPOINT.env[0]; +pub const AWS_SESSION_NAME: &str = AwsParams::SESSION_NAME.env[0]; +pub const AWS_PROFILE_NAME: &str = AwsParams::PROFILE_NAME.env[0]; +pub const AWS_ROLE_NAME: &str = AwsParams::ROLE_NAME.env[0]; +pub const AWS_WEB_IDENTITY_TOKEN: &str = AwsParams::WEB_IDENTITY_TOKEN.env[0]; pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN"; pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; -pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT"; -pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID"; +pub const AWS_STS_ENDPOINT: &str = AwsParams::STS_ENDPOINT.env[0]; +pub const AWS_EXTERNAL_ID: &str = AwsParams::EXTERNAL_ID.env[0]; pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK"; /// Headers SigV4 covers, beyond the `x-amz-` / `x-amzn-` prefixes. Mirrors /// Python's `_filter_headers_for_aws_signature` allowlist. diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index 8a3598234e1..24791ce7b53 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -12,9 +12,11 @@ google-sdk = ["dep:google-cloud-auth", "dep:http"] litellm-auth-types.workspace = true moka.workspace = true +serde.workspace = true serde_json.workspace = true sha2.workspace = true tokio.workspace = true +veil.workspace = true gcp_auth = "0.12.7" google-cloud-auth = { workspace = true, optional = true } @@ -22,3 +24,4 @@ http = { workspace = true, optional = true } [dev-dependencies] rstest.workspace = true +tempfile.workspace = true diff --git a/litellm-rust/crates/auth-gcp/src/auth.rs b/litellm-rust/crates/auth-gcp/src/auth.rs new file mode 100644 index 00000000000..d7c434fe26e --- /dev/null +++ b/litellm-rust/crates/auth-gcp/src/auth.rs @@ -0,0 +1,274 @@ +use std::{future::Future, path::Path, pin::Pin, sync::Arc}; + +use gcp_auth::{CustomServiceAccount, TokenProvider}; +use litellm_auth_types::{CredentialPlacement, Error, ErrorDetail, http::apply_credential}; +use moka::future::Cache; +use serde_json::Value; +use sha2::{Digest, Sha256}; + +use crate::{ + config::{ + CredentialSource, VertexConfig, credential_source, get_vertex_ai_project, non_empty_env, + }, + constants::{ + CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT, VERTEX_AI_API_KEY_ENV, + VERTEXAI_API_KEY_ENV, + }, +}; + +pub struct VertexEnvironment { + pub headers: Vec<(String, String)>, + pub project_id: String, +} + +struct VertexAccessToken { + token: String, + project_id: String, +} + +#[derive(Clone)] +pub struct VertexAuth { + providers: Cache>, + loader: Arc, +} + +impl Default for VertexAuth { + fn default() -> Self { + Self::new(Arc::new(GcpProviderLoader)) + } +} + +impl VertexAuth { + pub fn new(loader: Arc) -> Self { + Self { + providers: Cache::builder().max_capacity(64).build(), + loader, + } + } + + pub async fn access_token( + &self, + config: &VertexConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result { + self.load_provider(config, env_lookup).await?.token().await + } + + pub async fn validate_environment( + &self, + headers: Vec<(String, String)>, + api_key: Option<&str>, + config: &VertexConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result { + let has_authorization = headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case("Authorization")); + let static_token = api_key + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + .or_else(|| non_empty_env(env_lookup, VERTEX_AI_API_KEY_ENV)) + .or_else(|| non_empty_env(env_lookup, VERTEXAI_API_KEY_ENV)); + let project_id = get_vertex_ai_project(config, env_lookup); + + if !has_authorization && static_token.is_none() { + let access = self.get_access_token(config, env_lookup).await?; + return Ok(VertexEnvironment { + headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?, + project_id: project_id.unwrap_or(access.project_id), + }); + } + + let project_id = match project_id { + Some(project_id) => project_id, + None => { + self.load_provider(config, env_lookup) + .await? + .project_id() + .await? + } + }; + let headers = if has_authorization { + headers + } else { + apply_credential( + headers, + static_token.as_deref().expect("static token was checked"), + CredentialPlacement::Bearer, + )? + }; + Ok(VertexEnvironment { + headers, + project_id, + }) + } + + async fn get_access_token( + &self, + config: &VertexConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result { + let provider = self.load_provider(config, env_lookup).await?; + let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?; + Ok(VertexAccessToken { token, project_id }) + } + + async fn load_provider( + &self, + config: &VertexConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, Error> { + let source = credential_source(config, env_lookup); + let key = source.cache_key(); + self.providers + .try_get_with(key, self.loader.load(source)) + .await + .map_err(|error| (*error).clone()) + } +} + +pub trait VertexTokenSource: Send + Sync { + fn project_id(&self) -> VertexAuthFuture<'_, String>; + fn token(&self) -> VertexAuthFuture<'_, String>; +} + +pub trait VertexProviderLoader: Send + Sync { + fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc>; +} + +pub type VertexAuthFuture<'a, T> = Pin> + Send + 'a>>; + +struct GcpTokenSource(Arc); + +impl VertexTokenSource for GcpTokenSource { + fn project_id(&self) -> VertexAuthFuture<'_, String> { + Box::pin(async move { + self.0 + .project_id() + .await + .map(|project| project.to_string()) + .map_err(auth_acquisition_error) + }) + } + + fn token(&self) -> VertexAuthFuture<'_, String> { + Box::pin(async move { + self.0 + .token(&[CLOUD_PLATFORM_SCOPE]) + .await + .map(|token| token.as_str().to_string()) + .map_err(auth_acquisition_error) + }) + } +} + +struct GcpProviderLoader; + +impl VertexProviderLoader for GcpProviderLoader { + fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc> { + Box::pin(async move { + let provider: Arc = match source { + CredentialSource::Inline(configured) => Arc::new( + CustomServiceAccount::from_json(validate_request_credentials( + configured.expose(), + )?) + .map_err(auth_acquisition_error)?, + ), + CredentialSource::Trusted(configured) => { + let configured = configured.expose(); + let service_account = if Path::new(configured).is_file() { + CustomServiceAccount::from_file(configured) + } else { + CustomServiceAccount::from_json(configured) + } + .map_err(auth_acquisition_error)?; + Arc::new(service_account) + } + CredentialSource::ApplicationCredentials(path) => { + Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?) + } + CredentialSource::Adc => { + gcp_auth::provider().await.map_err(auth_acquisition_error)? + } + }; + Ok(Arc::new(GcpTokenSource(provider)) as Arc) + }) + } +} + +fn validate_request_credentials(configured: &str) -> Result<&str, Error> { + let token_uri = serde_json::from_str::(configured) + .ok() + .and_then(|credentials| { + credentials + .get("token_uri") + .and_then(Value::as_str) + .map(str::to_string) + }); + if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) { + return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { + field: "request-controlled vertex_credentials token_uri".into(), + expected: "the canonical Google OAuth token endpoint", + })); + } + Ok(configured) +} + +impl CredentialSource { + fn cache_key(&self) -> CredentialCacheKey { + match self { + Self::Inline(configured) => { + CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into()) + } + Self::Trusted(configured) => { + CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into()) + } + Self::ApplicationCredentials(path) => { + CredentialCacheKey::ApplicationCredentials(path.clone()) + } + Self::Adc => CredentialCacheKey::Adc, + } + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +enum CredentialCacheKey { + Inline([u8; 32]), + Trusted([u8; 32]), + ApplicationCredentials(String), + Adc, +} + +fn auth_acquisition_error(error: gcp_auth::Error) -> Error { + Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed( + "Vertex AI credentials", + error, + )) +} + +#[cfg(test)] +mod tests { + use litellm_auth_types::SecretValue; + + use super::*; + + #[rstest::rstest] + #[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)] + #[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)] + #[case::missing_endpoint("{}", false)] + fn request_credentials_require_canonical_token_endpoint( + #[case] credentials: &str, + #[case] accepted: bool, + ) { + assert_eq!(validate_request_credentials(credentials).is_ok(), accepted); + } + + #[rstest::rstest] + fn inline_and_trusted_credentials_never_share_a_cache_entry() { + assert_ne!( + CredentialSource::Inline(SecretValue::new("same-value")).cache_key(), + CredentialSource::Trusted(SecretValue::new("same-value")).cache_key() + ); + } +} diff --git a/litellm-rust/crates/auth-gcp/src/config.rs b/litellm-rust/crates/auth-gcp/src/config.rs new file mode 100644 index 00000000000..0a76bf93cf0 --- /dev/null +++ b/litellm-rust/crates/auth-gcp/src/config.rs @@ -0,0 +1,311 @@ +use std::{collections::BTreeMap, path::Path}; + +use litellm_auth_types::{Error, ErrorDetail, InputSource, SecretValue, Sourced}; +use serde_json::{Map, Value}; + +use litellm_auth_types::VertexParams; + +use crate::constants::GOOGLE_APPLICATION_CREDENTIALS_ENV; + +#[derive(Clone, Debug, Default)] +pub struct VertexConfig { + credentials: Option>, + project_id: Option, + location: Option, +} + +impl VertexConfig { + pub fn new( + credentials: Option>, + project_id: Option, + location: Option, + ) -> Self { + Self { + credentials: credentials.filter(|value| !value.value().expose().trim().is_empty()), + project_id: project_id.filter(|value| !value.trim().is_empty()), + location: location.filter(|value| !value.trim().is_empty()), + } + } + + pub fn from_sourced_optional_params( + params: &Map, + sources: &BTreeMap, + ) -> Result { + Ok(Self::new( + optional_credentials(params, sources, VertexParams::CREDENTIALS.wire)?, + optional_string(params, VertexParams::PROJECT.wire)?, + optional_string(params, VertexParams::LOCATION.wire)?, + )) + } + + /// The deployment's `litellm_params`, which a host already trusts the way it trusts its + /// own configuration. + pub fn from_params(params: &VertexParams) -> Self { + Self::new( + params + .credentials() + .map(|value| Sourced::new(SecretValue::new(value), InputSource::Deployment)), + params.project(), + params.location(), + ) + } + + pub fn project_id(&self) -> Option<&str> { + self.project_id.as_deref() + } + + pub fn location(&self) -> Option<&str> { + self.location.as_deref() + } +} + +pub fn get_vertex_ai_project( + config: &VertexConfig, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + config + .project_id() + .map(str::to_string) + .or_else(|| VertexParams::PROJECT.resolve(&|_| None, env_lookup)) +} + +/// The project a request is billed to when nothing names it: the `project_id` inside the +/// service account the call would authenticate with. Application default credentials carry +/// no project that can be read without a token exchange, so they resolve to `None`. +pub fn get_vertex_ai_project_from_credentials( + config: &VertexConfig, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + let text = match credential_source(config, env_lookup) { + CredentialSource::Inline(configured) => configured.expose().to_string(), + CredentialSource::Trusted(configured) => { + let configured = configured.expose(); + match std::fs::read_to_string(configured) { + Ok(contents) if Path::new(configured).is_file() => contents, + _ => configured.to_string(), + } + } + CredentialSource::ApplicationCredentials(path) => std::fs::read_to_string(path).ok()?, + CredentialSource::Adc => return None, + }; + serde_json::from_str::(&text) + .ok()? + .get("project_id") + .and_then(Value::as_str) + .map(str::to_string) +} + +pub fn get_vertex_ai_location( + config: &VertexConfig, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + config + .location() + .map(str::to_string) + .or_else(|| VertexParams::LOCATION.resolve(&|_| None, env_lookup)) +} + +#[derive(Clone, Debug)] +pub enum CredentialSource { + Inline(SecretValue), + Trusted(SecretValue), + ApplicationCredentials(String), + Adc, +} + +pub(crate) fn credential_source( + config: &VertexConfig, + env_lookup: &dyn Fn(&str) -> Option, +) -> CredentialSource { + if let Some(configured) = config.credentials.clone() { + return match configured.source() { + InputSource::Request => CredentialSource::Inline(configured.into_value()), + InputSource::Deployment | InputSource::Environment => { + CredentialSource::Trusted(configured.into_value()) + } + }; + } + if let Some(configured) = VertexParams::CREDENTIALS.resolve(&|_| None, env_lookup) { + return CredentialSource::Trusted(SecretValue::new(configured)); + } + non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV) + .map(CredentialSource::ApplicationCredentials) + .unwrap_or(CredentialSource::Adc) +} + +fn optional_credentials( + params: &Map, + sources: &BTreeMap, + names: &[&str], +) -> Result>, Error> { + for name in names { + let source = source_for(sources, name); + match params.get(*name) { + None | Some(Value::Null) => continue, + Some(Value::String(value)) if value.trim().is_empty() => continue, + Some(Value::String(value)) => { + return Ok(Some(Sourced::new(SecretValue::new(value), source))); + } + Some(Value::Object(value)) if value.is_empty() => continue, + Some(Value::Object(value)) => { + return serde_json::to_string(value) + .map(SecretValue::new) + .map(|value| Sourced::new(value, source)) + .map(Some) + .map_err(|error| { + Error::InvalidConfiguration(ErrorDetail::failed( + "credential serialization", + error, + )) + }); + } + Some(_) => { + return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { + field: names[0].into(), + expected: "a string or null", + })); + } + } + } + Ok(None) +} + +fn source_for(sources: &BTreeMap, name: &str) -> InputSource { + sources.get(name).copied().unwrap_or_default() +} + +fn optional_string(params: &Map, names: &[&str]) -> Result, Error> { + for name in names { + match params.get(*name) { + None | Some(Value::Null) => continue, + Some(Value::String(value)) if value.trim().is_empty() => continue, + Some(Value::String(value)) => return Ok(Some(value.clone())), + Some(_) => { + return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { + field: names[0].into(), + expected: "a string or null", + })); + } + } + } + Ok(None) +} + +pub(crate) fn non_empty_env( + env_lookup: &dyn Fn(&str) -> Option, + name: &str, +) -> Option { + env_lookup(name) + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::constants::VERTEXAI_CREDENTIALS_ENV; + + fn config(value: Value) -> VertexConfig { + VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new()) + .unwrap() + } + + #[rstest::rstest] + fn empty_primary_values_fall_back_to_python_aliases() { + let config = config(json!({ + "vertex_credentials": null, + "vertex_ai_credentials": "alias-credentials", + "vertex_project": " ", + "vertex_ai_project": "alias-project", + "vertex_location": null, + "vertex_ai_location": "alias-location" + })); + assert_eq!( + config.credentials.as_ref().unwrap().value().expose(), + "alias-credentials" + ); + assert_eq!(config.project_id(), Some("alias-project")); + assert_eq!(config.location(), Some("alias-location")); + } + + #[rstest::rstest] + fn typed_config_preserves_source_and_empty_value_fallback() { + let configured = VertexConfig::new( + Some(Sourced::new( + SecretValue::new("inline-json"), + InputSource::Request, + )), + Some("project".into()), + Some("location".into()), + ); + assert!(matches!( + credential_source(&configured, &|_| Some("environment-json".into())), + CredentialSource::Inline(value) if value.expose() == "inline-json" + )); + let empty = VertexConfig::new( + Some(Sourced::new(SecretValue::new(" "), InputSource::Request)), + Some(" ".into()), + Some(" ".into()), + ); + assert!(matches!( + credential_source(&empty, &|_| None), + CredentialSource::Adc + )); + assert_eq!( + get_vertex_ai_project(&empty, &|_| Some("env-project".into())).as_deref(), + Some("env-project") + ); + assert_eq!( + get_vertex_ai_location(&empty, &|_| Some("env-location".into())).as_deref(), + Some("env-location") + ); + } + + #[rstest::rstest] + fn credential_discovery_prefers_input_then_environment_then_adc() { + let params = json!({"vertex_credentials":"input-json"}); + let sources = BTreeMap::from([("vertex_credentials".to_string(), InputSource::Request)]); + let configured = + VertexConfig::from_sourced_optional_params(params.as_object().unwrap(), &sources) + .unwrap(); + assert!( + matches!(credential_source(&configured, &|_| Some("environment-value".into())), CredentialSource::Inline(value) if value.expose() == "input-json") + ); + let empty = VertexConfig::default(); + assert!( + matches!(credential_source(&empty, &|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into())), CredentialSource::Trusted(value) if value.expose() == "environment-json") + ); + assert!( + matches!(credential_source(&empty, &|name| (name == GOOGLE_APPLICATION_CREDENTIALS_ENV).then(|| "adc.json".into())), CredentialSource::ApplicationCredentials(path) if path == "adc.json") + ); + assert!(matches!( + credential_source(&empty, &|_| None), + CredentialSource::Adc + )); + } + + #[rstest::rstest] + fn params_become_a_deployment_sourced_config() { + let params = VertexParams { + vertex_credentials: Some("deployment-json".into()), + vertex_project: Some("project-1".into()), + vertex_ai_location: Some("europe-west4".into()), + ..VertexParams::default() + }; + let config = VertexConfig::from_params(¶ms); + assert_eq!(config.project_id(), Some("project-1")); + assert_eq!(config.location(), Some("europe-west4")); + assert!( + matches!(credential_source(&config, &|_| Some("environment-json".into())), CredentialSource::Trusted(value) if value.expose() == "deployment-json") + ); + assert!(matches!( + credential_source( + &VertexConfig::from_params(&VertexParams::default()), + &|_| None + ), + CredentialSource::Adc + )); + } +} diff --git a/litellm-rust/crates/auth-gcp/src/constants.rs b/litellm-rust/crates/auth-gcp/src/constants.rs new file mode 100644 index 00000000000..e76b8dfa8ff --- /dev/null +++ b/litellm-rust/crates/auth-gcp/src/constants.rs @@ -0,0 +1,36 @@ +use std::sync::LazyLock; + +use litellm_auth_types::VertexParams; + +pub const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; +pub const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token"; +pub const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS"; +pub const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY"; +pub const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY"; +pub const VERTEXAI_CREDENTIALS_ENV: &str = VertexParams::CREDENTIALS.env[0]; +pub const VERTEX_LOCATION_ENV: &str = VertexParams::LOCATION.env[1]; + +/// Every environment name the Vertex params and the token acquisition read, for hosts +/// that resolve secrets up front. +pub fn secret_names() -> &'static [&'static str] { + static NAMES: LazyLock> = LazyLock::new(|| { + [ + VERTEX_AI_API_KEY_ENV, + VERTEXAI_API_KEY_ENV, + GOOGLE_APPLICATION_CREDENTIALS_ENV, + ] + .into_iter() + .chain( + VertexParams::SPECS + .iter() + .flat_map(|spec| spec.env.iter().copied()), + ) + .fold(Vec::new(), |mut names, name| { + if !names.contains(&name) { + names.push(name); + } + names + }) + }); + &NAMES +} diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index 97bc2c482c3..2b389d45926 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -1,708 +1,16 @@ -use std::{collections::BTreeMap, future::Future, path::Path, pin::Pin, sync::Arc}; - -use gcp_auth::{CustomServiceAccount, TokenProvider}; -use litellm_auth_types::{ - CredentialPlacement, Error, InputSource, SecretValue, Sourced, http::apply_credential, -}; -use moka::future::Cache; -use serde_json::{Map, Value}; -use sha2::{Digest, Sha256}; - +mod auth; +mod config; +pub mod constants; #[cfg(feature = "google-sdk")] mod sdk; + +pub use auth::{ + VertexAuth, VertexAuthFuture, VertexEnvironment, VertexProviderLoader, VertexTokenSource, +}; +pub use config::{ + CredentialSource, VertexConfig, get_vertex_ai_location, get_vertex_ai_project, + get_vertex_ai_project_from_credentials, +}; +pub use constants::secret_names; #[cfg(feature = "google-sdk")] pub use sdk::GoogleCredentials; - -const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; -const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token"; -const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS"; -const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY"; -const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY"; -const VERTEXAI_CREDENTIALS_ENV: &str = "VERTEXAI_CREDENTIALS"; -const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT"; -const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION"; -const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION"; - -pub const SECRET_NAMES: &[&str] = &[ - VERTEX_AI_API_KEY_ENV, - VERTEXAI_API_KEY_ENV, - VERTEXAI_CREDENTIALS_ENV, - GOOGLE_APPLICATION_CREDENTIALS_ENV, - VERTEXAI_PROJECT_ENV, - VERTEXAI_LOCATION_ENV, - VERTEX_LOCATION_ENV, -]; - -#[derive(Clone, Debug, Default)] -pub struct VertexConfig { - credentials: Option>, - project_id: Option, - location: Option, -} - -impl VertexConfig { - pub fn new( - credentials: Option>, - project_id: Option, - location: Option, - ) -> Self { - Self { - credentials: credentials.filter(|value| !value.value().expose().trim().is_empty()), - project_id: project_id.filter(|value| !value.trim().is_empty()), - location: location.filter(|value| !value.trim().is_empty()), - } - } - - pub fn from_sourced_optional_params( - params: &Map, - sources: &BTreeMap, - ) -> Result { - Ok(Self::new( - optional_credentials( - params, - sources, - &["vertex_credentials", "vertex_ai_credentials"], - )?, - optional_string(params, &["vertex_project", "vertex_ai_project"])?, - optional_string(params, &["vertex_location", "vertex_ai_location"])?, - )) - } - - pub fn or_configured(self, project_id: Option<&str>, location: Option<&str>) -> Self { - let configured = - |value: Option<&str>| value.filter(|value| !value.is_empty()).map(str::to_string); - Self { - project_id: self.project_id.or_else(|| configured(project_id)), - location: self.location.or_else(|| configured(location)), - ..self - } - } - - pub fn project_id(&self) -> Option<&str> { - self.project_id.as_deref() - } - - pub fn location(&self) -> Option<&str> { - self.location.as_deref() - } -} - -pub struct VertexEnvironment { - pub headers: Vec<(String, String)>, - pub project_id: String, -} - -struct VertexAccessToken { - token: String, - project_id: String, -} - -pub fn get_vertex_ai_project( - config: &VertexConfig, - env_lookup: &dyn Fn(&str) -> Option, -) -> Option { - config - .project_id() - .map(str::to_string) - .or_else(|| non_empty_env(env_lookup, VERTEXAI_PROJECT_ENV)) -} - -pub fn get_vertex_ai_location( - config: &VertexConfig, - env_lookup: &dyn Fn(&str) -> Option, -) -> Option { - config - .location() - .map(str::to_string) - .or_else(|| non_empty_env(env_lookup, VERTEXAI_LOCATION_ENV)) - .or_else(|| non_empty_env(env_lookup, VERTEX_LOCATION_ENV)) -} - -#[derive(Clone)] -pub struct VertexAuth { - providers: Cache>, - loader: Arc, -} - -impl Default for VertexAuth { - fn default() -> Self { - Self::new(Arc::new(GcpProviderLoader)) - } -} - -impl VertexAuth { - pub fn new(loader: Arc) -> Self { - Self { - providers: Cache::builder().max_capacity(64).build(), - loader, - } - } - - pub async fn access_token( - &self, - config: &VertexConfig, - env_lookup: &(dyn Fn(&str) -> Option + Sync), - ) -> Result { - self.load_provider(config, env_lookup).await?.token().await - } - - pub async fn validate_environment( - &self, - headers: Vec<(String, String)>, - api_key: Option<&str>, - config: &VertexConfig, - env_lookup: &(dyn Fn(&str) -> Option + Sync), - ) -> Result { - let has_authorization = headers - .iter() - .any(|(name, _)| name.eq_ignore_ascii_case("Authorization")); - let static_token = api_key - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) - .or_else(|| non_empty_env(env_lookup, VERTEX_AI_API_KEY_ENV)) - .or_else(|| non_empty_env(env_lookup, VERTEXAI_API_KEY_ENV)); - let project_id = get_vertex_ai_project(config, env_lookup); - - if !has_authorization && static_token.is_none() { - let access = self.get_access_token(config, env_lookup).await?; - return Ok(VertexEnvironment { - headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?, - project_id: project_id.unwrap_or(access.project_id), - }); - } - - let project_id = match project_id { - Some(project_id) => project_id, - None => { - self.load_provider(config, env_lookup) - .await? - .project_id() - .await? - } - }; - let headers = if has_authorization { - headers - } else { - apply_credential( - headers, - static_token.as_deref().expect("static token was checked"), - CredentialPlacement::Bearer, - )? - }; - Ok(VertexEnvironment { - headers, - project_id, - }) - } - - async fn get_access_token( - &self, - config: &VertexConfig, - env_lookup: &(dyn Fn(&str) -> Option + Sync), - ) -> Result { - let provider = self.load_provider(config, env_lookup).await?; - let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?; - Ok(VertexAccessToken { token, project_id }) - } - - async fn load_provider( - &self, - config: &VertexConfig, - env_lookup: &(dyn Fn(&str) -> Option + Sync), - ) -> Result, Error> { - let source = credential_source(config, env_lookup); - let key = source.cache_key(); - self.providers - .try_get_with(key, self.loader.load(source)) - .await - .map_err(|error| (*error).clone()) - } -} - -pub trait VertexTokenSource: Send + Sync { - fn project_id(&self) -> VertexAuthFuture<'_, String>; - fn token(&self) -> VertexAuthFuture<'_, String>; -} - -pub trait VertexProviderLoader: Send + Sync { - fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc>; -} - -pub type VertexAuthFuture<'a, T> = Pin> + Send + 'a>>; - -struct GcpTokenSource(Arc); - -impl VertexTokenSource for GcpTokenSource { - fn project_id(&self) -> VertexAuthFuture<'_, String> { - Box::pin(async move { - self.0 - .project_id() - .await - .map(|project| project.to_string()) - .map_err(auth_acquisition_error) - }) - } - - fn token(&self) -> VertexAuthFuture<'_, String> { - Box::pin(async move { - self.0 - .token(&[CLOUD_PLATFORM_SCOPE]) - .await - .map(|token| token.as_str().to_string()) - .map_err(auth_acquisition_error) - }) - } -} - -struct GcpProviderLoader; - -impl VertexProviderLoader for GcpProviderLoader { - fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc> { - Box::pin(async move { - let provider: Arc = match source { - CredentialSource::Inline(configured) => Arc::new( - CustomServiceAccount::from_json(validate_request_credentials( - configured.expose(), - )?) - .map_err(auth_acquisition_error)?, - ), - CredentialSource::Trusted(configured) => { - let configured = configured.expose(); - let service_account = if Path::new(configured).is_file() { - CustomServiceAccount::from_file(configured) - } else { - CustomServiceAccount::from_json(configured) - } - .map_err(auth_acquisition_error)?; - Arc::new(service_account) - } - CredentialSource::ApplicationCredentials(path) => { - Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?) - } - CredentialSource::Adc => { - gcp_auth::provider().await.map_err(auth_acquisition_error)? - } - }; - Ok(Arc::new(GcpTokenSource(provider)) as Arc) - }) - } -} - -fn validate_request_credentials(configured: &str) -> Result<&str, Error> { - let token_uri = serde_json::from_str::(configured) - .ok() - .and_then(|credentials| { - credentials - .get("token_uri") - .and_then(Value::as_str) - .map(str::to_string) - }); - if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) { - return Err(Error::InvalidConfiguration("request-controlled Vertex credentials must use the canonical Google OAuth token endpoint".into())); - } - Ok(configured) -} - -#[derive(Clone, Debug)] -pub enum CredentialSource { - Inline(SecretValue), - Trusted(SecretValue), - ApplicationCredentials(String), - Adc, -} - -impl CredentialSource { - fn cache_key(&self) -> CredentialCacheKey { - match self { - Self::Inline(configured) => { - CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into()) - } - Self::Trusted(configured) => { - CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into()) - } - Self::ApplicationCredentials(path) => { - CredentialCacheKey::ApplicationCredentials(path.clone()) - } - Self::Adc => CredentialCacheKey::Adc, - } - } -} - -#[derive(Clone, Debug, Hash, PartialEq, Eq)] -enum CredentialCacheKey { - Inline([u8; 32]), - Trusted([u8; 32]), - ApplicationCredentials(String), - Adc, -} - -fn credential_source( - config: &VertexConfig, - env_lookup: &dyn Fn(&str) -> Option, -) -> CredentialSource { - if let Some(configured) = config.credentials.clone() { - return match configured.source() { - InputSource::Request => CredentialSource::Inline(configured.into_value()), - InputSource::Deployment | InputSource::Environment => { - CredentialSource::Trusted(configured.into_value()) - } - }; - } - if let Some(configured) = non_empty_env(env_lookup, VERTEXAI_CREDENTIALS_ENV) { - return CredentialSource::Trusted(SecretValue::new(configured)); - } - non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV) - .map(CredentialSource::ApplicationCredentials) - .unwrap_or(CredentialSource::Adc) -} - -fn optional_credentials( - params: &Map, - sources: &BTreeMap, - names: &[&str], -) -> Result>, Error> { - for name in names { - let source = source_for(sources, name); - match params.get(*name) { - None | Some(Value::Null) => continue, - Some(Value::String(value)) if value.trim().is_empty() => continue, - Some(Value::String(value)) => { - return Ok(Some(Sourced::new(SecretValue::new(value), source))); - } - Some(Value::Object(value)) if value.is_empty() => continue, - Some(Value::Object(value)) => { - return serde_json::to_string(value) - .map(SecretValue::new) - .map(|value| Sourced::new(value, source)) - .map(Some) - .map_err(|error| { - Error::InvalidConfiguration(litellm_auth_types::ErrorDetail::failed( - "credential serialization", - error, - )) - }); - } - Some(_) => { - return Err(Error::InvalidConfiguration( - litellm_auth_types::ErrorDetail::InvalidType { - field: names[0].into(), - expected: "a string or null", - }, - )); - } - } - } - Ok(None) -} - -fn source_for(sources: &BTreeMap, name: &str) -> InputSource { - sources.get(name).copied().unwrap_or_default() -} - -fn optional_string(params: &Map, names: &[&str]) -> Result, Error> { - for name in names { - match params.get(*name) { - None | Some(Value::Null) => continue, - Some(Value::String(value)) if value.trim().is_empty() => continue, - Some(Value::String(value)) => return Ok(Some(value.clone())), - Some(_) => { - return Err(Error::InvalidConfiguration( - litellm_auth_types::ErrorDetail::InvalidType { - field: names[0].into(), - expected: "a string or null", - }, - )); - } - } - } - Ok(None) -} - -fn non_empty_env(env_lookup: &dyn Fn(&str) -> Option, name: &str) -> Option { - env_lookup(name) - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) -} - -fn auth_acquisition_error(error: gcp_auth::Error) -> Error { - Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed( - "Vertex AI credentials", - error, - )) -} - -#[cfg(test)] -mod tests { - use std::collections::BTreeSet; - use std::sync::atomic::{AtomicUsize, Ordering}; - - use serde_json::json; - - use super::*; - - struct FakeProvider { - calls: Arc, - } - - impl VertexTokenSource for FakeProvider { - fn project_id(&self) -> VertexAuthFuture<'_, String> { - self.calls.fetch_add(1, Ordering::SeqCst); - Box::pin(async { Ok("adc-project".into()) }) - } - - fn token(&self) -> VertexAuthFuture<'_, String> { - self.calls.fetch_add(1, Ordering::SeqCst); - Box::pin(async { Ok("adc-token".into()) }) - } - } - - struct FakeLoader { - loads: Arc, - provider: Arc, - } - - impl VertexProviderLoader for FakeLoader { - fn load( - &self, - _source: CredentialSource, - ) -> VertexAuthFuture<'_, Arc> { - let loads = self.loads.clone(); - let provider = self.provider.clone(); - Box::pin(async move { - loads.fetch_add(1, Ordering::SeqCst); - Ok(provider) - }) - } - } - - fn config(value: Value) -> VertexConfig { - VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new()) - .unwrap() - } - - fn auth(calls: Arc, loads: Arc) -> VertexAuth { - let provider: Arc = Arc::new(FakeProvider { calls }); - VertexAuth::new(Arc::new(FakeLoader { loads, provider })) - } - - #[test] - fn config_is_typed_and_secrets_are_redacted() { - let config = config(json!({ - "vertex_credentials":{"private_key":"secret-key"}, - "vertex_project":"project-1", - "vertex_location":"europe-west4" - })); - assert_eq!(config.project_id(), Some("project-1")); - assert_eq!(config.location(), Some("europe-west4")); - assert!(!format!("{config:?}").contains("secret-key")); - assert!( - VertexConfig::from_sourced_optional_params( - json!({"vertex_credentials":true}).as_object().unwrap(), - &BTreeMap::new() - ) - .is_err() - ); - } - - #[tokio::test] - async fn secret_names_cover_environment_reads() { - let seen = Arc::new(std::sync::Mutex::new(BTreeSet::::new())); - let recorded = seen.clone(); - let env = |name: &str| { - recorded.lock().unwrap().insert(name.to_string()); - None - }; - let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))); - auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env) - .await - .unwrap(); - get_vertex_ai_location(&VertexConfig::default(), &env); - assert!( - seen.lock() - .unwrap() - .iter() - .all(|name| SECRET_NAMES.contains(&name.as_str())) - ); - } - - #[test] - fn empty_primary_values_fall_back_to_python_aliases() { - let config = config(json!({ - "vertex_credentials": null, - "vertex_ai_credentials": "alias-credentials", - "vertex_project": " ", - "vertex_ai_project": "alias-project", - "vertex_location": null, - "vertex_ai_location": "alias-location" - })); - assert_eq!( - config.credentials.as_ref().unwrap().value().expose(), - "alias-credentials" - ); - assert_eq!(config.project_id(), Some("alias-project")); - assert_eq!(config.location(), Some("alias-location")); - } - - #[test] - fn typed_config_preserves_source_and_empty_value_fallback() { - let configured = VertexConfig::new( - Some(Sourced::new( - SecretValue::new("inline-json"), - InputSource::Request, - )), - Some("project".into()), - Some("location".into()), - ); - assert!(matches!( - credential_source(&configured, &|_| Some("environment-json".into())), - CredentialSource::Inline(value) if value.expose() == "inline-json" - )); - let empty = VertexConfig::new( - Some(Sourced::new(SecretValue::new(" "), InputSource::Request)), - Some(" ".into()), - Some(" ".into()), - ); - assert!(matches!( - credential_source(&empty, &|_| None), - CredentialSource::Adc - )); - assert_eq!( - get_vertex_ai_project(&empty, &|_| Some("env-project".into())).as_deref(), - Some("env-project") - ); - assert_eq!( - get_vertex_ai_location(&empty, &|_| Some("env-location".into())).as_deref(), - Some("env-location") - ); - } - - #[test] - fn project_and_location_prefer_input_then_environment() { - let configured = - config(json!({"vertex_project":"input-project","vertex_location":"input-location"})); - let env = |name: &str| Some(format!("env-{name}")); - assert_eq!( - get_vertex_ai_project(&configured, &env).as_deref(), - Some("input-project") - ); - assert_eq!( - get_vertex_ai_location(&configured, &env).as_deref(), - Some("input-location") - ); - let empty = VertexConfig::default(); - assert_eq!( - get_vertex_ai_project(&empty, &|_| Some("env-project".into())).as_deref(), - Some("env-project") - ); - assert_eq!( - get_vertex_ai_location(&empty, &|name| (name == VERTEX_LOCATION_ENV) - .then(|| "fallback-location".into())) - .as_deref(), - Some("fallback-location") - ); - } - - #[test] - fn credential_discovery_prefers_input_then_environment_then_adc() { - let params = json!({"vertex_credentials":"input-json"}); - let sources = BTreeMap::from([("vertex_credentials".to_string(), InputSource::Request)]); - let configured = - VertexConfig::from_sourced_optional_params(params.as_object().unwrap(), &sources) - .unwrap(); - assert!( - matches!(credential_source(&configured, &|_| Some("environment-value".into())), CredentialSource::Inline(value) if value.expose() == "input-json") - ); - let empty = VertexConfig::default(); - assert!( - matches!(credential_source(&empty, &|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into())), CredentialSource::Trusted(value) if value.expose() == "environment-json") - ); - assert!( - matches!(credential_source(&empty, &|name| (name == GOOGLE_APPLICATION_CREDENTIALS_ENV).then(|| "adc.json".into())), CredentialSource::ApplicationCredentials(path) if path == "adc.json") - ); - assert!(matches!( - credential_source(&empty, &|_| None), - CredentialSource::Adc - )); - assert_ne!( - CredentialSource::Inline(SecretValue::new("same-value")).cache_key(), - CredentialSource::Trusted(SecretValue::new("same-value")).cache_key() - ); - } - - #[rstest::rstest] - #[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)] - #[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)] - #[case::missing_endpoint("{}", false)] - fn request_credentials_require_canonical_token_endpoint( - #[case] credentials: &str, - #[case] accepted: bool, - ) { - assert_eq!(validate_request_credentials(credentials).is_ok(), accepted); - } - - #[tokio::test] - async fn explicit_token_and_header_do_not_acquire_adc() { - let loads = Arc::new(AtomicUsize::new(0)); - let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone()); - let configured = config(json!({"vertex_project":"project-1"})); - let explicit = auth - .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None) - .await - .unwrap(); - assert_eq!(explicit.headers[0].1, "Bearer access-token"); - let existing = auth - .validate_environment( - vec![("authorization".into(), "Bearer existing".into())], - None, - &configured, - &|_| None, - ) - .await - .unwrap(); - assert_eq!(existing.headers[0].1, "Bearer existing"); - assert_eq!(loads.load(Ordering::SeqCst), 0); - } - - #[tokio::test] - async fn provider_is_reused_across_authentication_calls() { - let calls = Arc::new(AtomicUsize::new(0)); - let loads = Arc::new(AtomicUsize::new(0)); - let auth = auth(calls.clone(), loads.clone()); - for _ in 0..2 { - let environment = auth - .validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None) - .await - .unwrap(); - assert_eq!(environment.project_id, "adc-project"); - assert_eq!(environment.headers[0].1, "Bearer adc-token"); - } - assert_eq!(loads.load(Ordering::SeqCst), 1); - assert_eq!(calls.load(Ordering::SeqCst), 4); - } - - #[test] - fn configured_defaults_sit_between_call_params_and_the_environment() { - let env = |name: &str| Some(format!("env-{name}")); - let from_config = - VertexConfig::default().or_configured(Some("global-project"), Some("global-location")); - assert_eq!( - get_vertex_ai_project(&from_config, &env).as_deref(), - Some("global-project") - ); - assert_eq!( - get_vertex_ai_location(&from_config, &env).as_deref(), - Some("global-location") - ); - let from_call = - config(json!({"vertex_project":"call-project","vertex_location":"call-location"})) - .or_configured(Some("global-project"), Some("global-location")); - assert_eq!(from_call.project_id(), Some("call-project")); - assert_eq!(from_call.location(), Some("call-location")); - let empty_global = VertexConfig::default().or_configured(Some(""), None); - assert_eq!( - get_vertex_ai_project(&empty_global, &env).as_deref(), - Some("env-VERTEXAI_PROJECT") - ); - } -} diff --git a/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs b/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs new file mode 100644 index 00000000000..5991749ab8a --- /dev/null +++ b/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs @@ -0,0 +1,120 @@ +use std::{ + collections::{BTreeMap, BTreeSet}, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, +}; + +use litellm_auth_gcp::{ + CredentialSource, VertexAuth, VertexAuthFuture, VertexConfig, VertexProviderLoader, + VertexTokenSource, get_vertex_ai_location, secret_names, +}; +use rstest::rstest; +use serde_json::{Value, json}; + +struct FakeProvider { + calls: Arc, +} + +impl VertexTokenSource for FakeProvider { + fn project_id(&self) -> VertexAuthFuture<'_, String> { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Ok("adc-project".into()) }) + } + + fn token(&self) -> VertexAuthFuture<'_, String> { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Ok("adc-token".into()) }) + } +} + +struct FakeLoader { + loads: Arc, + provider: Arc, +} + +impl VertexProviderLoader for FakeLoader { + fn load(&self, _source: CredentialSource) -> VertexAuthFuture<'_, Arc> { + let loads = self.loads.clone(); + let provider = self.provider.clone(); + Box::pin(async move { + loads.fetch_add(1, Ordering::SeqCst); + Ok(provider) + }) + } +} + +fn config(value: Value) -> VertexConfig { + VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new()) + .unwrap() +} + +fn auth(calls: Arc, loads: Arc) -> VertexAuth { + let provider: Arc = Arc::new(FakeProvider { calls }); + VertexAuth::new(Arc::new(FakeLoader { loads, provider })) +} + +#[rstest] +#[tokio::test] +async fn secret_names_cover_environment_reads() { + let seen = Arc::new(std::sync::Mutex::new(BTreeSet::::new())); + let recorded = seen.clone(); + let env = |name: &str| { + recorded.lock().unwrap().insert(name.to_string()); + None + }; + let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))); + auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env) + .await + .unwrap(); + get_vertex_ai_location(&VertexConfig::default(), &env); + assert!( + seen.lock() + .unwrap() + .iter() + .all(|name| secret_names().contains(&name.as_str())) + ); +} + +#[rstest] +#[tokio::test] +async fn explicit_token_and_header_do_not_acquire_adc() { + let loads = Arc::new(AtomicUsize::new(0)); + let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone()); + let configured = config(json!({"vertex_project":"project-1"})); + let explicit = auth + .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None) + .await + .unwrap(); + assert_eq!(explicit.headers[0].1, "Bearer access-token"); + let existing = auth + .validate_environment( + vec![("authorization".into(), "Bearer existing".into())], + None, + &configured, + &|_| None, + ) + .await + .unwrap(); + assert_eq!(existing.headers[0].1, "Bearer existing"); + assert_eq!(loads.load(Ordering::SeqCst), 0); +} + +#[rstest] +#[tokio::test] +async fn provider_is_reused_across_authentication_calls() { + let calls = Arc::new(AtomicUsize::new(0)); + let loads = Arc::new(AtomicUsize::new(0)); + let auth = auth(calls.clone(), loads.clone()); + for _ in 0..2 { + let environment = auth + .validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None) + .await + .unwrap(); + assert_eq!(environment.project_id, "adc-project"); + assert_eq!(environment.headers[0].1, "Bearer adc-token"); + } + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!(calls.load(Ordering::SeqCst), 4); +} diff --git a/litellm-rust/crates/auth-gcp/tests/vertex_config.rs b/litellm-rust/crates/auth-gcp/tests/vertex_config.rs new file mode 100644 index 00000000000..3a589ab5266 --- /dev/null +++ b/litellm-rust/crates/auth-gcp/tests/vertex_config.rs @@ -0,0 +1,95 @@ +use std::collections::BTreeMap; + +use litellm_auth_gcp::{ + VertexConfig, + constants::{GOOGLE_APPLICATION_CREDENTIALS_ENV, VERTEX_LOCATION_ENV}, + get_vertex_ai_location, get_vertex_ai_project, get_vertex_ai_project_from_credentials, +}; +use litellm_auth_types::VertexParams; +use rstest::rstest; +use serde_json::{Value, json}; + +fn config(value: Value) -> VertexConfig { + VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new()) + .unwrap() +} + +#[rstest] +fn config_is_typed_and_secrets_are_redacted() { + let config = config(json!({ + "vertex_credentials":{"private_key":"secret-key"}, + "vertex_project":"project-1", + "vertex_location":"europe-west4" + })); + assert_eq!(config.project_id(), Some("project-1")); + assert_eq!(config.location(), Some("europe-west4")); + assert!(!format!("{config:?}").contains("secret-key")); + assert!( + VertexConfig::from_sourced_optional_params( + json!({"vertex_credentials":true}).as_object().unwrap(), + &BTreeMap::new() + ) + .is_err() + ); +} + +#[rstest] +fn project_and_location_prefer_input_then_environment() { + let configured = + config(json!({"vertex_project":"input-project","vertex_location":"input-location"})); + let env = |name: &str| Some(format!("env-{name}")); + assert_eq!( + get_vertex_ai_project(&configured, &env).as_deref(), + Some("input-project") + ); + assert_eq!( + get_vertex_ai_location(&configured, &env).as_deref(), + Some("input-location") + ); + let empty = VertexConfig::default(); + assert_eq!( + get_vertex_ai_project(&empty, &|_| Some("env-project".into())).as_deref(), + Some("env-project") + ); + assert_eq!( + get_vertex_ai_location(&empty, &|name| (name == VERTEX_LOCATION_ENV) + .then(|| "fallback-location".into())) + .as_deref(), + Some("fallback-location") + ); +} + +#[rstest] +fn project_is_read_from_the_credentials_that_would_authenticate() { + let key = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap(); + let path = key.path().to_str().unwrap().to_string(); + let inline = VertexConfig::from_params(&VertexParams { + vertex_credentials: Some(r#"{"project_id":"inline-project"}"#.into()), + ..VertexParams::default() + }); + assert_eq!( + get_vertex_ai_project_from_credentials(&inline, &|_| None).as_deref(), + Some("inline-project") + ); + let from_file = VertexConfig::from_params(&VertexParams { + vertex_credentials: Some(path.clone()), + ..VertexParams::default() + }); + assert_eq!( + get_vertex_ai_project_from_credentials(&from_file, &|_| None).as_deref(), + Some("file-project") + ); + let empty = VertexConfig::default(); + assert_eq!( + get_vertex_ai_project_from_credentials(&empty, &|name| (name + == GOOGLE_APPLICATION_CREDENTIALS_ENV) + .then(|| path.clone())) + .as_deref(), + Some("file-project") + ); + assert_eq!( + get_vertex_ai_project_from_credentials(&empty, &|_| None), + None + ); +} diff --git a/litellm-rust/crates/auth-types/src/lib.rs b/litellm-rust/crates/auth-types/src/lib.rs index 38b13988448..c875bda7979 100644 --- a/litellm-rust/crates/auth-types/src/lib.rs +++ b/litellm-rust/crates/auth-types/src/lib.rs @@ -53,7 +53,7 @@ pub use credential::{ }; pub use error::{Error, ErrorDetail, ErrorSource}; pub use http::CredentialPlacement; -pub use params::{AwsParams, ParamSpec}; +pub use params::{AwsParams, ParamSpec, VertexParams}; pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; pub use secret::SecretValue; pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; diff --git a/litellm-rust/crates/auth-types/src/params/mod.rs b/litellm-rust/crates/auth-types/src/params/mod.rs index 093cdb42d64..88bc64cc2fc 100644 --- a/litellm-rust/crates/auth-types/src/params/mod.rs +++ b/litellm-rust/crates/auth-types/src/params/mod.rs @@ -1,6 +1,8 @@ mod aws; +mod vertex; pub use aws::AwsParams; +pub use vertex::VertexParams; /// One connection param of Python's `CredentialLiteLLMParams` and every place Python reads /// it from, in order: the wire names on a call or a deployment, the `litellm.` module diff --git a/litellm-rust/crates/auth-types/src/params/vertex.rs b/litellm-rust/crates/auth-types/src/params/vertex.rs new file mode 100644 index 00000000000..06ec3008dfb --- /dev/null +++ b/litellm-rust/crates/auth-types/src/params/vertex.rs @@ -0,0 +1,118 @@ +use serde::{Deserialize, Deserializer, Serialize, de::Unexpected}; +use serde_json::Value; +use veil::Redact; + +use super::ParamSpec; + +/// The Vertex AI fields of Python's `GenericLiteLLMParams`, both the current names and the +/// `vertex_ai_*` spellings `VertexBase.safe_get_vertex_ai_*` still read. +#[derive(Redact, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct VertexParams { + #[serde(default, deserialize_with = "optional_credential_text")] + #[redact(with = "[REDACTED]")] + pub vertex_credentials: Option, + #[serde(default)] + pub vertex_project: Option, + #[serde(default)] + pub vertex_location: Option, + #[serde(default, deserialize_with = "optional_credential_text")] + #[redact(with = "[REDACTED]")] + pub vertex_ai_credentials: Option, + #[serde(default)] + pub vertex_ai_project: Option, + #[serde(default)] + pub vertex_ai_location: Option, +} + +impl VertexParams { + pub const CREDENTIALS: ParamSpec = ParamSpec { + setting: "credentials", + wire: &["vertex_credentials", "vertex_ai_credentials"], + module_global: None, + env: &["VERTEXAI_CREDENTIALS"], + }; + pub const PROJECT: ParamSpec = ParamSpec { + setting: "project", + wire: &["vertex_project", "vertex_ai_project"], + module_global: Some("vertex_project"), + env: &["VERTEXAI_PROJECT"], + }; + pub const LOCATION: ParamSpec = ParamSpec { + setting: "location", + wire: &["vertex_location", "vertex_ai_location"], + module_global: Some("vertex_location"), + env: &["VERTEXAI_LOCATION", "VERTEX_LOCATION"], + }; + + /// Every Vertex AI param Python reads, with both spellings, the module global + /// `VertexBase.get_vertex_ai_*` consults, and the environment names it falls back to. + pub const SPECS: [ParamSpec; 3] = [Self::CREDENTIALS, Self::PROJECT, Self::LOCATION]; + + /// The wire names a host projects out of a caller's kwargs, derived from [`Self::SPECS`]. + pub fn fields() -> impl Iterator { + Self::SPECS + .iter() + .flat_map(|spec| spec.wire.iter().copied()) + } + + /// The value under a wire name, read through serde so the names can never drift from + /// the struct. + pub fn get(&self, wire: &str) -> Option { + serde_json::to_value(self) + .ok()? + .get(wire)? + .as_str() + .map(str::to_string) + } + + /// The spec's value from these params, then the environment. + pub fn resolve( + &self, + spec: &ParamSpec, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Option { + spec.resolve(&|name| self.get(name), env_lookup) + } + + pub fn credentials(&self) -> Option { + self.resolve(&Self::CREDENTIALS, &|_| None) + } + + pub fn project(&self) -> Option { + self.resolve(&Self::PROJECT, &|_| None) + } + + pub fn location(&self) -> Option { + self.resolve(&Self::LOCATION, &|_| None) + } +} + +/// A credential is a service account JSON text or a path to one; a YAML or kwargs mapping is +/// accepted as the JSON text it spells, the way Python passes a dict through. +fn optional_credential_text<'de, D: Deserializer<'de>>( + deserializer: D, +) -> Result, D::Error> { + match Option::::deserialize(deserializer)? { + None | Some(Value::Null) => Ok(None), + Some(Value::String(text)) => Ok(Some(text)), + Some(Value::Object(object)) if object.is_empty() => Ok(None), + Some(Value::Object(object)) => serde_json::to_string(&object) + .map(Some) + .map_err(serde::de::Error::custom), + Some(other) => Err(serde::de::Error::invalid_type( + Unexpected::Other(value_kind(&other)), + &"a string or an object", + )), + } +} + +fn value_kind(value: &Value) -> &'static str { + match value { + Value::Null => "null", + Value::Bool(_) => "a boolean", + Value::Number(_) => "a number", + Value::String(_) => "a string", + Value::Array(_) => "a list", + Value::Object(_) => "an object", + } +} diff --git a/litellm-rust/crates/auth-types/tests/vertex_params.rs b/litellm-rust/crates/auth-types/tests/vertex_params.rs new file mode 100644 index 00000000000..cf0b7f824d1 --- /dev/null +++ b/litellm-rust/crates/auth-types/tests/vertex_params.rs @@ -0,0 +1,126 @@ +use std::collections::BTreeSet; + +use litellm_auth_types::VertexParams; +use rstest::rstest; +use serde_json::{Map, Value, json}; + +#[rstest] +#[case::current_spelling( + VertexParams { vertex_project: Some("p".into()), vertex_ai_project: Some("old".into()), ..VertexParams::default() }, + Some("p") +)] +#[case::legacy_spelling( + VertexParams { vertex_ai_project: Some("old".into()), ..VertexParams::default() }, + Some("old") +)] +#[case::blank_falls_through( + VertexParams { vertex_project: Some(" ".into()), vertex_ai_project: Some("old".into()), ..VertexParams::default() }, + Some("old") +)] +#[case::unset(VertexParams::default(), None)] +fn params_prefer_the_current_spelling_over_the_legacy_one( + #[case] params: VertexParams, + #[case] project: Option<&str>, +) { + assert_eq!(params.project().as_deref(), project); + let mirrored = VertexParams { + vertex_credentials: params.vertex_project.clone(), + vertex_ai_credentials: params.vertex_ai_project.clone(), + vertex_location: params.vertex_project.clone(), + vertex_ai_location: params.vertex_ai_project.clone(), + ..VertexParams::default() + }; + assert_eq!(mirrored.credentials().as_deref(), project); + assert_eq!(mirrored.location().as_deref(), project); +} + +#[rstest] +#[case::text(json!({"vertex_credentials": "/path/to/key.json"}), Some("/path/to/key.json"))] +#[case::object(json!({"vertex_credentials": {"type": "service_account"}}), Some(r#"{"type":"service_account"}"#))] +#[case::null(json!({"vertex_credentials": null}), None)] +#[case::empty_object_falls_through_to_the_legacy_spelling( + json!({"vertex_credentials": {}, "vertex_ai_credentials": "/legacy/key.json"}), + Some("/legacy/key.json") +)] +#[case::absent(json!({}), None)] +fn credentials_deserialize_from_text_or_an_object( + #[case] params: Value, + #[case] credentials: Option<&str>, +) { + let params: VertexParams = serde_json::from_value(params).unwrap(); + assert_eq!(params.credentials().as_deref(), credentials); +} + +#[rstest] +#[case::number(json!({"vertex_credentials": 7}))] +#[case::list(json!({"vertex_ai_credentials": ["a"]}))] +#[case::project(json!({"vertex_project": ["p"]}))] +fn a_param_of_the_wrong_type_is_rejected(#[case] params: Value) { + assert!(serde_json::from_value::(params).is_err()); +} + +#[rstest] +fn fields_name_every_param_once() { + let filled: Value = VertexParams::fields() + .map(|name| (name.to_string(), Value::from(name))) + .collect::>() + .into(); + let params: VertexParams = serde_json::from_value(filled).unwrap(); + let expected = VertexParams { + vertex_credentials: Some("vertex_credentials".into()), + vertex_project: Some("vertex_project".into()), + vertex_location: Some("vertex_location".into()), + vertex_ai_credentials: Some("vertex_ai_credentials".into()), + vertex_ai_project: Some("vertex_ai_project".into()), + vertex_ai_location: Some("vertex_ai_location".into()), + }; + assert_eq!(params, expected); + assert_eq!( + VertexParams::fields().collect::>().len(), + VertexParams::fields().count() + ); +} + +#[rstest] +fn debug_output_does_not_expose_credentials() { + let params = VertexParams { + vertex_credentials: Some("current-secret".into()), + vertex_ai_credentials: Some("legacy-secret".into()), + vertex_project: Some("project-1".into()), + ..VertexParams::default() + }; + let debug = format!("{params:?}"); + + assert!(!debug.contains("current-secret")); + assert!(!debug.contains("legacy-secret")); + assert!(debug.contains("project-1")); +} + +#[rstest] +#[case::current_spelling(Some("p"), Some("legacy"), &[("VERTEXAI_PROJECT", "env")], Some("p"))] +#[case::legacy_spelling(None, Some("legacy"), &[("VERTEXAI_PROJECT", "env")], Some("legacy"))] +#[case::blank_params_fall_to_the_environment(Some(" "), None, &[("VERTEXAI_PROJECT", "env")], Some("env"))] +#[case::nothing(None, None, &[], None)] +fn a_spec_resolves_from_the_params_then_the_environment( + #[case] vertex_project: Option<&str>, + #[case] vertex_ai_project: Option<&str>, + #[case] environment: &[(&str, &str)], + #[case] expected: Option<&str>, +) { + let params = VertexParams { + vertex_project: vertex_project.map(str::to_string), + vertex_ai_project: vertex_ai_project.map(str::to_string), + ..VertexParams::default() + }; + let env = |name: &str| { + environment + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + }; + + assert_eq!( + params.resolve(&VertexParams::PROJECT, &env).as_deref(), + expected + ); +} diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/vertex_ai.rs b/litellm-rust/crates/inference-ocr/tests/ocr/vertex_ai.rs index 47ef3847d1f..81b52d076a5 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/vertex_ai.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/vertex_ai.rs @@ -1,6 +1,5 @@ use litellm_auth::{InputSource, Sourced}; use litellm_inference_ocr::arguments::is_supported_request; -use litellm_llms::base_llm::ocr::settings::OcrSettings; use rstest::rstest; use super::*; @@ -42,30 +41,6 @@ async fn mistral_is_served_at_the_resolved_project_and_location() { ); } -#[tokio::test] -async fn configured_project_and_location_apply_when_the_call_sets_neither() { - let upstream = upstream([pages_response()]).await; - let route = ocr_route_with(OcrSettings { - vertex_project: Some("configured-project".into()), - vertex_location: Some("europe-west4".into()), - ..OcrSettings::default() - }); - - route - .execute( - ocr_request("vertex_ai/mistral-ocr-maas", &upstream.uri(), json!({})), - &(), - None, - ) - .await - .unwrap(); - - assert_eq!( - only_request(&upstream).await.url.path(), - "/v1/projects/configured-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); -} - #[tokio::test] async fn a_supplied_authorization_is_forwarded_without_a_static_token() { let upstream = upstream([pages_response()]).await; diff --git a/litellm-rust/crates/inference-ocr/tests/resources.rs b/litellm-rust/crates/inference-ocr/tests/resources.rs index 8a5090bf54c..d618647e89b 100644 --- a/litellm-rust/crates/inference-ocr/tests/resources.rs +++ b/litellm-rust/crates/inference-ocr/tests/resources.rs @@ -105,10 +105,7 @@ async fn auth_survives_per_call_clients_without_freezing_settings_or_secrets( owner, &http, Default::default(), - OcrSettings { - vertex_location: Some(location.into()), - ..OcrSettings::default() - }, + OcrSettings::default(), Arc::new(RecordingSecrets::new([("VERTEXAI_CREDENTIALS", identity)])), ); let request = decode_request(OcrWireRequest { @@ -118,7 +115,7 @@ async fn auth_survives_per_call_clients_without_freezing_settings_or_secrets( api_base: Some(upstream.uri()), custom_llm_provider: None, extra_headers: None, - optional_params: Default::default(), + optional_params: json!({"vertex_location": location}).as_object().cloned().unwrap(), input_sources: Default::default(), timeout_seconds: Some(5.0), }).unwrap(); diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs b/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs index 87461cd36aa..af4037ddc0b 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs @@ -9,8 +9,6 @@ pub struct OcrSettings { pub poll_timeout: Duration, pub document_intelligence_api_version: String, pub document_intelligence_dpi: i64, - pub vertex_project: Option, - pub vertex_location: Option, pub enable_azure_ad_token_refresh: bool, } @@ -22,8 +20,6 @@ impl Default for OcrSettings { poll_timeout: Duration::from_secs(120), document_intelligence_api_version: "2024-11-30".into(), document_intelligence_dpi: 96, - vertex_project: None, - vertex_location: None, enable_azure_ad_token_refresh: false, } } diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs index 93f8a25605e..810462c57f8 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs @@ -7,15 +7,10 @@ use crate::base_llm::ocr::{ }; pub(super) fn vertex_config(request: &PreparedOcrRequest) -> Result { - let settings = &request.connection.settings; Ok(VertexConfig::from_sourced_optional_params( &request.optional_params, &request.input_sources, - )? - .or_configured( - settings.vertex_project.as_deref(), - settings.vertex_location.as_deref(), - )) + )?) } pub(super) fn validate_destination(connection: &OcrConnection) -> Result<(), Error> { diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index 4656b2534b6..bfd7b530ffa 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -35,7 +35,7 @@ impl BaseOcrConfig for VertexAiOcrConfig { } fn secret_names(&self) -> Vec<&'static str> { - litellm_auth_gcp::SECRET_NAMES.to_vec() + litellm_auth_gcp::secret_names().to_vec() } fn map_ocr_params( diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 99afc772b99..f0fea5adb06 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -68,7 +68,7 @@ fn missing_module(py: Python<'_>, error: &PyErr, expected: &str) -> PyResult>()? - .is_some_and(|name| name == expected)) + .is_some_and(|name| name == expected || name.starts_with(&format!("{expected}.")))) } #[cfg(test)] diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index ca683600d71..5ae72acd104 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -13,7 +13,7 @@ use pyo3::{ use super::{ errors::to_pyerr as ocr_error_to_pyerr, - project::{OcrHostHandles, project_request}, + project::{OcrHostHandles, SETTINGS_ERROR_MARKER, project_request}, }; use crate::marshal::public_response; @@ -68,7 +68,7 @@ impl OcrPythonHost { } fn map_failure(&self, py: Python<'_>, error: PyErr) -> PyErr { - if !error.is_instance_of::(py) { + if !error.is_instance_of::(py) || is_settings_error(py, &error) { return error; } let provider = match &self.data { @@ -87,6 +87,15 @@ impl OcrPythonHost { } } +fn is_settings_error(py: Python<'_>, error: &PyErr) -> bool { + error + .value(py) + .getattr_opt(SETTINGS_ERROR_MARKER) + .ok() + .flatten() + .is_some() +} + impl PythonBinding for OcrPythonHost { type Protocol = Ocr; type Failure = PyErr; diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index 854cdb25f7b..fab9ecc13c1 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -19,10 +19,6 @@ use crate::{ python_settings::{PythonSettings, Snapshot}, }; -const VERTEX_PROJECT: FieldSpec> = - FieldSpec::new("vertex_project", |field| field.falsy_optional_string()); -const VERTEX_LOCATION: FieldSpec> = - FieldSpec::new("vertex_location", |field| field.falsy_optional_string()); const ENABLE_AZURE_AD_TOKEN_REFRESH: FieldSpec = FieldSpec::new("enable_azure_ad_token_refresh", |field| { Ok(field.exact_true()) @@ -60,8 +56,6 @@ fn ocr_settings(py: Python<'_>) -> PyResult { fn project_provider_defaults(snapshot: &Snapshot<'_>) -> PyResult { Ok(OcrSettings { - vertex_project: snapshot.read(&VERTEX_PROJECT)?, - vertex_location: snapshot.read(&VERTEX_LOCATION)?, enable_azure_ad_token_refresh: snapshot.read(&ENABLE_AZURE_AD_TOKEN_REFRESH)?, ..OcrSettings::from_environment(&ProcessEnvironment) }) @@ -108,31 +102,29 @@ mod tests { use crate::python_settings::PythonSettings; #[rstest::rstest] - fn provider_defaults_distinguish_falsey_values_and_exact_true() { + fn provider_defaults_read_exact_true_only() { Python::initialize(); Python::attach(|py| { - let value = py.eval(c"__import__('types').SimpleNamespace(vertex_project=[], vertex_location=0, enable_azure_ad_token_refresh=1)", None, None).unwrap(); + let value = py + .eval( + c"__import__('types').SimpleNamespace(enable_azure_ad_token_refresh=1)", + None, + None, + ) + .unwrap(); let snapshot = PythonSettings::ProviderDefaults.snapshot(value.clone()); - let projected = super::project_provider_defaults(&snapshot).unwrap(); - assert_eq!(projected.vertex_project, None); - assert_eq!(projected.vertex_location, None); - assert!(!projected.enable_azure_ad_token_refresh); - value.setattr("vertex_project", "project").unwrap(); - value.setattr("vertex_location", "region").unwrap(); + assert!( + !super::project_provider_defaults(&snapshot) + .unwrap() + .enable_azure_ad_token_refresh + ); value .setattr("enable_azure_ad_token_refresh", true) .unwrap(); - let next = super::project_provider_defaults(&snapshot).unwrap(); - assert_eq!(next.vertex_project.as_deref(), Some("project")); - assert_eq!(next.vertex_location.as_deref(), Some("region")); - assert!(next.enable_azure_ad_token_refresh); - value.setattr("vertex_project", 1).unwrap(); - let error = super::project_provider_defaults(&snapshot).err().unwrap(); - assert!(error.is_instance_of::(py)); assert!( - error - .to_string() - .contains("provider_defaults.vertex_project") + super::project_provider_defaults(&snapshot) + .unwrap() + .enable_azure_ad_token_refresh ); }); } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index 7ae5a1f3ce2..93f04966f08 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -1,4 +1,5 @@ use litellm_auth::SecretValue; +use litellm_auth::VertexParams; use litellm_host_python::{from_py, present}; use litellm_inference_ocr::{ types::{LiteLLMOcrRequest, OcrDocumentInput}, @@ -10,8 +11,10 @@ use serde_json::{Map, Value}; use super::{document::FileDocumentInput, errors::to_pyerr as ocr_error_to_pyerr}; use crate::{ + coercion::FieldSpec, credentials::{self, CallerTokenProvider}, marshal::{project_optional_fields, python_timeout_seconds, request_input_sources}, + python_settings::{PythonSettings, Snapshot}, }; /// What the host keeps after projection: the caller's token callable that answers the @@ -123,6 +126,16 @@ pub(super) fn project_request( let names = specs.iter().map(|spec| spec.name).collect::>(); let optional_params = project_optional_fields(names.iter().copied(), |name| present(kwargs, bound, name))?; + let globals = module_globals_to_read(&optional_params, &names); + let defaults = if globals.is_empty() { + None + } else { + PythonSettings::ProviderDefaults.read_or_unset(bound.py())? + }; + let optional_params = match defaults { + Some(defaults) => fold_module_globals(bound.py(), optional_params, &globals, &defaults)?, + None => optional_params, + }; let input_sources = request_input_sources( kwargs, names @@ -156,6 +169,59 @@ pub(super) fn project_request( )) } +/// Python's `VertexBase.get_vertex_ai_project` and `get_vertex_ai_location` fall back to a +/// `litellm.` module global when a call names neither spelling: the globals of the +/// specs this route consumes whose wire names the call left out. +fn module_globals_to_read( + optional_params: &Map, + consumed: &[&str], +) -> Vec<&'static str> { + VertexParams::SPECS + .iter() + .filter(|spec| spec.wire.iter().all(|name| consumed.contains(name))) + .filter(|spec| { + !spec + .wire + .iter() + .any(|name| optional_params.contains_key(*name)) + }) + .filter_map(|spec| spec.module_global) + .collect() +} + +/// Folds each module global in under the name its spec declares, skipping falsy values the +/// way Python's `or` chain does. +fn fold_module_globals( + py: Python<'_>, + optional_params: Map, + globals: &[&'static str], + defaults: &Snapshot<'_>, +) -> PyResult> { + let read = globals + .iter() + .map(|name| { + let spec = FieldSpec::new(name, |field| field.falsy_optional_string()); + Ok(defaults + .read(&spec) + .map_err(|error| settings_error(py, error.into()))? + .map(|value| (name.to_string(), Value::String(value)))) + }) + .collect::>>()?; + Ok(optional_params + .into_iter() + .chain(read.into_iter().flatten()) + .collect()) +} + +/// A misconfigured litellm setting is the operator's error, not the request's or the +/// provider's, so the host raises it as is instead of mapping it onto a provider failure. +pub(super) const SETTINGS_ERROR_MARKER: &str = "native_settings_error"; + +fn settings_error(py: Python<'_>, error: PyErr) -> PyErr { + error.value(py).setattr(SETTINGS_ERROR_MARKER, true).ok(); + error +} + #[cfg(test)] mod tests { use litellm_llms_types::formats::ocr::OcrDocument; @@ -163,6 +229,105 @@ mod tests { use super::*; + #[rstest::rstest] + #[case::global_fills_a_missing_project( + serde_json::json!({"vertex_location": "us-east5"}), + "vertex_project='from-global', vertex_location='ignored'", + serde_json::json!({"vertex_location": "us-east5", "vertex_project": "from-global"}), + )] + #[case::either_spelling_on_the_call_wins( + serde_json::json!({"vertex_ai_project": "from-call"}), + "vertex_project='from-global', vertex_location=None", + serde_json::json!({"vertex_ai_project": "from-call"}), + )] + #[case::falsy_globals_are_absent( + serde_json::json!({}), + "vertex_project=[], vertex_location=''", + serde_json::json!({}), + )] + fn module_globals_fill_the_vertex_params_the_call_leaves_out( + #[case] params: Value, + #[case] defaults: &str, + #[case] expected: Value, + ) { + Python::initialize(); + Python::attach(|py| { + let namespace = py + .eval( + &std::ffi::CString::new(format!( + "__import__('types').SimpleNamespace({defaults})" + )) + .unwrap(), + None, + None, + ) + .unwrap(); + let snapshot = PythonSettings::ProviderDefaults.snapshot(namespace); + let consumed = VertexParams::fields().collect::>(); + let params = params.as_object().cloned().unwrap(); + let globals = module_globals_to_read(¶ms, &consumed); + + let folded = fold_module_globals(py, params, &globals, &snapshot).unwrap(); + + assert_eq!(Value::Object(folded), expected); + }); + } + + #[rstest::rstest] + fn a_bad_global_is_a_marked_settings_error() { + Python::initialize(); + Python::attach(|py| { + let namespace = py + .eval( + c"__import__('types').SimpleNamespace(vertex_project=1)", + None, + None, + ) + .unwrap(); + let snapshot = PythonSettings::ProviderDefaults.snapshot(namespace); + + let error = + fold_module_globals(py, Map::new(), &["vertex_project"], &snapshot).unwrap_err(); + + assert!(error.is_instance_of::(py)); + assert!( + error + .to_string() + .contains("provider_defaults.vertex_project") + ); + assert!( + error + .value(py) + .getattr_opt(SETTINGS_ERROR_MARKER) + .unwrap() + .is_some() + ); + }); + } + + #[rstest::rstest] + #[case::not_consumed(&["pages"], serde_json::json!({}), &[])] + #[case::consumed_and_absent( + &["vertex_project", "vertex_ai_project", "vertex_location", "vertex_ai_location"], + serde_json::json!({}), + &["vertex_project", "vertex_location"] + )] + #[case::consumed_and_present_in_either_spelling( + &["vertex_project", "vertex_ai_project", "vertex_location", "vertex_ai_location"], + serde_json::json!({"vertex_ai_project": "p"}), + &["vertex_location"] + )] + fn only_the_globals_of_consumed_absent_params_are_read( + #[case] consumed: &[&str], + #[case] params: Value, + #[case] expected: &[&str], + ) { + assert_eq!( + module_globals_to_read(params.as_object().unwrap(), consumed), + expected + ); + } + fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { let locals = PyDict::new(py); py.run(source, Some(&locals), Some(&locals)).unwrap(); diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 8767b273b68..88cdd9ff5c1 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -527,8 +527,6 @@ async def test_native_failures_raise_the_public_exception_class( ("ssl_verify", object()), ("ssl_certificate", 1), ("ssl_certificate", ""), - ("vertex_project", 1), - ("vertex_location", ["region"]), ("user_url_allowed_hosts", ["example.test", 1]), ], ) @@ -541,11 +539,35 @@ async def test_native_settings_fail_before_provider_io( ) -> None: ocr_server.expected_requests = 0 monkeypatch.setattr(litellm, name, value) - with pytest.raises(ValueError, match=r"http_settings|provider_defaults|url_policy"): + with pytest.raises(ValueError, match=r"http_settings|url_policy"): await call_native(ocr_server, asynchronous, num_retries=0) assert ocr_server.requests == [] +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + "name,value", + [ + ("vertex_project", 1), + ("vertex_location", ["region"]), + ], +) +async def test_native_vertex_globals_fail_before_provider_io_only_for_vertex_calls( + ocr_server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + name: str, + value: object, +) -> None: + ocr_server.expected_requests = 1 + monkeypatch.setattr(litellm, name, value) + await call_native(ocr_server, asynchronous, num_retries=0) + with pytest.raises(ValueError, match=r"provider_defaults"): + await call_native(ocr_server, asynchronous, model="vertex_ai/mistral-ocr-latest", num_retries=0) + assert len(ocr_server.requests) == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) async def test_native_ssl_context_is_terminal_configuration(ocr_server: RecordingServer, asynchronous: bool) -> None: From 3d2abf9a441babc9d6646fe2f824aa269b01cb98 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 9 Oct 2026 16:29:45 -0700 Subject: [PATCH 003/104] feat(rust): one LitellmParams type for a config and a caller (#45642) litellm-router-types mirrors litellm/types/router.py: LitellmParams holds the model, the credentials, one flattened group per credential family from auth-types (aws, vertex), the deployment settings a config spells, and every other key in extra, as Python's extra="allow" keeps it. Config parses it directly, so config::LiteLlmParams and its untyped additional_fields are gone and a mistyped provider param is a parse error. fields() and specs() derive from the families' param specs for hosts that project kwargs and fold module globals. Spelled is the one untagged shape for a typed value or the text a config spells, organization takes the list Python's router expands, and the callback shorthand keeps a value-level one-or-many. No credentials wrapper group: serde's flatten only consumes keys for struct-shaped children. --- litellm-rust/Cargo.lock | 13 ++ litellm-rust/Cargo.toml | 1 + litellm-rust/crates/auth-aws/Cargo.toml | 1 + .../crates/auth-types/src/params/aws.rs | 7 +- litellm-rust/crates/config/Cargo.toml | 1 + litellm-rust/crates/config/src/lib.rs | 5 +- litellm-rust/crates/config/src/model.rs | 71 +------- litellm-rust/crates/config/src/settings.rs | 19 ++- litellm-rust/crates/config/src/value.rs | 47 ++---- litellm-rust/crates/config/tests/config.rs | 47 ++++-- litellm-rust/crates/gateway/src/lib.rs | 6 +- litellm-rust/crates/router-types/Cargo.toml | 16 ++ litellm-rust/crates/router-types/src/lib.rs | 5 + .../crates/router-types/src/litellm_params.rs | 92 ++++++++++ litellm-rust/crates/router-types/src/value.rs | 10 ++ .../router-types/tests/litellm_params.rs | 158 ++++++++++++++++++ 16 files changed, 364 insertions(+), 135 deletions(-) create mode 100644 litellm-rust/crates/router-types/Cargo.toml create mode 100644 litellm-rust/crates/router-types/src/lib.rs create mode 100644 litellm-rust/crates/router-types/src/litellm_params.rs create mode 100644 litellm-rust/crates/router-types/src/value.rs create mode 100644 litellm-rust/crates/router-types/tests/litellm_params.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 4cbde047a71..ec8d10b58de 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3383,6 +3383,7 @@ dependencies = [ "sha2 0.10.9", "thiserror 2.0.19", "tokio", + "veil", ] [[package]] @@ -3645,6 +3646,7 @@ name = "litellm-config" version = "0.1.0" dependencies = [ "litellm-auth-types", + "litellm-router-types", "rstest", "serde", "serde_json", @@ -4271,6 +4273,17 @@ dependencies = [ "rstest", ] +[[package]] +name = "litellm-router-types" +version = "0.1.0" +dependencies = [ + "litellm-auth-types", + "rstest", + "serde", + "serde_json", + "serde_with", +] + [[package]] name = "litellm-secrets" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 794070366af..b94988eb8e1 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -11,6 +11,7 @@ repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] litellm-config = { path = "crates/config" } litellm-router = { path = "crates/router" } +litellm-router-types = { path = "crates/router-types" } litellm-tracing = { path = "crates/tracing" } litellm-spend-clickhouse = { path = "crates/spend-clickhouse" } litellm-storage-clickhouse = { path = "crates/storage-clickhouse" } diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index bf757c6744e..65b411ffaef 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -14,6 +14,7 @@ serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } sha2.workspace = true thiserror.workspace = true +veil.workspace = true aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"] } aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"] } diff --git a/litellm-rust/crates/auth-types/src/params/aws.rs b/litellm-rust/crates/auth-types/src/params/aws.rs index f3653e969d6..25fe52b2ce1 100644 --- a/litellm-rust/crates/auth-types/src/params/aws.rs +++ b/litellm-rust/crates/auth-types/src/params/aws.rs @@ -2,6 +2,7 @@ use std::sync::LazyLock; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; +use veil::Redact; use super::ParamSpec; @@ -19,13 +20,15 @@ const fn spec( } /// The `aws_*` fields of Python's `GenericLiteLLMParams`. -#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Redact, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct AwsParams { #[serde(default)] pub aws_access_key_id: Option, #[serde(default)] + #[redact(with = "[REDACTED]")] pub aws_secret_access_key: Option, #[serde(default)] + #[redact(with = "[REDACTED]")] pub aws_session_token: Option, #[serde(default)] pub aws_region_name: Option, @@ -36,10 +39,12 @@ pub struct AwsParams { #[serde(default)] pub aws_role_name: Option, #[serde(default)] + #[redact(with = "[REDACTED]")] pub aws_web_identity_token: Option, #[serde(default)] pub aws_sts_endpoint: Option, #[serde(default)] + #[redact(with = "[REDACTED]")] pub aws_external_id: Option, #[serde(default)] pub aws_bedrock_runtime_endpoint: Option, diff --git a/litellm-rust/crates/config/Cargo.toml b/litellm-rust/crates/config/Cargo.toml index e23739997e2..23d53923362 100644 --- a/litellm-rust/crates/config/Cargo.toml +++ b/litellm-rust/crates/config/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] litellm-auth-types.workspace = true +litellm-router-types.workspace = true serde.workspace = true serde_yaml_ng = "0.10.0" strum.workspace = true diff --git a/litellm-rust/crates/config/src/lib.rs b/litellm-rust/crates/config/src/lib.rs index fed60ab1a4f..8603ef3279c 100644 --- a/litellm-rust/crates/config/src/lib.rs +++ b/litellm-rust/crates/config/src/lib.rs @@ -10,13 +10,14 @@ use std::{fmt, path::Path}; use serde::Deserialize; pub use error::Error; +pub use litellm_router_types::{LitellmParams, Spelled}; pub use mcp::{McpAuth, McpServer, McpTransport}; -pub use model::{LiteLlmParams, Model}; +pub use model::Model; pub use settings::{ ClickHouseStoreSettings, GeneralSettings, LiteLlmSettings, RouterSettings, TracingSettings, TracingStoreSettings, }; -pub use value::{AdditionalFields, Flag, NumberOrString, Object, OneOrMany, Value}; +pub use value::{AdditionalFields, Object, Value}; #[derive(Clone, Default, Deserialize)] #[serde(default)] diff --git a/litellm-rust/crates/config/src/model.rs b/litellm-rust/crates/config/src/model.rs index cd05bb4b7c8..c0f8516103d 100644 --- a/litellm-rust/crates/config/src/model.rs +++ b/litellm-rust/crates/config/src/model.rs @@ -1,14 +1,14 @@ use std::fmt; -use litellm_auth_types::SecretValue; +use litellm_router_types::LitellmParams; use serde::Deserialize; -use crate::{AdditionalFields, Flag, NumberOrString, Object}; +use crate::{AdditionalFields, Object}; #[derive(Clone, Deserialize)] pub struct Model { pub model_name: String, - pub litellm_params: LiteLlmParams, + pub litellm_params: LitellmParams, #[serde(default)] pub model_info: Object, pub blocked: Option, @@ -28,68 +28,3 @@ impl fmt::Debug for Model { .finish() } } - -#[derive(Clone, Deserialize)] -pub struct LiteLlmParams { - pub model: String, - pub api_key: Option, - pub api_base: Option, - pub api_version: Option, - pub custom_llm_provider: Option, - pub timeout: Option, - pub stream_timeout: Option, - pub max_retries: Option, - pub tpm: Option, - pub rpm: Option, - pub itpm: Option, - pub otpm: Option, - pub max_parallel_requests: Option, - pub organization: Option, - pub drop_params: Option, - pub tags: Option>, - pub tag_regex: Option>, - pub max_budget: Option, - pub budget_duration: Option, - pub default_api_key_tpm_limit: Option, - pub default_api_key_rpm_limit: Option, - pub use_in_pass_through: Option, - pub use_chat_completions_api: Option, - pub litellm_credential_name: Option, - pub provider_affinity_header: Option, - #[serde(flatten)] - pub additional_fields: AdditionalFields, -} - -impl fmt::Debug for LiteLlmParams { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("LiteLlmParams") - .field("model", &self.model) - .field("api_key", &self.api_key) - .field("api_base", &self.api_base) - .field("api_version", &self.api_version) - .field("custom_llm_provider", &self.custom_llm_provider) - .field("timeout", &self.timeout) - .field("stream_timeout", &self.stream_timeout) - .field("max_retries", &self.max_retries) - .field("tpm", &self.tpm) - .field("rpm", &self.rpm) - .field("itpm", &self.itpm) - .field("otpm", &self.otpm) - .field("max_parallel_requests", &self.max_parallel_requests) - .field("organization", &self.organization) - .field("drop_params", &self.drop_params) - .field("tags", &self.tags) - .field("tag_regex", &self.tag_regex) - .field("max_budget", &self.max_budget) - .field("budget_duration", &self.budget_duration) - .field("default_api_key_tpm_limit", &self.default_api_key_tpm_limit) - .field("default_api_key_rpm_limit", &self.default_api_key_rpm_limit) - .field("use_in_pass_through", &self.use_in_pass_through) - .field("use_chat_completions_api", &self.use_chat_completions_api) - .field("litellm_credential_name", &self.litellm_credential_name) - .field("provider_affinity_header", &self.provider_affinity_header) - .field("additional_fields", &self.additional_fields.keys()) - .finish() - } -} diff --git a/litellm-rust/crates/config/src/settings.rs b/litellm-rust/crates/config/src/settings.rs index b1ead35e55b..1094de1638e 100644 --- a/litellm-rust/crates/config/src/settings.rs +++ b/litellm-rust/crates/config/src/settings.rs @@ -3,7 +3,7 @@ use std::fmt; use litellm_auth_types::SecretValue; use serde::Deserialize; -use crate::{AdditionalFields, Flag, NumberOrString, Object, OneOrMany, Value}; +use crate::{AdditionalFields, Object, Spelled, Value, value::one_or_many}; #[derive(Clone, Debug, Deserialize)] #[serde(rename_all = "lowercase")] @@ -18,7 +18,7 @@ pub struct ClickHouseStoreSettings { pub kind: TracingStoreKind, pub url: Option, pub database: Option, - pub retention_days: Option, + pub retention_days: Option>, } impl fmt::Debug for ClickHouseStoreSettings { @@ -207,7 +207,7 @@ impl fmt::Debug for RouterSettings { #[derive(Clone, Default, Deserialize)] #[serde(default)] pub struct LiteLlmSettings { - pub ssl_verify: Option, + pub ssl_verify: Option>, pub ssl_certificate: Option, pub ssl_security_level: Option, pub ssl_ecdh_curve: Option, @@ -216,14 +216,17 @@ pub struct LiteLlmSettings { pub aiohttp_trust_env: Option, pub disable_aiohttp_trust_env: Option, pub disable_aiohttp_transport: Option, - pub drop_params: Option, - pub request_timeout: Option, + pub drop_params: Option>, + pub request_timeout: Option>, pub num_retries: Option, pub cache: Option, pub cache_params: Option, - pub callbacks: Option>, - pub success_callback: Option>, - pub failure_callback: Option>, + #[serde(deserialize_with = "one_or_many")] + pub callbacks: Option>, + #[serde(deserialize_with = "one_or_many")] + pub success_callback: Option>, + #[serde(deserialize_with = "one_or_many")] + pub failure_callback: Option>, pub json_logs: Option, pub set_verbose: Option, #[serde(flatten)] diff --git a/litellm-rust/crates/config/src/value.rs b/litellm-rust/crates/config/src/value.rs index 6b68b5bc2d0..5becdc87b70 100644 --- a/litellm-rust/crates/config/src/value.rs +++ b/litellm-rust/crates/config/src/value.rs @@ -1,6 +1,6 @@ use std::{collections::BTreeMap, fmt, ops::Deref}; -use serde::Deserialize; +use serde::{Deserialize, Deserializer}; pub type Value = serde_yaml_ng::Value; pub type AdditionalFields = BTreeMap; @@ -40,39 +40,14 @@ impl fmt::Debug for Object { } } -#[derive(Clone, Debug, Deserialize, PartialEq)] -#[serde(untagged)] -pub enum NumberOrString { - Number(f64), - String(String), -} - -#[derive(Clone, Debug, Deserialize, PartialEq, Eq)] -#[serde(untagged)] -pub enum Flag { - Boolean(bool), - String(String), -} - -#[derive(Clone, Debug, Deserialize)] -#[serde(untagged)] -pub enum OneOrMany { - Many(Box<[T]>), - One(T), -} - -impl OneOrMany { - pub fn len(&self) -> usize { - match self { - Self::Many(values) => values.len(), - Self::One(_) => 1, - } - } - - pub fn is_empty(&self) -> bool { - match self { - Self::Many(values) => values.is_empty(), - Self::One(_) => false, - } - } +/// Python's callback shorthand: one entry or a list of them, kept as the values written. +pub(crate) fn one_or_many<'de, D: Deserializer<'de>>( + deserializer: D, +) -> Result>, D::Error> { + Ok( + Option::::deserialize(deserializer)?.map(|value| match value { + Value::Sequence(values) => values, + one => vec![one], + }), + ) } diff --git a/litellm-rust/crates/config/tests/config.rs b/litellm-rust/crates/config/tests/config.rs index ab9f3403a01..d425b62efc2 100644 --- a/litellm-rust/crates/config/tests/config.rs +++ b/litellm-rust/crates/config/tests/config.rs @@ -1,4 +1,4 @@ -use litellm_config::{Config, Error, Flag, NumberOrString, TracingStoreSettings}; +use litellm_config::{Config, Error, Spelled, TracingStoreSettings}; use rstest::{fixture, rstest}; use tempfile::TempDir; @@ -75,6 +75,9 @@ fn config_debug_redacts_api_keys() { #[case::malformed_yaml("model_list: [")] #[case::missing_params("model_list: [{model_name: assistant}]")] #[case::missing_model("model_list: [{model_name: assistant, litellm_params: {api_key: key}}]")] +#[case::provider_param_of_the_wrong_type( + "model_list: [{model_name: assistant, litellm_params: {model: test, aws_region_name: [eu-central-1]}}]" +)] fn rejects_malformed_and_incomplete_config(#[case] yaml: &str) { assert!(matches!(Config::from_yaml(yaml), Err(Error::Parse(_)))); } @@ -128,7 +131,7 @@ fn tracing_settings_are_typed_and_redact_the_url() { "https://writer:password@example.com" ); assert_eq!(store.database.as_deref(), Some("analytics")); - assert_eq!(store.retention_days, Some(NumberOrString::Number(7.0))); + assert_eq!(store.retention_days, Some(Spelled::Value(7.0))); assert!(!format!("{config:?}").contains("password")); } @@ -145,9 +148,7 @@ fn tracing_settings_accept_environment_references() { }; assert_eq!( store.retention_days, - Some(NumberOrString::String( - "os.environ/RETENTION_DAYS".to_owned() - )) + Some(Spelled::Text("os.environ/RETENTION_DAYS".to_owned())) ); } @@ -190,6 +191,9 @@ model_list: rpm: 5 drop_params: "true" vertex_project: test-project + vertex_credentials: + type: service_account + future_param: 1 model_info: mode: chat access_groups: [internal] @@ -214,25 +218,26 @@ future_section: let model = &config.model_list[0]; assert_eq!( model.litellm_params.timeout, - Some(NumberOrString::String( - "os.environ/REQUEST_TIMEOUT".to_string() - )) + Some(Spelled::Text("os.environ/REQUEST_TIMEOUT".to_string())) ); assert_eq!( model.litellm_params.tpm, - Some(NumberOrString::String("os.environ/TPM_LIMIT".to_string())) + Some(Spelled::Text("os.environ/TPM_LIMIT".to_string())) ); - assert_eq!(model.litellm_params.rpm, Some(NumberOrString::Number(5.0))); + assert_eq!(model.litellm_params.rpm, Some(Spelled::Value(5.0))); assert_eq!( model.litellm_params.drop_params, - Some(Flag::String("true".to_string())) + Some(Spelled::Text("true".to_string())) ); - assert!( - model - .litellm_params - .additional_fields - .contains_key("vertex_project") + assert_eq!( + model.litellm_params.vertex.vertex_project.as_deref(), + Some("test-project") ); + assert_eq!( + model.litellm_params.vertex.vertex_credentials.as_deref(), + Some(r#"{"type":"service_account"}"#) + ); + assert!(model.litellm_params.extra.contains_key("future_param")); assert!(model.additional_fields.contains_key("access_groups")); assert_eq!(config.router_settings.allowed_fails, Some(2)); assert!( @@ -345,7 +350,9 @@ fn accepts_python_callback_shorthand(#[case] setting: &str, #[case] expected_len #[rstest] #[case::api_key("api_key", "provider-secret")] -#[case::provider_extension("aws_secret_access_key", "aws-secret")] +#[case::aws_secret("aws_secret_access_key", "aws-secret")] +#[case::vertex_credentials("vertex_credentials", "vertex-secret")] +#[case::provider_extension("azure_ad_token", "azure-secret")] #[case::general_extension("custom_auth_secret", "auth-secret")] #[case::router_extension("redis_password", "redis-secret")] #[case::litellm_extension("callback_token", "callback-secret")] @@ -358,6 +365,12 @@ fn debug_output_does_not_expose_config_values(#[case] field: &str, #[case] secre "aws_secret_access_key" => format!( "model_list: [{{model_name: assistant, litellm_params: {{model: test, aws_secret_access_key: {secret}}}}}]" ), + "vertex_credentials" => format!( + "model_list: [{{model_name: assistant, litellm_params: {{model: test, vertex_credentials: {secret}}}}}]" + ), + "azure_ad_token" => format!( + "model_list: [{{model_name: assistant, litellm_params: {{model: test, azure_ad_token: {secret}}}}}]" + ), "custom_auth_secret" => format!("general_settings: {{{field}: {secret}}}"), "redis_password" => format!("router_settings: {{{field}: {secret}}}"), "callback_token" => format!("litellm_settings: {{{field}: {secret}}}"), diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index d711eb78caa..0b3aba81181 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -40,9 +40,9 @@ pub fn build_inference(config: &Config) -> Result, Error> { HttpSettingsLayer::from_environment(&lookup), HttpSettingsLayer { ssl_verify: settings.ssl_verify.as_ref().map(|value| match value { - litellm_config::Flag::Boolean(true) => SslVerify::Enabled, - litellm_config::Flag::Boolean(false) => SslVerify::Disabled, - litellm_config::Flag::String(value) => SslVerify::parse(value), + litellm_config::Spelled::Value(true) => SslVerify::Enabled, + litellm_config::Spelled::Value(false) => SslVerify::Disabled, + litellm_config::Spelled::Text(value) => SslVerify::parse(value), }), ssl_certificate: settings.ssl_certificate.as_ref().map(Into::into), ssl_security_level: settings.ssl_security_level.clone(), diff --git a/litellm-rust/crates/router-types/Cargo.toml b/litellm-rust/crates/router-types/Cargo.toml new file mode 100644 index 00000000000..2e9fec28685 --- /dev/null +++ b/litellm-rust/crates/router-types/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "litellm-router-types" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true +description = "The types of litellm/types/router.py: a deployment's litellm_params as a config or a caller spells them" + +[dependencies] +litellm-auth-types.workspace = true +serde.workspace = true +serde_json.workspace = true +serde_with.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/router-types/src/lib.rs b/litellm-rust/crates/router-types/src/lib.rs new file mode 100644 index 00000000000..2235e40a891 --- /dev/null +++ b/litellm-rust/crates/router-types/src/lib.rs @@ -0,0 +1,5 @@ +mod litellm_params; +mod value; + +pub use litellm_params::{ExtraParams, LitellmParams}; +pub use value::Spelled; diff --git a/litellm-rust/crates/router-types/src/litellm_params.rs b/litellm-rust/crates/router-types/src/litellm_params.rs new file mode 100644 index 00000000000..186a1d98036 --- /dev/null +++ b/litellm-rust/crates/router-types/src/litellm_params.rs @@ -0,0 +1,92 @@ +use std::{collections::BTreeMap, fmt, ops::Deref}; + +use litellm_auth_types::{AwsParams, ParamSpec, SecretValue, VertexParams}; +use serde::Deserialize; +use serde_json::Value; +use serde_with::{OneOrMany, formats::PreferMany, serde_as}; + +use crate::Spelled; + +/// Python's `LiteLLM_Params`: a deployment's `litellm_params`, as a config spells them and as +/// a caller passes them. Each credential family is one flattened group owned by its auth +/// crate, so a config reaches for `litellm_params.aws` or `litellm_params.vertex`, never a +/// map. Every key no field names lands in [`Self::extra`], as Python's `extra="allow"` keeps it. +/// +/// A key a field names but cannot hold is an error, as it is for Python's pydantic model; +/// the coercions Python's validators apply (`max_retries` from text, `drop_params` from a flag +/// word) are left to the readers, so the shapes here are the ones the YAML carries. An +/// `organization` list is what Python's router expands into one deployment per entry. +#[serde_as] +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +pub struct LitellmParams { + pub model: String, + pub api_key: Option, + pub api_base: Option, + pub api_version: Option, + #[serde(flatten)] + pub aws: AwsParams, + #[serde(flatten)] + pub vertex: VertexParams, + pub custom_llm_provider: Option, + pub timeout: Option>, + pub stream_timeout: Option>, + pub max_retries: Option>, + pub tpm: Option>, + pub rpm: Option>, + pub itpm: Option>, + pub otpm: Option>, + pub max_parallel_requests: Option, + #[serde_as(as = "Option>")] + pub organization: Option>, + pub drop_params: Option>, + pub tags: Option>, + pub tag_regex: Option>, + pub max_budget: Option, + pub budget_duration: Option, + pub default_api_key_tpm_limit: Option, + pub default_api_key_rpm_limit: Option, + pub use_in_pass_through: Option, + pub use_chat_completions_api: Option, + pub litellm_credential_name: Option, + pub provider_affinity_header: Option, + #[serde(flatten)] + pub extra: ExtraParams, +} + +impl LitellmParams { + /// Every connection param spec the credential families declare, for hosts that fold a + /// spec's module global in when a call names none of its spellings. + pub fn specs() -> impl Iterator { + AwsParams::SPECS.iter().chain(VertexParams::SPECS.iter()) + } + + /// The wire names a host projects out of a caller's kwargs: the model, the credentials + /// and each credential family. The routing and limit fields are a deployment's + /// settings, which only a config spells. + pub fn fields() -> impl Iterator { + ["model", "api_key", "api_base", "api_version"] + .into_iter() + .chain(AwsParams::fields()) + .chain(VertexParams::fields()) + } +} + +/// The `litellm_params` keys no field names, kept as written. Its `Debug` lists the keys +/// only, since a provider's secret may sit among the values. +#[derive(Clone, Default, PartialEq, Deserialize)] +#[serde(transparent)] +pub struct ExtraParams(BTreeMap); + +impl Deref for ExtraParams { + type Target = BTreeMap; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl fmt::Debug for ExtraParams { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.debug_list().entries(self.0.keys()).finish() + } +} diff --git a/litellm-rust/crates/router-types/src/value.rs b/litellm-rust/crates/router-types/src/value.rs new file mode 100644 index 00000000000..e52cf60c0ef --- /dev/null +++ b/litellm-rust/crates/router-types/src/value.rs @@ -0,0 +1,10 @@ +use serde::Deserialize; + +/// A typed value, or the text a config spells in its place: an `os.environ/NAME` reference +/// the secrets layer resolves later, or a word Python's coercion reads the value from. +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(untagged)] +pub enum Spelled { + Value(T), + Text(String), +} diff --git a/litellm-rust/crates/router-types/tests/litellm_params.rs b/litellm-rust/crates/router-types/tests/litellm_params.rs new file mode 100644 index 00000000000..7d021fd857a --- /dev/null +++ b/litellm-rust/crates/router-types/tests/litellm_params.rs @@ -0,0 +1,158 @@ +use litellm_auth_types::{AwsParams, SecretValue, VertexParams}; +use litellm_router_types::{LitellmParams, Spelled}; +use rstest::rstest; +use serde_json::{Map, Value, json}; + +fn params(value: Value) -> LitellmParams { + serde_json::from_value(value).unwrap() +} + +#[rstest] +#[case::aws_group( + json!({"model": "m", "aws_region_name": "eu-west-1", "aws_access_key_id": "AKIA"}), + LitellmParams { + model: "m".into(), + aws: AwsParams { + aws_region_name: Some("eu-west-1".into()), + aws_access_key_id: Some("AKIA".into()), + ..AwsParams::default() + }, + ..LitellmParams::default() + }, +)] +#[case::vertex_group( + json!({"model": "m", "vertex_project": "p", "vertex_ai_location": "eu", "vertex_credentials": {"type": "service_account"}}), + LitellmParams { + model: "m".into(), + vertex: VertexParams { + vertex_project: Some("p".into()), + vertex_ai_location: Some("eu".into()), + vertex_credentials: Some(r#"{"type":"service_account"}"#.into()), + ..VertexParams::default() + }, + ..LitellmParams::default() + }, +)] +#[case::both_groups_and_the_credentials( + json!({"model": "m", "api_key": "k", "api_base": "b", "aws_region_name": "eu-west-1", "vertex_location": "us-east5"}), + LitellmParams { + model: "m".into(), + api_key: Some(SecretValue::new("k")), + api_base: Some("b".into()), + aws: AwsParams { + aws_region_name: Some("eu-west-1".into()), + ..AwsParams::default() + }, + vertex: VertexParams { + vertex_location: Some("us-east5".into()), + ..VertexParams::default() + }, + ..LitellmParams::default() + }, +)] +#[case::deployment_settings( + json!({"model": "m", "timeout": "os.environ/T", "rpm": 5, "drop_params": "true", "tags": ["a"], "max_budget": 1.5}), + LitellmParams { + model: "m".into(), + timeout: Some(Spelled::Text("os.environ/T".into())), + rpm: Some(Spelled::Value(5.0)), + drop_params: Some(Spelled::Text("true".into())), + tags: Some(Box::from(["a".to_string()])), + max_budget: Some(1.5), + ..LitellmParams::default() + }, +)] +#[case::explicit_null_is_absent( + json!({"model": "m", "aws_region_name": null, "vertex_project": null, "api_key": null, "tpm": null}), + LitellmParams { model: "m".into(), ..LitellmParams::default() }, +)] +#[case::one_organization( + json!({"model": "m", "organization": "org-a"}), + LitellmParams { + model: "m".into(), + organization: Some(vec!["org-a".into()]), + ..LitellmParams::default() + }, +)] +#[case::organizations_the_router_expands( + json!({"model": "m", "organization": ["org-a", "org-b"]}), + LitellmParams { + model: "m".into(), + organization: Some(vec!["org-a".into(), "org-b".into()]), + ..LitellmParams::default() + }, +)] +fn deserializes_each_group_from_a_config_or_a_call( + #[case] value: Value, + #[case] expected: LitellmParams, +) { + assert_eq!(params(value), expected); +} + +#[rstest] +fn keys_no_field_names_land_in_extra_and_nothing_else_does() { + let typed = params(json!({ + "model": "m", + "api_key": "k", + "aws_region_name": "eu-west-1", + "vertex_project": "p", + "rpm": 5, + "azure_ad_token": "t", + "messages": [], + })); + + assert_eq!( + typed.extra.keys().collect::>(), + ["azure_ad_token", "messages"] + ); + assert_eq!(typed.extra["azure_ad_token"], "t"); + assert_eq!(typed.aws.aws_region_name.as_deref(), Some("eu-west-1")); + assert_eq!(typed.vertex.vertex_project.as_deref(), Some("p")); +} + +#[rstest] +#[case::model_missing(json!({"api_key": "k"}))] +#[case::aws(json!({"model": "m", "aws_region_name": 7}))] +#[case::vertex(json!({"model": "m", "vertex_project": ["p"]}))] +#[case::vertex_credentials(json!({"model": "m", "vertex_credentials": 7}))] +#[case::rate_limit(json!({"model": "m", "max_parallel_requests": "many"}))] +#[case::tags(json!({"model": "m", "tags": "not-a-list"}))] +fn a_param_of_the_wrong_type_is_rejected(#[case] value: Value) { + assert!(serde_json::from_value::(value).is_err()); +} + +#[rstest] +fn fields_fill_the_model_the_credentials_and_every_group() { + let filled: Value = LitellmParams::fields() + .map(|name| (name.to_string(), Value::from(name))) + .collect::>() + .into(); + let typed = params(filled); + let nulls = |group: Value| { + group + .as_object() + .unwrap() + .values() + .filter(|value| value.is_null()) + .count() + }; + + assert_eq!(typed.model, "model"); + assert_eq!(typed.api_key, Some(SecretValue::new("api_key"))); + assert_eq!(typed.api_base.as_deref(), Some("api_base")); + assert_eq!(typed.api_version.as_deref(), Some("api_version")); + assert_eq!(nulls(serde_json::to_value(&typed.aws).unwrap()), 0); + assert_eq!(nulls(serde_json::to_value(&typed.vertex).unwrap()), 0); + assert!(typed.extra.is_empty()); +} + +#[rstest] +fn debug_shows_neither_the_api_key_nor_the_extra_values() { + let typed = + params(json!({"model": "m", "api_key": "sk-secret", "azure_ad_token": "token-secret"})); + let debug = format!("{typed:?}"); + + assert!(!debug.contains("sk-secret")); + assert!(!debug.contains("token-secret")); + assert!(debug.contains("azure_ad_token")); +} From ee295aea686ae41a9ccc05c63504c4bd94d9a351 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 9 Oct 2026 16:29:45 -0700 Subject: [PATCH 004/104] refactor(rust): give Messages configs typed litellm params (#45611) * refactor(rust): give Messages configs typed litellm params * feat(rust): port _normalize_system_role_messages for Messages hosts mid_conversation_system mirrors Python's module function for function: a leading run of system turns is hoisted into the top-level system, a later turn stays in place when the model map flags supports_mid_conversation_system and otherwise becomes a user turn where it was, never between a tool_use and its tool_result, and billing blocks are stripped from the hoisted system. The capability joins MessagesModelCapabilities and the bridge projection. Azure AI uses it in place of fold_system_role_messages, which hoisted every turn. * fix(rust): keep the Bedrock runtime endpoint behind a blank api_base * refactor(rust): import LitellmParams from its crate instead of a re-export * refactor(rust): template the Messages request decode errors Both decode failures go through ErrorDetail::invalid with the subject and the serde error as source instead of a preformatted sentence. * refactor(rust_bridge): read the module globals the param specs name The Messages host no longer keeps its own list of litellm globals; a spec whose spellings the call leaves out names the global to read, so the next family that has one needs no host change. * refactor(rust): give each route lookup Python's implicit None ProviderConfigManager's per-route methods list only the providers the route serves and fall through to None; the Rust lookups now do the same, so a new LlmProviders variant touches only the route that serves it. --- litellm-rust/Cargo.lock | 4 + .../crates/gateway-inference/src/messages.rs | 1 + .../crates/inference-chat/src/common_utils.rs | 8 +- .../crates/inference-messages/Cargo.toml | 1 + .../inference-messages/src/common_utils.rs | 8 +- .../crates/inference-messages/src/handler.rs | 4 +- .../crates/inference-messages/src/lib.rs | 3 +- .../crates/inference-messages/src/prepare.rs | 17 +- .../crates/inference-messages/src/types.rs | 15 +- .../inference-messages/tests/caching.rs | 2 +- .../inference-messages/tests/messages/main.rs | 1 + .../tests/messages/request.rs | 10 + .../tests/messages/response.rs | 1 + .../tests/messages/secrets.rs | 2 + .../tests/messages/stream.rs | 2 + .../crates/inference-ocr/src/prepare.rs | 9 +- .../inference-ocr/src/provider_config.rs | 5 +- .../inference-transcription/src/prepare.rs | 10 +- litellm-rust/crates/llms/Cargo.toml | 1 + .../crates/llms/src/anthropic/common_utils.rs | 29 +++ .../src/anthropic/messages/transformation.rs | 40 +++- .../src/azure_ai/messages/transformation.rs | 207 ++++++++++-------- litellm-rust/crates/llms/src/base_llm/auth.rs | 11 + .../llms/src/base_llm/messages/context.rs | 3 + .../messages/mid_conversation_system.rs | 201 +++++++++++++++++ .../crates/llms/src/base_llm/messages/mod.rs | 1 + .../src/base_llm/messages/normalization.rs | 128 ++++++++--- .../src/base_llm/messages/transformation.rs | 23 +- .../anthropic_claude3_transformation.rs | 180 ++++++++++++--- .../llms/tests/messages_normalization.rs | 23 +- litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../python-bridge/src/python_settings.rs | 2 +- .../python-bridge/src/routes/messages/host.rs | 153 ++++++++++++- litellm-rust/crates/router/Cargo.toml | 1 + litellm-rust/crates/router/src/deployment.rs | 2 + litellm-rust/crates/router/src/lib.rs | 1 + litellm-rust/crates/router/tests/router.rs | 27 +++ litellm/rust_bridge/model_capabilities.py | 1 + .../rust_bridge/test_model_capabilities.py | 3 + 39 files changed, 901 insertions(+), 240 deletions(-) create mode 100644 litellm-rust/crates/llms/src/base_llm/messages/mid_conversation_system.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ec8d10b58de..b0c5b02a245 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4031,6 +4031,7 @@ dependencies = [ "litellm-inference-testing", "litellm-llms", "litellm-llms-types", + "litellm-router-types", "litellm-secrets", "litellm-tracing", "reqwest 0.12.28", @@ -4151,6 +4152,7 @@ dependencies = [ "litellm-http", "litellm-llms-types", "litellm-python-compat", + "litellm-router-types", "litellm-secrets", "reqwest 0.12.28", "rstest", @@ -4222,6 +4224,7 @@ dependencies = [ "litellm-inference-transcription", "litellm-llms", "litellm-llms-types", + "litellm-router-types", "litellm-secrets", "litellm-secrets-aws", "litellm-secrets-types", @@ -4270,6 +4273,7 @@ dependencies = [ "litellm-config", "litellm-inference", "litellm-inference-messages", + "litellm-router-types", "rstest", ] diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index 8969ad80934..e56308b7ae1 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -79,6 +79,7 @@ fn project( api_key: deployment.api_key.clone(), api_base: deployment.api_base.clone(), custom_llm_provider: deployment.custom_llm_provider.clone(), + litellm_params: deployment.litellm_params.clone(), extra_headers: None, provider_specific_header: anthropic_api_headers(headers), timeout: deployment.timeout, diff --git a/litellm-rust/crates/inference-chat/src/common_utils.rs b/litellm-rust/crates/inference-chat/src/common_utils.rs index 14abd42f6b7..90b75dcddaa 100644 --- a/litellm-rust/crates/inference-chat/src/common_utils.rs +++ b/litellm-rust/crates/inference-chat/src/common_utils.rs @@ -33,13 +33,7 @@ pub(super) fn chat_completions_provider(provider: LlmProviders) -> Option Some(ChatProvider::Anthropic), LlmProviders::Bedrock => Some(ChatProvider::Bedrock), LlmProviders::OpenaiLike => Some(ChatProvider::OpenaiLike), - LlmProviders::AwsTextract - | LlmProviders::AzureAi - | LlmProviders::Cohere - | LlmProviders::Mistral - | LlmProviders::Openai - | LlmProviders::Reducto - | LlmProviders::VertexAi => None, + _ => None, } } diff --git a/litellm-rust/crates/inference-messages/Cargo.toml b/litellm-rust/crates/inference-messages/Cargo.toml index 8dd516732d0..6801c758e54 100644 --- a/litellm-rust/crates/inference-messages/Cargo.toml +++ b/litellm-rust/crates/inference-messages/Cargo.toml @@ -17,6 +17,7 @@ litellm-http.workspace = true litellm-inference.workspace = true litellm-llms.workspace = true litellm-llms-types.workspace = true +litellm-router-types.workspace = true litellm-secrets.workspace = true litellm-tracing.workspace = true reqwest.workspace = true diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index 145192e672b..5102d063a8a 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -44,13 +44,7 @@ pub(crate) fn messages_provider(provider: LlmProviders) -> Option Some(MessagesProvider::Anthropic), LlmProviders::AzureAi => Some(MessagesProvider::AzureAi), LlmProviders::Bedrock => Some(MessagesProvider::Bedrock), - LlmProviders::AwsTextract - | LlmProviders::Cohere - | LlmProviders::Mistral - | LlmProviders::Openai - | LlmProviders::OpenaiLike - | LlmProviders::Reducto - | LlmProviders::VertexAi => None, + _ => None, } } diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index ee056a533b0..4e883320303 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -70,7 +70,9 @@ impl MessagesRoute { WireRequest { url, headers: authenticated.headers, - body: serde_json::to_value(&body).map_err(serialize_failure)?, + body: provider + .config() + .wire_body(serde_json::to_value(&body).map_err(serialize_failure)?), }, request_context, ) diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index cc48a2a39b1..6cf7e34e52a 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -15,7 +15,8 @@ use std::sync::Arc; pub use litellm_inference::RouteError as Error; pub use types::{ - MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, messages_body, + MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params, + messages_body, }; #[derive(Clone)] diff --git a/litellm-rust/crates/inference-messages/src/prepare.rs b/litellm-rust/crates/inference-messages/src/prepare.rs index b2a09f3baf9..0c189d0a309 100644 --- a/litellm-rust/crates/inference-messages/src/prepare.rs +++ b/litellm-rust/crates/inference-messages/src/prepare.rs @@ -69,6 +69,7 @@ fn prepare_provider_request( body, api_key, api_base, + litellm_params, extra_headers, provider_specific_header, timeout, @@ -98,6 +99,7 @@ fn prepare_provider_request( forwarded, api_key.as_deref(), &transformed.model, + &litellm_params, &env_lookup, )?; let environment = ValidatedEnvironment { @@ -108,11 +110,13 @@ fn prepare_provider_request( auth: validated.auth, }; - let url = if transformed.params.stream == Some(true) { - config.complete_stream_url(api_base.as_deref(), &transformed.model, &env_lookup)? - } else { - config.get_complete_url(api_base.as_deref(), &transformed.model, &env_lookup)? - }; + let url = config.get_complete_url( + api_base.as_deref(), + &transformed.model, + &litellm_params, + transformed.params.stream == Some(true), + &env_lookup, + )?; Ok(ProviderMessagesRequest { provider, @@ -227,6 +231,7 @@ mod tests { api_key: None, api_base: None, custom_llm_provider: Some("anthropic".into()), + litellm_params: Default::default(), extra_headers: None, provider_specific_header: None, timeout: None, @@ -253,6 +258,7 @@ mod tests { api_key: Some("sk-test".into()), api_base: Some("https://anthropic.test".into()), custom_llm_provider: Some("anthropic".into()), + litellm_params: Default::default(), extra_headers: None, provider_specific_header: None, timeout: None, @@ -360,6 +366,7 @@ mod tests { api_key: Some("sk-test".into()), api_base: Some("https://resource.services.ai.azure.com".into()), custom_llm_provider: custom_llm_provider.map(Into::into), + litellm_params: Default::default(), extra_headers: Some(Map::from_iter([("x-priority".into(), json!("extra"))])), provider_specific_header: Some(configured), timeout: None, diff --git a/litellm-rust/crates/inference-messages/src/types.rs b/litellm-rust/crates/inference-messages/src/types.rs index 8521345b00e..669f7b283a1 100644 --- a/litellm-rust/crates/inference-messages/src/types.rs +++ b/litellm-rust/crates/inference-messages/src/types.rs @@ -2,11 +2,12 @@ use std::time::Duration; use bytes::Bytes; use litellm_host::call::CallOutput; -use litellm_llms::base_llm::messages::context::MessagesModelCapabilities; +use litellm_llms::{ErrorDetail, base_llm::messages::context::MessagesModelCapabilities}; use litellm_llms_types::{ formats::messages::{MessagesRequest, MessagesResponse}, headers::ProviderSpecificHeaders, }; +use litellm_router_types::LitellmParams; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; @@ -17,6 +18,7 @@ pub struct MessagesCall { pub api_key: Option, pub api_base: Option, pub custom_llm_provider: Option, + pub litellm_params: LitellmParams, pub extra_headers: Option>, pub provider_specific_header: Option, pub timeout: Option, @@ -27,8 +29,16 @@ pub fn messages_body(body: Map) -> Result serde_json::from_value(Value::Object(body)).map_err(invalid_request) } +/// The caller's litellm params, projected by a host from the keys [`LitellmParams::fields`] +/// names. A key present with a value of the wrong type is a request error, as it is for +/// Python's `GenericLiteLLMParams(**kwargs)`. +pub fn litellm_params(fields: Map) -> Result { + serde_json::from_value(Value::Object(fields)) + .map_err(|err| Error::InvalidRequest(ErrorDetail::invalid("litellm params", err))) +} + pub(super) fn invalid_request(err: serde_json::Error) -> Error { - Error::InvalidRequest(format!("invalid Anthropic messages request: {err}").into()) + Error::InvalidRequest(ErrorDetail::invalid("Anthropic messages request", err)) } pub type MessagesCallResponse = @@ -126,6 +136,7 @@ mod tests { supports_output_config: true, supports_sampling_params: false, supports_speed: true, + supports_mid_conversation_system: false, effort_tiers: SupportedEffortTiers { minimal: false, low: true, diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index dbec5b87d81..e9005123119 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -206,7 +206,7 @@ async fn messages_cache_identity_follows_resolved_configuration_and_request_call let cache = ScopedCache::new(cache.clone(), CacheScope::Shared); support::messages_route(secrets.clone()).with_cache(cache).execute(MessagesCall { body: serde_json::from_value(json!({"model":"anthropic/cache-test-model","messages":[{"role":"user","content":"hello"}],"max_tokens":32})).unwrap(), - api_key:None,api_base:None,custom_llm_provider:None,extra_headers:None,provider_specific_header:None,timeout:None,shaping:Default::default(), + api_key:None,api_base:None,custom_llm_provider:None,litellm_params:Default::default(),extra_headers:None,provider_specific_header:None,timeout:None,shaping:Default::default(), }, &hooks, None).await.unwrap(); } assert_eq!(hooks.calls.load(Ordering::SeqCst), 4); diff --git a/litellm-rust/crates/inference-messages/tests/messages/main.rs b/litellm-rust/crates/inference-messages/tests/messages/main.rs index 509a8b19847..8f92b31c999 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/main.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/main.rs @@ -81,6 +81,7 @@ fn call() -> MessagesCall { api_key: None, api_base: None, custom_llm_provider: Some("anthropic".into()), + litellm_params: Default::default(), extra_headers: None, provider_specific_header: None, timeout: Some(Duration::from_secs(5)), diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index 707e91c2d0d..523b25f81d6 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -58,6 +58,7 @@ async fn credentials_become_exactly_one_auth_header( run_message(MessagesCall { custom_llm_provider: Some(provider.into()), + litellm_params: Default::default(), api_key: api_key.map(Into::into), api_base: Some(upstream.uri()), extra_headers: headers(extra_headers.iter().copied()), @@ -85,6 +86,7 @@ async fn a_call_without_credentials_fails_before_sending( let error = run(MessagesCall { custom_llm_provider: Some(provider.into()), + litellm_params: Default::default(), api_base: Some(upstream.uri()), ..call }) @@ -124,6 +126,7 @@ async fn each_provider_posts_to_its_messages_endpoint( run_message(MessagesCall { custom_llm_provider: provider.map(Into::into), + litellm_params: Default::default(), api_key: Some("sk".into()), api_base: Some(format!("{}{base_suffix}", upstream.uri())), ..with_model(call, model) @@ -154,6 +157,7 @@ async fn unsupported_providers_are_rejected_before_sending( ) { let error = run(MessagesCall { custom_llm_provider: provider.map(Into::into), + litellm_params: Default::default(), api_key: Some("sk".into()), api_base: Some(UNREACHABLE_BASE.into()), ..with_model(call, model) @@ -206,6 +210,7 @@ async fn cache_scope_removal_is_selected_by_the_provider( run_message(MessagesCall { custom_llm_provider: Some(provider.into()), + litellm_params: Default::default(), api_key: Some("sk-azure".into()), api_base: Some(upstream.uri()), body: body(json!({ @@ -349,6 +354,7 @@ async fn caller_protocol_headers_win_over_the_defaults(call: MessagesCall, #[cas run_message(MessagesCall { custom_llm_provider: Some(provider.into()), + litellm_params: Default::default(), api_key: Some("sk".into()), api_base: Some(upstream.uri()), extra_headers: headers([ @@ -400,6 +406,7 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i }, body: call.body.clone(), custom_llm_provider: call.custom_llm_provider.clone(), + litellm_params: Default::default(), extra_headers: None, provider_specific_header: None, timeout: call.timeout, @@ -577,6 +584,7 @@ async fn metadata_is_reduced_to_the_user_id(call: MessagesCall, #[case] provider run_message(with_fields( MessagesCall { custom_llm_provider: Some(provider.into()), + litellm_params: Default::default(), api_key: Some("sk".into()), api_base: Some(upstream.uri()), ..call @@ -638,6 +646,7 @@ async fn system_message_folding_is_selected_by_the_provider( run_message(with_fields( MessagesCall { custom_llm_provider: Some(provider.into()), + litellm_params: Default::default(), api_key: Some("sk-azure".into()), api_base: Some(upstream.uri()), shaping: MessagesShaping { @@ -702,6 +711,7 @@ async fn provider_validation_runs_before_caller_parameter_removal( let result = run(with_fields( MessagesCall { custom_llm_provider: Some(provider.into()), + litellm_params: Default::default(), api_key: Some("sk-test".into()), api_base: Some(upstream.uri()), shaping: MessagesShaping { diff --git a/litellm-rust/crates/inference-messages/tests/messages/response.rs b/litellm-rust/crates/inference-messages/tests/messages/response.rs index 13b99d45241..b5d623a5582 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/response.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/response.rs @@ -104,6 +104,7 @@ async fn the_provider_message_is_returned(call: MessagesCall, #[case] provider: let message = run_message(MessagesCall { custom_llm_provider: Some(provider.into()), + litellm_params: Default::default(), api_key: Some("sk".into()), api_base: Some(upstream.uri()), ..call diff --git a/litellm-rust/crates/inference-messages/tests/messages/secrets.rs b/litellm-rust/crates/inference-messages/tests/messages/secrets.rs index c565edac9dc..a33fdd130f1 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/secrets.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/secrets.rs @@ -38,6 +38,7 @@ async fn the_credential_and_base_come_from_the_secret_source( secrets.clone(), MessagesCall { custom_llm_provider: Some(provider.into()), + litellm_params: Default::default(), ..call }, ) @@ -189,6 +190,7 @@ async fn azure_without_a_base_anywhere_fails_before_sending(call: MessagesCall) Arc::new(RecordingSecrets::new([("AZURE_API_KEY", "sk-azure")])), MessagesCall { custom_llm_provider: Some("azure_ai".into()), + litellm_params: Default::default(), ..call }, ) diff --git a/litellm-rust/crates/inference-messages/tests/messages/stream.rs b/litellm-rust/crates/inference-messages/tests/messages/stream.rs index 4dec4de2c02..049a620e454 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/stream.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/stream.rs @@ -346,6 +346,7 @@ async fn the_sdk_returns_stream_headers_and_every_sse_byte( .execute( MessagesCall { custom_llm_provider: Some(provider.into()), + litellm_params: Default::default(), ..streaming(call, upstream.uri()) }, &(), @@ -470,6 +471,7 @@ async fn a_host_on_anthropic_sse_is_relayed_byte_for_byte(call: MessagesCall) { let host = RecordingStreamHost::new( MessagesCall { custom_llm_provider: Some("azure_ai".into()), + litellm_params: Default::default(), ..streaming(call, upstream.uri()) }, usize::MAX, diff --git a/litellm-rust/crates/inference-ocr/src/prepare.rs b/litellm-rust/crates/inference-ocr/src/prepare.rs index e3b0e819fb2..14120606864 100644 --- a/litellm-rust/crates/inference-ocr/src/prepare.rs +++ b/litellm-rust/crates/inference-ocr/src/prepare.rs @@ -21,14 +21,7 @@ pub(crate) fn prepare_request( Some("MISTRAL_AZURE_API_BASE"), ), LlmProviders::AzureAi => (None, Some("AZURE_AI_API_BASE")), - LlmProviders::Anthropic - | LlmProviders::AwsTextract - | LlmProviders::Bedrock - | LlmProviders::Cohere - | LlmProviders::Openai - | LlmProviders::OpenaiLike - | LlmProviders::Reducto - | LlmProviders::VertexAi => (None, None), + _ => (None, None), }; let secret = |name: &str| secrets.truthy(name); let dynamic_api_key = credentials.dynamic_api_key.or_else(|| { diff --git a/litellm-rust/crates/inference-ocr/src/provider_config.rs b/litellm-rust/crates/inference-ocr/src/provider_config.rs index 28cb96238a3..da708d17b86 100644 --- a/litellm-rust/crates/inference-ocr/src/provider_config.rs +++ b/litellm-rust/crates/inference-ocr/src/provider_config.rs @@ -225,10 +225,7 @@ pub(crate) fn resolve_provider_config( OcrConfigKind::VertexDeepSeek } LlmProviders::VertexAi => OcrConfigKind::VertexAi, - LlmProviders::Anthropic - | LlmProviders::Bedrock - | LlmProviders::Openai - | LlmProviders::OpenaiLike => { + _ => { return Err(Error::InvalidProvider( provider.custom_llm_provider.to_string(), )); diff --git a/litellm-rust/crates/inference-transcription/src/prepare.rs b/litellm-rust/crates/inference-transcription/src/prepare.rs index 91b2d11e43a..15ca7627339 100644 --- a/litellm-rust/crates/inference-transcription/src/prepare.rs +++ b/litellm-rust/crates/inference-transcription/src/prepare.rs @@ -17,15 +17,7 @@ use litellm_inference::provider::resolve_llm_provider; fn provider_config(provider: LlmProviders) -> Option<&'static dyn BaseAudioTranscriptionConfig> { match provider { LlmProviders::Bedrock => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG), - LlmProviders::Anthropic - | LlmProviders::AwsTextract - | LlmProviders::AzureAi - | LlmProviders::Cohere - | LlmProviders::Mistral - | LlmProviders::Openai - | LlmProviders::OpenaiLike - | LlmProviders::Reducto - | LlmProviders::VertexAi => None, + _ => None, } } diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index ab9c8bc5646..2d3ea85d1fc 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -20,6 +20,7 @@ litellm-framer.workspace = true litellm-http.workspace = true litellm-secrets.workspace = true litellm-python-compat.workspace = true +litellm-router-types.workspace = true base64.workspace = true bytes.workspace = true data-url = "0.3.2" diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index ccf3a6d15f2..335731fa001 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -7,6 +7,7 @@ use litellm_http::request::{ use litellm_llms_types::{ formats::messages::{ ContentBlock, ContentBlockType, EffortLevel, Message, MessageContent, MessagesTool, + SystemPrompt, }, providers::anthropic::{AnthropicBeta, BetaSet}, recognized::Recognized, @@ -385,6 +386,33 @@ pub fn is_encrypted_reasoning_block(block: &ContentBlock) -> bool { field.is_some_and(|value| value.starts_with(ENCRYPTED_REASONING_SIGNATURE_PREFIX)) } +const BILLING_HEADER_PREFIX: &str = "x-anthropic-billing-header:"; + +fn is_billing_header_block(block: &ContentBlock) -> bool { + block.block_type == Some(ContentBlockType::Text) + && block + .text + .as_deref() + .is_some_and(|text| text.starts_with(BILLING_HEADER_PREFIX)) +} + +/// Python's `AnthropicMessagesConfig._filter_billing_headers_from_system`: the Claude Code +/// attribution blocks the first-party API reads, dropped for hosts that reject them. `None` +/// when nothing else was in the system prompt. +pub fn filter_billing_headers_from_system(system: SystemPrompt) -> Option { + match system { + SystemPrompt::Text(text) if text.starts_with(BILLING_HEADER_PREFIX) => None, + SystemPrompt::Text(text) => Some(SystemPrompt::Text(text)), + SystemPrompt::Blocks(blocks) => { + let kept: Vec = blocks + .into_iter() + .filter(|block| !is_billing_header_block(block)) + .collect(); + (!kept.is_empty()).then_some(SystemPrompt::Blocks(kept)) + } + } +} + pub fn strip_encrypted_reasoning_blocks(messages: Vec) -> Vec { retain_blocks(messages, |block| !is_encrypted_reasoning_block(block)) } @@ -1858,6 +1886,7 @@ mod tests { supports_output_config: false, supports_sampling_params: true, supports_speed: false, + supports_mid_conversation_system: false, effort_tiers: tiers(false, false, false, false, false, false), } ); diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 74ac35752ab..7021a4199df 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -3,9 +3,10 @@ use litellm_llms_types::{ formats::messages::{ ContextEdit, ContextManagement, Message, MessagesOptionalParams, MessagesRequest, Speed, }, - providers::anthropic::{AnthropicBeta, BetaSet}, + providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet}, recognized::Recognized, }; +use litellm_router_types::LitellmParams; use serde_json::{Map, Value, json}; use super::{handler::shape_anthropic_messages_request, thinking::translate_thinking}; @@ -47,6 +48,8 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { &self, api_base: Option<&str>, _model: &str, + _litellm_params: &LitellmParams, + _stream: bool, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { Ok(complete_anthropic_url(api_base, env_lookup)) @@ -76,6 +79,7 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { headers: Headers, api_key: Option<&str>, _model: &str, + _litellm_params: &LitellmParams, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { let headers = match optionally_handle_anthropic_oauth(headers, api_key) { @@ -150,6 +154,15 @@ pub(crate) fn update_headers_with_anthropic_beta( merge_beta_headers(headers, feature_betas(request)) } +/// The betas a request needs on a host other than the first-party API: the Anthropic set +/// filtered and renamed through the host's `BetaProvider` policy. +pub fn provider_feature_betas(request: &MessagesRequest, provider: BetaProvider) -> BetaSet { + feature_betas(request) + .iter() + .filter_map(|beta| beta.on(provider)) + .collect() +} + fn feature_betas(request: &MessagesRequest) -> BetaSet { let params = &request.params; let tools = params.tools.as_deref(); @@ -815,7 +828,13 @@ mod tests { #[case] expected: &str, ) { assert_eq!( - ANTHROPIC_MESSAGES_CONFIG.get_complete_url(api_base, "claude", &env(vars)), + ANTHROPIC_MESSAGES_CONFIG.get_complete_url( + api_base, + "claude", + &LitellmParams::default(), + false, + &env(vars), + ), Ok(expected.to_string()) ); } @@ -833,6 +852,7 @@ mod tests { headers(forwarded), api_key, "claude", + &LitellmParams::default(), &env(vars), ) } @@ -1108,8 +1128,20 @@ mod tests { requested.borrow_mut().push(name.to_string()); None }; - let _ = ANTHROPIC_MESSAGES_CONFIG.validate_environment(Vec::new(), None, "claude", &record); - let _ = ANTHROPIC_MESSAGES_CONFIG.get_complete_url(None, "claude", &record); + let _ = ANTHROPIC_MESSAGES_CONFIG.validate_environment( + Vec::new(), + None, + "claude", + &LitellmParams::default(), + &record, + ); + let _ = ANTHROPIC_MESSAGES_CONFIG.get_complete_url( + None, + "claude", + &LitellmParams::default(), + false, + &record, + ); let requested = requested.into_inner(); assert!(!requested.is_empty()); let undeclared: Vec<&String> = requested diff --git a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs index ffd079b2d04..3e6abf2d14e 100644 --- a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs @@ -1,9 +1,7 @@ use litellm_auth::{CredentialPlacement, SecretValue}; use litellm_http::request::{has_bearer_auth, has_header}; -use litellm_llms_types::formats::messages::{ - CacheControl, ContentBlock, Message, MessageContent, MessagesOptionalParams, MessagesRequest, - SystemPrompt, -}; +use litellm_llms_types::formats::messages::MessagesRequest; +use litellm_router_types::LitellmParams; use crate::{ Error, @@ -20,7 +18,7 @@ use crate::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, messages::{ context::MessagesTransformContext, - normalization::fold_system_role_messages, + normalization::{normalize_system_role_messages, strip_cache_control_scope}, transformation::{BaseMessagesConfig, MESSAGES_PATH_SUFFIX}, }, }, @@ -47,6 +45,8 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig { &self, api_base: Option<&str>, _model: &str, + _litellm_params: &LitellmParams, + _stream: bool, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { complete_azure_anthropic_url(api_base, env_lookup) @@ -57,20 +57,14 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig { request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { - let request = fold_system_role_messages(request); transform_messages_request( - MessagesRequest { - messages: request - .messages - .into_iter() - .map(strip_scope_from_message) - .collect(), - params: MessagesOptionalParams { - system: request.params.system.map(strip_scope_from_system), - ..request.params - }, - ..request - }, + strip_cache_control_scope(normalize_system_role_messages( + request, + context + .thinking + .capabilities + .supports_mid_conversation_system, + )), context, ) } @@ -86,6 +80,7 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig { headers: Headers, api_key: Option<&str>, _model: &str, + _litellm_params: &LitellmParams, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { if has_header(&headers, API_KEY_PLACEMENT.header_name()) || has_bearer_auth(&headers) { @@ -129,37 +124,6 @@ pub fn complete_azure_anthropic_url( Ok(format!("{with_anthropic}{MESSAGES_PATH_SUFFIX}")) } -fn strip_scope_from_block(block: ContentBlock) -> ContentBlock { - ContentBlock { - cache_control: block.cache_control.map(|cache_control| CacheControl { - scope: None, - ..cache_control - }), - ..block - } -} - -fn strip_scope_from_system(system: SystemPrompt) -> SystemPrompt { - match system { - SystemPrompt::Blocks(blocks) => { - SystemPrompt::Blocks(blocks.into_iter().map(strip_scope_from_block).collect()) - } - text => text, - } -} - -fn strip_scope_from_message(message: Message) -> Message { - Message { - content: match message.content { - MessageContent::Blocks(blocks) => { - MessageContent::Blocks(blocks.into_iter().map(strip_scope_from_block).collect()) - } - text => text, - }, - ..message - } -} - #[cfg(test)] mod tests { use litellm_llms_types::formats::messages::MessagesResponse; @@ -169,7 +133,9 @@ mod tests { use litellm_auth::CredentialPlacement; use super::*; - use crate::base_llm::messages::context::MessagesModelCapabilities; + use crate::base_llm::messages::{ + context::MessagesModelCapabilities, mid_conversation_system::CONVERTED_SYSTEM_NOTE, + }; fn request_from(value: serde_json::Value) -> MessagesRequest { serde_json::from_value(value).expect("valid request") @@ -268,6 +234,7 @@ mod tests { .collect(), api_key, "claude", + &LitellmParams::default(), &|_| None, ) .unwrap() @@ -427,62 +394,101 @@ mod tests { assert_eq!(transformed, body); } - #[test] - fn transform_request_folds_system_role_message_into_top_level_system() { - let request = request_from(json!({ - "model": "claude-sonnet-4-5", - "max_tokens": 256, - "system": [{"type": "text", "text": "base system"}], - "messages": [ - {"role": "user", "content": "fix the bug"}, - {"role": "system", "content": "Available agent types: claude"} - ] - })); - - let transformed = to_value( - AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_anthropic_messages_request(request, &MessagesTransformContext::default()) - .expect("request transforms"), - ); - - assert_eq!( - transformed["messages"], - json!([{"role": "user", "content": "fix the bug"}]) - ); - assert_eq!( - transformed["system"], - json!([ - {"type": "text", "text": "base system"}, - {"type": "text", "text": "Available agent types: claude"} - ]) - ); + fn context_with(supports_mid_conversation_system: bool) -> MessagesTransformContext { + MessagesTransformContext::new( + MessagesModelCapabilities { + supports_mid_conversation_system, + ..MessagesModelCapabilities::default() + }, + false, + ) } - #[test] - fn transform_request_folds_system_role_when_no_top_level_system() { - let request = request_from(json!({ - "model": "claude-sonnet-4-5", - "max_tokens": 256, + #[rstest] + #[case::leading_turn_joins_the_top_level_system( + json!({ + "system": [{"type": "text", "text": "base system"}], "messages": [ - {"role": "user", "content": [{"type": "text", "text": "hi"}]}, - {"role": "system", "content": [{"type": "text", "text": "sys block"}]} + {"role": "system", "content": "Available agent types: claude"}, + {"role": "user", "content": "fix the bug"} ] - })); + }), + false, + json!([{"type": "text", "text": "base system"}, {"type": "text", "text": "Available agent types: claude"}]), + json!([{"role": "user", "content": "fix the bug"}]), + )] + #[case::leading_turn_becomes_the_system_when_there_is_none( + json!({"messages": [ + {"role": "system", "content": [{"type": "text", "text": "sys block"}]}, + {"role": "user", "content": [{"type": "text", "text": "hi"}]} + ]}), + false, + json!([{"type": "text", "text": "sys block"}]), + json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]), + )] + #[case::billing_block_in_a_leading_turn_is_dropped( + json!({"messages": [ + {"role": "system", "content": [{"type": "text", "text": "x-anthropic-billing-header: cc_version=1"}, {"type": "text", "text": "keep"}]}, + {"role": "user", "content": "hi"} + ]}), + false, + json!([{"type": "text", "text": "keep"}]), + json!([{"role": "user", "content": "hi"}]), + )] + #[case::later_turn_is_converted_in_place_for_a_model_without_the_capability( + json!({ + "system": "be terse", + "messages": [ + {"role": "user", "content": "fix the bug"}, + {"role": "system", "content": "reminder"} + ] + }), + false, + json!("be terse"), + json!([ + {"role": "user", "content": "fix the bug"}, + {"role": "user", "content": [ + {"type": "text", "text": CONVERTED_SYSTEM_NOTE}, + {"type": "text", "text": "reminder"} + ]} + ]), + )] + #[case::later_turn_stays_for_a_model_with_the_capability( + json!({ + "system": "be terse", + "messages": [ + {"role": "user", "content": "fix the bug"}, + {"role": "system", "content": "reminder"} + ] + }), + true, + json!("be terse"), + json!([ + {"role": "user", "content": "fix the bug"}, + {"role": "system", "content": "reminder"} + ]), + )] + fn transform_request_normalizes_system_role_messages_like_python( + #[case] body: serde_json::Value, + #[case] supports_mid_conversation_system: bool, + #[case] system: serde_json::Value, + #[case] messages: serde_json::Value, + ) { + let mut body = body; + body["model"] = json!("claude-sonnet-4-5"); + body["max_tokens"] = json!(256); let transformed = to_value( AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_anthropic_messages_request(request, &MessagesTransformContext::default()) + .transform_anthropic_messages_request( + request_from(body), + &context_with(supports_mid_conversation_system), + ) .expect("request transforms"), ); - assert_eq!( - transformed["messages"], - json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) - ); - assert_eq!( - transformed["system"], - json!([{"type": "text", "text": "sys block"}]) - ); + assert_eq!(transformed["system"], system); + assert_eq!(transformed["messages"], messages); } #[test] @@ -607,9 +613,16 @@ mod tests { Vec::new(), None, "claude", + &LitellmParams::default(), + &record, + ); + let _ = AZURE_ANTHROPIC_MESSAGES_CONFIG.get_complete_url( + None, + "claude", + &LitellmParams::default(), + false, &record, ); - let _ = AZURE_ANTHROPIC_MESSAGES_CONFIG.get_complete_url(None, "claude", &record); let requested = requested.into_inner(); assert!(!requested.is_empty()); let undeclared: Vec<&String> = requested diff --git a/litellm-rust/crates/llms/src/base_llm/auth.rs b/litellm-rust/crates/llms/src/base_llm/auth.rs index 92a19e3ef2a..de271a269d3 100644 --- a/litellm-rust/crates/llms/src/base_llm/auth.rs +++ b/litellm-rust/crates/llms/src/base_llm/auth.rs @@ -7,6 +7,7 @@ use litellm_auth::{AuthServices, CredentialPlacement, SecretValue, TokenProviderHandle}; use litellm_auth_aws::{AwsCredentialSource, SigV4Signer}; +use litellm_auth_gcp::VertexConfig; use litellm_http::request::with_header; pub type Headers = Vec<(String, String)>; @@ -24,6 +25,9 @@ pub enum AuthScheme { /// A bearer acquired when the request is sent, from a token source such as a cloud SDK /// or a caller-supplied callable. Token { provider: TokenProviderHandle }, + /// A Google access token for the Vertex AI project the config names, acquired when the + /// request is sent through the shared GCP token cache. + GcpAccessToken { config: Box }, /// AWS SigV4 over the bytes that go on the wire, so the handler signs after the body is /// serialized. AwsSigV4 { @@ -73,6 +77,13 @@ pub async fn resolve_auth( signer: None, }) } + AuthScheme::GcpAccessToken { config } => { + let token = services.gcp.access_token(&config, env_lookup).await?; + Ok(Authenticated { + headers: with_credential(headers, CredentialPlacement::Bearer, &token), + signer: None, + }) + } AuthScheme::AwsSigV4 { region, service, diff --git a/litellm-rust/crates/llms/src/base_llm/messages/context.rs b/litellm-rust/crates/llms/src/base_llm/messages/context.rs index a20a98f56b1..e59a63813b0 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/context.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/context.rs @@ -40,6 +40,8 @@ pub struct MessagesModelCapabilities { #[serde(default)] pub supports_speed: bool, #[serde(default)] + pub supports_mid_conversation_system: bool, + #[serde(default)] pub effort_tiers: SupportedEffortTiers, } @@ -57,6 +59,7 @@ impl Default for MessagesModelCapabilities { supports_output_config: false, supports_sampling_params: true, supports_speed: false, + supports_mid_conversation_system: false, effort_tiers: SupportedEffortTiers::default(), } } diff --git a/litellm-rust/crates/llms/src/base_llm/messages/mid_conversation_system.rs b/litellm-rust/crates/llms/src/base_llm/messages/mid_conversation_system.rs new file mode 100644 index 00000000000..ec1a9863786 --- /dev/null +++ b/litellm-rust/crates/llms/src/base_llm/messages/mid_conversation_system.rs @@ -0,0 +1,201 @@ +use litellm_llms_types::formats::messages::{ + ContentBlock, ContentBlockType, Message, MessageContent, SystemPrompt, +}; +use serde_json::Map; + +/// Python's `CONVERTED_SYSTEM_NOTE`. +pub const CONVERTED_SYSTEM_NOTE: &str = "Operator note (not from the user): the following was originally a mid-conversation system-role reminder."; + +const SYSTEM_ROLE: &str = "system"; +const USER_ROLE: &str = "user"; + +pub fn as_system_content_blocks(value: Option) -> Vec { + match value { + None => Vec::new(), + Some(SystemPrompt::Text(text)) => vec![ContentBlock::text(text)], + Some(SystemPrompt::Blocks(blocks)) => blocks, + } +} + +pub fn message_content_blocks(content: MessageContent) -> Vec { + match content { + MessageContent::Text(text) => vec![ContentBlock::text(text)], + MessageContent::Blocks(blocks) => blocks, + } +} + +pub fn is_system_role_message(message: &Message) -> bool { + message.role == SYSTEM_ROLE +} + +pub fn system_role_message_as_user(message: Message) -> Message { + Message { + role: USER_ROLE.into(), + content: MessageContent::Blocks( + [ContentBlock::text(CONVERTED_SYSTEM_NOTE)] + .into_iter() + .chain(message_content_blocks(message.content)) + .collect(), + ), + extra: Map::new(), + } +} + +pub fn opens_with_tool_results(message: &Message) -> bool { + message.role == USER_ROLE + && matches!( + &message.content, + MessageContent::Blocks(blocks) + if blocks.first().and_then(|block| block.block_type.as_ref()) == Some(&ContentBlockType::ToolResult) + ) +} + +fn system_run_placed_after_tool_results( + system_run: Vec, + follower_run: Vec, +) -> Vec { + let mut follower_run = follower_run.into_iter(); + match follower_run.next() { + Some(first) if opens_with_tool_results(&first) => [first] + .into_iter() + .chain(system_run) + .chain(follower_run) + .collect(), + first => system_run + .into_iter() + .chain(first) + .chain(follower_run) + .collect(), + } +} + +fn runs(messages: Vec) -> Vec> { + messages + .into_iter() + .fold(Vec::new(), |mut runs: Vec>, message| { + match runs.last_mut() { + Some(run) + if is_system_role_message(&run[0]) == is_system_role_message(&message) => + { + run.push(message); + } + _ => runs.push(vec![message]), + } + runs + }) +} + +fn system_turns_after_tool_results(messages: Vec) -> Vec { + let mut runs = runs(messages).into_iter(); + let mut placed = Vec::new(); + while let Some(run) = runs.next() { + if !is_system_role_message(&run[0]) { + placed.extend(run); + continue; + } + let follower_run = runs.next().unwrap_or_default(); + placed.extend(system_run_placed_after_tool_results(run, follower_run)); + } + placed +} + +pub fn convert_mid_conversation_system_turns(messages: Vec) -> Vec { + system_turns_after_tool_results(messages) + .into_iter() + .map(|message| { + if is_system_role_message(&message) { + system_role_message_as_user(message) + } else { + message + } + }) + .collect() +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::{Value, json}; + + use super::*; + + fn messages(value: Value) -> Vec { + serde_json::from_value(value).unwrap() + } + + fn roles(messages: &[Message]) -> Vec<&str> { + messages + .iter() + .map(|message| message.role.as_str()) + .collect() + } + + #[rstest] + #[case::no_system_turns( + json!([{"role": "user", "content": "a"}, {"role": "assistant", "content": "b"}]), + vec!["user", "assistant"], + )] + #[case::a_turn_between_messages_becomes_a_user_turn_in_place( + json!([ + {"role": "user", "content": "a"}, + {"role": "assistant", "content": "b"}, + {"role": "system", "content": "reminder"}, + {"role": "user", "content": "c"} + ]), + vec!["user", "assistant", "user", "user"], + )] + #[case::a_run_between_a_tool_use_and_its_result_moves_after_the_result( + json!([ + {"role": "user", "content": "a"}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "t", "name": "n", "input": {}}]}, + {"role": "system", "content": "reminder"}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t", "content": "ok"}, {"type": "text", "text": "next"}]}, + {"role": "user", "content": "d"} + ]), + vec!["user", "assistant", "user", "user", "user"], + )] + fn turns_are_converted_in_place_and_never_split_a_tool_call_from_its_result( + #[case] input: Value, + #[case] expected_roles: Vec<&str>, + ) { + let converted = convert_mid_conversation_system_turns(messages(input)); + + assert_eq!(roles(&converted), expected_roles); + assert!( + converted + .iter() + .all(|message| !is_system_role_message(message)) + ); + } + + #[rstest] + fn a_converted_turn_carries_the_note_then_the_original_content_only() { + let converted = convert_mid_conversation_system_turns(messages(json!([ + {"role": "user", "content": "a"}, + {"role": "system", "content": "reminder", "name": "ignored"} + ]))); + + assert_eq!( + serde_json::to_value(&converted[1]).unwrap(), + json!({"role": "user", "content": [ + {"type": "text", "text": CONVERTED_SYSTEM_NOTE}, + {"type": "text", "text": "reminder"} + ]}) + ); + } + + #[rstest] + fn the_tool_result_turn_keeps_its_place_before_the_moved_run() { + let converted = convert_mid_conversation_system_turns(messages(json!([ + {"role": "assistant", "content": [{"type": "tool_use", "id": "t", "name": "n", "input": {}}]}, + {"role": "system", "content": "reminder"}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t", "content": "ok"}]} + ]))); + + assert!(opens_with_tool_results(&converted[1])); + assert_eq!( + serde_json::to_value(&converted[2]).unwrap()["content"][1], + json!({"type": "text", "text": "reminder"}) + ); + } +} diff --git a/litellm-rust/crates/llms/src/base_llm/messages/mod.rs b/litellm-rust/crates/llms/src/base_llm/messages/mod.rs index ac2112176f9..e9591daa9da 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/mod.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/mod.rs @@ -1,4 +1,5 @@ pub mod context; +pub mod mid_conversation_system; pub mod normalization; pub mod streaming; pub mod transformation; diff --git a/litellm-rust/crates/llms/src/base_llm/messages/normalization.rs b/litellm-rust/crates/llms/src/base_llm/messages/normalization.rs index bddd829304e..3ae4145ae32 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/normalization.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/normalization.rs @@ -1,47 +1,103 @@ use litellm_llms_types::formats::messages::{ - ContentBlock, Message, MessageContent, MessagesOptionalParams, MessagesRequest, SystemPrompt, + CacheControl, ContentBlock, Message, MessageContent, MessagesOptionalParams, MessagesRequest, + SystemPrompt, }; -const SYSTEM_ROLE: &str = "system"; +use crate::{ + anthropic::common_utils::filter_billing_headers_from_system, + base_llm::messages::mid_conversation_system::{ + as_system_content_blocks, convert_mid_conversation_system_turns, is_system_role_message, + message_content_blocks, + }, +}; -fn content_into_blocks(content: MessageContent) -> Vec { - match content { - MessageContent::Text(text) => vec![ContentBlock::text(text)], - MessageContent::Blocks(blocks) => blocks, - } -} - -fn system_into_blocks(system: Option) -> Vec { - match system { - None => Vec::new(), - Some(SystemPrompt::Text(text)) => vec![ContentBlock::text(text)], - Some(SystemPrompt::Blocks(blocks)) => blocks, - } -} - -pub fn fold_system_role_messages(request: MessagesRequest) -> MessagesRequest { - if !request.messages.iter().any(|msg| msg.role == SYSTEM_ROLE) { - return request; - } - - let (system_messages, chat_messages): (Vec, Vec) = request +/// Python's `_normalize_system_role_messages`: a leading run of `role: "system"` entries is +/// hoisted into the top-level `system`, a later one stays in place when the model accepts it +/// and otherwise becomes a user turn where it was, and billing blocks are stripped from the +/// top-level `system` either way. +pub fn normalize_system_role_messages( + request: MessagesRequest, + supports_mid_conversation_system: bool, +) -> MessagesRequest { + let leading_count = request .messages - .into_iter() - .partition(|msg| msg.role == SYSTEM_ROLE); - - let folded_system: Vec = system_into_blocks(request.params.system) - .into_iter() - .chain( - system_messages + .iter() + .take_while(|message| is_system_role_message(message)) + .count(); + let mut messages = request.messages.into_iter(); + let hoisted: Vec = messages.by_ref().take(leading_count).collect(); + let remaining: Vec = messages.collect(); + let remaining = if supports_mid_conversation_system { + remaining + } else { + convert_mid_conversation_system_turns(remaining) + }; + let system = if hoisted.is_empty() { + request.params.system + } else { + Some(SystemPrompt::Blocks( + as_system_content_blocks(request.params.system) .into_iter() - .flat_map(|msg| content_into_blocks(msg.content)), - ) - .collect(); - + .chain( + hoisted + .into_iter() + .flat_map(|message| message_content_blocks(message.content)), + ) + .collect(), + )) + }; MessagesRequest { - messages: chat_messages, + messages: remaining, params: MessagesOptionalParams { - system: (!folded_system.is_empty()).then_some(SystemPrompt::Blocks(folded_system)), + system: system.and_then(filter_billing_headers_from_system), + ..request.params + }, + ..request + } +} + +fn strip_scope_from_block(block: ContentBlock) -> ContentBlock { + ContentBlock { + cache_control: block.cache_control.map(|cache_control| CacheControl { + scope: None, + ..cache_control + }), + ..block + } +} + +fn strip_scope_from_system(system: SystemPrompt) -> SystemPrompt { + match system { + SystemPrompt::Blocks(blocks) => { + SystemPrompt::Blocks(blocks.into_iter().map(strip_scope_from_block).collect()) + } + text => text, + } +} + +fn strip_scope_from_message(message: Message) -> Message { + Message { + content: match message.content { + MessageContent::Blocks(blocks) => { + MessageContent::Blocks(blocks.into_iter().map(strip_scope_from_block).collect()) + } + text => text, + }, + ..message + } +} + +/// Python's `_remove_scope_from_cache_control`: hosts other than the first-party API reject +/// `cache_control.scope`, so it is dropped from the system prompt and every message block. +pub fn strip_cache_control_scope(request: MessagesRequest) -> MessagesRequest { + MessagesRequest { + messages: request + .messages + .into_iter() + .map(strip_scope_from_message) + .collect(), + params: MessagesOptionalParams { + system: request.params.system.map(strip_scope_from_system), ..request.params }, ..request diff --git a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs index 2f2d3bcf909..f42adcdcfe6 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs @@ -1,4 +1,6 @@ use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse}; +use litellm_router_types::LitellmParams; +use serde_json::Value; use super::context::MessagesTransformContext; @@ -20,18 +22,11 @@ pub trait BaseMessagesConfig: Sync { &self, api_base: Option<&str>, model: &str, + litellm_params: &LitellmParams, + stream: bool, env_lookup: &dyn Fn(&str) -> Option, ) -> Result; - fn complete_stream_url( - &self, - api_base: Option<&str>, - model: &str, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - self.get_complete_url(api_base, model, env_lookup) - } - fn transform_anthropic_messages_request( &self, request: MessagesRequest, @@ -58,6 +53,7 @@ pub trait BaseMessagesConfig: Sync { headers: Headers, api_key: Option<&str>, model: &str, + litellm_params: &LitellmParams, env_lookup: &dyn Fn(&str) -> Option, ) -> Result; @@ -75,6 +71,12 @@ pub trait BaseMessagesConfig: Sync { fn request_headers(&self, headers: Headers, _request: &MessagesRequest) -> Headers { headers } + + /// The JSON that goes on the wire, for a host whose body differs from the typed request: + /// Python's configs `pop("model")` when the model is addressed by the URL. + fn wire_body(&self, body: Value) -> Value { + body + } } #[cfg(test)] @@ -94,6 +96,8 @@ mod tests { &self, _api_base: Option<&str>, _model: &str, + _litellm_params: &LitellmParams, + _stream: bool, _env_lookup: &dyn Fn(&str) -> Option, ) -> Result { Ok(String::new()) @@ -104,6 +108,7 @@ mod tests { headers: Headers, _api_key: Option<&str>, _model: &str, + _litellm_params: &LitellmParams, _env_lookup: &dyn Fn(&str) -> Option, ) -> Result { Ok(ValidatedEnvironment { diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs index c5bf7d8fc95..59452e0b693 100644 --- a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs @@ -19,6 +19,7 @@ use litellm_llms_types::formats::messages::{ MessagesRequest, streaming::{MessagesStreamEvent, MessagesStreamUsage}, }; +use litellm_router_types::LitellmParams; use serde_json::{Map, Value}; use crate::{ @@ -74,16 +75,26 @@ fn bearer_token( fn invoke_url( api_base: Option<&str>, model: &str, + aws: &AwsParams, + stream: bool, env_lookup: &dyn Fn(&str) -> Option, - path: &str, ) -> String { let (model_id, model_region) = bedrock_model_id_and_region(model.strip_prefix(INVOKE_MODEL_PREFIX).unwrap_or(model)); - let region = resolve_bedrock_region(model_region.as_deref(), &AwsParams::default(), env_lookup); - let endpoint = api_base - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) + let region = resolve_bedrock_region(model_region.as_deref(), aws, env_lookup); + let path = if stream { + INVOKE_STREAM_PATH + } else { + INVOKE_PATH + }; + let configured = |value: Option<&str>| { + value + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + }; + let endpoint = configured(api_base) + .or_else(|| configured(aws.aws_bedrock_runtime_endpoint.as_deref())) .or_else(|| env_lookup(AWS_BEDROCK_RUNTIME_ENDPOINT)) .unwrap_or_else(|| BEDROCK_RUNTIME_ENDPOINT_TEMPLATE.replace("{region}", ®ion)); format!("{}/model/{model_id}/{path}", endpoint.trim_end_matches('/')) @@ -102,18 +113,17 @@ impl BaseMessagesConfig for AmazonAnthropicClaudeMessagesConfig { &self, api_base: Option<&str>, model: &str, + litellm_params: &LitellmParams, + stream: bool, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { - Ok(invoke_url(api_base, model, env_lookup, INVOKE_PATH)) - } - - fn complete_stream_url( - &self, - api_base: Option<&str>, - model: &str, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - Ok(invoke_url(api_base, model, env_lookup, INVOKE_STREAM_PATH)) + Ok(invoke_url( + api_base, + model, + &litellm_params.aws, + stream, + env_lookup, + )) } fn transform_anthropic_messages_request( @@ -137,6 +147,7 @@ impl BaseMessagesConfig for AmazonAnthropicClaudeMessagesConfig { headers: Headers, api_key: Option<&str>, model: &str, + litellm_params: &LitellmParams, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { if let Some(token) = bearer_token(api_key, env_lookup) { @@ -150,13 +161,13 @@ impl BaseMessagesConfig for AmazonAnthropicClaudeMessagesConfig { } let (_, model_region) = bedrock_model_id_and_region(model.strip_prefix(INVOKE_MODEL_PREFIX).unwrap_or(model)); - let params = AwsParams::default(); + let params = &litellm_params.aws; Ok(ValidatedEnvironment { headers, auth: AuthScheme::AwsSigV4 { - region: resolve_bedrock_region(model_region.as_deref(), ¶ms, env_lookup), + region: resolve_bedrock_region(model_region.as_deref(), params, env_lookup), service: BEDROCK_SERVICE, - credentials: Box::new(AwsCredentialSource::from_params(¶ms, env_lookup)), + credentials: Box::new(AwsCredentialSource::from_params(params, env_lookup)), }, }) } @@ -503,19 +514,129 @@ mod tests { assert_eq!(sse, expected); } - #[test] - fn config_uses_the_streaming_url_only_for_streams() { - let env = |_: &str| -> Option { None }; - let config = AmazonAnthropicClaudeMessagesConfig; + fn in_region(region: &str) -> LitellmParams { + LitellmParams { + aws: AwsParams { + aws_region_name: Some(region.into()), + ..AwsParams::default() + }, + ..LitellmParams::default() + } + } + fn with_runtime_endpoint(endpoint: &str) -> LitellmParams { + LitellmParams { + aws: AwsParams { + aws_bedrock_runtime_endpoint: Some(endpoint.into()), + ..in_region("us-west-2").aws + }, + ..LitellmParams::default() + } + } + + #[rstest] + #[case::invoke_in_the_params_region( + None, + in_region("us-west-2"), + false, + None, + "https://bedrock-runtime.us-west-2.amazonaws.com/model/anthropic.claude-3/invoke" + )] + #[case::stream_in_the_params_region( + None, + in_region("us-west-2"), + true, + None, + "https://bedrock-runtime.us-west-2.amazonaws.com/model/anthropic.claude-3/invoke-with-response-stream" + )] + #[case::api_base_outranks_the_params_endpoint( + Some("https://base.test/"), + with_runtime_endpoint("https://params.test"), + false, + None, + "https://base.test/model/anthropic.claude-3/invoke" + )] + #[case::a_blank_api_base_does_not_hide_the_params_endpoint( + Some(" "), + with_runtime_endpoint("https://params.test"), + false, + Some("https://env.test"), + "https://params.test/model/anthropic.claude-3/invoke" + )] + #[case::params_endpoint_outranks_the_environment( + None, + with_runtime_endpoint("https://params.test"), + false, + Some("https://env.test"), + "https://params.test/model/anthropic.claude-3/invoke" + )] + #[case::environment_outranks_the_region_template( + None, + in_region("us-west-2"), + false, + Some("https://env.test"), + "https://env.test/model/anthropic.claude-3/invoke" + )] + fn url_follows_python_endpoint_precedence_and_the_stream_path( + #[case] api_base: Option<&str>, + #[case] litellm_params: LitellmParams, + #[case] stream: bool, + #[case] env_endpoint: Option<&str>, + #[case] expected: &str, + ) { + let env = |name: &str| { + (name == AWS_BEDROCK_RUNTIME_ENDPOINT) + .then_some(env_endpoint) + .flatten() + .map(str::to_string) + }; assert_eq!( - config - .get_complete_url(None, "anthropic.claude-3", &env) + AmazonAnthropicClaudeMessagesConfig + .get_complete_url( + api_base, + "anthropic.claude-3", + &litellm_params, + stream, + &env + ) .unwrap(), - config - .complete_stream_url(None, "anthropic.claude-3", &env) - .unwrap() - .replace(INVOKE_STREAM_PATH, INVOKE_PATH) + expected + ); + } + + #[test] + fn sigv4_scope_and_credentials_come_from_the_litellm_params() { + let litellm_params = LitellmParams { + aws: AwsParams { + aws_access_key_id: Some("AKIAPARAMS".into()), + aws_secret_access_key: Some("params-secret".into()), + ..in_region("eu-central-1").aws + }, + ..LitellmParams::default() + }; + let validated = AmazonAnthropicClaudeMessagesConfig + .validate_environment( + Vec::new(), + None, + "anthropic.claude-3", + &litellm_params, + &|_| None, + ) + .unwrap(); + let AuthScheme::AwsSigV4 { + region, + credentials, + .. + } = validated.auth + else { + panic!("expected SigV4, got {:?}", validated.auth); + }; + let AwsCredentialSource::HostSupplied(credentials) = *credentials else { + panic!("expected the params' static keys, got {credentials:?}"); + }; + assert_eq!( + (region.as_str(), credentials.access_key_id()), + ("eu-central-1", "AKIAPARAMS") ); } @@ -538,6 +659,7 @@ mod tests { vec![("authorization".into(), "Bearer forwarded".into())], api_key, "anthropic.claude-3", + &LitellmParams::default(), &env, ) .unwrap(); diff --git a/litellm-rust/crates/llms/tests/messages_normalization.rs b/litellm-rust/crates/llms/tests/messages_normalization.rs index 27dba22c662..bd54618aef0 100644 --- a/litellm-rust/crates/llms/tests/messages_normalization.rs +++ b/litellm-rust/crates/llms/tests/messages_normalization.rs @@ -1,4 +1,4 @@ -use litellm_llms::base_llm::messages::normalization::fold_system_role_messages; +use litellm_llms::base_llm::messages::normalization::normalize_system_role_messages; use litellm_llms_types::formats::messages::MessagesRequest; use rstest::rstest; use serde_json::{Value, json}; @@ -7,42 +7,47 @@ use serde_json::{Value, json}; #[case::text_system(json!("existing"), json!([{"type": "text", "text": "existing"}]))] #[case::block_system(json!([{ "type": "future", "payload": 7 }]), json!([{ "type": "future", "payload": 7 }]))] #[case::no_system(Value::Null, json!([]))] -fn folding_preserves_block_fields_order_and_unrelated_request_fields( +fn hoisting_preserves_block_fields_order_and_unrelated_request_fields( #[case] system: Value, #[case] initial_blocks: Value, ) { let cache_control = json!({"type": "ephemeral", "scope": "global", "future": true}); - let folded_block = json!({"type": "text", "text": "second", "cache_control": cache_control}); + let leading_block = json!({"type": "text", "text": "second", "cache_control": cache_control}); let user = json!({"role": "user", "content": "hello", "future_message": 42}); + let later_turn = json!({"role": "system", "content": "mid-conversation reminder"}); let request: MessagesRequest = serde_json::from_value(json!({ "model": "test-model", "max_tokens": 64, "system": system, "messages": [ {"role": "system", "content": "first"}, + {"role": "system", "content": [leading_block]}, user, - {"role": "system", "content": [folded_block]} + later_turn ], "future_request": {"nested": true} })) .unwrap(); - let folded = fold_system_role_messages(request); + let normalized = normalize_system_role_messages(request, true); let expected_blocks: Vec = initial_blocks .as_array() .unwrap() .iter() .cloned() - .chain([json!({"type": "text", "text": "first"}), folded_block]) + .chain([json!({"type": "text", "text": "first"}), leading_block]) .collect(); assert_eq!( - serde_json::to_value(&folded).unwrap(), + serde_json::to_value(&normalized).unwrap(), json!({ "model": "test-model", "max_tokens": 64, "system": expected_blocks, - "messages": [user], + "messages": [user, later_turn], "future_request": {"nested": true} }) ); - assert_eq!(fold_system_role_messages(folded.clone()), folded); + assert_eq!( + normalize_system_role_messages(normalized.clone(), true), + normalized + ); } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 282f5b8544c..6fc05c43146 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -46,6 +46,7 @@ litellm-llms.workspace = true litellm-secrets = { workspace = true, features = ["aws", "azure", "google", "hashicorp", "cyberark"] } litellm-secrets-types.workspace = true litellm-llms-types.workspace = true +litellm-router-types.workspace = true litellm-host-python.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index f0fea5adb06..88900544779 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -60,7 +60,7 @@ impl PythonSettings { } } -fn missing_module(py: Python<'_>, error: &PyErr, expected: &str) -> PyResult { +pub(crate) fn missing_module(py: Python<'_>, error: &PyErr, expected: &str) -> PyResult { if !error.is_instance_of::(py) { return Ok(false); } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 841ef3081e3..d34047c735d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -5,11 +5,12 @@ use bytes::Bytes; use litellm_host_python::{InvokeError, PythonBinding, from_py, present, to_py}; use litellm_http::transport::Error as TransportError; use litellm_inference_messages::{ - Error, MessagesCall, MessagesSettings, MessagesShaping, messages_body, + Error, MessagesCall, MessagesSettings, MessagesShaping, litellm_params, messages_body, route::{Messages, MessagesStreamHead}, }; use litellm_llms::base_llm::messages::context::MessagesModelCapabilities; use litellm_llms_types::headers::ProviderSpecificHeaders; +use litellm_router_types::LitellmParams; use pyo3::{ exceptions::{PyException, PyValueError}, gc::{PyTraverseError, PyVisit}, @@ -21,6 +22,7 @@ use serde_json::{Map, Value}; use crate::{ errors::{RustUpstreamError, route_error_to_pyerr}, marshal::{optional_timeout, project_optional_fields, public_response, python_timeout_seconds}, + python_settings::missing_module, }; const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.messages.route_host"; @@ -63,6 +65,33 @@ fn merge_headers( (!merged.is_empty()).then_some(merged) } +/// Reads `litellm.`, the module global a param spec names, the way +/// `VertexBase.get_vertex_ai_project` reads `litellm.vertex_project`. +fn module_global<'py>(py: Python<'py>, name: &str) -> PyResult>> { + let module = match py.import("litellm") { + Ok(module) => module, + Err(error) if missing_module(py, &error, "litellm")? => return Ok(None), + Err(error) => return Err(error), + }; + let value = module.getattr(name)?; + Ok((!value.is_none()).then_some(value)) +} + +/// The caller's litellm params, read from the kwargs by the names the typed params declare, +/// so a key the configs do not read is never converted; a spec whose spellings the call all +/// leaves out falls back to the module global it names, which stays the host's concern. +fn project_litellm_params<'py>( + argument: impl Fn(&str) -> PyResult>>, + global: impl Fn(&str) -> PyResult>>, +) -> PyResult> { + let fields = project_optional_fields(LitellmParams::fields(), &argument)?; + let globals = LitellmParams::specs() + .filter(|spec| !spec.wire.iter().any(|name| fields.contains_key(*name))) + .filter_map(|spec| spec.module_global); + let folded = project_optional_fields(globals, &global)?; + Ok(litellm_params(fields.into_iter().chain(folded).collect())) +} + fn native_error(py: Python<'_>, error: Error) -> PyResult { match error { Error::Transport(TransportError::Http { status, body }) => { @@ -132,15 +161,19 @@ impl MessagesPythonHost { let api_base = string("api_base")?; let extra_headers = self.merged_headers(py, arguments)?; let provider_specific_header = self.provider_specific_header(py, arguments)?; - Ok(messages_body(body).map(|body| MessagesCall { - body, - api_key, - api_base, - extra_headers, - provider_specific_header, - custom_llm_provider, - timeout: optional_timeout(timeout), - shaping, + let litellm_params = project_litellm_params(argument, |name| module_global(py, name))?; + Ok(messages_body(body).and_then(|body| { + Ok(MessagesCall { + body, + api_key, + api_base, + extra_headers, + provider_specific_header, + custom_llm_provider, + litellm_params: litellm_params?, + timeout: optional_timeout(timeout), + shaping, + }) })) } @@ -351,6 +384,106 @@ mod tests { ); } + #[rstest] + #[case::model_and_credentials_are_projected( + json!({"model": "m", "api_key": "k", "api_base": "b", "api_version": "v", "timeout": 5, "messages": []}), + Ok(json!({"model": "m", "api_key": "k", "api_base": "b", "api_version": "v"})), + )] + #[case::aws_keys_are_projected( + json!({"model": "m", "aws_region_name": "eu-central-1", "aws_access_key_id": "AKIA", "messages": [{"role": "user"}]}), + Ok(json!({"model": "m", "aws_region_name": "eu-central-1", "aws_access_key_id": "AKIA"})), + )] + #[case::vertex_keys_are_projected_in_both_spellings( + json!({"model": "m", "vertex_project": "p", "vertex_ai_location": "us-east5", "vertex_credentials": {"type": "service_account"}}), + Ok(json!({"model": "m", "vertex_project": "p", "vertex_ai_location": "us-east5", "vertex_credentials": {"type": "service_account"}})), + )] + #[case::an_explicit_none_is_absent(json!({"model": "m", "aws_region_name": null}), Ok(json!({"model": "m"})))] + #[case::a_wrong_type_is_a_request_error(json!({"model": "m", "aws_region_name": 7}), Err(()))] + #[case::a_missing_model_is_a_request_error(json!({"aws_region_name": "eu-central-1"}), Err(()))] + fn litellm_params_are_projected_by_their_declared_names( + #[case] kwargs: Value, + #[case] expected: Result, + ) { + assert_projection(kwargs, json!({}), expected); + } + + #[rstest] + #[case::global_fills_a_missing_vertex_project( + json!({"model": "m"}), + json!({"vertex_project": "from-global", "vertex_location": "us-east5"}), + Ok(json!({"model": "m", "vertex_project": "from-global", "vertex_location": "us-east5"})), + )] + #[case::the_call_wins_over_the_global( + json!({"model": "m", "vertex_project": "from-call"}), + json!({"vertex_project": "from-global"}), + Ok(json!({"model": "m", "vertex_project": "from-call"})), + )] + #[case::an_explicit_none_in_the_call_still_falls_back( + json!({"model": "m", "vertex_project": null}), + json!({"vertex_project": "from-global"}), + Ok(json!({"model": "m", "vertex_project": "from-global"})), + )] + #[case::a_global_of_the_wrong_type_is_a_request_error( + json!({"model": "m"}), + json!({"vertex_location": 5}), + Err(()), + )] + fn module_globals_fill_the_litellm_params_the_call_leaves_out( + #[case] kwargs: Value, + #[case] globals: Value, + #[case] expected: Result, + ) { + assert_projection(kwargs, globals, expected); + } + + #[rstest] + #[case::a_global_without_a_spec_is_never_read( + json!({"model": "m"}), + json!({"aws_region_name": "eu-central-1", "api_key": "k"}), + Ok(json!({"model": "m"})), + )] + #[case::only_the_spec_named_global_is_read( + json!({"model": "m"}), + json!({"vertex_ai_project": "legacy-global", "vertex_project": "from-global"}), + Ok(json!({"model": "m", "vertex_project": "from-global"})), + )] + fn only_the_globals_the_specs_name_are_consulted( + #[case] kwargs: Value, + #[case] globals: Value, + #[case] expected: Result, + ) { + assert_projection(kwargs, globals, expected); + } + + fn assert_projection(kwargs: Value, globals: Value, expected: Result) { + Python::initialize(); + Python::attach(|py| { + let dict = |value: &Value| { + to_py(py, value) + .unwrap() + .into_bound(py) + .cast_into::() + .unwrap() + }; + let kwargs = dict(&kwargs); + let globals = dict(&globals); + let bound = PyDict::new(py); + let projected = project_litellm_params( + |name| present(&kwargs, &bound, name), + |name| present(&globals, &bound, name), + ) + .unwrap(); + match (projected, expected) { + (Ok(projected), Ok(fields)) => assert_eq!( + projected, + litellm_params(serde_json::from_value(fields).unwrap()).unwrap() + ), + (Err(error), Err(())) => assert!(matches!(error, Error::InvalidRequest(_))), + (projected, expected) => panic!("got {projected:?}, expected {expected:?}"), + } + }); + } + #[rstest] #[case::rejected_request(Error::InvalidRequest("does not support top_k=5".into()), true)] #[case::missing_field(Error::MissingField("max_tokens"), true)] diff --git a/litellm-rust/crates/router/Cargo.toml b/litellm-rust/crates/router/Cargo.toml index 8c01c1e5a99..a356bedcf16 100644 --- a/litellm-rust/crates/router/Cargo.toml +++ b/litellm-rust/crates/router/Cargo.toml @@ -9,6 +9,7 @@ repository.workspace = true litellm-config.workspace = true litellm-inference.workspace = true litellm-inference-messages.workspace = true +litellm-router-types.workspace = true [dev-dependencies] rstest.workspace = true diff --git a/litellm-rust/crates/router/src/deployment.rs b/litellm-rust/crates/router/src/deployment.rs index 4a091f8845c..d57284bae5a 100644 --- a/litellm-rust/crates/router/src/deployment.rs +++ b/litellm-rust/crates/router/src/deployment.rs @@ -1,6 +1,7 @@ use std::time::Duration; use litellm_inference_messages::MessagesShaping; +use litellm_router_types::LitellmParams; #[derive(Clone, Debug, Default)] pub struct Deployment { @@ -8,6 +9,7 @@ pub struct Deployment { pub api_key: Option, pub api_base: Option, pub custom_llm_provider: Option, + pub litellm_params: LitellmParams, pub timeout: Option, pub shaping: MessagesShaping, } diff --git a/litellm-rust/crates/router/src/lib.rs b/litellm-rust/crates/router/src/lib.rs index da33bfb04bd..fb48213e8a7 100644 --- a/litellm-rust/crates/router/src/lib.rs +++ b/litellm-rust/crates/router/src/lib.rs @@ -25,6 +25,7 @@ impl Router { .map(|value| value.expose().to_string()), api_base: model.litellm_params.api_base.clone(), custom_llm_provider: model.litellm_params.custom_llm_provider.clone(), + litellm_params: model.litellm_params.clone(), ..Deployment::default() }, ) diff --git a/litellm-rust/crates/router/tests/router.rs b/litellm-rust/crates/router/tests/router.rs index 97ee544f120..bbaeb806b96 100644 --- a/litellm-rust/crates/router/tests/router.rs +++ b/litellm-rust/crates/router/tests/router.rs @@ -29,6 +29,32 @@ fn configuration_preserves_deployment_parameters(#[case] parameters: &str) { assert_eq!(deployment.custom_llm_provider, params.custom_llm_provider); assert_eq!(deployment.timeout, Deployment::default().timeout); assert_eq!(deployment.shaping, Deployment::default().shaping); + assert_eq!(deployment.litellm_params, *params); +} + +#[rstest] +fn configuration_types_the_provider_params_of_a_deployment() { + let config = Config::from_yaml( + "model_list: + - model_name: public-model + litellm_params: + model: bedrock/anthropic.claude-3 + aws_region_name: eu-central-1 + aws_bedrock_runtime_endpoint: https://runtime.example + azure_ad_token: ignored-by-every-group", + ) + .unwrap(); + let router = Router::from_model_list(&config.model_list); + let params = &router.get("public-model").unwrap().litellm_params; + + assert_eq!( + ( + params.aws.aws_region_name.as_deref(), + params.aws.aws_bedrock_runtime_endpoint.as_deref(), + params.aws.aws_access_key_id.as_deref(), + ), + (Some("eu-central-1"), Some("https://runtime.example"), None) + ); } #[rstest] @@ -72,6 +98,7 @@ fn programmatic_deployments_preserve_overrides_and_last_entry_wins() { api_key: Some("test-key".into()), api_base: Some("https://provider.example/v1".into()), custom_llm_provider: Some("test-provider".into()), + litellm_params: Default::default(), timeout: Some(Duration::from_secs(7)), shaping: MessagesShaping { settings: MessagesSettings { diff --git a/litellm/rust_bridge/model_capabilities.py b/litellm/rust_bridge/model_capabilities.py index eb8496c722f..cbf59dfb97b 100644 --- a/litellm/rust_bridge/model_capabilities.py +++ b/litellm/rust_bridge/model_capabilities.py @@ -31,5 +31,6 @@ def anthropic_model_capabilities(model: str, custom_llm_provider: str | None) -> "supports_output_config": supports("supports_output_config"), "supports_sampling_params": AnthropicModelInfo._supports_sampling_params(resolved_model), # pyright: ignore[reportPrivateUsage] # same gate the handler applies "supports_speed": AnthropicConfig._model_supports_speed_param(resolved_model, provider), # pyright: ignore[reportPrivateUsage] # same gate the handler applies + "supports_mid_conversation_system": supports("supports_mid_conversation_system"), "effort_tiers": {level: tier(level) for level in ("minimal", "low", "medium", "high", "xhigh", "max")}, } diff --git a/tests/unit/rust_bridge/test_model_capabilities.py b/tests/unit/rust_bridge/test_model_capabilities.py index 9f02660a98f..93f7dfdc04a 100644 --- a/tests/unit/rust_bridge/test_model_capabilities.py +++ b/tests/unit/rust_bridge/test_model_capabilities.py @@ -32,12 +32,14 @@ def test_capabilities_come_from_the_model_map_under_the_callers_provider(monkeyp supports_output_config=True, supports_xhigh_reasoning_effort=True, supports_sampling_params=False, + supports_mid_conversation_system=True, ) capabilities: Final = anthropic_model_capabilities("anthropic/claude-test-adaptive", None) assert capabilities["supports_adaptive_thinking"] assert capabilities["supports_output_config"] + assert capabilities["supports_mid_conversation_system"] assert not capabilities["supports_legacy_thinking"] assert not capabilities["supports_sampling_params"] assert capabilities["effort_tiers"] == { @@ -56,6 +58,7 @@ def test_unmapped_model_keeps_sampling_params_and_no_reasoning_features() -> Non assert capabilities["supports_sampling_params"] assert not capabilities["supports_reasoning"] assert not capabilities["supports_adaptive_thinking"] + assert not capabilities["supports_mid_conversation_system"] assert capabilities["effort_tiers"] == dict.fromkeys(("minimal", "low", "medium", "high", "xhigh", "max"), False) From a0b3c46cf0b289fd766f4912643f99c6f5f809db Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 9 Oct 2026 16:29:45 -0700 Subject: [PATCH 005/104] feat(rust): add the Vertex AI Anthropic Messages config (#45626) * feat(rust): add the Vertex AI Anthropic Messages config * fix(rust): normalize Vertex system turns the way Python does Vertex uses the shared normalize_system_role_messages: a leading run is hoisted, a later turn stays or is converted in place by the model's supports_mid_conversation_system, and billing blocks are filtered after the hoist so one inside a system-role message no longer reaches Vertex. * fix(rust_bridge): keep non-Claude Vertex Messages calls on Python The catalog admits every vertex_ai Messages call, but Rust serves only Claude there and reports anything else as a terminal InvalidProvider instead of a decline. The dispatcher now bypasses Rust for those models, as get_provider_anthropic_messages_config does. * refactor(rust): template the Vertex location and project errors The location failure is the shared InvalidType detail and the missing project is MissingSetting with the provider, setting, param and variable name, so no sentence is spelled at the return site. * test(rust): use rstest in the Vertex AI Messages tests * refactor(rust): derive the missing Vertex project error from its spec --- .../inference-messages/src/common_utils.rs | 25 +- .../crates/inference-messages/src/prepare.rs | 2 +- .../tests/messages/request.rs | 51 ++ .../crates/llms/src/vertex_ai/common_utils.rs | 71 ++ .../crates/llms/src/vertex_ai/messages/mod.rs | 1 + .../src/vertex_ai/messages/transformation.rs | 682 ++++++++++++++++++ litellm-rust/crates/llms/src/vertex_ai/mod.rs | 2 + litellm/messages/dispatch.py | 10 +- litellm/rust_bridge/catalog.py | 2 +- tests/unit/messages/test_dispatch.py | 47 ++ tests/unit/rust_bridge/test_catalog.py | 6 +- 11 files changed, 889 insertions(+), 10 deletions(-) create mode 100644 litellm-rust/crates/llms/src/vertex_ai/common_utils.rs create mode 100644 litellm-rust/crates/llms/src/vertex_ai/messages/mod.rs create mode 100644 litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index 5102d063a8a..d61761b9348 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -6,6 +6,7 @@ use litellm_llms::{ azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG, base_llm::messages::transformation::BaseMessagesConfig, bedrock::messages::invoke_transformations::anthropic_claude3_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG, + vertex_ai::messages::transformation::VERTEX_ANTHROPIC_MESSAGES_CONFIG, }; use serde_json::{Map, Value}; @@ -18,6 +19,7 @@ pub(crate) enum MessagesProvider { Anthropic, AzureAi, Bedrock, + VertexAi, } impl MessagesProvider { @@ -26,6 +28,7 @@ impl MessagesProvider { Self::Anthropic => LlmProviders::Anthropic, Self::AzureAi => LlmProviders::AzureAi, Self::Bedrock => LlmProviders::Bedrock, + Self::VertexAi => LlmProviders::VertexAi, } .into() } @@ -35,15 +38,21 @@ impl MessagesProvider { Self::Anthropic => &ANTHROPIC_MESSAGES_CONFIG, Self::AzureAi => &AZURE_ANTHROPIC_MESSAGES_CONFIG, Self::Bedrock => &BEDROCK_ANTHROPIC_MESSAGES_CONFIG, + Self::VertexAi => &VERTEX_ANTHROPIC_MESSAGES_CONFIG, } } } -pub(crate) fn messages_provider(provider: LlmProviders) -> Option { +/// Python's `get_provider_anthropic_messages_config`: Vertex AI serves only its Claude +/// partner models on this route. +pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option { match provider { LlmProviders::Anthropic => Some(MessagesProvider::Anthropic), LlmProviders::AzureAi => Some(MessagesProvider::AzureAi), LlmProviders::Bedrock => Some(MessagesProvider::Bedrock), + LlmProviders::VertexAi if model.to_ascii_lowercase().contains("claude") => { + Some(MessagesProvider::VertexAi) + } _ => None, } } @@ -68,20 +77,26 @@ mod tests { #[case::anthropic("anthropic", MessagesProvider::Anthropic)] #[case::azure_ai("azure_ai", MessagesProvider::AzureAi)] #[case::bedrock("bedrock", MessagesProvider::Bedrock)] + #[case::vertex_ai("vertex_ai", MessagesProvider::VertexAi)] fn provider_round_trips_through_its_python_name( #[case] name: &str, #[case] provider: MessagesProvider, ) { assert_eq!( - messages_provider(name.parse::().unwrap()), + messages_provider(name.parse::().unwrap(), "claude-sonnet-4-5"), Some(provider) ); assert_eq!(provider.as_str(), name); } - #[test] - fn provider_without_a_messages_config_is_rejected() { - assert_eq!(messages_provider(LlmProviders::Openai), None); + #[rstest] + #[case::openai(LlmProviders::Openai, "gpt-5")] + #[case::vertex_gemini(LlmProviders::VertexAi, "gemini-2.5-pro")] + fn provider_without_a_messages_config_is_rejected( + #[case] provider: LlmProviders, + #[case] model: &str, + ) { + assert_eq!(messages_provider(provider, model), None); } #[test] diff --git a/litellm-rust/crates/inference-messages/src/prepare.rs b/litellm-rust/crates/inference-messages/src/prepare.rs index 0c189d0a309..50a079d5eb3 100644 --- a/litellm-rust/crates/inference-messages/src/prepare.rs +++ b/litellm-rust/crates/inference-messages/src/prepare.rs @@ -51,7 +51,7 @@ fn resolve_provider( custom_llm_provider: Option<&str>, ) -> Result { let resolved = resolve_llm_provider(model, custom_llm_provider, "messages")?; - let provider = messages_provider(resolved.provider) + let provider = messages_provider(resolved.provider, resolved.model) .ok_or_else(|| Error::InvalidProvider(<&str>::from(resolved.provider).to_string()))?; Ok(ResolvedProvider { model: resolved.model.to_string(), diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index 523b25f81d6..b1a989713ec 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -1,8 +1,10 @@ +use litellm_auth::VertexParams; use litellm_llms::base_llm::messages::context::{MessagesModelCapabilities, SupportedEffortTiers}; use litellm_llms_types::{ headers::{ProviderSpecificHeader, ProviderSpecificHeaders}, providers::anthropic::{AnthropicBeta, BetaSet}, }; +use litellm_router_types::LitellmParams; use rstest::rstest; use super::*; @@ -599,6 +601,55 @@ async fn metadata_is_reduced_to_the_user_id(call: MessagesCall, #[case] provider ); } +#[rstest] +#[case::not_streaming(false, ":rawPredict")] +#[case::streaming(true, ":streamRawPredict?alt=sse")] +#[tokio::test] +async fn vertex_ai_addresses_the_model_in_the_url_and_not_in_the_body( + call: MessagesCall, + #[case] stream: bool, + #[case] suffix: &str, +) { + let streamed = ResponseTemplate::new(200).set_body_raw( + "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + "text/event-stream", + ); + let upstream = upstream([if stream { streamed } else { message_response() }]).await; + + run(with_fields( + MessagesCall { + custom_llm_provider: Some("vertex_ai".into()), + litellm_params: LitellmParams { + vertex: VertexParams { + vertex_project: Some("proj".into()), + vertex_location: Some("us-east5".into()), + ..VertexParams::default() + }, + ..LitellmParams::default() + }, + api_base: Some(upstream.uri()), + extra_headers: headers([("Authorization", "Bearer caller-token")]), + ..with_model(call, "claude-sonnet-4-5@20250929") + }, + json!({"stream": stream}), + )) + .await + .expect("messages call succeeds"); + + let request = only_request(&upstream).await; + assert_eq!( + request.url.path().to_string() + request.url.query().map_or("", |_| "?alt=sse"), + format!( + "/v1/projects/proj/locations/us-east5/publishers/anthropic/models/claude-sonnet-4-5@20250929{suffix}" + ) + ); + assert_eq!(request.header("authorization"), Some("Bearer caller-token")); + assert_eq!(request.header("anthropic-version"), None); + let body = request.json(); + assert_eq!(body.get("model"), None); + assert_eq!(body["anthropic_version"], json!("vertex-2023-10-16")); +} + #[rstest] #[case::numeric_user_id(json!({"metadata": {"user_id": 7}}))] #[case::missing_max_tokens(json!({"max_tokens": null}))] diff --git a/litellm-rust/crates/llms/src/vertex_ai/common_utils.rs b/litellm-rust/crates/llms/src/vertex_ai/common_utils.rs new file mode 100644 index 00000000000..45abf5cd9f2 --- /dev/null +++ b/litellm-rust/crates/llms/src/vertex_ai/common_utils.rs @@ -0,0 +1,71 @@ +use crate::Error; + +const GLOBAL_LOCATION: &str = "global"; +pub const DEFAULT_VERTEX_LOCATION: &str = "us-central1"; + +fn is_location_token(location: &str) -> bool { + !location.is_empty() + && location + .bytes() + .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-') + && location + .as_bytes() + .first() + .is_some_and(u8::is_ascii_alphanumeric) + && location + .as_bytes() + .last() + .is_some_and(u8::is_ascii_alphanumeric) +} + +/// Python's `validate_vertex_location`: a location reaches a hostname, so only `global` or a +/// lowercase alphanumeric token with hyphens is allowed. +pub fn validate_vertex_location(location: &str) -> Result<&str, Error> { + if location == GLOBAL_LOCATION || is_location_token(location) { + return Ok(location); + } + Err(Error::Auth(litellm_auth::Error::InvalidConfiguration( + litellm_auth::ErrorDetail::InvalidType { + field: "vertex_location".into(), + expected: "global or a lowercase alphanumeric token with hyphens", + }, + ))) +} + +/// Python's `get_vertex_base_url`: the global control plane, a multi-region geography, or a +/// regional host. +pub fn get_vertex_base_url(location: &str) -> Result { + let location = validate_vertex_location(location)?; + if location == GLOBAL_LOCATION { + return Ok("https://aiplatform.googleapis.com".to_string()); + } + if !location.contains('-') { + return Ok(format!("https://aiplatform.{location}.rep.googleapis.com")); + } + Ok(format!("https://{location}-aiplatform.googleapis.com")) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::regional("us-east5", "https://us-east5-aiplatform.googleapis.com")] + #[case::geography("eu", "https://aiplatform.eu.rep.googleapis.com")] + #[case::global("global", "https://aiplatform.googleapis.com")] + fn base_url_follows_the_location_kind(#[case] location: &str, #[case] expected: &str) { + assert_eq!(get_vertex_base_url(location).unwrap(), expected); + } + + #[rstest] + #[case::host_injection("attacker.example/")] + #[case::fragment("evil.com#")] + #[case::uppercase("US-EAST5")] + #[case::leading_hyphen("-us")] + #[case::empty("")] + fn a_location_that_is_not_a_token_is_rejected(#[case] location: &str) { + assert!(get_vertex_base_url(location).is_err()); + } +} diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/mod.rs b/litellm-rust/crates/llms/src/vertex_ai/messages/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs new file mode 100644 index 00000000000..8bedea4d100 --- /dev/null +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs @@ -0,0 +1,682 @@ +use litellm_auth::VertexParams; +use litellm_auth_gcp::{ + VertexConfig, get_vertex_ai_location, get_vertex_ai_project, + get_vertex_ai_project_from_credentials, secret_names, +}; +use litellm_http::request::has_header; +use litellm_llms_types::{ + formats::messages::{MessagesOptionalParams, MessagesRequest, MessagesTool, OutputConfig}, + providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet}, + recognized::Recognized, +}; +use litellm_router_types::LitellmParams; +use serde_json::{Value, json}; + +use crate::{ + Error, + anthropic::{ + common_utils::{merge_beta_headers, supports_effort_param}, + messages::{ + handler::shape_anthropic_messages_request, + transformation::{provider_feature_betas, transform_messages_request}, + }, + }, + base_llm::{ + auth::{AuthScheme, Headers, ValidatedEnvironment}, + messages::{ + context::MessagesTransformContext, + normalization::{normalize_system_role_messages, strip_cache_control_scope}, + transformation::BaseMessagesConfig, + }, + }, + vertex_ai::common_utils::{DEFAULT_VERTEX_LOCATION, get_vertex_base_url}, +}; + +pub const VERTEX_ANTHROPIC_VERSION: &str = "vertex-2023-10-16"; +const RAW_PREDICT: &str = "rawPredict"; +const STREAM_RAW_PREDICT: &str = "streamRawPredict"; +const WEB_SEARCH_TOOL_PREFIX: &str = "web_search"; + +/// Claude on Vertex AI, Python's `VertexAIPartnerModelsAnthropicMessagesConfig`: the Anthropic +/// payload addressed to a project and location, authenticated with a Google access token. +pub struct VertexAiPartnerModelsAnthropicMessagesConfig; + +pub const VERTEX_ANTHROPIC_MESSAGES_CONFIG: VertexAiPartnerModelsAnthropicMessagesConfig = + VertexAiPartnerModelsAnthropicMessagesConfig; + +impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig { + fn shape_request( + &self, + request: MessagesRequest, + reasoning_auto_summary: bool, + ) -> Result { + shape_anthropic_messages_request(request, reasoning_auto_summary) + } + + fn get_complete_url( + &self, + api_base: Option<&str>, + model: &str, + litellm_params: &LitellmParams, + stream: bool, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + let config = VertexConfig::from_params(&litellm_params.vertex); + let project = get_vertex_ai_project(&config, env_lookup) + .or_else(|| get_vertex_ai_project_from_credentials(&config, env_lookup)) + .ok_or_else(|| Error::Auth(VertexParams::PROJECT.missing("Vertex AI")))?; + let location = get_vertex_ai_location(&config, env_lookup) + .unwrap_or_else(|| DEFAULT_VERTEX_LOCATION.to_string()); + complete_vertex_anthropic_url(api_base, &project, &location, model, stream) + } + + fn transform_anthropic_messages_request( + &self, + request: MessagesRequest, + context: &MessagesTransformContext, + ) -> Result { + let request = transform_messages_request(request, context)?; + let request = strip_cache_control_scope(normalize_system_role_messages( + request, + context + .thinking + .capabilities + .supports_mid_conversation_system, + )); + let accepts_effort = supports_effort_param(&context.thinking.capabilities); + Ok(MessagesRequest { + params: MessagesOptionalParams { + output_config: request + .params + .output_config + .and_then(|config| sanitize_output_config(config, accepts_effort)), + extra: request + .params + .extra + .into_iter() + .chain([( + "anthropic_version".to_string(), + json!(VERTEX_ANTHROPIC_VERSION), + )]) + .collect(), + ..request.params + }, + ..request + }) + } + + fn secret_names(&self) -> &'static [&'static str] { + secret_names() + } + + /// A forwarded bearer is sent as is; otherwise a Google access token for the configured + /// credentials (`vertex_credentials`, `VERTEXAI_CREDENTIALS`, `GOOGLE_APPLICATION_CREDENTIALS`, + /// then application default credentials) is acquired when the request is sent. + fn validate_environment( + &self, + headers: Headers, + _api_key: Option<&str>, + _model: &str, + litellm_params: &LitellmParams, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + if has_header(&headers, "authorization") { + return Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Forwarded, + }); + } + Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::GcpAccessToken { + config: Box::new(VertexConfig::from_params(&litellm_params.vertex)), + }, + }) + } + + fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers { + merge_beta_headers(headers, vertex_feature_betas(request)) + } + + /// Vertex addresses the model in the URL and rejects it in the body. + fn wire_body(&self, mut body: Value) -> Value { + if let Value::Object(fields) = &mut body { + fields.remove("model"); + } + body + } +} + +/// Python's `create_vertex_url` plus `_check_custom_proxy` for the Claude partner: the +/// publisher path under the location host, or under a custom base that has no path of its +/// own, else the custom base with the predict verb appended. A custom base also gets +/// `alt=sse` when streaming. +pub fn complete_vertex_anthropic_url( + api_base: Option<&str>, + project: &str, + location: &str, + model: &str, + stream: bool, +) -> Result { + let verb = if stream { + STREAM_RAW_PREDICT + } else { + RAW_PREDICT + }; + let path = format!( + "/v1/projects/{project}/locations/{location}/publishers/anthropic/models/{model}:{verb}" + ); + let custom = api_base.map(str::trim).filter(|base| !base.is_empty()); + let Some(custom) = custom else { + return Ok(format!("{}{path}", get_vertex_base_url(location)?)); + }; + let custom = custom.trim_end_matches('/'); + let authority_and_path = custom.split_once("://").map_or(custom, |(_, rest)| rest); + let base_path = authority_and_path + .find('/') + .map_or("", |index| &authority_and_path[index..]); + let url = match base_path { + "" => format!("{custom}{path}"), + "/v1" | "/v1beta1" => format!("{custom}{}", path.trim_start_matches("/v1")), + _ => format!("{custom}:{verb}"), + }; + if !stream { + return Ok(url); + } + Ok(match url.contains('?') { + true => format!("{url}&alt=sse"), + false => format!("{url}?alt=sse"), + }) +} + +/// Python's `sanitize_vertex_anthropic_output_params`: `effort` only reaches models whose +/// capabilities accept it, a malformed `output_config` is dropped, and an emptied one is omitted. +fn sanitize_output_config( + config: Recognized, + accepts_effort: bool, +) -> Option> { + let Recognized::Known(config) = config else { + return None; + }; + let sanitized = OutputConfig { + effort: config.effort.filter(|_| accepts_effort), + ..config + }; + (!sanitized.is_empty()).then_some(Recognized::Known(sanitized)) +} + +fn has_web_search_tool(tools: Option<&[Recognized]>) -> bool { + tools.unwrap_or_default().iter().any(|tool| match tool { + Recognized::Unrecognized(value) => value + .get("type") + .and_then(Value::as_str) + .is_some_and(|kind| kind.starts_with(WEB_SEARCH_TOOL_PREFIX)), + Recognized::Known(_) => false, + }) +} + +/// Python's Vertex `validate_anthropic_messages_environment` beta set: the shared feature +/// betas under Vertex's policy, plus the web search and safeguards betas Vertex needs spelled +/// out. +fn vertex_feature_betas(request: &MessagesRequest) -> BetaSet { + let params = &request.params; + provider_feature_betas(request, BetaProvider::VertexAi).union( + [ + has_web_search_tool(params.tools.as_deref()) + .then_some(AnthropicBeta::WebSearch20250305), + params + .extra + .get("safeguards") + .is_some_and(|value| !value.is_null()) + .then_some(AnthropicBeta::DangerousToolUse20260903), + ] + .into_iter() + .flatten() + .collect(), + ) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + use crate::base_llm::messages::context::{MessagesModelCapabilities, ThinkingContext}; + + fn no_env(_: &str) -> Option { + None + } + + fn params(project: Option<&str>, location: Option<&str>) -> LitellmParams { + LitellmParams { + vertex: VertexParams { + vertex_project: project.map(Into::into), + vertex_location: location.map(Into::into), + ..VertexParams::default() + }, + ..LitellmParams::default() + } + } + + fn url( + api_base: Option<&str>, + litellm_params: &LitellmParams, + stream: bool, + env: &dyn Fn(&str) -> Option, + ) -> Result { + VERTEX_ANTHROPIC_MESSAGES_CONFIG.get_complete_url( + api_base, + "claude-sonnet-4-5@20250929", + litellm_params, + stream, + env, + ) + } + + const PATH: &str = "/v1/projects/proj/locations/us-east5/publishers/anthropic/models/claude-sonnet-4-5@20250929"; + + #[rstest] + #[case::regional( + None, + false, + "https://us-east5-aiplatform.googleapis.com", + ":rawPredict" + )] + #[case::streaming( + None, + true, + "https://us-east5-aiplatform.googleapis.com", + ":streamRawPredict" + )] + #[case::custom_host_gets_the_default_path( + Some("https://proxy.example/"), + false, + "https://proxy.example", + ":rawPredict" + )] + #[case::custom_host_streaming_adds_alt_sse( + Some("https://proxy.example"), + true, + "https://proxy.example", + ":streamRawPredict?alt=sse" + )] + #[case::custom_versioned_base_is_grafted( + Some("https://proxy.example/v1"), + false, + "https://proxy.example", + ":rawPredict" + )] + fn the_url_is_the_publisher_path_under_the_host( + #[case] api_base: Option<&str>, + #[case] stream: bool, + #[case] host: &str, + #[case] suffix: &str, + ) { + assert_eq!( + url( + api_base, + ¶ms(Some("proj"), Some("us-east5")), + stream, + &no_env + ) + .unwrap(), + format!("{host}{PATH}{suffix}") + ); + } + + #[rstest] + #[case::not_streaming(false, "https://proxy.example/custom/path:rawPredict")] + #[case::streaming(true, "https://proxy.example/custom/path:streamRawPredict?alt=sse")] + fn a_custom_base_with_its_own_path_only_gets_the_verb( + #[case] stream: bool, + #[case] expected: &str, + ) { + assert_eq!( + url( + Some("https://proxy.example/custom/path"), + ¶ms(Some("proj"), Some("us-east5")), + stream, + &no_env + ) + .unwrap(), + expected + ); + } + + #[rstest] + #[case::global( + "global", + "https://aiplatform.googleapis.com/v1/projects/proj/locations/global/" + )] + #[case::geography( + "eu", + "https://aiplatform.eu.rep.googleapis.com/v1/projects/proj/locations/eu/" + )] + fn the_host_follows_the_location_kind(#[case] location: &str, #[case] prefix: &str) { + assert!( + url(None, ¶ms(Some("proj"), Some(location)), false, &no_env) + .unwrap() + .starts_with(prefix) + ); + } + + #[rstest] + fn a_location_that_could_name_another_host_is_rejected() { + assert!(matches!( + url( + None, + ¶ms(Some("proj"), Some("attacker.example/")), + false, + &no_env + ), + Err(Error::Auth(_)) + )); + } + + #[rstest] + #[case::param_wins( + params(Some("from-param"), None), + &[("VERTEXAI_PROJECT", "from-env"), ("VERTEXAI_LOCATION", "europe-west4")], + "https://europe-west4-aiplatform.googleapis.com/v1/projects/from-param/locations/europe-west4/" + )] + #[case::legacy_spelling( + LitellmParams { + vertex: VertexParams { + vertex_ai_project: Some("legacy".into()), + vertex_ai_location: Some("us-east5".into()), + ..VertexParams::default() + }, + ..LitellmParams::default() + }, + &[], + "https://us-east5-aiplatform.googleapis.com/v1/projects/legacy/locations/us-east5/" + )] + #[case::env_then_default_location( + LitellmParams::default(), + &[("VERTEXAI_PROJECT", "from-env")], + "https://us-central1-aiplatform.googleapis.com/v1/projects/from-env/locations/us-central1/" + )] + #[case::project_from_the_credentials( + LitellmParams { + vertex: VertexParams { + vertex_credentials: Some(r#"{"type": "service_account", "project_id": "from-creds"}"#.into()), + ..VertexParams::default() + }, + ..LitellmParams::default() + }, + &[("VERTEX_LOCATION", "us-east5")], + "https://us-east5-aiplatform.googleapis.com/v1/projects/from-creds/locations/us-east5/" + )] + fn project_and_location_follow_the_python_precedence( + #[case] litellm_params: LitellmParams, + #[case] env: &[(&str, &str)], + #[case] prefix: &str, + ) { + let lookup = |name: &str| { + env.iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + }; + assert!( + url(None, &litellm_params, false, &lookup) + .unwrap() + .starts_with(prefix) + ); + } + + #[rstest] + fn without_a_project_anywhere_the_call_fails_before_sending() { + let error = url(None, &LitellmParams::default(), false, &no_env).unwrap_err(); + + assert!(matches!( + error, + Error::Auth(litellm_auth::Error::MissingParam { + provider: "Vertex AI", + spec: &VertexParams::PROJECT, + }) + )); + assert_eq!( + error.to_string(), + "Missing Vertex AI project - pass vertex_project or vertex_ai_project, set litellm.vertex_project, or set VERTEXAI_PROJECT" + ); + } + + fn validated(forwarded: &[(&str, &str)]) -> ValidatedEnvironment { + VERTEX_ANTHROPIC_MESSAGES_CONFIG + .validate_environment( + forwarded + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(), + Some("ignored-api-key"), + "claude-sonnet-4-5", + ¶ms(Some("proj"), None), + &no_env, + ) + .unwrap() + } + + #[rstest] + fn the_credential_is_a_google_access_token_for_the_configured_project() { + assert!(matches!( + validated(&[]).auth, + AuthScheme::GcpAccessToken { config } if config.project_id() == Some("proj") + )); + } + + #[rstest] + fn a_forwarded_bearer_is_sent_as_is() { + assert!(matches!( + validated(&[("Authorization", "Bearer caller")]).auth, + AuthScheme::Forwarded + )); + } + + #[rstest] + fn default_headers_carry_no_anthropic_version() { + assert_eq!( + VERTEX_ANTHROPIC_MESSAGES_CONFIG.default_headers(), + &[("content-type", "application/json")] + ); + } + + fn transformed(value: Value, context: &MessagesTransformContext) -> Value { + let request: MessagesRequest = serde_json::from_value(value).unwrap(); + serde_json::to_value( + VERTEX_ANTHROPIC_MESSAGES_CONFIG + .transform_anthropic_messages_request(request, context) + .unwrap(), + ) + .unwrap() + } + + #[rstest] + fn the_body_carries_the_vertex_version_and_the_wire_drops_the_model() { + let body = transformed( + json!({ + "model": "claude-sonnet-4-5@20250929", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}] + }), + &MessagesTransformContext::default(), + ); + assert_eq!(body["anthropic_version"], json!("vertex-2023-10-16")); + assert_eq!(body["model"], json!("claude-sonnet-4-5@20250929")); + let wire = VERTEX_ANTHROPIC_MESSAGES_CONFIG.wire_body(body); + assert_eq!(wire.get("model"), None); + assert_eq!(wire["anthropic_version"], json!("vertex-2023-10-16")); + } + + #[rstest] + fn leading_system_roles_are_hoisted_billing_blocks_drop_and_scope_is_stripped() { + let body = transformed( + json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "system": [ + {"type": "text", "text": "x-anthropic-billing-header: cc_version=1"}, + {"type": "text", "text": "be terse", "cache_control": {"type": "ephemeral", "scope": "global"}} + ], + "messages": [ + {"role": "system", "content": [ + {"type": "text", "text": "x-anthropic-billing-header: cc_version=2"}, + {"type": "text", "text": "leading system"} + ]}, + {"role": "user", "content": [ + {"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h", "scope": "global"}} + ]} + ] + }), + &MessagesTransformContext::default(), + ); + assert_eq!( + body["system"], + json!([ + {"type": "text", "text": "be terse", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "leading system"} + ]) + ); + assert_eq!( + body["messages"], + json!([{"role": "user", "content": [ + {"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ]}]) + ); + } + + #[rstest] + #[case::converted_in_place_without_the_capability(false, "user")] + #[case::kept_in_place_with_the_capability(true, "system")] + fn a_later_system_turn_never_moves_into_the_prompt_prefix( + #[case] supports_mid_conversation_system: bool, + #[case] role: &str, + ) { + let body = transformed( + json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "system": "be terse", + "messages": [ + {"role": "user", "content": "a"}, + {"role": "assistant", "content": "b"}, + {"role": "system", "content": "reminder"}, + {"role": "user", "content": "c"} + ] + }), + &MessagesTransformContext { + thinking: ThinkingContext { + capabilities: MessagesModelCapabilities { + supports_mid_conversation_system, + ..MessagesModelCapabilities::default() + }, + ..ThinkingContext::default() + }, + drop_params: false, + }, + ); + assert_eq!(body["system"], json!("be terse")); + assert_eq!(body["messages"][2]["role"], json!(role)); + assert_eq!(body["messages"].as_array().unwrap().len(), 4); + } + + fn with_output_config() -> MessagesTransformContext { + MessagesTransformContext { + thinking: ThinkingContext { + capabilities: MessagesModelCapabilities { + supports_output_config: true, + ..MessagesModelCapabilities::default() + }, + ..ThinkingContext::default() + }, + drop_params: false, + } + } + + #[rstest] + #[case::effort_kept_when_the_model_accepts_it( + with_output_config(), + json!({"effort": "high", "format": {"type": "json_schema"}}), + Some(json!({"effort": "high", "format": {"type": "json_schema"}})) + )] + #[case::effort_dropped_otherwise( + MessagesTransformContext::default(), + json!({"effort": "high", "format": {"type": "json_schema"}}), + Some(json!({"format": {"type": "json_schema"}})) + )] + #[case::emptied_config_is_omitted( + MessagesTransformContext::default(), + json!({"effort": "high"}), + None + )] + #[case::malformed_config_is_dropped( + with_output_config(), + json!("not an object"), + None + )] + fn output_config_is_sanitized_for_vertex( + #[case] context: MessagesTransformContext, + #[case] output_config: Value, + #[case] expected: Option, + ) { + let body = transformed( + json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "output_config": output_config, + "messages": [{"role": "user", "content": "hi"}] + }), + &context, + ); + assert_eq!(body.get("output_config").cloned(), expected); + } + + fn betas(fields: Value) -> Vec { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + })) + .unwrap(); + let request: MessagesRequest = serde_json::from_value({ + let mut base = serde_json::to_value(request).unwrap(); + base.as_object_mut() + .unwrap() + .extend(fields.as_object().unwrap().clone()); + base + }) + .unwrap(); + VERTEX_ANTHROPIC_MESSAGES_CONFIG + .request_headers(Vec::new(), &request) + .into_iter() + .filter(|(name, _)| name == "anthropic-beta") + .flat_map(|(_, value)| value.split(',').map(str::to_string).collect::>()) + .collect() + } + + #[rstest] + #[case::web_search_tool( + json!({"tools": [{"type": "web_search_20250305", "name": "web_search"}]}), + &["web-search-2025-03-05"] + )] + #[case::tool_search_uses_the_vertex_header( + json!({"tools": [{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}]}), + &["tool-search-tool-2025-10-19"] + )] + #[case::safeguards(json!({"safeguards": {"mode": "strict"}}), &["dangerous-tool-use-2026-09-03"])] + #[case::context_management_edits( + json!({"context_management": {"edits": [ + {"type": "compact_20260112"}, + {"type": "clear_tool_uses_20250919"} + ]}}), + &["compact-2026-01-12", "context-management-2025-06-27"] + )] + #[case::advisor_tool_is_not_a_vertex_beta( + json!({"tools": [{"type": "advisor_20260301", "name": "advisor"}]}), + &[] + )] + #[case::fast_mode_is_not_a_vertex_beta(json!({"speed": "fast"}), &[])] + fn beta_headers_follow_the_vertex_policy(#[case] fields: Value, #[case] expected: &[&str]) { + assert_eq!(betas(fields), expected); + } +} diff --git a/litellm-rust/crates/llms/src/vertex_ai/mod.rs b/litellm-rust/crates/llms/src/vertex_ai/mod.rs index 3621ff6a2fd..138b893e4ca 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/mod.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/mod.rs @@ -1 +1,3 @@ +pub mod common_utils; +pub mod messages; pub mod ocr; diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index c7865132840..e895f0e72d1 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -72,6 +72,13 @@ def _resolved_provider(request: NativeCall) -> str | None: return optional_str(request.resolved.get("custom_llm_provider")) +def _python_serves(request: NativeCall) -> bool: + """Vertex AI serves only Claude through the Anthropic Messages config, as + ``get_provider_anthropic_messages_config`` does; every other Vertex model keeps Python's + translation through Chat Completions.""" + return _resolved_provider(request) == "vertex_ai" and "claude" not in str(request.resolved["model"]).lower() + + def _context(request: NativeCall) -> RouteContext: return RouteContext( Route.MESSAGES, @@ -84,13 +91,14 @@ _DISPATCH: Final = PublicDispatch( route=Route.MESSAGES, request=lambda args, kwargs: _public_request(_MESSAGES, args, kwargs), context=_context, - bypass=lambda request: request.kwargs.get("is_async") is True, + bypass=lambda request: request.kwargs.get("is_async") is True or _python_serves(request), ) _ADISPATCH: Final = PublicDispatch( route=Route.MESSAGES, request=lambda args, kwargs: _public_request(_AMESSAGES, args, kwargs), context=_context, + bypass=_python_serves, ) diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 73e95c24de4..5024e58b4ca 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -70,7 +70,7 @@ RULES: Final[Rules] = ( RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY), RouteRule(Route.EMBEDDINGS, Rollout.PYTHON_ONLY), RouteRule(Route.OCR, Rollout.RUST_REQUIRED), - RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic"})), + RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic", "vertex_ai"})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), diff --git a/tests/unit/messages/test_dispatch.py b/tests/unit/messages/test_dispatch.py index faaefa1437e..23400bf6440 100644 --- a/tests/unit/messages/test_dispatch.py +++ b/tests/unit/messages/test_dispatch.py @@ -196,6 +196,53 @@ def test_internal_async_marker_bypasses_native() -> None: assert captured == [(args, kwargs)] +@pytest.mark.parametrize( + ("model", "custom_llm_provider"), + [ + ("vertex_ai/gemini-2.5-pro", None), + ("gemini-2.5-pro", "vertex_ai"), + ], +) +def test_vertex_models_other_than_claude_stay_on_python(model: str, custom_llm_provider: str | None) -> None: + args: Final[tuple[object, ...]] = (16, MESSAGES, model) + kwargs: Final[Mapping[str, object]] = {"custom_llm_provider": custom_llm_provider} + expected: Final = response(model) + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records call shape + return expected + + def native(request: NativeCall) -> AnthropicMessagesResponse: + pytest.fail("Rust serves only Claude on Vertex AI") + + result: Final = _DISPATCH.run( + args, + kwargs, + python=python, + binding=messages_binding(native), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=RUST_RULES, + ) + assert result is expected + + +def test_vertex_claude_reaches_native() -> None: + args: Final[tuple[object, ...]] = (16, MESSAGES, "vertex_ai/claude-sonnet-4-5@20250929") + expected: Final = response("vertex_ai/claude-sonnet-4-5@20250929") + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records call shape + pytest.fail("Claude on Vertex AI is served natively") + + result: Final = _DISPATCH.run( + args, + {}, + python=python, + binding=messages_binding(lambda request: expected), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=RUST_RULES, + ) + assert result is expected + + @pytest.mark.parametrize( ("args", "kwargs"), ( diff --git a/tests/unit/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py index 4242294a33a..ea3e61ecc3d 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/tests/unit/rust_bridge/test_catalog.py @@ -27,7 +27,9 @@ def isolated_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: @pytest.mark.parametrize("route", tuple(Route)) -@pytest.mark.parametrize("provider", (None, "bedrock", "mistral", "anthropic", "openai", "azure_ai", "unknown")) +@pytest.mark.parametrize( + "provider", (None, "bedrock", "mistral", "anthropic", "vertex_ai", "openai", "azure_ai", "unknown") +) @pytest.mark.parametrize("process", (None, False, True)) @pytest.mark.parametrize("environment", (None, "0", "1")) def test_shipped_decisions( @@ -45,7 +47,7 @@ def test_shipped_decisions( if route is Route.OCR or (route is Route.TRANSCRIPTION and provider == "bedrock"): assert catalog.rollout(context) is Rollout.RUST_REQUIRED assert catalog.decision(context) is Decision.RUST_REQUIRED - elif route is Route.MESSAGES and provider == "anthropic": + elif route is Route.MESSAGES and provider in ("anthropic", "vertex_ai"): assert catalog.rollout(context) is Rollout.RUST_OPT_IN opted_in: Final = environment == "1" or (environment is None and process is True) assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if opted_in else Decision.PYTHON) From 4cb47edd2a2d6d56c4dc6dbe99fa8585eeb790dd Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 9 Oct 2026 16:29:46 -0700 Subject: [PATCH 006/104] feat(rust): add the DeepSeek Anthropic Messages config (#45614) * feat(rust): add the DeepSeek Anthropic Messages config * refactor(rust): drop the tool discriminator without mutation and use rstest in DeepSeek tests --- .../core-utils/src/get_llm_provider_logic.rs | 2 + .../gateway-inference/tests/messages.rs | 1 + .../inference-messages/src/common_utils.rs | 6 + .../tests/messages/request.rs | 49 ++ .../tests/messages/secrets.rs | 7 + .../src/anthropic/messages/transformation.rs | 2 +- .../crates/llms/src/deepseek/messages/mod.rs | 1 + .../src/deepseek/messages/transformation.rs | 458 ++++++++++++++++++ litellm-rust/crates/llms/src/deepseek/mod.rs | 1 + litellm-rust/crates/llms/src/lib.rs | 1 + litellm/rust_bridge/catalog.py | 2 +- tests/unit/rust_bridge/test_catalog.py | 4 +- 12 files changed, 530 insertions(+), 4 deletions(-) create mode 100644 litellm-rust/crates/llms/src/deepseek/messages/mod.rs create mode 100644 litellm-rust/crates/llms/src/deepseek/messages/transformation.rs create mode 100644 litellm-rust/crates/llms/src/deepseek/mod.rs diff --git a/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs b/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs index a2bb462c2a9..6a82511db5c 100644 --- a/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs +++ b/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs @@ -18,6 +18,8 @@ pub enum LlmProviders { Bedrock, #[strum(serialize = "cohere")] Cohere, + #[strum(serialize = "deepseek")] + Deepseek, #[strum(serialize = "mistral")] Mistral, #[strum(serialize = "openai")] diff --git a/litellm-rust/crates/gateway-inference/tests/messages.rs b/litellm-rust/crates/gateway-inference/tests/messages.rs index 818fdab27f2..be500bad86c 100644 --- a/litellm-rust/crates/gateway-inference/tests/messages.rs +++ b/litellm-rust/crates/gateway-inference/tests/messages.rs @@ -16,6 +16,7 @@ use wiremock::{ #[rstest] #[case::anthropic("anthropic/test-model", "/v1/messages", true)] #[case::azure("azure_ai/test-model", "/anthropic/v1/messages", false)] +#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", false)] #[tokio::test] async fn messages_reaches_the_provider_and_preserves_json_or_sse( #[case] model: &str, diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index d61761b9348..ea8c7cca4f7 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -6,6 +6,7 @@ use litellm_llms::{ azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG, base_llm::messages::transformation::BaseMessagesConfig, bedrock::messages::invoke_transformations::anthropic_claude3_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG, + deepseek::messages::transformation::DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG, vertex_ai::messages::transformation::VERTEX_ANTHROPIC_MESSAGES_CONFIG, }; use serde_json::{Map, Value}; @@ -19,6 +20,7 @@ pub(crate) enum MessagesProvider { Anthropic, AzureAi, Bedrock, + Deepseek, VertexAi, } @@ -28,6 +30,7 @@ impl MessagesProvider { Self::Anthropic => LlmProviders::Anthropic, Self::AzureAi => LlmProviders::AzureAi, Self::Bedrock => LlmProviders::Bedrock, + Self::Deepseek => LlmProviders::Deepseek, Self::VertexAi => LlmProviders::VertexAi, } .into() @@ -38,6 +41,7 @@ impl MessagesProvider { Self::Anthropic => &ANTHROPIC_MESSAGES_CONFIG, Self::AzureAi => &AZURE_ANTHROPIC_MESSAGES_CONFIG, Self::Bedrock => &BEDROCK_ANTHROPIC_MESSAGES_CONFIG, + Self::Deepseek => &DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG, Self::VertexAi => &VERTEX_ANTHROPIC_MESSAGES_CONFIG, } } @@ -50,6 +54,7 @@ pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option Some(MessagesProvider::Anthropic), LlmProviders::AzureAi => Some(MessagesProvider::AzureAi), LlmProviders::Bedrock => Some(MessagesProvider::Bedrock), + LlmProviders::Deepseek => Some(MessagesProvider::Deepseek), LlmProviders::VertexAi if model.to_ascii_lowercase().contains("claude") => { Some(MessagesProvider::VertexAi) } @@ -77,6 +82,7 @@ mod tests { #[case::anthropic("anthropic", MessagesProvider::Anthropic)] #[case::azure_ai("azure_ai", MessagesProvider::AzureAi)] #[case::bedrock("bedrock", MessagesProvider::Bedrock)] + #[case::deepseek("deepseek", MessagesProvider::Deepseek)] #[case::vertex_ai("vertex_ai", MessagesProvider::VertexAi)] fn provider_round_trips_through_its_python_name( #[case] name: &str, diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index b1a989713ec..25a318ecad5 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -12,6 +12,14 @@ use super::*; #[rstest] #[case::anthropic_key("anthropic", Some("sk-ant"), &[], ("x-api-key", "sk-ant"), &["authorization"])] #[case::azure_key("azure_ai", Some("sk-azure"), &[], ("x-api-key", "sk-azure"), &["authorization"])] +#[case::deepseek_key("deepseek", Some("sk-deepseek"), &[], ("x-api-key", "sk-deepseek"), &["authorization"])] +#[case::deepseek_forwards_caller_authorization( + "deepseek", + Some("sk-deepseek"), + &[("Authorization", "Bearer caller")], + ("authorization", "Bearer caller"), + &["x-api-key"] +)] #[case::caller_x_api_key_wins( "azure_ai", Some("rust-fallback-key"), @@ -79,6 +87,7 @@ async fn credentials_become_exactly_one_auth_header( #[rstest] #[case::anthropic("anthropic")] #[case::azure_ai("azure_ai")] +#[case::deepseek("deepseek")] #[tokio::test] async fn a_call_without_credentials_fails_before_sending( call: MessagesCall, @@ -115,6 +124,8 @@ async fn a_call_without_credentials_fails_before_sending( "/v1/messages" )] #[case::azure_ai(MODEL, Some("azure_ai"), "", "/anthropic/v1/messages")] +#[case::deepseek(MODEL, Some("deepseek"), "", "/anthropic/v1/messages")] +#[case::deepseek_openai_compatible_base(MODEL, Some("deepseek"), "/beta", "/anthropic/v1/messages")] #[case::provider_from_model_prefix("anthropic/claude-sonnet-4-5", None, "", "/v1/messages")] #[tokio::test] async fn each_provider_posts_to_its_messages_endpoint( @@ -350,6 +361,7 @@ async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: M #[rstest] #[case::anthropic("anthropic")] #[case::azure_ai("azure_ai")] +#[case::deepseek("deepseek")] #[tokio::test] async fn caller_protocol_headers_win_over_the_defaults(call: MessagesCall, #[case] provider: &str) { let upstream = upstream([message_response()]).await; @@ -579,6 +591,7 @@ async fn replayed_history_is_cleaned_before_sending( #[rstest] #[case::anthropic("anthropic")] #[case::azure("azure_ai")] +#[case::deepseek("deepseek")] #[tokio::test] async fn metadata_is_reduced_to_the_user_id(call: MessagesCall, #[case] provider: &str) { let upstream = upstream([message_response()]).await; @@ -601,6 +614,42 @@ async fn metadata_is_reduced_to_the_user_id(call: MessagesCall, #[case] provider ); } +#[rstest] +#[tokio::test] +async fn deepseek_sends_neither_billing_blocks_nor_the_custom_tool_discriminator( + call: MessagesCall, +) { + let upstream = upstream([message_response()]).await; + + run_message(with_fields( + MessagesCall { + custom_llm_provider: Some("deepseek".into()), + litellm_params: Default::default(), + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }, + json!({ + "system": [ + {"type": "text", "text": "x-anthropic-billing-header: cc_version=1"}, + {"type": "text", "text": "be terse"} + ], + "tools": [{"type": "custom", "name": "get_weather", "input_schema": {"type": "object"}}] + }), + )) + .await; + + let body = only_request(&upstream).await.json(); + assert_eq!( + body["system"], + json!([{"type": "text", "text": "be terse"}]) + ); + assert_eq!( + body["tools"], + json!([{"name": "get_weather", "input_schema": {"type": "object"}}]) + ); +} + #[rstest] #[case::not_streaming(false, ":rawPredict")] #[case::streaming(true, ":streamRawPredict?alt=sse")] diff --git a/litellm-rust/crates/inference-messages/tests/messages/secrets.rs b/litellm-rust/crates/inference-messages/tests/messages/secrets.rs index a33fdd130f1..203cc620d35 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/secrets.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/secrets.rs @@ -18,6 +18,13 @@ use super::*; "/anthropic/v1/messages", &["AZURE_API_KEY", "AZURE_API_BASE"] )] +#[case::deepseek( + "deepseek", + "DEEPSEEK_API_KEY", + "DEEPSEEK_API_BASE", + "/anthropic/v1/messages", + &["DEEPSEEK_API_KEY", "DEEPSEEK_ANTHROPIC_API_BASE", "DEEPSEEK_API_BASE"] +)] #[tokio::test] async fn the_credential_and_base_come_from_the_secret_source( call: MessagesCall, diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 7021a4199df..f83b477471b 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -156,7 +156,7 @@ pub(crate) fn update_headers_with_anthropic_beta( /// The betas a request needs on a host other than the first-party API: the Anthropic set /// filtered and renamed through the host's `BetaProvider` policy. -pub fn provider_feature_betas(request: &MessagesRequest, provider: BetaProvider) -> BetaSet { +pub(crate) fn provider_feature_betas(request: &MessagesRequest, provider: BetaProvider) -> BetaSet { feature_betas(request) .iter() .filter_map(|beta| beta.on(provider)) diff --git a/litellm-rust/crates/llms/src/deepseek/messages/mod.rs b/litellm-rust/crates/llms/src/deepseek/messages/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/llms/src/deepseek/messages/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs new file mode 100644 index 00000000000..43db1529cc8 --- /dev/null +++ b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs @@ -0,0 +1,458 @@ +use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_core_utils::settings::resolve_non_empty; +use litellm_llms_types::{ + formats::messages::{MessagesOptionalParams, MessagesRequest, MessagesTool}, + recognized::Recognized, +}; +use litellm_router_types::LitellmParams; +use serde_json::Value; + +use crate::{ + Error, + anthropic::{ + common_utils::{filter_billing_headers_from_system, has_anthropic_credential}, + messages::{ + handler::shape_anthropic_messages_request, + transformation::{ + DEFAULT_HEADERS, transform_messages_request, update_headers_with_anthropic_beta, + }, + }, + }, + base_llm::{ + auth::{AuthScheme, Headers, ValidatedEnvironment}, + messages::{ + context::MessagesTransformContext, + transformation::{BaseMessagesConfig, MESSAGES_PATH_SUFFIX}, + }, + }, +}; + +pub const DEEPSEEK_API_KEY_ENV: &str = "DEEPSEEK_API_KEY"; +pub const DEEPSEEK_ANTHROPIC_API_BASE_ENV: &str = "DEEPSEEK_ANTHROPIC_API_BASE"; +pub const DEEPSEEK_API_BASE_ENV: &str = "DEEPSEEK_API_BASE"; +pub const DEFAULT_DEEPSEEK_ANTHROPIC_API_BASE: &str = "https://api.deepseek.com/anthropic"; + +const ANTHROPIC_PATH_SEGMENT: &str = "/anthropic"; +const CUSTOM_TOOL_TYPE: &str = "custom"; + +/// DeepSeek's Anthropic-compatible Messages API: the Anthropic payload shaping, a key in +/// `x-api-key`, and a host that rejects the explicit `custom` tool discriminator and the +/// first-party billing system blocks. +pub struct DeepSeekAnthropicMessagesConfig; + +pub const DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG: DeepSeekAnthropicMessagesConfig = + DeepSeekAnthropicMessagesConfig; + +impl BaseMessagesConfig for DeepSeekAnthropicMessagesConfig { + fn shape_request( + &self, + request: MessagesRequest, + reasoning_auto_summary: bool, + ) -> Result { + shape_anthropic_messages_request(request, reasoning_auto_summary) + } + + fn get_complete_url( + &self, + api_base: Option<&str>, + _model: &str, + _litellm_params: &LitellmParams, + _stream: bool, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + Ok(complete_deepseek_anthropic_url(api_base, env_lookup)) + } + + fn transform_anthropic_messages_request( + &self, + request: MessagesRequest, + context: &MessagesTransformContext, + ) -> Result { + let request = transform_messages_request( + MessagesRequest { + params: MessagesOptionalParams { + system: request + .params + .system + .and_then(filter_billing_headers_from_system), + ..request.params + }, + ..request + }, + context, + )?; + Ok(MessagesRequest { + params: MessagesOptionalParams { + tools: request + .params + .tools + .map(|tools| tools.into_iter().map(sanitize_tool).collect()), + ..request.params + }, + ..request + }) + } + + fn secret_names(&self) -> &'static [&'static str] { + &[ + DEEPSEEK_API_KEY_ENV, + DEEPSEEK_ANTHROPIC_API_BASE_ENV, + DEEPSEEK_API_BASE_ENV, + ] + } + + /// A forwarded `x-api-key` or `authorization` is the credential; otherwise the DeepSeek + /// key goes in `x-api-key`. + fn validate_environment( + &self, + headers: Headers, + api_key: Option<&str>, + _model: &str, + _litellm_params: &LitellmParams, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + if has_anthropic_credential(&headers) { + return Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Forwarded, + }); + } + let key = get_api_key(api_key, env_lookup).ok_or(Error::Auth( + litellm_auth::Error::MissingApiKey { + provider: "DeepSeek", + environment_variable: DEEPSEEK_API_KEY_ENV, + }, + ))?; + Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Credential { + placement: CredentialPlacement::Header("x-api-key"), + secret: SecretValue::new(key), + }, + }) + } + + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { + DEFAULT_HEADERS + } + + fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers { + update_headers_with_anthropic_beta(headers, request) + } +} + +pub fn get_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + resolve_non_empty(api_key, env_lookup, &[DEEPSEEK_API_KEY_ENV]) +} + +pub fn get_api_base(api_base: Option<&str>, env_lookup: &dyn Fn(&str) -> Option) -> String { + resolve_non_empty( + api_base, + env_lookup, + &[DEEPSEEK_ANTHROPIC_API_BASE_ENV, DEEPSEEK_API_BASE_ENV], + ) + .unwrap_or_else(|| DEFAULT_DEEPSEEK_ANTHROPIC_API_BASE.to_string()) +} + +/// Python's `get_complete_url`: a base already ending in the Anthropic messages path is kept, +/// otherwise the OpenAI compatible suffixes are peeled off and `/anthropic/v1/messages` is +/// appended once. +pub fn complete_deepseek_anthropic_url( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> String { + let api_base = get_api_base(api_base, env_lookup); + let api_base = api_base.trim_end_matches('/'); + if api_base.ends_with(MESSAGES_PATH_SUFFIX) && api_base.contains("/anthropic/") { + return api_base.to_string(); + } + let api_base = [MESSAGES_PATH_SUFFIX, "/v1", "/beta"] + .into_iter() + .fold(api_base, |base, suffix| { + base.strip_suffix(suffix).unwrap_or(base) + }); + if api_base.ends_with(ANTHROPIC_PATH_SEGMENT) || api_base.contains("/anthropic/") { + return format!("{api_base}{MESSAGES_PATH_SUFFIX}"); + } + format!("{api_base}{ANTHROPIC_PATH_SEGMENT}{MESSAGES_PATH_SUFFIX}") +} + +/// DeepSeek rejects Anthropic's explicit `{"type": "custom"}` tool discriminator, so it is +/// dropped; every other tool is sent as given. +fn sanitize_tool(tool: Recognized) -> Recognized { + match tool { + Recognized::Unrecognized(Value::Object(fields)) + if fields.get("type").and_then(Value::as_str) == Some(CUSTOM_TOOL_TYPE) => + { + Recognized::Unrecognized(Value::Object( + fields + .into_iter() + .filter(|(name, _)| name != "type") + .collect(), + )) + } + other => other, + } +} + +#[cfg(test)] +mod tests { + use litellm_auth::CredentialPlacement; + use rstest::rstest; + use serde_json::{Value, json}; + + use super::*; + + fn no_env(_: &str) -> Option { + None + } + + fn request_from(value: Value) -> MessagesRequest { + serde_json::from_value(value).unwrap() + } + + fn transformed(value: Value) -> Value { + serde_json::to_value( + DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG + .transform_anthropic_messages_request( + request_from(value), + &MessagesTransformContext::default(), + ) + .unwrap(), + ) + .unwrap() + } + + #[rstest] + #[case::no_base(None)] + #[case::anthropic_v1("https://api.deepseek.com/anthropic/v1")] + #[case::anthropic("https://api.deepseek.com/anthropic")] + #[case::anthropic_with_trailing_slash("https://api.deepseek.com/anthropic/")] + #[case::host_only("https://api.deepseek.com")] + #[case::openai_v1("https://api.deepseek.com/v1")] + #[case::openai_messages("https://api.deepseek.com/v1/messages")] + #[case::openai_beta("https://api.deepseek.com/beta")] + #[case::complete("https://api.deepseek.com/anthropic/v1/messages")] + fn every_base_spelling_reaches_the_anthropic_messages_endpoint( + #[case] api_base: impl Into>, + ) { + assert_eq!( + complete_deepseek_anthropic_url(api_base.into(), &no_env), + "https://api.deepseek.com/anthropic/v1/messages" + ); + } + + #[rstest] + #[case::proxy_under_anthropic_prefix( + "https://proxy.example/anthropic/v1/messages", + "https://proxy.example/anthropic/v1/messages" + )] + #[case::messages_path_without_the_anthropic_segment_is_rebuilt( + "https://proxy.example/v1/messages", + "https://proxy.example/anthropic/v1/messages" + )] + #[case::anthropic_segment_in_the_middle_is_kept( + "https://proxy.example/anthropic/tenant", + "https://proxy.example/anthropic/tenant/v1/messages" + )] + fn custom_bases_follow_the_python_url_rules(#[case] api_base: &str, #[case] expected: &str) { + assert_eq!( + complete_deepseek_anthropic_url(Some(api_base), &no_env), + expected + ); + } + + #[rstest] + #[case::anthropic_base_wins( + &[("DEEPSEEK_ANTHROPIC_API_BASE", "https://anthropic.example"), ("DEEPSEEK_API_BASE", "https://openai.example/v1")], + "https://anthropic.example/anthropic/v1/messages" + )] + #[case::openai_base_is_the_fallback( + &[("DEEPSEEK_API_BASE", "https://openai.example/v1")], + "https://openai.example/anthropic/v1/messages" + )] + #[case::blank_env_is_absent( + &[("DEEPSEEK_ANTHROPIC_API_BASE", " "), ("DEEPSEEK_API_BASE", "")], + "https://api.deepseek.com/anthropic/v1/messages" + )] + fn the_base_comes_from_the_environment_in_python_order( + #[case] env: &[(&str, &str)], + #[case] expected: &str, + ) { + let lookup = |name: &str| { + env.iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + }; + assert_eq!(complete_deepseek_anthropic_url(None, &lookup), expected); + } + + fn validated( + forwarded: &[(&str, &str)], + api_key: Option<&str>, + env: &dyn Fn(&str) -> Option, + ) -> Result { + DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG.validate_environment( + forwarded + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(), + api_key, + "deepseek-v4-pro", + &LitellmParams::default(), + env, + ) + } + + #[rstest] + #[case::param(Some("sk-deepseek"), &[], "sk-deepseek")] + #[case::env(None, &[("DEEPSEEK_API_KEY", "sk-env")], "sk-env")] + #[case::blank_param_falls_back_to_env(Some(" "), &[("DEEPSEEK_API_KEY", "sk-env")], "sk-env")] + fn the_deepseek_key_goes_in_x_api_key( + #[case] api_key: Option<&str>, + #[case] env: &[(&str, &str)], + #[case] expected: &str, + ) { + let lookup = |name: &str| { + env.iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + }; + let validated = validated(&[], api_key, &lookup).unwrap(); + assert!(matches!( + validated.auth, + AuthScheme::Credential { + placement: CredentialPlacement::Header("x-api-key"), + ref secret + } if secret.expose() == expected + )); + assert!(validated.headers.is_empty()); + } + + #[rstest] + #[case::x_api_key(&[("X-Api-Key", "caller")])] + #[case::bearer(&[("Authorization", "Bearer caller")])] + fn a_forwarded_credential_is_sent_as_is(#[case] forwarded: &[(&str, &str)]) { + assert!(matches!( + validated(forwarded, Some("sk-deepseek"), &no_env) + .unwrap() + .auth, + AuthScheme::Forwarded + )); + } + + #[rstest] + fn a_call_without_a_key_names_the_deepseek_variable() { + assert!(matches!( + validated(&[], None, &no_env).unwrap_err(), + Error::Auth(litellm_auth::Error::MissingApiKey { + provider: "DeepSeek", + environment_variable: "DEEPSEEK_API_KEY", + }) + )); + } + + #[rstest] + fn secret_names_cover_the_key_and_both_bases() { + assert_eq!( + DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG.secret_names(), + &[ + "DEEPSEEK_API_KEY", + "DEEPSEEK_ANTHROPIC_API_BASE", + "DEEPSEEK_API_BASE" + ] + ); + } + + #[rstest] + fn default_headers_are_anthropic_version_and_json() { + assert_eq!( + DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG.default_headers(), + &[ + ("anthropic-version", "2023-06-01"), + ("content-type", "application/json"), + ] + ); + } + + #[rstest] + fn the_custom_discriminator_is_dropped_and_other_tools_are_kept() { + let body = transformed(json!({ + "model": "deepseek-v4-pro", + "max_tokens": 100, + "messages": [{"role": "user", "content": "Use the tool."}], + "tools": [ + {"type": "custom", "name": "get_weather", "description": "Get weather", "input_schema": {"type": "object"}}, + {"name": "untyped", "input_schema": {"type": "object"}}, + {"type": "web_search_20260209", "name": "web_search", "max_uses": 1}, + {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"} + ] + })); + assert_eq!( + body["tools"], + json!([ + {"name": "get_weather", "description": "Get weather", "input_schema": {"type": "object"}}, + {"name": "untyped", "input_schema": {"type": "object"}}, + {"type": "web_search_20260209", "name": "web_search", "max_uses": 1}, + {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"} + ]) + ); + } + + #[rstest] + #[case::billing_block_among_others( + json!([ + {"type": "text", "text": "x-anthropic-billing-header: cc_version=1"}, + {"type": "text", "text": "real system prompt"} + ]), + Some(json!([{"type": "text", "text": "real system prompt"}])) + )] + #[case::only_billing_blocks( + json!([{"type": "text", "text": "x-anthropic-billing-header: cc_version=1"}]), + None + )] + #[case::billing_string(json!("x-anthropic-billing-header: cc_version=1"), None)] + #[case::ordinary_string(json!("be terse"), Some(json!("be terse")))] + #[case::billing_text_in_a_non_text_block_is_kept( + json!([{"type": "other", "text": "x-anthropic-billing-header: cc_version=1"}]), + Some(json!([{"type": "other", "text": "x-anthropic-billing-header: cc_version=1"}])) + )] + fn billing_metadata_is_stripped_from_the_system_prompt( + #[case] system: Value, + #[case] expected: Option, + ) { + let body = transformed(json!({ + "model": "deepseek-v4-pro", + "max_tokens": 100, + "system": system, + "messages": [{"role": "user", "content": "hi"}] + })); + assert_eq!(body.get("system").cloned(), expected); + } + + #[rstest] + fn thinking_history_and_params_pass_through() { + let messages = json!([ + {"role": "user", "content": "Use the tool."}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "I should call the tool.", "signature": "sig"}, + {"type": "tool_use", "id": "toolu_123", "name": "get_weather", "input": {"city": "Sao Paulo"}} + ]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_123", "content": "Sunny"}]} + ]); + let body = transformed(json!({ + "model": "deepseek-v4-pro", + "max_tokens": 100, + "thinking": {"type": "enabled", "budget_tokens": 1024}, + "messages": messages + })); + assert_eq!(body["messages"], messages); + assert_eq!( + body["thinking"], + json!({"type": "enabled", "budget_tokens": 1024}) + ); + } +} diff --git a/litellm-rust/crates/llms/src/deepseek/mod.rs b/litellm-rust/crates/llms/src/deepseek/mod.rs new file mode 100644 index 00000000000..ba63992f3cb --- /dev/null +++ b/litellm-rust/crates/llms/src/deepseek/mod.rs @@ -0,0 +1 @@ +pub mod messages; diff --git a/litellm-rust/crates/llms/src/lib.rs b/litellm-rust/crates/llms/src/lib.rs index 957cae95df6..9c62be44fc2 100644 --- a/litellm-rust/crates/llms/src/lib.rs +++ b/litellm-rust/crates/llms/src/lib.rs @@ -4,6 +4,7 @@ pub mod azure_ai; pub mod base_llm; pub mod bedrock; pub mod cohere; +pub mod deepseek; mod error; pub mod mistral; pub mod openai; diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 5024e58b4ca..dda9ed0b566 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -70,7 +70,7 @@ RULES: Final[Rules] = ( RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY), RouteRule(Route.EMBEDDINGS, Rollout.PYTHON_ONLY), RouteRule(Route.OCR, Rollout.RUST_REQUIRED), - RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic", "vertex_ai"})), + RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic", "deepseek", "vertex_ai"})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), diff --git a/tests/unit/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py index ea3e61ecc3d..1ba95507d03 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/tests/unit/rust_bridge/test_catalog.py @@ -28,7 +28,7 @@ def isolated_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: @pytest.mark.parametrize("route", tuple(Route)) @pytest.mark.parametrize( - "provider", (None, "bedrock", "mistral", "anthropic", "vertex_ai", "openai", "azure_ai", "unknown") + "provider", (None, "bedrock", "mistral", "anthropic", "deepseek", "vertex_ai", "openai", "azure_ai", "unknown") ) @pytest.mark.parametrize("process", (None, False, True)) @pytest.mark.parametrize("environment", (None, "0", "1")) @@ -47,7 +47,7 @@ def test_shipped_decisions( if route is Route.OCR or (route is Route.TRANSCRIPTION and provider == "bedrock"): assert catalog.rollout(context) is Rollout.RUST_REQUIRED assert catalog.decision(context) is Decision.RUST_REQUIRED - elif route is Route.MESSAGES and provider in ("anthropic", "vertex_ai"): + elif route is Route.MESSAGES and provider in ("anthropic", "deepseek", "vertex_ai"): assert catalog.rollout(context) is Rollout.RUST_OPT_IN opted_in: Final = environment == "1" or (environment is None and process is True) assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if opted_in else Decision.PYTHON) From 3069a467d0deb9f6cdf9c65e4b56ca93d57ff774 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:36:32 -0700 Subject: [PATCH 007/104] docs(agents): ban semicolons and colons in human-facing text (#45684) * docs(agents): ban semicolons and colons in human-facing text Tighten the writing rule so ";" and ":" are only allowed where the sentence would be nonsensical without them, rewrite the existing semicolons in AGENTS.md to follow it, and note that docs (litellm-docs) keep their trailing "." * docs(agents): rewrite semicolons in nested AGENTS.md files * docs(agents): hand-maintain the dashboard Next.js warning and sync e2e CONTRIBUTING lines * docs(agents): point at the PR template path that exists * docs(agents): restore the generated Next.js block in the dashboard AGENTS.md --- AGENTS.md | 34 ++++----- litellm-rust/AGENTS.md | 10 +-- litellm-rust/crates/cache-response/AGENTS.md | 2 +- .../crates/callbacks-legacy-python/AGENTS.md | 20 +++--- litellm-rust/crates/coroutine/AGENTS.md | 6 +- litellm-rust/crates/gateway-auth/AGENTS.md | 8 +-- .../crates/gateway-inference/AGENTS.md | 12 ++-- litellm-rust/crates/gateway/AGENTS.md | 6 +- litellm-rust/crates/host-native/AGENTS.md | 8 +-- litellm-rust/crates/host-python/AGENTS.md | 30 ++++---- litellm-rust/crates/host/AGENTS.md | 14 ++-- .../crates/host/src/machine/AGENTS.md | 4 +- litellm-rust/crates/inference-chat/AGENTS.md | 2 +- .../crates/inference-messages/AGENTS.md | 4 +- litellm-rust/crates/inference-ocr/AGENTS.md | 2 +- .../crates/inference-responses/AGENTS.md | 2 +- .../crates/inference-transcription/AGENTS.md | 4 +- litellm-rust/crates/inference/AGENTS.md | 32 ++++----- litellm-rust/crates/llms-types/AGENTS.md | 2 +- litellm-rust/crates/llms/AGENTS.md | 4 +- litellm-rust/crates/python-bridge/AGENTS.md | 58 +++++++-------- .../crates/python-bridge/src/cache/AGENTS.md | 2 +- .../crates/python-bridge/src/routes/AGENTS.md | 4 +- litellm-rust/crates/python-compat/AGENTS.md | 18 ++--- .../crates/storage-clickhouse/AGENTS.md | 2 +- .../proxy/_experimental/mcp_server/AGENTS.md | 12 ++-- litellm/rust_bridge/AGENTS.md | 4 +- litellm/tracing/AGENTS.md | 2 +- tests/AGENTS.md | 6 +- tests/e2e/AGENTS.md | 72 +++++++++---------- tests/e2e/CONTRIBUTING.md | 8 +-- tests/e2e_harness/AGENTS.md | 4 +- tests/integration/AGENTS.md | 4 +- tests/rust-python-harness/AGENTS.md | 18 ++--- .../strategies/trace_parity/AGENTS.md | 2 +- tests/unit/AGENTS.md | 2 +- ui/litellm-dashboard/AGENTS.md | 4 +- .../src/components/chat/AGENTS.md | 4 +- 38 files changed, 216 insertions(+), 216 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index f0d1a083abb..0254893554d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -25,7 +25,7 @@ Same thing for bug fixes. The tests should make it so that this specific bug can Never test structure of code only function of it -A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent"; assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken +A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent". Assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken `tests/unit/` mirrors `litellm/` in a parallel path (see `tests/unit/AGENTS.md`). Name tests `test_.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_.py` if you're the first test there). One focused regression test beats many shallow ones @@ -33,20 +33,20 @@ End-to-end tests belong in `tests/e2e/` and must follow the harness conventions When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD` -When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule. A section you have nothing to put in (Relevant issues, Affected release, Linear ticket, Caveats, QA runbook, and so on) is removed entirely, heading included, never left as an empty title +When writing a PR body, treat the comments and imperative instructions inside .github/PULL_REQUEST_TEMPLATE/general.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/PULL_REQUEST_TEMPLATE/general.md from disk before writing a PR body to make sure you see every comment rule. A section you have nothing to put in (Relevant issues, Affected release, Linear ticket, Caveats, QA runbook, and so on) is removed entirely, heading included, never left as an empty title Same applies for filing bug reports and feature requests, with .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml, respectively If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just drop the section -Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it +Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`. The Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into. I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y: - don't use emojis -- don't use "—". 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, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message +- 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 "." +- 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 @@ -58,17 +58,17 @@ The four lint gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gat `make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice -`make check`, `make lint`, `scripts/pre_commit_lint.sh`, and the standalone lint gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gate.py`, `scripts/type_check_gate.py`) each hold one of 2 machine-wide slots, so when other sessions or worktrees on the same box are already running heavy work, yours prints "all N machine-wide slots are busy; queueing" and then stays quiet until a slot frees. Give the command a long timeout and let it wait rather than killing it, retrying it, or assuming it hung. Don't change the # of machine-wide slots or make it unlimited by setting `LITELLM_GATE_SLOTS=0` +`make check`, `make lint`, `scripts/pre_commit_lint.sh`, and the standalone lint gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gate.py`, `scripts/type_check_gate.py`) each hold one of 2 machine-wide slots, so when other sessions or worktrees on the same box are already running heavy work, yours prints "all N machine-wide slots are busy" followed by "queueing" and then stays quiet until a slot frees. Give the command a long timeout and let it wait rather than killing it, retrying it, or assuming it hung. Don't change the # of machine-wide slots or make it unlimited by setting `LITELLM_GATE_SLOTS=0` If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their cap in `ANY_CAPS`, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in -If you get an LIT001 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason +If you get an LIT001 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used. Reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # `. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing Commit and push your work when you're done without asking -When referencing or running models (coding, QA'ing, writing docs, writing tests, etc.), use the latest model in that model family unless otherwise specified; treat your training knowledge, memories, configs, and tests as stale, and determine the family's latest with model_prices_and_context_window.json or the web +When referencing or running models (coding, QA'ing, writing docs, writing tests, etc.), use the latest model in that model family unless otherwise specified. Treat your training knowledge, memories, configs, and tests as stale, and determine the family's latest with model_prices_and_context_window.json or the web Always pull before starting any work. The checkout or worktree may be sitting on a stale branch @@ -84,7 +84,7 @@ Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer de Do not put names of customers or customer company names in code, PR descriptions, issue bodies, etc. This means never mention literally any company name. Especially if you're about to say a sentence mentioning that the reason the PR exists was a feature/model/bug fix/etc. requested by a company. That's the indication that you should replace that company name with "the customer". e.g. not "Model request from Acme (Pylon #1234)" but "Model request from a customer (Pylon #1234)". This is because the codebase is public. The only exception is for publicly known providers or vendors such as OpenAI, Anthropic, AWS Bedrock, etc. only IF we're adding support for that provider/vendor in general and NOT if that PR or whatnot was a request by one of them, and they're actually one of our customers. -CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI +CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`). Download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: ` @@ -92,19 +92,19 @@ Follow these coding conventions for new/updated code (a three-line fix in a lega - Composition over inheritance - Never-nester: early returns over deep nesting -- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never) -- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `MappingProxyType`, etc. - - Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: ` +- Don't throw. Model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never) +- No mutation. Don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `MappingProxyType`, etc. + - Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011). `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: ` - Qualify every TypedDict field with `ReadOnly[...]` (LIT012), which nests freely with `Required` / `NotRequired` / `Annotated` in any order. If making the key writable is truly unavoidable, suppress with `# writable-ok: ` - - Comprehensions take at most one `for` clause and one `if` clause (LIT014); split stacked clauses into a helper generator, a named intermediate, or a plain loop. Suppress with `# comprehension-ok: ` only when unavoidable + - Comprehensions take at most one `for` clause and one `if` clause (LIT014). Split stacked clauses into a helper generator, a named intermediate, or a plain loop. Suppress with `# comprehension-ok: ` only when unavoidable - Every pydantic model must be frozen (LIT015), set via `model_config = ConfigDict(frozen=True)`, a dict-literal `model_config`, an inner `class Config`, or the class keywords. Subclasses inherit it unless they override it. Replace in-place field writes with `model_copy(update=...)`. If making the model mutable is truly unavoidable, suppress with `# frozen-ok: ` on the `class` line - Use dependency injection -- Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed +- Fully typed. No `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed - Use tagged unions + match - No monster files or god objects - No file sprawl: deliberate file and folder structure -- Standard over hand-rolled: use the official SDK or a library where one exists; where none does, follow industry standards instead of inventing local conventions -- API-fragmentation-aware: when logic must branch on which API surface produced or consumes data (e.g. chat completions vs Anthropic Messages vs Responses API shapes), proactively look for an existing shared helper (e.g. `litellm_core_utils/prompt_templates/factory.py`) before writing per-surface parsing in the new module; if none exists, add one there instead of duplicating the same format-detection logic in every new guardrail/integration +- Standard over hand-rolled: use the official SDK or a library where one exists. Where none does, follow industry standards instead of inventing local conventions +- API-fragmentation-aware: when logic must branch on which API surface produced or consumes data (e.g. chat completions vs Anthropic Messages vs Responses API shapes), proactively look for an existing shared helper (e.g. `litellm_core_utils/prompt_templates/factory.py`) before writing per-surface parsing in the new module. If none exists, add one there instead of duplicating the same format-detection logic in every new guardrail/integration Follow conventional commits for commit names and PR titles @@ -130,4 +130,4 @@ Before implementing: Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify -Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review; never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or raise `ANY_CAPS` to satisfy a check +Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review. Never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or raise `ANY_CAPS` to satisfy a check diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index c22f9575f74..b028425c7c1 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -4,7 +4,7 @@ For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.a For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md) -For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of a hand-listed `ALL` array; derive Strum string conversions instead of hand-written variant-to-string matches +For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of a hand-listed `ALL` array. Derive Strum string conversions instead of hand-written variant-to-string matches Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all` @@ -15,20 +15,20 @@ Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `st - A test that only uses the crate's public API lives in `crates//tests/.rs`, next to `src/` - Split a mixed test file along that line instead of widening visibility to move it - A test for another crate's item belongs in that crate, not in a downstream one -- Never set `autotests = false` or hand-list `[[test]]` targets; every file directly under `tests/` is discovered by cargo, and a shared helper goes in `tests//mod.rs` or `tests//support.rs` so it is not picked up as a test crate of its own +- Never set `autotests = false` or hand-list `[[test]]` targets. Every file directly under `tests/` is discovered by cargo, and a shared helper goes in `tests//mod.rs` or `tests//support.rs` so it is not picked up as a test crate of its own ## Test fixtures and cases Use [`#[rstest]`](https://docs.rs/rstest/latest/rstest/attr.rstest.html) for new and updated tests and [`#[fixture]`](https://docs.rs/rstest/latest/rstest/attr.fixture.html) for reusable setup, injected through typed test arguments. Express input variations as named `#[case::name(...)]` cases instead of loops or duplicated tests so each failure identifies its case. Keep behavior assertions in the test body and fixtures focused on setup. Use the workspace `rstest` dependency -- Never loop over inputs (`for`, `.iter().for_each`, `.all`) inside a test body; give each input its own `#[case::name(...)]`, or use `#[values(...)]` for a cross product +- Never loop over inputs (`for`, `.iter().for_each`, `.all`) inside a test body. Give each input its own `#[case::name(...)]`, or use `#[values(...)]` for a cross product - Exception: a test pinning a Rust table against a repo-owned data file (for example `include_str!` of a JSON config) may iterate that file's entries ## Error definitions - A crate's errors live in `src/error.rs`, defined with `thiserror`, and re-exported from `lib.rs` -- Put message templates in the variant's `#[error(...)]` declaration. Callers pass only the small typed arguments needed to fill them, never `Error::Variant(format!(...))` or a preformatted message. Keep the smallest set of neutral variants that callers need to distinguish; different wording or providers do not justify new variants -- Default to one top-level `Error` enum per crate, with one variant per failure mode and a `#[error(...)]` message on each. A failure mode is something a caller handles differently (phase, status code, retry, a message Python parity pins exactly); failures no caller tells apart share one variant and differ only in its message +- Put message templates in the variant's `#[error(...)]` declaration. Callers pass only the small typed arguments needed to fill them, never `Error::Variant(format!(...))` or a preformatted message. Keep the smallest set of neutral variants that callers need to distinguish. Different wording or providers do not justify new variants +- Default to one top-level `Error` enum per crate, with one variant per failure mode and a `#[error(...)]` message on each. A failure mode is something a caller handles differently (phase, status code, retry, a message Python parity pins exactly). Failures no caller tells apart share one variant and differ only in its message - Keep shared error enums minimal and provider-neutral. Provider names, credential types, configuration fields, and setup guidance belong in caller-supplied data, not dedicated variants or hardcoded shared messages. Reuse a variant for the same failure mode across providers, such as `MissingApiBase { provider: "Azure", guidance: "..." }`. An exact parity message does not justify a provider-specific variant when caller-supplied context can preserve it - Wrap a lower-level error as a variant with `#[from]` or `#[source]` instead of flattening it to a string - Exception: split into separate types when different functions fail in disjoint ways, especially when different callers see them. A shared enum would force every caller to match variants its function can never return diff --git a/litellm-rust/crates/cache-response/AGENTS.md b/litellm-rust/crates/cache-response/AGENTS.md index 4dbcc65d403..edd6a2f4ea9 100644 --- a/litellm-rust/crates/cache-response/AGENTS.md +++ b/litellm-rust/crates/cache-response/AGENTS.md @@ -24,6 +24,6 @@ Keep unary caching independent of stream-only methods. Store streams only after Test each contract in its owner: storage capabilities in backend tests, envelopes and freshness here, reuse and replay in core, Python callback and fallback behavior at the bridge, and HTTP behavior at the gateway. Run backend contract checks and Python response-codec fixtures before exposing a new backend -`ScopedCache` requires an explicit shared or isolated scope at construction. Per-call `CachePolicy` controls reads, writes, expiry, and freshness without replacing the attached scope or service. `CacheOptions` binds that policy to an explicit scope for storage requests and has no default sharing policy. Versioned native envelopes reject incompatible API surfaces and versions as misses; this envelope is distinct from the legacy Python response codec +`ScopedCache` requires an explicit shared or isolated scope at construction. Per-call `CachePolicy` controls reads, writes, expiry, and freshness without replacing the attached scope or service. `CacheOptions` binds that policy to an explicit scope for storage requests and has no default sharing policy. Versioned native envelopes reject incompatible API surfaces and versions as misses. This envelope is distinct from the legacy Python response codec Response storage is not the source of budget or rate-limit coordination dependencies. Keep counters, reservations, and atomic admission operations out of `ResponseCacheService`, including when both services happen to use Redis diff --git a/litellm-rust/crates/callbacks-legacy-python/AGENTS.md b/litellm-rust/crates/callbacks-legacy-python/AGENTS.md index de99fe17a4b..1dd09575c08 100644 --- a/litellm-rust/crates/callbacks-legacy-python/AGENTS.md +++ b/litellm-rust/crates/callbacks-legacy-python/AGENTS.md @@ -1,20 +1,20 @@ - Target invariants, not completion claims - This crate owns compatibility for all existing Python callbacks and loggers, including `CustomLogger`. `mapping.rs` owns the executable call bindings and the inventory of Python-owned hooks. A Python-owned entry records an existing path, never permission to invoke it a second time. The native call adapter preserves the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on) - Smell test: if a future callback host (`callbacks-v1-python`, WASM, in-process Rust) could share a piece of this crate, it does not belong here - - SDK request policy (credential inheritance, the budget and retry-count limits) is a separate hook supplied by `python-bridge`; compose it after this adapter so logging adopts the final keyword view before policy mutates or rejects it - - The driver in `litellm-host-python`, the routes and core see one `PythonCallHooks` using the shared `CallEvent`; they never learn which Python objects consume a call + - SDK request policy (credential inheritance, the budget and retry-count limits) is a separate hook supplied by `python-bridge`. Compose it after this adapter so logging adopts the final keyword view before policy mutates or rejects it + - The driver in `litellm-host-python`, the routes and core see one `PythonCallHooks` using the shared `CallEvent`. They never learn which Python objects consume a call - Every litellm Python internal Rust still borrows is a variant of `LegacyPython`, grouped by subsystem, with its signature pinned in `python_contract.json` - The enum only shrinks: when Rust owns a subsystem, delete its group rather than adding a Rust path beside it - Calling a user's own callback directly is permanent Python surface and gets its own type outside `LegacyPython` -- `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the call rewrites it (setup, deployment hook, preflight) and the bound request object backing omitted keywords; shared bridge composition hands it to `LegacyLogging`; routes use the neutral call boundary +- `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the call rewrites it (setup, deployment hook, preflight) and the bound request object backing omitted keywords. Shared bridge composition hands it to `LegacyLogging`. Routes use the neutral call boundary - `LoggingOperation` selects legacy logging entrypoints and response handling. It belongs here rather than in shared inference data contracts -- `setup` reuses a `Logging` passed as `litellm_logging_obj` (the proxy and Router) and otherwise builds one through `function_setup`; which callbacks run is `Logging`'s decision, never this crate's -- Callbacks receive the caller's own objects and may mutate them; this crate alone carries that obligation - - Retain complete boundary arguments, opaque values, aliases, omitted/default distinctions and deliberate copies; preserve the deployment-hook kwargs view +- `setup` reuses a `Logging` passed as `litellm_logging_obj` (the proxy and Router) and otherwise builds one through `function_setup`. Which callbacks run is `Logging`'s decision, never this crate's +- Callbacks receive the caller's own objects and may mutate them. This crate alone carries that obligation + - Retain complete boundary arguments, opaque values, aliases, omitted/default distinctions and deliberate copies. Preserve the deployment-hook kwargs view - Before `pre_call`, re-alias every body key whose value equals the caller's argument to the caller's own object, resolved through `litellm_host_python::lookup` - - Retain body/header roots from `pre_call` to `post_call`; in-place mutation reaches the wire, envelope field replacement is visible to later callbacks only + - Retain body/header roots from `pre_call` to `post_call`. In-place mutation reaches the wire, envelope field replacement is visible to later callbacks only - Success and failure handlers receive the exact selected public response or exception - - A failure-handler error cannot suppress the other eligible family or replace the mapped provider error; a cancellation ends the call with no further dispatch - - Dispatch errors never replay provider work or trigger the opposite outcome; the proxy releases deferred success at most once + - A failure-handler error cannot suppress the other eligible family or replace the mapped provider error. A cancellation ends the call with no further dispatch + - Dispatch errors never replay provider work or trigger the opposite outcome. The proxy releases deferred success at most once - Delivery follows the registry, not the callable's type: direct, awaited, executor-submitted, logging-worker and deferred paths stay distinct -- Traverse every retained Python edge; `close` is idempotent and restores the correlation context once +- Traverse every retained Python edge. `close` is idempotent and restores the correlation context once diff --git a/litellm-rust/crates/coroutine/AGENTS.md b/litellm-rust/crates/coroutine/AGENTS.md index fcb4f4df47f..403d53a0a6d 100644 --- a/litellm-rust/crates/coroutine/AGENTS.md +++ b/litellm-rust/crates/coroutine/AGENTS.md @@ -1,9 +1,9 @@ # Requirements -Core must pause mid-call to ask the host for things it cannot do itself (Python callbacks, secret and token reads, `before_send` rewrites, stream demand), then continue where it stopped. Any change to this crate must keep every requirement below; the alternatives section says which one each rejected design breaks +Core must pause mid-call to ask the host for things it cannot do itself (Python callbacks, secret and token reads, `before_send` rewrites, stream demand), then continue where it stopped. Any change to this crate must keep every requirement below. The alternatives section says which one each rejected design breaks - R1 Core never calls the host: it names an op and waits for the answer, so it stays free of PyO3 and of any other host runtime -- R2 Async host work is awaited by the host's own driver in the caller's asyncio task (`litellm/rust_bridge/lifecycle.py`), so `contextvars` writes reach the caller; a Rust-side `into_future` would run it in a copied context +- R2 Async host work is awaited by the host's own driver in the caller's asyncio task (`litellm/rust_bridge/lifecycle.py`), so `contextvars` writes reach the caller. A Rust-side `into_future` would run it in a copied context - R3 The body awaits real I/O (HTTP, `spawn_blocking`, timers) between yields, so `resume` is itself a future driven by the caller's runtime - R4 Route code stays straight-line async (`host.route(OcrOp::ReadDocument).await?`) instead of hand-written states - R5 Each op fixes its answer type at compile time: a host cannot answer `ReadDocument` with a token, and core never matches a result variant it did not ask for @@ -20,7 +20,7 @@ Core must pause mid-call to ask the host for things it cannot do itself (Python - `corosensei` and other stackful coroutines: sync bodies on their own stack, no async I/O inside (R3) - A hand-written phase enum with an `advance` match (the old `HostPhase`): every await point becomes a state (R4) - An injected host trait with `async fn`s: core would call the host itself (R1, R2) -- Sans-IO, where core does no I/O and HTTP becomes one more host op: keeps every requirement and makes `resume` a pure step function, but HTTP, streaming, retries and timeouts would move out of core into every bridge; the one real alternative, not taken +- Sans-IO, where core does no I/O and HTTP becomes one more host op: keeps every requirement and makes `resume` a pure step function, but HTTP, streaming, retries and timeouts would move out of core into every bridge. The one real alternative, not taken - Temporal's Rust workflow SDK (`WorkflowFuture`, `WfContext`) is the closest precedent: an `async fn` polled in place, commands sent over a channel with a oneshot to unblock them. Roles are inverted there (the language SDK owns the program, core answers), and its workflow body may not do real I/O # Tradeoffs accepted diff --git a/litellm-rust/crates/gateway-auth/AGENTS.md b/litellm-rust/crates/gateway-auth/AGENTS.md index 916e96bd358..58457a3151b 100644 --- a/litellm-rust/crates/gateway-auth/AGENTS.md +++ b/litellm-rust/crates/gateway-auth/AGENTS.md @@ -2,9 +2,9 @@ Inbound authentication uses a verifier, identity resolver, and authorizer injected into `Auth`. `Auth::from_config` composes the current master-key implementation. `Auth::new` accepts alternate implementations and an injected clock -`CredentialExtractor` selects a token or transport credential from HTTP request parts. The default `Bearer` extractor is strict; a host can supply another extractor with `Auth::with_extractor` for a separate route profile. A transport credential is only a selection marker, never proof of identity +`CredentialExtractor` selects a token or transport credential from HTTP request parts. The default `Bearer` extractor is strict. A host can supply another extractor with `Auth::with_extractor` for a separate route profile. A transport credential is only a selection marker, never proof of identity -`Authenticator` receives that selected credential and HTTP request parts. Transport verifiers must obtain proof from trusted server extensions, not client-supplied identity headers. It returns `VerifiedIdentity` only after verification. Request parts provide headers and trusted transport extensions for integrations; the verifier must validate any client-controlled claims before treating them as identity. There is no fallback chain after verification failure +`Authenticator` receives that selected credential and HTTP request parts. Transport verifiers must obtain proof from trusted server extensions, not client-supplied identity headers. It returns `VerifiedIdentity` only after verification. Request parts provide headers and trusted transport extensions for integrations. The verifier must validate any client-controlled claims before treating them as identity. There is no fallback chain after verification failure `IdentityResolver` maps verified identity into an internal principal and permissions. The caller retains the original authentication evidence and credential restrictions independently of that mapping. Principals include an authority, subject, and kind. External subjects must remain authority-scoped unless the resolver explicitly maps them to a shared internal identity @@ -14,13 +14,13 @@ Inbound authentication uses a verifier, identity resolver, and authorizer inject The Axum `authenticate` middleware currently accepts one Authorization header using the existing Bearer format. Missing, duplicate, empty, or malformed credentials fail. Authentication replaces any preexisting caller extension and checks the method and matched route before dispatch. Handlers extract `AuthenticatedRequest` and authorize their parsed operation before calling a provider. Missing authenticated context fails closed -The gateway also inserts an MCP `Authorization` extension and a `SessionOwner` derived from the verified principal and credential reference. The MCP host checks the injected policy at initialization and before sending resolved upstream operations. Standalone MCP library hosts retain their existing explicitly trusted registry behavior when no policy is installed; production mounts must install authentication and the policy extension together +The gateway also inserts an MCP `Authorization` extension and a `SessionOwner` derived from the verified principal and credential reference. The MCP host checks the injected policy at initialization and before sending resolved upstream operations. Standalone MCP library hosts retain their existing explicitly trusted registry behavior when no policy is installed. Production mounts must install authentication and the policy extension together Local UI login continues using axum-login and tower-sessions. Only after session and CSRF validation does `UiSession` expose an authenticated caller. Its credentials are restricted to session-info and logout operations, so the UI CSRF bearer cannot authorize inference. Future UI management routes must extend that explicit scope Authentication evidence and principals contain no raw token, password, request body, or mutable accounting state. Session ownership is scoped by principal authority, subject, verifier, and credential ID, so separate credentials do not silently share an MCP session. Credential rotation may retain ownership when the verifier preserves a stable credential ID. Scope and expiry checks still run for each operation -Failures distinguish invalid or expired credentials, forbidden operations, unavailable authentication services, and missing server configuration/context. HTTP adapters map those outcomes to status codes; inference keeps its API-specific error envelopes and MCP translates policy failures to its own protocol +Failures distinguish invalid or expired credentials, forbidden operations, unavailable authentication services, and missing server configuration/context. HTTP adapters map those outcomes to status codes. Inference keeps its API-specific error envelopes and MCP translates policy failures to its own protocol Virtual-key storage, JWT verification, OAuth2 introspection, trusted-proxy validation, SSO, and custom Python hook adapters are not implemented here yet. They should supply these contracts rather than bypassing the shared authorization boundary. Request-body-dependent custom hooks will need a bounded endpoint adapter after parsing diff --git a/litellm-rust/crates/gateway-inference/AGENTS.md b/litellm-rust/crates/gateway-inference/AGENTS.md index 73340982811..601f8d8ed63 100644 --- a/litellm-rust/crates/gateway-inference/AGENTS.md +++ b/litellm-rust/crates/gateway-inference/AGENTS.md @@ -1,6 +1,6 @@ -- Expose a mountable Axum router; listener binding, server lifecycle, and shared inbound middleware belong to `gateway` -- Own endpoint paths, request parsing, model alias resolution, and API-specific response and SSE error formats; delegate hosted call execution and HTTP body delivery to host-http -- Delegate inference execution to the `inference-` crates and provider transformations and authentication to `llms` and the auth crates; do not duplicate them in handlers -- Let the `inference-` crates validate inference fields and supported features, then map their errors to HTTP responses; do not add gateway checks for temporary inference limitations -- Use injected deployments, HTTP pools, settings, and secret sources; do not load process configuration or construct independent clients in handlers -- Test HTTP contracts here, including status codes, forwarded headers, error envelopes, and streaming behavior; keep inference and provider tests in their owning crates +- Expose a mountable Axum router. Listener binding, server lifecycle, and shared inbound middleware belong to `gateway` +- Own endpoint paths, request parsing, model alias resolution, and API-specific response and SSE error formats. Delegate hosted call execution and HTTP body delivery to host-http +- Delegate inference execution to the `inference-` crates and provider transformations and authentication to `llms` and the auth crates. Do not duplicate them in handlers +- Let the `inference-` crates validate inference fields and supported features, then map their errors to HTTP responses. Do not add gateway checks for temporary inference limitations +- Use injected deployments, HTTP pools, settings, and secret sources. Do not load process configuration or construct independent clients in handlers +- Test HTTP contracts here, including status codes, forwarded headers, error envelopes, and streaming behavior. Keep inference and provider tests in their owning crates diff --git a/litellm-rust/crates/gateway/AGENTS.md b/litellm-rust/crates/gateway/AGENTS.md index 3c417111233..d54fa062f2b 100644 --- a/litellm-rust/crates/gateway/AGENTS.md +++ b/litellm-rust/crates/gateway/AGENTS.md @@ -1,5 +1,5 @@ - Keep this crate a thin composition layer: mount endpoint routers and serve the supplied listener - Server lifecycle and shared inbound middleware belong here, including client authentication, rate limiting, and request logging -- Endpoint paths, request handling, model resolution, and response encoding belong to the mounted crates; provider execution belongs to the `inference-` crates and `llms` -- Inject shared state and infrastructure; avoid global runtimes, duplicate client pools, and abstractions for hypothetical endpoint groups -- Test mounting and server lifecycle through public HTTP behavior; test endpoint semantics in the owning crate +- Endpoint paths, request handling, model resolution, and response encoding belong to the mounted crates. Provider execution belongs to the `inference-` crates and `llms` +- Inject shared state and infrastructure. Avoid global runtimes, duplicate client pools, and abstractions for hypothetical endpoint groups +- Test mounting and server lifecycle through public HTTP behavior. Test endpoint semantics in the owning crate diff --git a/litellm-rust/crates/host-native/AGENTS.md b/litellm-rust/crates/host-native/AGENTS.md index 6882d3aaca2..316427976f5 100644 --- a/litellm-rust/crates/host-native/AGENTS.md +++ b/litellm-rust/crates/host-native/AGENTS.md @@ -1,7 +1,7 @@ -`litellm-host-native` is the Rust driver for hosted calls. `Driver` owns the machine, a `HostCallHandler` and a `Interceptors`; `advance()` answers services and hooks inline and returns at completion or at the next stream boundary, holding the `Reply>` until the consumer calls `advance()` or `detach()` again. Dropping the driver drops the machine and so cancels the call +`litellm-host-native` is the Rust driver for hosted calls. `Driver` owns the machine, a `HostCallHandler` and a `Interceptors`. `advance()` answers services and hooks inline and returns at completion or at the next stream boundary, holding the `Reply>` until the consumer calls `advance()` or `detach()` again. Dropping the driver drops the machine and so cancels the call -The consumer decides demand, so the driver never spawns a producer task and never buffers chunks ahead of demand. `litellm-host-http` polls it from the response body; `in_process::run_hosted` polls it on behalf of a `StreamConsumer`. Both observe lifecycle terminals themselves, the driver reports none +The consumer decides demand, so the driver never spawns a producer task and never buffers chunks ahead of demand. `litellm-host-http` polls it from the response body. `in_process::run_hosted` polls it on behalf of a `StreamConsumer`. Both observe lifecycle terminals themselves, the driver reports none -Depend on `litellm-host` only. HTTP encoding stays in `litellm-host-http`; `litellm-host-python` drives the machine directly so Python callbacks stay in the caller's asyncio task +Depend on `litellm-host` only. HTTP encoding stays in `litellm-host-http`. `litellm-host-python` drives the machine directly so Python callbacks stay in the caller's asyncio task -`services.rs` owns `HostCallHandler` and its borrowed and no-service implementations. This is the Rust driver's handler contract; the shared host crate owns the service request protocol +`services.rs` owns `HostCallHandler` and its borrowed and no-service implementations. This is the Rust driver's handler contract. The shared host crate owns the service request protocol diff --git a/litellm-rust/crates/host-python/AGENTS.md b/litellm-rust/crates/host-python/AGENTS.md index 9cd65c4d0a6..7776df32de4 100644 --- a/litellm-rust/crates/host-python/AGENTS.md +++ b/litellm-rust/crates/host-python/AGENTS.md @@ -1,4 +1,4 @@ -- Target invariants; implementation and runtime validation may lag these rules +- Target invariants. Implementation and runtime validation may lag these rules ## Boundary with Python consumers @@ -6,15 +6,15 @@ This crate owns CPython execution mechanics for generic `litellm-host` machines `src/native.rs` belongs here: it runs a generic machine through the Python runtime and owns its pending execution and abort handle. Keep provider selection, request projection and public exception policy out of it. A rename to `machine_runner.rs` is optional and must not change behavior -The execution handle receives its Python lifecycle binding from its consumer through `PythonLifecycle` rather than import a fixed `litellm.rust_bridge` module. Generic suspension and execution state validation belong here; public stream wrappers and `_hidden_params` conventions belong to the consumer +The execution handle receives its Python lifecycle binding from its consumer through `PythonLifecycle` rather than import a fixed `litellm.rust_bridge` module. Generic suspension and execution state validation belong here. Public stream wrappers and `_hidden_params` conventions belong to the consumer -Creating a resolved asyncio Future from an already constructed Python value belongs here, alongside runtime waiting, interpreter detachment and panic containment. Choosing which callable exceptions become a public `RuntimeError` belongs to the consumer; `python-bridge::callable::wrap_failure` owns that policy +Creating a resolved asyncio Future from an already constructed Python value belongs here, alongside runtime waiting, interpreter detachment and panic containment. Choosing which callable exceptions become a public `RuntimeError` belongs to the consumer. `python-bridge::callable::wrap_failure` owns that policy The driver owns ordering: start, argument preparation, prepared-argument hooks, binding decode and machine start. Fallible per-call resource setup supplied by the consumer runs after all argument hooks, using the prepared argument view, and before provider work. Setup failure follows the existing terminal failure path. Creating or discarding an unstarted coroutine must not initialize clients, acquire credentials or capture execution context Boundary tests exercise behavior with a supplied lifecycle binding without importing the LiteLLM Python package. Pin inline awaiting, awaitable final values, exception identity, cancellation, re-entry and release of retained objects, rather than module names or source layout -`HookChain` composes Python runtime hooks in order. Each argument, wire-request and response transformation feeds its result to the next hook. After all argument transformations, the driver calls `arguments_prepared` on every hook in order. Retained callback views must adopt that dictionary before later policy hooks can mutate or reject it. SDK policy is supplied by bridge composition as a hook, never a separate driver phase or parameter. Hooks implement only the stages they need; default stages preserve the supplied values +`HookChain` composes Python runtime hooks in order. Each argument, wire-request and response transformation feeds its result to the next hook. After all argument transformations, the driver calls `arguments_prepared` on every hook in order. Retained callback views must adopt that dictionary before later policy hooks can mutate or reject it. SDK policy is supplied by bridge composition as a hook, never a separate driver phase or parameter. Hooks implement only the stages they need. Default stages preserve the supplied values Terminal notifications share the selected response or exception. An ordinary notification error is reported as unraisable and does not skip the next hook or replace the selected outcome. Preparation, interception and transformation errors stop the chain. Cancellation stops all further hook dispatch. Suspensions stay inline in the existing driver, and the chain traverses retained event values for GC @@ -23,21 +23,21 @@ Terminal notifications share the selected response or exception. An ordinary not - Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonBinding`, `PythonHostCalls` and `PythonOwned` traits, and the `PythonRuntime` specialization of `host::hooks::CallHooks` - No LiteLLM domain dependencies beyond `litellm-host`: no route types, no `Logging` policy, no public API registration, no cdylib build features - `PythonCallHooks` only constrains the shared call-stage interface to `PythonRuntime` and Python ownership. It must not redeclare the stages - - `PythonCallEvent` is a specialization of the shared `CallEvent`, never a separately defined lifecycle. The driver emits `Succeeded` or `Failed` exactly once and never dispatches callbacks after a cancellation; which Python objects consume those events is the legacy adapter's business + - `PythonCallEvent` is a specialization of the shared `CallEvent`, never a separately defined lifecycle. The driver emits `Succeeded` or `Failed` exactly once and never dispatches callbacks after a cancellation. Which Python objects consume those events is the legacy adapter's business - `CallOptions` can publish snapshots independently of callback delivery. Terminal snapshots follow completed hook dispatch, and cancellation never calls a Python callback. For Python-driven calls, leave the machine's observation publisher unset so the driver is the sole publisher of intercepted provider-response snapshots - - `PythonBinding::decode_request` receives the keyword view returned by `prepare_arguments` and updated by `arguments_prepared`, not the caller's dict; a binding that decodes from it inherits all composed argument rewrites - - A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the binding's `map_error`; a Python exception raised inside the call, and a failure in `prepare_arguments` or `transform_response`, is raised as is + - `PythonBinding::decode_request` receives the keyword view returned by `prepare_arguments` and updated by `arguments_prepared`, not the caller's dict. A binding that decodes from it inherits all composed argument rewrites + - A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the binding's `map_error`. A Python exception raised inside the call, and a failure in `prepare_arguments` or `transform_response`, is raised as is - A failing `map_error` is raised with the native error's text as its `__context__`, never swallowed - Use standard PyO3 ownership and conversion APIs - - Prefer `Bound<'py, T>` for attached operations/results, `Py` for retention; binding/unbinding does not copy payloads - - Use `pythonize` for selected Serde data, never a JSON-text round trip; share conversion with `Pythonized` - - Preserve `PythonizeError`'s standard conversion into `PyErr`; do not stringify original Python exceptions into new `ValueError`s + - Prefer `Bound<'py, T>` for attached operations/results, `Py` for retention. Binding/unbinding does not copy payloads + - Use `pythonize` for selected Serde data, never a JSON-text round trip. Share conversion with `Pythonized` + - Preserve `PythonizeError`'s standard conversion into `PyErr`. Do not stringify original Python exceptions into new `ValueError`s - Keep serializer-panic containment in `Pythonized`: async output conversion can run in an unjoined blocking task and otherwise strand delivery -- Use `Python::detach` for Rust-only work; Python operations require attachment - - Keep diagnostic counters in the consumer; wrapper invocations do not measure every interpreter release - - Release exclusive class borrows/locks before Python calls or decrements that can invoke finalizers; expose retained Python edges to GC without calling Python during traversal +- Use `Python::detach` for Rust-only work. Python operations require attachment + - Keep diagnostic counters in the consumer. Wrapper invocations do not measure every interpreter release + - Release exclusive class borrows/locks before Python calls or decrements that can invoke finalizers. Expose retained Python edges to GC without calling Python during traversal - Keep coroutine driving in the shared Python driver and the native handle - - Shared driver implementation: `litellm/rust_bridge/lifecycle.py`; handle: `src/handle.rs`; call driver: `src/driver.rs`; native-backed behavior tests: `tests/lifecycle.py`. The consumer supplies the lifecycle binding - - Every lifecycle suspension is awaited inline in the caller's task; `into_future` creates a separate task and cannot satisfy this contract + - Shared driver implementation: `litellm/rust_bridge/lifecycle.py`, handle: `src/handle.rs`, call driver: `src/driver.rs`, native-backed behavior tests: `tests/lifecycle.py`. The consumer supplies the lifecycle binding + - Every lifecycle suspension is awaited inline in the caller's task. `into_future` creates a separate task and cannot satisfy this contract - References: [ownership](https://pyo3.rs/v0.29.2/types.html), [conversions](https://pyo3.rs/v0.29.2/conversions/traits.html), [pythonize errors](https://docs.rs/pythonize/0.29.0/src/pythonize/error.rs.html) - [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [re-entry](https://pyo3.rs/v0.29.2/class/call.html), [parallelism](https://pyo3.rs/v0.29.2/parallelism.html), [async delivery source](https://docs.rs/pyo3-async-runtimes/0.29.0/src/pyo3_async_runtimes/generic.rs.html) diff --git a/litellm-rust/crates/host/AGENTS.md b/litellm-rust/crates/host/AGENTS.md index 317a43516c7..0f06f0e57e1 100644 --- a/litellm-rust/crates/host/AGENTS.md +++ b/litellm-rust/crates/host/AGENTS.md @@ -2,28 +2,28 @@ | Responsibility | Shared contract | HTTP | Python | | --- | --- | --- | --- | -| Input and output conversion | `Protocol::{Request, Response, StreamHead, Chunk, Error}` | Typed input; `ResponseEncoder` and `StreamEncoder` produce HTTP values | `PythonBinding` decodes prepared arguments, encodes public values and maps native errors | -| Host services | `Protocol::HostCall`, `HostServices::call` | `host-native::services::HostCallHandler` answers typed calls; `()` handles protocols without host calls | `PythonHostCalls` invokes retained Python objects; it may share an owner with the binding | +| Input and output conversion | `Protocol::{Request, Response, StreamHead, Chunk, Error}` | Typed input. `ResponseEncoder` and `StreamEncoder` produce HTTP values | `PythonBinding` decodes prepared arguments, encodes public values and maps native errors | +| Host services | `Protocol::HostCall`, `HostServices::call` | `host-native::services::HostCallHandler` answers typed calls. `()` handles protocols without host calls | `PythonHostCalls` invokes retained Python objects. It may share an owner with the binding | | Active hooks | `Interceptors::{before_provider_request, after_provider_response}` and `hooks::CallHooks` | Request interception and fallible execution callbacks | `PythonCallHooks` also prepares arguments, transforms public responses and receives stream callbacks | | Passive observation | `observation::ObservationSender` | Queued execution and lifecycle snapshots, retained by the response body | Public Python callbacks remain active hooks with their existing failure policy | | Runtime driving | `Machine`, `HostRequest::{HostCall, Intercept, Stream}` | Body polling controls demand | Native polling and inline caller-task Python awaits control progress | -`hosted_call(request, observers, execute)` starts with a typed request. Its route closure receives separate `HostServices`, `ChannelInterceptors` and optional observation publisher; it returns `CallOutput`. Hosted-call plumbing alone forwards the returned stream through demand replies. Lower-level `CallMachine` users receive a `CallContext` containing separately named services, interceptors, observers and stream delivery +`hosted_call(request, observers, execute)` starts with a typed request. Its route closure receives separate `HostServices`, `ChannelInterceptors` and optional observation publisher. It returns `CallOutput`. Hosted-call plumbing alone forwards the returned stream through demand replies. Lower-level `CallMachine` users receive a `CallContext` containing separately named services, interceptors, observers and stream delivery Core route constructors prepare their dependencies and return a closure accepting the typed request. Python starts that closure only after argument preparation, preflight and decoding succeed. These steps remain inside the driver's terminal and error handling. Decoding may retain objects for subsequent host service calls -Each driver owns terminal dispatch. Hooks can change or fail execution; passive observers return no result. HTTP observes success after response conversion or stream exhaustion, failure on errors, and cancellation on body drop. Python preserves exception identity and maps native failures once. Explicit Python stream close reports success for delivered chunks; cancellation stops further callback dispatch +Each driver owns terminal dispatch. Hooks can change or fail execution. Passive observers return no result. HTTP observes success after response conversion or stream exhaustion, failure on errors, and cancellation on body drop. Python preserves exception identity and maps native failures once. Explicit Python stream close reports success for delivered chunks. Cancellation stops further callback dispatch `interceptors.rs` owns `Interceptors` and its request/response payload types. `lifecycle.rs` owns `CallObserver`, `CallEvent`, `ExecutionEvent`, timing, failure origin, and the observation wrappers. Event payloads are generic so a runtime can retain its own response, exception and raw-response references without introducing a language dependency. `snapshot()` projects them into the owned observation contract without retaining runtime objects. Pass interceptors and observers separately at direct route and HTTP entrypoints. Routes publish execution events independently of interception `hooks.rs` owns the call-stage interface and its runtime-associated types. It contains no Python types or legacy callback policy. A runtime supplies its context and continuation representation through `HookRuntime` -`protocol.rs` owns `Protocol` and suspension messages, including `InterceptRequest` and `StreamDelivery`. `call.rs` owns route outputs and their adaptation into a hosted machine. Rust service handling belongs in `host-native::services`; coroutine channel handles stay in `machine/context.rs` +`protocol.rs` owns `Protocol` and suspension messages, including `InterceptRequest` and `StreamDelivery`. `call.rs` owns route outputs and their adaptation into a hosted machine. Rust service handling belongs in `host-native::services`. Coroutine channel handles stay in `machine/context.rs` Rust handlers answer suspensions through `litellm-host-native::Driver`, which `litellm-host-http` and `litellm_host_native::in_process` share. `in_process::Host` is an assembly of services, interceptors, stream consumer and optional observation publisher. It is not a trait mirroring every suspension. Use `run_hosted` to preserve the distinction between stream completion and detachment Keep API policy in gateway-inference and python-bridge, and legacy callback policy in callbacks-legacy-python. Python bindings and hooks expose retained references through `PythonOwned`, with idempotent close and GC traversal. Runtime machinery stays in driver, native, handle and runtime modules -Interceptors run inline and can rewrite values or fail execution. Observers consume owned `CallEvent` snapshots from `observation_channel`; its bounded `ObservationSender` never waits for delivery and counts events dropped when the queue is full or closed. The host owns receiver processing and draining. Pass the same publisher to machine construction and the driver when one receiver should collect execution and lifecycle events. Legacy Python callbacks retain their existing awaited, fallible behavior through the Python adapter +Interceptors run inline and can rewrite values or fail execution. Observers consume owned `CallEvent` snapshots from `observation_channel`. Its bounded `ObservationSender` never waits for delivery and counts events dropped when the queue is full or closed. The host owns receiver processing and draining. Pass the same publisher to machine construction and the driver when one receiver should collect execution and lifecycle events. Legacy Python callbacks retain their existing awaited, fallible behavior through the Python adapter -`ExecutionFacts` and `ResultSource` describe execution without pricing or budget policy. `Interceptors::result_ready` delivers these facts through an awaited `InterceptRequest::ResultReady`; hosts receive them before response transformation or stream delivery. `ExecutionEvent::ResultReady` is the matching lifecycle event and can also be published as a passive snapshot. Accounting must consume the awaited path rather than a lossy observation queue +`ExecutionFacts` and `ResultSource` describe execution without pricing or budget policy. `Interceptors::result_ready` delivers these facts through an awaited `InterceptRequest::ResultReady`. Hosts receive them before response transformation or stream delivery. `ExecutionEvent::ResultReady` is the matching lifecycle event and can also be published as a passive snapshot. Accounting must consume the awaited path rather than a lossy observation queue diff --git a/litellm-rust/crates/host/src/machine/AGENTS.md b/litellm-rust/crates/host/src/machine/AGENTS.md index 048f9aacf3a..89ad82b07dd 100644 --- a/litellm-rust/crates/host/src/machine/AGENTS.md +++ b/litellm-rust/crates/host/src/machine/AGENTS.md @@ -19,7 +19,7 @@ Credential acquisition contracts and reusable adapters belong in `litellm-auth-t The coroutine polls the route future until it completes or yields a `HostRequest`. Each request carries a typed `Reply` that its driver must answer before resuming, or abandon when interrupting or dropping the execution. A pending network future is an ordinary async wait, not a host suspension -`CallContext` gives the route separate capabilities: `HostServices` requests host operations, `ChannelInterceptors` requests interception, `ObservationSender` publishes events, and `StreamSender` delivers stream values. Keep their yield-and-reply mechanics in `context.rs`. The actual service and hook implementations belong to the host. This follows the effect-handler pattern: the route requests an operation, the driver handles it, and the route continues with the reply. The suspended computation stays in the coroutine; `Reply` only supplies its result +`CallContext` gives the route separate capabilities: `HostServices` requests host operations, `ChannelInterceptors` requests interception, `ObservationSender` publishes events, and `StreamSender` delivers stream values. Keep their yield-and-reply mechanics in `context.rs`. The actual service and hook implementations belong to the host. This follows the effect-handler pattern: the route requests an operation, the driver handles it, and the route continues with the reply. The suspended computation stays in the coroutine. `Reply` only supplies its result Stream replies use `std::ops::ControlFlow<()>`. `ControlFlow::Continue(())` permits stream execution to continue. `ControlFlow::Break(())` tells it that the consumer stopped reading. Holding the reply applies backpressure until the consumer advances. Keep stream forwarding and the distinction between stream exhaustion and detachment in `crate::call::hosted_call` @@ -29,4 +29,4 @@ Stream replies use `std::ops::ControlFlow<()>`. `ControlFlow::Continue(())` perm Rust drivers live in `litellm-host-native` and `litellm-host-http`. The Python driver lives in `litellm-host-python` and awaits Python hooks in the caller's task. Keep runtime scheduling, encoding, terminal observation, and callback policy in those layers and their adapters -Public behavior tests belong in `crates/host/tests`; private behavior tests stay inline with their owning implementation. Test suspension answers, interruption, resource release, and stream demand through behavior, rather than asserting file layout +Public behavior tests belong in `crates/host/tests`. Private behavior tests stay inline with their owning implementation. Test suspension answers, interruption, resource release, and stream demand through behavior, rather than asserting file layout diff --git a/litellm-rust/crates/inference-chat/AGENTS.md b/litellm-rust/crates/inference-chat/AGENTS.md index 36dd72fc011..a73880c2858 100644 --- a/litellm-rust/crates/inference-chat/AGENTS.md +++ b/litellm-rust/crates/inference-chat/AGENTS.md @@ -1,3 +1,3 @@ - Shared base and layering rules: [`../inference/AGENTS.md`](../inference/AGENTS.md) -- This crate owns Chat Completions call orchestration; payloads belong in `litellm-llms-types`, transformations in `llms/src/base_llm/chat` and `llms/src//chat` +- This crate owns Chat Completions call orchestration. Payloads belong in `litellm-llms-types`, transformations in `llms/src/base_llm/chat` and `llms/src//chat` - Chat Completions is currently non-streaming: the route returns its completed response, cached through `litellm_inference::caching::execute_unary` diff --git a/litellm-rust/crates/inference-messages/AGENTS.md b/litellm-rust/crates/inference-messages/AGENTS.md index c992fe60fe9..860eb83c1f3 100644 --- a/litellm-rust/crates/inference-messages/AGENTS.md +++ b/litellm-rust/crates/inference-messages/AGENTS.md @@ -7,8 +7,8 @@ - delegate authentication policy, beta selection, payload rewriting, and response interpretation to those adapters - keep provider policy out of request preparation and transport handlers - calling a concrete provider helper for every provider is still a policy dependency -- Route types such as `MessagesCall`, prepared requests, and response wrappers containing live streams describe execution; reuse the shared Messages payload types inside them instead of defining another request or response schema here +- Route types such as `MessagesCall`, prepared requests, and response wrappers containing live streams describe execution. Reuse the shared Messages payload types inside them instead of defining another request or response schema here - The route returns `litellm_host::call::CallOutput`: a completed response, or a stream head and chunks -- Per-call dependencies are grouped in `litellm_inference::context::CallContext`; `src/lib.rs` explicitly sequences cache lookup, provider execution, result acceptance (`CallContext::result_ready`), and cache storage +- Per-call dependencies are grouped in `litellm_inference::context::CallContext`. `src/lib.rs` explicitly sequences cache lookup, provider execution, result acceptance (`CallContext::result_ready`), and cache storage - Preserve the order of validation, normalization, caller-requested parameter removal, and provider transformation when that order affects observable behavior - Test provider dispatch, auth precedence, header handling, transformations, and responses through behavior, not source structure diff --git a/litellm-rust/crates/inference-ocr/AGENTS.md b/litellm-rust/crates/inference-ocr/AGENTS.md index c3644e850e0..d74b7e7af06 100644 --- a/litellm-rust/crates/inference-ocr/AGENTS.md +++ b/litellm-rust/crates/inference-ocr/AGENTS.md @@ -1,5 +1,5 @@ - Shared base and layering rules: [`../inference/AGENTS.md`](../inference/AGENTS.md) -- This crate owns OCR call orchestration; payloads belong in `litellm-llms-types`, transformations and the request handler in `llms/src/base_llm/ocr` and `llms/src//ocr` +- This crate owns OCR call orchestration. Payloads belong in `litellm-llms-types`, transformations and the request handler in `llms/src/base_llm/ocr` and `llms/src//ocr` - The route returns its completed response directly - Route closures supply OCR-specific host capabilities, such as the caller's Azure AD token provider (`src/route.rs`) - Provider code reaches the caller's hooks mid-call only through `litellm_llms::base_llm::ocr::handler::CallHooks`, which this crate implements over its host until it folds into `litellm_host::interceptors::Interceptors` diff --git a/litellm-rust/crates/inference-responses/AGENTS.md b/litellm-rust/crates/inference-responses/AGENTS.md index d3e49a1579a..1dd199024b0 100644 --- a/litellm-rust/crates/inference-responses/AGENTS.md +++ b/litellm-rust/crates/inference-responses/AGENTS.md @@ -1,4 +1,4 @@ - Shared base and layering rules: [`../inference/AGENTS.md`](../inference/AGENTS.md) -- This crate owns Responses API call orchestration; payloads belong in `litellm-llms-types`, transformations in `llms/src/base_llm/responses` and `llms/src//responses` +- This crate owns Responses API call orchestration. Payloads belong in `litellm-llms-types`, transformations in `llms/src/base_llm/responses` and `llms/src//responses` - The HTTP route returns `litellm_host::call::CallOutput`: a completed response, or a stream head and chunks, cached through `litellm_inference::caching::execute_streaming` - WebSocket sessions (`src/websocket.rs`) stay separate from the HTTP call driver because a connection can accept multiple requests while receiving events diff --git a/litellm-rust/crates/inference-transcription/AGENTS.md b/litellm-rust/crates/inference-transcription/AGENTS.md index 632e5cea92c..a18039f9e0b 100644 --- a/litellm-rust/crates/inference-transcription/AGENTS.md +++ b/litellm-rust/crates/inference-transcription/AGENTS.md @@ -1,3 +1,3 @@ - Shared base and layering rules: [`../inference/AGENTS.md`](../inference/AGENTS.md) -- This crate owns audio transcription call orchestration; transformations belong in `llms/src/base_llm/audio_transcription` and `llms/src//audio_transcription` -- `AudioTranscriptionRoute::execute` returns the completed response directly; there is no hosted route machine or response cache for transcription yet +- This crate owns audio transcription call orchestration. Transformations belong in `llms/src/base_llm/audio_transcription` and `llms/src//audio_transcription` +- `AudioTranscriptionRoute::execute` returns the completed response directly. There is no hosted route machine or response cache for transcription yet diff --git a/litellm-rust/crates/inference/AGENTS.md b/litellm-rust/crates/inference/AGENTS.md index e36faad2905..2c4bd00108b 100644 --- a/litellm-rust/crates/inference/AGENTS.md +++ b/litellm-rust/crates/inference/AGENTS.md @@ -6,7 +6,7 @@ - diagnostic spans (`src/diagnostic.rs`), outbound send and signing (`src/outbound.rs`), provider resolution (`src/provider.rs`), `CoreResources` (`src/resources.rs`) - response caching (`src/caching.rs`): `Cachable`, `StreamCachable`, `CacheRequest`, `CallCache`, `execute_unary`, `execute_streaming`, stream capture - Shared test helpers belong in litellm-inference-testing, used only as a dev dependency -- Nothing here names an API format; format-specific code, constants, tests and test builders live in their `inference-` crate +- Nothing here names an API format. Format-specific code, constants, tests and test builders live in their `inference-` crate - Never depend on an `inference-*` crate from here ## Format crates @@ -26,22 +26,22 @@ - Hosts assemble route objects from shared `CoreResources`, HTTP settings, and secret sources - each route owns its provider client and authentication dependencies - - gateway routes live for the gateway lifetime; Python assembles routes per call from its settings snapshot + - gateway routes live for the gateway lifetime. Python assembles routes per call from its settings snapshot - Calls pass `Interceptors` and an optional `ObservationSender` separately - use `&()` for no hooks and `None` for no observer - handlers accept `Interceptors`, never a concrete `ChannelInterceptors` -- Construction does no work; preparation and lifecycle observation begin when the future is polled -- Native observers receive start and terminal events through the shared call runner; a stream retains its lifecycle until exhaustion, error, or drop. Hosted routes leave terminal observation to their driver -- A hosted format's `route.rs` declares the concrete `Protocol` and a route method that takes a typed request and builds a `litellm_host::call::HostedMachine` with `hosted_call`; audio transcription has no hosted machine and returns its response directly +- Construction does no work. Preparation and lifecycle observation begin when the future is polled +- Native observers receive start and terminal events through the shared call runner. A stream retains its lifecycle until exhaustion, error, or drop. Hosted routes leave terminal observation to their driver +- A hosted format's `route.rs` declares the concrete `Protocol` and a route method that takes a typed request and builds a `litellm_host::call::HostedMachine` with `hosted_call`. Audio transcription has no hosted machine and returns its response directly - the shared call plumbing owns stream opening, delivery, backpressure, and detachment - request decoding belongs to the boundary before the machine starts - route closures only supply execution dependencies and route-specific host capabilities - Python uses its own shared driver and preserves caller-task callback execution -- Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback execution of any kind. Routes run as machines that yield host operations and call events; which integrations consume those events is the host's business +- Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback execution of any kind. Routes run as machines that yield host operations and call events. Which integrations consume those events is the host's business ## Crate boundaries -- Dependencies only point down; Python package names identify counterparts, not ownership +- Dependencies only point down. Python package names identify counterparts, not ownership - `litellm-llms-types` owns shared inference API contracts, grouped by format: pure serde data and shape validation, no I/O - `litellm-core-utils` mirrors `litellm/litellm_core_utils/`: pure helpers (provider resolution, prompt factory, call arguments, settings lookup and layer merge), no network I/O - `litellm-http` is Rust-only and route-neutral: settings resolution, pooled `reqwest` clients, TLS, proxies, the SSRF-safe media fetcher, request and header helpers, and transport errors @@ -49,35 +49,35 @@ - `litellm-inference-` mirrors the route packages (`litellm/ocr/`, `litellm/messages/`, ...) - Provider code never imports from `inference` or `inference-*` - Handlers belong in `inference-*` or `llms`, never in a host crate -- Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only allowed re-exports are a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`) and each format crate's `pub use litellm_inference::RouteError as Error;`, which keeps the pre-split `::Error` name +- Import every item from its canonical path. Never re-export another crate's items or give an item a second public path. The only allowed re-exports are a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`) and each format crate's `pub use litellm_inference::RouteError as Error;`, which keeps the pre-split `::Error` name ## Error placement -- The workspace `Error definitions` rules shape each crate's error; this section decides which crate and module a failure belongs to +- The workspace `Error definitions` rules shape each crate's error. This section decides which crate and module a failure belongs to - A failure is declared once, by the lowest crate that raises it - every crate above nests it unchanged (`#[error(transparent)] Auth(#[from] litellm_auth::Error)`) or maps it once at its boundary, as `src/error.rs` does for `litellm_llms::Error` - `RouteError` collects route failures for every format crate and never re-declares a variant a lower crate raises - Scope follows the concept, not the first caller - an error type under `litellm-llms`'s `/` is private to that provider: no other provider and nothing in `base_llm` may import it - a failure two providers or two routes can hit (wire framing, stream event decoding, a malformed provider response) belongs to the crate that owns the concept: `litellm-framer` for framing, `litellm_llms::Error` for the transformation layer -- `litellm_llms::Error` (`crates/llms/src/error.rs`) is the one transformation error for every provider and API; `base_llm/ocr/error.rs` is the recorded exception until OCR folds into it +- `litellm_llms::Error` (`crates/llms/src/error.rs`) is the one transformation error for every provider and API. `base_llm/ocr/error.rs` is the recorded exception until OCR folds into it ## Response caching and accounting -- Attach a `litellm_cache_response::ScopedCache` with `route.with_cache(cache)`; cached and uncached routes use the same `execute` and `machine` methods -- `CallOptions` carries a scope-free `CachePolicy` and observation; per-call policy never replaces the attached scope or service -- Provider transport does not own cache orchestration; stream capture stays in `src/caching.rs` +- Attach a `litellm_cache_response::ScopedCache` with `route.with_cache(cache)`. Cached and uncached routes use the same `execute` and `machine` methods +- `CallOptions` carries a scope-free `CachePolicy` and observation. Per-call policy never replaces the attached scope or service +- Provider transport does not own cache orchestration. Stream capture stays in `src/caching.rs` - This crate and the format crates own request identity, typed response reconstruction, and stream capture and replay - `cache-response` owns cache policy, namespacing, scope encoding, versioned envelopes, and freshness -- The SDK explicitly chooses shared scope; the gateway derives isolated scope from authenticated identity before attaching its service +- The SDK explicitly chooses shared scope. The gateway derives isolated scope from authenticated identity before attaching its service - `ExecutionFacts` are delivered through the awaited `ResultReady` host operation for both provider and cached results, before public response processing or stream opening - facts carry resolved model/provider and result source, including the hit key - usage remains in the typed response or delivered stream, where completion and cancellation determine what was actually reported - passive observation is not an accounting delivery mechanism - Inference does not calculate prices, charge budgets, or update rate-limit counters - - the legacy Python callback adapter translates execution facts into the existing Python logging contract; Python remains the accounting owner on that path + - the legacy Python callback adapter translates execution facts into the existing Python logging contract. Python remains the accounting owner on that path - native gateway accounting belongs to gateway dependencies, independently of `host-python` - - response-cache services expose no coordination counters or reservation APIs; a shared Redis deployment does not make response storage and accounting coordination the same dependency + - response-cache services expose no coordination counters or reservation APIs. A shared Redis deployment does not make response storage and accounting coordination the same dependency - Cache lookup follows provider preparation, credential resolution, and the request interceptor - keys describe the effective provider URL, authenticated headers, and rewritten body - signed requests bypass caching until the signing identity has a stable cache representation diff --git a/litellm-rust/crates/llms-types/AGENTS.md b/litellm-rust/crates/llms-types/AGENTS.md index e5bd83b1d6e..c8e590845d4 100644 --- a/litellm-rust/crates/llms-types/AGENTS.md +++ b/litellm-rust/crates/llms-types/AGENTS.md @@ -10,7 +10,7 @@ The same ownership rule applies to Messages, Responses, Chat Completions, OCR, a - Keep one canonical definition and import path when moving a contract, updating consumers together instead of adding duplicate models or compatibility re-exports - Keep shared provider-specific wire types and extensions under `providers` - - Provider types may reuse format types; format types must not depend on provider types + - Provider types may reuse format types. Format types must not depend on provider types - A field belonging to an API format stays under `formats` even when provider support varies. Including it in a type does not promise provider support - Add a typed provider extension when a consumer needs to interpret or construct it. Keep adapter-only projections in `llms` until a shared public data contract is needed - Keep one authoritative representation of each field, preserving unknown fields without duplicating typed values in an extension map diff --git a/litellm-rust/crates/llms/AGENTS.md b/litellm-rust/crates/llms/AGENTS.md index a34a962101d..494c3fa2d39 100644 --- a/litellm-rust/crates/llms/AGENTS.md +++ b/litellm-rust/crates/llms/AGENTS.md @@ -14,11 +14,11 @@ Use named `#[rstest]` cases for independent input/output scenarios instead of lo Shared OCR document and response contracts live in `litellm-llms-types::formats::ocr`. `BaseOcrConfig` and decoding into adapter errors remain in `src/base_llm/ocr/transformation.rs`. Rust context/environment types support the runtime. `BaseOcrConfig::prepare_request` corresponds to Python's HTTP-handler preparation rather than a `BaseOCRConfig` method, and `validate_request_body` is a Rust-only hook. `src/base_llm/ocr/error.rs` and `src/base_llm/ocr/document.rs` are Rust-only: the OCR error taxonomy shared with the route, and inline-document helpers shared by several providers -For Mistral, `async_transform_ocr_request` uses the base default in both languages. `resolve_headers` and `build_ocr_url` implement the respective environment and URL operations, and `normalize_response` implements the typed part of response transformation. Existing auth key/header handling and top-level response-extra preservation differ between languages; layout refactors must preserve those behaviors and verify them with the existing tests +For Mistral, `async_transform_ocr_request` uses the base default in both languages. `resolve_headers` and `build_ocr_url` implement the respective environment and URL operations, and `normalize_response` implements the typed part of response transformation. Existing auth key/header handling and top-level response-extra preservation differ between languages. Layout refactors must preserve those behaviors and verify them with the existing tests For non-OCR pairs, order corresponding methods as parameter support/mapping, environment validation, URL construction, request transformation, and response transformation, followed by Rust-only runtime hooks. Auth resolution remains split between configs and route preparation in the `litellm-inference-` crates. Chat `supported_openai_param_mappings` describes accepted OpenAI/provider name pairs, unlike Python's `get_supported_openai_params` name list. Audio `map_transcription_params` remains a Rust filtering helper -Azure Messages maps to `llms/azure_ai/anthropic/messages_transformation.py`; Bedrock Converse maps to `llms/bedrock/chat/converse_transformation.py`. `AnthropicConfig`, `AmazonConverseConfig`, and the non-OCR base traits are partial ports. `OpenAiResponsesApiConfig` implements WebSocket transformations and a direct HTTP Responses path. Its HTTP path does not implement Python model-specific parameter rewriting or Responses-to-Chat emulation. Preserve their acceptance gates, passthrough behavior, and host fallback contracts when aligning layout +Azure Messages maps to `llms/azure_ai/anthropic/messages_transformation.py`. Bedrock Converse maps to `llms/bedrock/chat/converse_transformation.py`. `AnthropicConfig`, `AmazonConverseConfig`, and the non-OCR base traits are partial ports. `OpenAiResponsesApiConfig` implements WebSocket transformations and a direct HTTP Responses path. Its HTTP path does not implement Python model-specific parameter rewriting or Responses-to-Chat emulation. Preserve their acceptance gates, passthrough behavior, and host fallback contracts when aligning layout ## Provider and format boundaries diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 55837c4f32f..254a493d1ea 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,10 +1,10 @@ -- Target invariants, not completion claims; these supersede the crate guidance below where they conflict +- Target invariants, not completion claims. These supersede the crate guidance below where they conflict ## Boundary migration Keep domain composition here and execution mechanics in `litellm-host-python`. A helper does not belong in the runtime adapter merely because it uses PyO3. LiteLLM argument rules, provider defaults, public responses, public exception policy and cache or secret-manager compatibility remain product responsibilities -The bridge supplies the Python lifecycle binding and public stream construction to the runtime adapter. Preserve the single inline coroutine driver; removing the host's hardcoded import must not introduce another driver or a separate asyncio task for caller hooks +The bridge supplies the Python lifecycle binding and public stream construction to the runtime adapter. Preserve the single inline coroutine driver. Removing the host's hardcoded import must not introduce another driver or a separate asyncio task for caller hooks `src/callable.rs::wrap_failure` owns callable exception policy here. Resolved-Future construction uses `litellm-host-python::ready_future`, passing an already constructed Python value. Keep cache-specific serialization and disabled-cache results here @@ -16,18 +16,18 @@ Each step needs focused regression tests in the owning crate and Python integrat - Keep this crate the product-specific PyO3 consumer of `litellm-host-python` - Own registration, input projection, the route host and the caller callables it answers operations with (file readers, token providers), public response/error construction and the per-call composition of machine, route host and callback contract - - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, re-aliasing unchanged body keys) lives in `litellm-callbacks-legacy-python` behind `PublicCall` and `LegacyLogging`; the bridge hands the public call over and keeps no copy - - Value-oriented execution, sync waiting, nested-runtime checks, signal polling and panic containment live in `litellm-host-python`; native async work uses `pyo3-async-runtimes`, Serde output uses `Pythonized` - - Core owns typed native state, the route machine, provider preparation/I/O and normalization; the host driver owns terminal events; the legacy adapter in `litellm-callbacks-legacy-python` owns `Logging` dispatch policy - - Python, Rust SDK and gateway use one lifecycle-bearing core route entrypoint; provider helpers stay private, never bridge-accessible transport drivers - - Built-in provider/config/secret/auth/document preparation stays in Rust; caller-authored callbacks and focused Python-file reads run only at core-selected points -- Target GIL-enabled CPython explicitly with `#[pymodule(gil_used = true)]`; detach Rust-only work + - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, re-aliasing unchanged body keys) lives in `litellm-callbacks-legacy-python` behind `PublicCall` and `LegacyLogging`. The bridge hands the public call over and keeps no copy + - Value-oriented execution, sync waiting, nested-runtime checks, signal polling and panic containment live in `litellm-host-python`. Native async work uses `pyo3-async-runtimes`, Serde output uses `Pythonized` + - Core owns typed native state, the route machine, provider preparation/I/O and normalization. The host driver owns terminal events. The legacy adapter in `litellm-callbacks-legacy-python` owns `Logging` dispatch policy + - Python, Rust SDK and gateway use one lifecycle-bearing core route entrypoint. Provider helpers stay private, never bridge-accessible transport drivers + - Built-in provider/config/secret/auth/document preparation stays in Rust. Caller-authored callbacks and focused Python-file reads run only at core-selected points +- Target GIL-enabled CPython explicitly with `#[pymodule(gil_used = true)]`. Detach Rust-only work - GIL and tokio invariants, each pinned by a test in `host-python` (`runtime.rs`, `gil.rs`) so a regression fails there before it deadlocks a proxy: - Never hold the GIL while waiting on the runtime. A sync entrypoint releases it with `release_gil` around `block_on`, because every task that attaches would otherwise wait on the thread that is waiting on them ([pyo3 parallelism](https://pyo3.rs/v0.29.2/parallelism.html)) - - Never `block_on` from a tokio worker; the sync entrypoints refuse with "cannot run from + - Never `block_on` from a tokio worker. The sync entrypoints refuse with "cannot run from a Tokio context" instead of panicking inside the runtime ([tokio `Runtime::block_on`](https://docs.rs/tokio/latest/tokio/runtime/struct.Runtime.html#method.block_on)) - Inside a future, `Python::attach` only for GIL-cheap work: cloning a `Py`, building a small value, reading a settings snapshot. Anything that can block (a secret manager @@ -39,33 +39,33 @@ Each step needs focused regression tests in the owning crate and Python integrat `threading.get_ident()` test). Like any foreign-thread attach it therefore has no running asyncio loop and a fresh `contextvars` context: do not hand it a coroutine or anything bound to the caller's loop - - Dropping the await (an asyncio cancel) does not interrupt the Python call; it runs to + - Dropping the await (an asyncio cancel) does not interrupt the Python call. It runs to completion and its result is discarded. A panic in it reaches the awaiting task as a panic - - Add a case to `gil.rs` when a new seam changes any of these; the tests are the spec - - Free-threading requires separate runtime/concurrency validation; omitting the attribute does not opt out on PyO3 0.28+ + - Add a case to `gil.rs` when a new seam changes any of these. The tests are the spec + - Free-threading requires separate runtime/concurrency validation. Omitting the attribute does not opt out on PyO3 0.28+ - Preserve public argument binding and Python object provenance - - Project only consumed fields at reference read points; no eager whole-graph serialization or equality-based alias reconstruction - - Preserve provider-specific upload/submission/poll observation and encoding boundaries; signed/build-captured bytes must not be silently reserialized + - Project only consumed fields at reference read points. No eager whole-graph serialization or equality-based alias reconstruction + - Preserve provider-specific upload/submission/poll observation and encoding boundaries. Signed/build-captured bytes must not be silently reserialized - Conversion errors and every failure after the call starts are terminal - - Disabled/unavailable native execution may select legacy once; callback exceptions never authorize fallback or replay + - Disabled/unavailable native execution may select legacy once. Callback exceptions never authorize fallback or replay - Use one ordinary inline `async def` driver in `litellm/rust_bridge/lifecycle.py`, with the native handle and call driver in `litellm-host-python` - - Contract: `start`, `resume_value`, `resume_error`, idempotent `close`; explicitly tagged `Await`/`Complete` preserve awaitable final values - - Validate Created/Running/Suspended/Closed protocol states; the machine yields ops, the driver emits one terminal event, the adapter chooses dispatch policy - - Defer effectful setup/context reads/timestamps until start; unstarted-handle destruction releases inputs independently of Python `finally` - - Catch only the selected await's errors; start/resume errors propagate, `GeneratorExit` closes without further awaits - - Inline hooks preserve caller task/thread/loop and context writes; `into_future` creates a separate task and cannot satisfy this contract + - Contract: `start`, `resume_value`, `resume_error`, idempotent `close`. Explicitly tagged `Await`/`Complete` preserve awaitable final values + - Validate Created/Running/Suspended/Closed protocol states. The machine yields ops, the driver emits one terminal event, the adapter chooses dispatch policy + - Defer effectful setup/context reads/timestamps until start. Unstarted-handle destruction releases inputs independently of Python `finally` + - Catch only the selected await's errors. Start/resume errors propagate, `GeneratorExit` closes without further awaits + - Inline hooks preserve caller task/thread/loop and context writes. `into_future` creates a separate task and cannot satisfy this contract - Finalize fallible public response/error construction, replacements and metadata before terminal dispatch - Make ownership safe across suspension, re-entry, cancellation and GC - - Keep native provider state typed in core; do not shuttle it through opaque Python transport/response classes - - Prefer one retained `Py` via `PyErr::into_value(py)`; reconstruct transient `PyErr`s, preserving identity, traceback, cause and context - - Traverse every owned Python edge, including duplicate references; traversal cannot call Python + - Keep native provider state typed in core. Do not shuttle it through opaque Python transport/response classes + - Prefer one retained `Py` via `PyErr::into_value(py)`. Reconstruct transient `PyErr`s, preserving identity, traceback, cause and context + - Traverse every owned Python edge, including duplicate references. Traversal cannot call Python - Take state out and mark Running under a short borrow, release borrows/locks before Python invocation, publish terminal state before finalizer-capable drops - - Close/GC/deferred release are idempotent and re-entry-safe, including during Rust unwinding; release only owned references, never clear caller containers or mask the selected error - - The machine owns its in-flight provider future; `interrupt` drops it synchronously, so provider captures are released before the driver returns and no task outlives the call + - Close/GC/deferred release are idempotent and re-entry-safe, including during Rust unwinding. Release only owned references, never clear caller containers or mask the selected error + - The machine owns its in-flight provider future. `interrupt` drops it synchronously, so provider captures are released before the driver returns and no task outlives the call - Verify behavior through a fresh, provenance-checked installed extension and positive native execution evidence before replacing the custom coroutine - Cover admitted provider workflows, binding/read-point/identity behavior, failure continuation, finalization, no replay, deferred gates, re-entry, GC and cancellation termination - - Measure real conversion/copy costs before optimizing; preserve input contracts and capture lifetimes with `PyBackedBytes`, and lookup timing when interning names - - Ship accurate `_native.pyi` declarations and typing markers; distinguish Future-returning bindings from coroutine-returning bindings + - Measure real conversion/copy costs before optimizing. Preserve input contracts and capture lifetimes with `PyBackedBytes`, and lookup timing when interning names + - Ship accurate `_native.pyi` declarations and typing markers. Distinguish Future-returning bindings from coroutine-returning bindings - References: [ownership](https://pyo3.rs/v0.29.2/types.html), [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [exception transfer](https://docs.rs/pyo3/0.29.2/pyo3/struct.PyErr.html#method.into_value), [re-entry](https://pyo3.rs/v0.29.2/class/call.html) - [GIL policy](https://pyo3.rs/v0.29.2/free-threading.html), [experimental async limits](https://pyo3.rs/v0.29.2/async-await.html), [task conversion](https://docs.rs/pyo3-async-runtimes/0.29.0/pyo3_async_runtimes/fn.into_future_with_locals.html), [native cancellation/delivery](https://docs.rs/pyo3-async-runtimes/0.29.0/pyo3_async_runtimes/tokio/fn.future_into_py.html) - [performance](https://pyo3.rs/v0.29.2/performance.html), [PyBackedBytes](https://docs.rs/pyo3/0.29.2/pyo3/pybacked/struct.PyBackedBytes.html), [typing](https://pyo3.rs/v0.29.2/python-typing-hints.html) @@ -87,7 +87,7 @@ GIL handling to `litellm-host-python`. measured reason. - Provider dispatch belongs in the `litellm-inference-*` route crate (e.g. `litellm_inference_messages`), not in this PyO3 crate. -- Python owns rollout state and fallback. Rust should return errors; Python +- Python owns rollout state and fallback. Rust should return errors. Python decides whether to raise or fall back. For a rust-only provider/route (no Python reference), the Python side is a thin dispatch that calls Rust and raises when the bridge is unavailable, with no fallback. @@ -101,7 +101,7 @@ GIL handling to `litellm-host-python`. Python implementation, switch its dispatch to `NO_PYTHON` in the same change - Keep the Python interface minimal (well under 100 lines per route): it only marshals inputs and calls Rust. Do not add per-route feature flags, and do - not put provider dispatch in `litellm/main.py`; it lives in a thin dispatch + not put provider dispatch in `litellm/main.py`. It lives in a thin dispatch class under `litellm/llms///`. ## Data Handling diff --git a/litellm-rust/crates/python-bridge/src/cache/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/AGENTS.md index 7a692670eba..3f1097f4bc2 100644 --- a/litellm-rust/crates/python-bridge/src/cache/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/cache/AGENTS.md @@ -2,6 +2,6 @@ This folder owns how Rust inference reaches the selected cache: global cache selection, route admission, delegation to a Python cache and the experimental V2 native handles. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates -`mod.rs` exposes the cache boundary to routes and module registration; adapter directories remain private. `selection.rs` owns global cache selection, route admission and inference protocol composition for both adapters +`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `selection.rs` owns global cache selection, route admission and inference protocol composition for both adapters `python/` delegates operations to the selected Python cache without discovering configuration. `native/` owns the experimental V2 native cache handles. Neither adapter depends on shared selection or the other adapter. Shared composition depends on the adapters, and routes use only the parent module's exports diff --git a/litellm-rust/crates/python-bridge/src/routes/AGENTS.md b/litellm-rust/crates/python-bridge/src/routes/AGENTS.md index 948c67eb647..310c25ea955 100644 --- a/litellm-rust/crates/python-bridge/src/routes/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/routes/AGENTS.md @@ -1,12 +1,12 @@ # Route boundary -These invariants apply to lifecycle-bearing public calls; value-oriented APIs that return an already running Future retain their explicit contract +These invariants apply to lifecycle-bearing public calls. Value-oriented APIs that return an already running Future retain their explicit contract Route modules own public argument projection, route-specific host operations, response and error construction, and composition of the core route with the Python host. Provider dispatch, transport execution and normalization belong to core and provider crates. Runtime waiting, cancellation mechanics and execution state validation belong to `litellm-host-python` Before execution starts, perform only admission checks needed to select native execution or legacy fallback. Do not fully project a request just to decide admission. Keep effectful settings reads, HTTP client acquisition, secret-source construction and caller context capture inside execution, after argument preparation and SDK preflight. Configure resources from that prepared view, never the original entrypoint kwargs -The host driver owns sequencing and terminal events; the bridge supplies fallible resource composition without exposing route types to the driver. An unstarted async call performs no resource setup. Setup errors after start follow the terminal failure contract and never authorize fallback or provider replay +The host driver owns sequencing and terminal events. The bridge supplies fallible resource composition without exposing route types to the driver. An unstarted async call performs no resource setup. Setup errors after start follow the terminal failure contract and never authorize fallback or provider replay Use the shared `run_public_call` boundary with hooks supplied by bridge composition. `callbacks-legacy-python` owns legacy argument sharing and `Logging` dispatch behind `PublicCall` and `LegacyLogging`. Route bindings supply `callbacks-legacy-python::LoggingOperation` when composing legacy logging and may retain the request needed for projection, but must not duplicate the legacy callback contract diff --git a/litellm-rust/crates/python-compat/AGENTS.md b/litellm-rust/crates/python-compat/AGENTS.md index 0869568d33b..55049daf612 100644 --- a/litellm-rust/crates/python-compat/AGENTS.md +++ b/litellm-rust/crates/python-compat/AGENTS.md @@ -7,17 +7,17 @@ - Cache-key derivation: sha256 over `str(value)` must match Python byte for byte (`repr::to_str`) - Byte-identical writes where Python compares raw values (`json::dumps`, `repr`) - Relation to [`py_literal`](https://docs.rs/py_literal/latest/py_literal/): replaced, do not reintroduce - - Its pest grammar backtracks: parse time doubles per nested `[`/`{` (105 ms at depth 16); ours is linear (19 µs at depth 128) - - Its formatter is not `repr` (`2e-1`, always single quotes, escapes non-ASCII); it also lost `-0.0`, `(1+2j)`, `set()` - - `cache-response` and `cache-disk` still depend on it; migrate them here + - Its pest grammar backtracks: parse time doubles per nested `[`/`{` (105 ms at depth 16). Ours is linear (19 µs at depth 128) + - Its formatter is not `repr` (`2e-1`, always single quotes, escapes non-ASCII). It also lost `-0.0`, `(1+2j)`, `set()` + - `cache-response` and `cache-disk` still depend on it. Migrate them here - Relation to [`serde-pickle`](https://docs.rs/serde-pickle/latest/serde_pickle/): the pickle codec, used only through its serde interface - Never `serde_pickle::Value`: its `BTreeMap` dicts reorder keys - - Accepted limits: ints beyond i64, `tuple`/`set`/`frozenset` decode as lists, class references (`GLOBAL`/`REDUCE`) fail; writes protocol 3 -- Every behavior is pinned by CPython output, not by reasoning; everything under `generated/` is script output, never hand-edited - - Regenerate `generated/values.json` with `scripts/generate_fixtures.py`; add a corpus row before changing behavior - - Divergences go in `KNOWN` in `tests/fixtures.rs` with a reason; an entry that starts matching fails until deleted + - Accepted limits: ints beyond i64, `tuple`/`set`/`frozenset` decode as lists, class references (`GLOBAL`/`REDUCE`) fail. Writes protocol 3 +- Every behavior is pinned by CPython output, not by reasoning. Everything under `generated/` is script output, never hand-edited + - Regenerate `generated/values.json` with `scripts/generate_fixtures.py`. Add a corpus row before changing behavior + - Divergences go in `KNOWN` in `tests/fixtures.rs` with a reason. An entry that starts matching fails until deleted - Regenerate `generated/nonprintable.rs` with `scripts/generate_nonprintable.py` when the target Python's Unicode version changes - - `scripts/verify_rust_pickles.py` checks CPython reads Rust pickles, with class resolution disabled; CI does not run Python + - `scripts/verify_rust_pickles.py` checks CPython reads Rust pickles, with class resolution disabled. CI does not run Python - Decoders recurse, so they reject nesting beyond `MAX_DEPTH` for stack safety: deliberately stricter than CPython, whose parser takes ~200 levels and whose unpickler has no limit (pinned as `nested_150`) - - Formatters (`repr`, `json`) are unbounded; values from the decoders are already capped, a hand-built `Value` is the caller's responsibility + - Formatters (`repr`, `json`) are unbounded. Values from the decoders are already capped, a hand-built `Value` is the caller's responsibility - `literal_eval` must stay linear in depth: `tests/limits.rs` times the deepest parse, `benches/formats.rs` measures the curve but is manual, since CI runs no Rust bench diff --git a/litellm-rust/crates/storage-clickhouse/AGENTS.md b/litellm-rust/crates/storage-clickhouse/AGENTS.md index 569e1402859..594d7a65da2 100644 --- a/litellm-rust/crates/storage-clickhouse/AGENTS.md +++ b/litellm-rust/crates/storage-clickhouse/AGENTS.md @@ -6,4 +6,4 @@ The crate has no product tables, OTLP types, or named trace queries. `litellm-sp It also applies embedded SQLx migrations through the `_sqlx_migrations` ledger -Migration files are append-only, and changed applied files are rejected by their checksums. Startup migrations must be replay-safe schema changes because the runner records success after execution without dirty states or locks. Backfills belong in coordinated jobs outside proxy startup. The replay policy lives in `ClickHouseMigrate::apply`, `dirty_version`, and `lock`; a Keeper-backed or deploy-time runner changes only those methods +Migration files are append-only, and changed applied files are rejected by their checksums. Startup migrations must be replay-safe schema changes because the runner records success after execution without dirty states or locks. Backfills belong in coordinated jobs outside proxy startup. The replay policy lives in `ClickHouseMigrate::apply`, `dirty_version`, and `lock`. A Keeper-backed or deploy-time runner changes only those methods diff --git a/litellm/proxy/_experimental/mcp_server/AGENTS.md b/litellm/proxy/_experimental/mcp_server/AGENTS.md index 626646c4c5c..8c53216f77a 100644 --- a/litellm/proxy/_experimental/mcp_server/AGENTS.md +++ b/litellm/proxy/_experimental/mcp_server/AGENTS.md @@ -18,9 +18,9 @@ litellm/proxy/_experimental/mcp_server/ mcp_server_manager.py # upstream server registry, clients, tool routing [PR7: _create_mcp_client swaps resolve_mcp_auth -> resolve_credentials] auth/ user_api_key_auth_mcp.py # LiteLLM admission auth and MCP request headers - token_exchange.py # OAuth token exchange handling [unchanged; V1TokenExchangeAdapter delegates here] + token_exchange.py # OAuth token exchange handling [unchanged, V1TokenExchangeAdapter delegates here] litellm_auth_handler.py # authenticated-user adapter for MCP sessions - client_allowlist.py # gateway-level client application allowlist (mcp_allowed_clients); leaf module, no litellm.proxy imports + client_allowlist.py # gateway-level client application allowlist (mcp_allowed_clients), a leaf module with no litellm.proxy imports outbound_credentials/ # NEW — typed upstream-credential resolution (resolve_credentials + arms) __init__.py # public surface: resolve_credentials, the configs, CredError result.py # Ok | Error union (pure stdlib) @@ -28,13 +28,13 @@ litellm/proxy/_experimental/mcp_server/ httpx_auth.py # NoOpAuth, StaticHeaderAuth (every mode -> one httpx.Auth) resolver.py # resolve_credentials(): exhaustive per-mode match + assert_never seams.py # injected Protocols (one per cache-touching mode) - v1_adapters.py # v1-backed seam bodies; delegate to auth/oauth2/db owners + v1_adapters.py # v1-backed seam bodies that delegate to auth/oauth2/db owners adapter.py # to_subject / to_server_spec / raise_public (v1 <-> v2 boundary) discoverable_endpoints.py # MCP OAuth metadata, authorize, token, callback byok_oauth_endpoints.py # BYOK OAuth UI/API flow oauth_utils.py # redirect URI and proxy base URL validation - oauth2_token_cache.py # OAuth2 and per-user token resolution/cache [PR7: resolve_mcp_auth removed; cache class stays, V1OAuth2CacheAdapter delegates to async_get_token] - db.py # MCP server, credential, env var, submission DB access [unchanged; V1ByokStore delegates to _get_byok_credential / get_user_credential] + oauth2_token_cache.py # OAuth2 and per-user token resolution/cache [PR7: resolve_mcp_auth removed, cache class stays, V1OAuth2CacheAdapter delegates to async_get_token] + db.py # MCP server, credential, env var, submission DB access [unchanged, V1ByokStore delegates to _get_byok_credential / get_user_credential] toolset_db.py # MCP toolset DB access rest_endpoints.py # proxy REST facade for listing/calling MCP tools [PR7: 7-arm only — pass identity + inbound token down instead of mcp_auth_header] openapi_to_mcp_generator.py# OpenAPI spec to MCP tool generation @@ -60,7 +60,7 @@ module materially harder to understand. ## Implementation Rules - Preserve the boundary between LiteLLM admission auth and upstream MCP auth. - Admission belongs in `auth/user_api_key_auth_mcp.py`; upstream token exchange, + Admission belongs in `auth/user_api_key_auth_mcp.py`. Upstream token exchange, delegated auth, per-user OAuth, BYOK, and raw header forwarding belong in the dedicated OAuth/header modules. - Treat `none`, bearer/API key, OAuth, OAuth token exchange, delegated upstream diff --git a/litellm/rust_bridge/AGENTS.md b/litellm/rust_bridge/AGENTS.md index a3fa17217fc..0c313b097cf 100644 --- a/litellm/rust_bridge/AGENTS.md +++ b/litellm/rust_bridge/AGENTS.md @@ -1,8 +1,8 @@ # Python boundary -This package owns native rollout and fallback selection, Python public API compatibility, settings projection and the Python bindings supplied to the Rust bridge. Rust core owns provider execution; `litellm-host-python` owns CPython runtime mechanics; `callbacks-legacy-python` owns legacy callback sharing and dispatch policy +This package owns native rollout and fallback selection, Python public API compatibility, settings projection and the Python bindings supplied to the Rust bridge. Rust core owns provider execution. `litellm-host-python` owns CPython runtime mechanics. `callbacks-legacy-python` owns legacy callback sharing and dispatch policy -`lifecycle.py` owns generic inline execution and stream iteration. `streams.py` supplies the product binding and public stream wrappers through the bridge rather than let the generic Rust host import this package by name. Keep one driver implementing `start`, `resume_value`, `resume_error` and idempotent `close`; do not create a second implementation +`lifecycle.py` owns generic inline execution and stream iteration. `streams.py` supplies the product binding and public stream wrappers through the bridge rather than let the generic Rust host import this package by name. Keep one driver implementing `start`, `resume_value`, `resume_error` and idempotent `close`. Do not create a second implementation Generic execution steps and inline suspension handling must not depend on LiteLLM response metadata. Public stream construction, `_hidden_params` and header compatibility remain product responsibilities. Keep generic stream heads opaque and preserve public header behavior in the product wrappers diff --git a/litellm/tracing/AGENTS.md b/litellm/tracing/AGENTS.md index a901fe9acaa..397252be7df 100644 --- a/litellm/tracing/AGENTS.md +++ b/litellm/tracing/AGENTS.md @@ -3,7 +3,7 @@ - Lens owns trace normalization, schemas, storage, graph assembly and querying in `BerriAI/lens` - This package owns the HTTP client, gateway response validation and asynchronous gateway spend export. Keep the trace reader remote-only - Gateway endpoints preserve authenticated user/team scope, trace references, pagination and public error categories -- `generated/` comes from `scripts/generate_trace_types.py` using the Lens schemas pinned in `scripts/lens_assets/source.json`. Update the canonical Lens Rust contracts and import their schema outputs before regenerating; never edit generated Python +- `generated/` comes from `scripts/generate_trace_types.py` using the Lens schemas pinned in `scripts/lens_assets/source.json`. Update the canonical Lens Rust contracts and import their schema outputs before regenerating. Never edit generated Python - Keep the pinned schema and fixture checksums current. CI verifies assets and regenerated Python without a Rust trace implementation in this repository - Independent gateway ClickHouse spend logging lives in `litellm/integrations/clickhouse`, using `litellm/rust_bridge/clickhouse.py` and the native spend writer - OTLP upload routes only return setup guidance. Keep their JSON/protobuf error encoding without restoring ingestion or a native trace fallback diff --git a/tests/AGENTS.md b/tests/AGENTS.md index accd552cf8e..ac21776bdcc 100644 --- a/tests/AGENTS.md +++ b/tests/AGENTS.md @@ -6,7 +6,7 @@ itself, no proxy at all: `tests/e2e_harness`. Two fit, split it ## What good looks like -Red when the claim in the name is broken. Prove it: mutate the behaviour, red; restore, green. Put the +Red when the claim in the name is broken. Prove it: mutate the behaviour, red. Restore, green. Put the mutation in the PR body ```python @@ -21,7 +21,7 @@ def test_custom_price_is_reported_and_charged(gateway: Gateway) -> None: Rates in the test, expected computed by hand, one call, `response.text` in the assert Assert the whole value. Iterating `expected_body.items()` (`test_responses_api_request_body.py`) cannot -see an extra key; that is the shape of `stream_options.include_usage` (#19777, #28553) +see an extra key. That is the shape of `stream_options.include_usage` (#19777, #28553) The linter catches no-assert, mock-echo and credential skips. It cannot see an assert behind an `if` (a poll that ends in `pytest.fail` is fine), `except Exception` around the call @@ -29,7 +29,7 @@ behind an `if` (a poll that ends in `pytest.fail` is fine), `except Exception` a ## Where it goes -What the assertion depends on goes in the test; everything else in conftest. A rate in a fixture three +What the assertion depends on goes in the test. Everything else in conftest. A rate in a fixture three directories up makes a failed assertion unreadable. Extend the file that already covers the behaviour ## Writing it so a human can read it diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 279fe076c50..3fd2fb9b9a6 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -1,6 +1,6 @@ # e2e harness conventions -Code-style rules for writing tests under `tests/e2e/`. The harness already encodes the plumbing; your job is the feature-specific behavior, not reinventing it. For what a complete test must do (the lifecycle contract, asserting both recorded state and enforced behavior) and how to run a suite, see `CONTRIBUTING.md` in this directory. Repo-wide conventions live in the root `AGENTS.md` +Code-style rules for writing tests under `tests/e2e/`. The harness already encodes the plumbing. Your job is the feature-specific behavior, not reinventing it. For what a complete test must do (the lifecycle contract, asserting both recorded state and enforced behavior) and how to run a suite, see `CONTRIBUTING.md` in this directory. Repo-wide conventions live in the root `AGENTS.md` ## What good looks like @@ -24,13 +24,13 @@ assert still tears down. Assert what the caller receives ## Where it goes By the surface a customer would name: `guardrails`, `llm_translation`, `management`. Mutation check -deferred; it needs credentials +deferred. It needs credentials ## Suite folders Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family or behavior area. If you add a new folder, you must add a line here describing what kind of tests belong in it, so the layout stays self-describing. `gateway/` is the exception: it holds proxy configuration only and never tests -- `migrations/` - isolated Docker startup, concurrent migration, crash recovery, and legacy database compatibility. The CircleCI migration workflow enables `LITELLM_MIGRATION_TESTS=1`; these tests own their proxy containers and databases, so they do not use the shared proxy preflight or shared database cleanup +- `migrations/` - isolated Docker startup, concurrent migration, crash recovery, and legacy database compatibility. The CircleCI migration workflow enables `LITELLM_MIGRATION_TESTS=1`. These tests own their proxy containers and databases, so they do not use the shared proxy preflight or shared database cleanup - `llm_translation/` - LLM endpoint and provider-translation behavior: passthrough, custom pricing, OCR, and the non-chat inference endpoints (`/v1/responses`, `/v1/messages`, `/embeddings`, `/v1/rerank`, `/v1/audio/speech`, `/v1/images/generations`), each against a deployment the test creates via `/model/new` and deletes on teardown - `access_control/` - the gateway's authorization and error-shape contract: per-key model allow-lists, route-group permissions (`allowed_routes`), and unknown-model validation @@ -38,33 +38,33 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `batches/` - the `/batches` endpoint (placeholder until the first test lands) - `realtime/` - realtime websocket sessions, including the pipecat audio path - `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) -- `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface; not Playwright) +- `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface, not Playwright) - `a2a/` - the A2A (agent-to-agent) surface: admin registration via `/v1/agents`, proxy-fronted card discovery at `/.well-known/agent-card.json`, and JSON-RPC `message/send` invocation, driving agents backed by the litellm completion bridge (a real provider) and asserting protocol-version normalization (0.3 vs 1.0) -- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below); plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions` in `test_mcp_chat_completion_oauth_e2e.py` and direct MCP protocol operations in `test_mcp_oauth_happy_path_e2e.py`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion or protocol call lists and executes the server's tools with the stored per-user token +- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below), plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions` in `test_mcp_chat_completion_oauth_e2e.py` and direct MCP protocol operations in `test_mcp_oauth_happy_path_e2e.py`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion or protocol call lists and executes the server's tools with the stored per-user token - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection -- `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did; and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a fixed budget nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) -- `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). Its aggregation logic (locust, process usage, session anomaly) is covered by `tests/e2e_harness/load/` -- `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite +- `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a fixed budget nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) +- `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). Its aggregation logic (locust, process usage, session anomaly) is covered by `tests/e2e_harness/load/` +- `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) -- `gateway/` - proxy configuration only (`litellm-config.yml`); no tests -- `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher, covered by `tests/e2e_harness/claude_code/`. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke -- `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` +- `gateway/` - proxy configuration only (`litellm-config.yml`). No tests +- `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture. Ships its own driver/builder/publisher, covered by `tests/e2e_harness/claude_code/`. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`). The CLI-driving path stays bespoke +- `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream, see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport. The Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` ## MCP suite: real Datadog only Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datadog remote MCP server, except the two Linear OAuth tests `test_mcp_chat_completion_oauth_e2e.py` and `test_mcp_oauth_happy_path_e2e.py`. Do not add a compose service, FastMCP fixture, mock upstream, or any other fake MCP host for this suite - Register via `register_datadog_mcp` in `tests/e2e/mcp/datadog_mcp.py` (or extend that helper if you need a different `toolsets=` / `allowed_tools` slice of the same Datadog endpoint). That posts `/v1/mcp/server` with `url=datadog_mcp_url(...)` and static headers `DD-API-KEY` / `DD-APPLICATION-KEY` from the process env -- Auth is Datadog's documented CI/header path, not a browser OAuth authorize/token dance. Hard-fail when `DD_API_KEY` or `DD_APP_KEY` is missing (`assert_dd_mcp_creds`); never skip for a missing fake upstream -- Prefer calling real Datadog tools that prove the product path (e.g. `search_datadog_logs` for list/call and permission denials). Seed a unique marker (`e2e-datadog-mcp-*`) in a chat completion when you need a log the tool can find; dual-read with `dd_logs` from conftest when delivery matters +- Auth is Datadog's documented CI/header path, not a browser OAuth authorize/token dance. Hard-fail when `DD_API_KEY` or `DD_APP_KEY` is missing (`assert_dd_mcp_creds`). Never skip for a missing fake upstream +- Prefer calling real Datadog tools that prove the product path (e.g. `search_datadog_logs` for list/call and permission denials). Seed a unique marker (`e2e-datadog-mcp-*`) in a chat completion when you need a log the tool can find. Dual-read with `dd_logs` from conftest when delivery matters - Delete the MCP server (and any keys) through `resources.defer` the same way every other suite tears down -- If a new MCP behavior cannot be covered with Datadog's tool surface, say so in the PR and get agreement before inventing another upstream; the default is always Datadog -- The two standing exceptions are `test_mcp_chat_completion_oauth_e2e.py` and `test_mcp_oauth_happy_path_e2e.py`. Datadog authenticates with the static `DD-API-KEY` / `DD-APPLICATION-KEY` headers and exposes no authorize/token dance at all, so these tests drive a real Linear MCP server instead; they are still real remote upstreams, so the no-mock, no-fixture rule above holds unchanged. The direct OAuth test also uses the existing live provider edge to inspect forwarded headers without replay, and owns a separate source-built gateway for cold restarts +- If a new MCP behavior cannot be covered with Datadog's tool surface, say so in the PR and get agreement before inventing another upstream. The default is always Datadog +- The two standing exceptions are `test_mcp_chat_completion_oauth_e2e.py` and `test_mcp_oauth_happy_path_e2e.py`. Datadog authenticates with the static `DD-API-KEY` / `DD-APPLICATION-KEY` headers and exposes no authorize/token dance at all, so these tests drive a real Linear MCP server instead. They are still real remote upstreams, so the no-mock, no-fixture rule above holds unchanged. The direct OAuth test also uses the existing live provider edge to inspect forwarded headers without replay, and owns a separate source-built gateway for cold restarts ## Lay the pattern down in a class -Keep the cases for one feature inside a class so the file reads as a spec for how that feature behaves in production. The class name says what is under test; each method is one behavior. Think of it as documenting the contract, with the rough intent being +Keep the cases for one feature inside a class so the file reads as a spec for how that feature behaves in production. The class name says what is under test. Each method is one behavior. Think of it as documenting the contract, with the rough intent being ```python # pseudo-code to convey intent, not the real API @@ -80,13 +80,13 @@ class TestPromptCompression: assert response.cost == compressed_value # the cost was actually reduced ``` -That snippet only conveys intent. What you actually write uses the real harness: the `client` fixture for your suite, the `scoped_key` fixture for an auto-deleted key, typed pydantic bodies from `models.py`, and `unwrap(...)` on the tagged-union result. `tests/e2e/llm_translation/test_custom_pricing_e2e.py` is the reference to copy from; it creates a scoped key, drives a real gemini call, polls `/spend/logs` to a deadline for the cost-breakdown row, then asserts the input and output costs match the configured custom rates and that a sibling deployment kept its own price. Read it before writing yours +That snippet only conveys intent. What you actually write uses the real harness: the `client` fixture for your suite, the `scoped_key` fixture for an auto-deleted key, typed pydantic bodies from `models.py`, and `unwrap(...)` on the tagged-union result. `tests/e2e/llm_translation/test_custom_pricing_e2e.py` is the reference to copy from. It creates a scoped key, drives a real gemini call, polls `/spend/logs` to a deadline for the cost-breakdown row, then asserts the input and output costs match the configured custom rates and that a sibling deployment kept its own price. Read it before writing yours -## Use the shared transport; never touch requests directly +## Use the shared transport, never touch requests directly Every HTTP call goes through the shared transport, never through `requests.*` in a test. `e2e_http.py` is the only module permitted to call `requests.*`, and that is enforced in CI by `tests/code_coverage_tests/check_e2e_no_raw_requests.py`. A test that imports requests will fail the check -One deliberate exception: LLM-endpoint calls in `llm_translation/` go through the real provider SDKs (OpenAI, Anthropic) via the suite's `sdk` fixture (`llm_translation/sdk_clients.py`), because that is what customers actually run against the proxy (LIT-4577). The SDKs raise their own typed exceptions on failure, which is exactly the customer-observable contract; management routes (model/key CRUD, spend read-back) and endpoints no official SDK covers (e.g. `/v1/rerank`, `/v1/ocr`, custom passthrough paths) stay on the shared transport. Raw HTTP client imports remain banned either way +One deliberate exception: LLM-endpoint calls in `llm_translation/` go through the real provider SDKs (OpenAI, Anthropic) via the suite's `sdk` fixture (`llm_translation/sdk_clients.py`), because that is what customers actually run against the proxy (LIT-4577). The SDKs raise their own typed exceptions on failure, which is exactly the customer-observable contract. Management routes (model/key CRUD, spend read-back) and endpoints no official SDK covers (e.g. `/v1/rerank`, `/v1/ocr`, custom passthrough paths) stay on the shared transport. Raw HTTP client imports remain banned either way The shape is layered so tests stay declarative @@ -96,7 +96,7 @@ The shape is layered so tests stay declarative Each suite provides its own `client` fixture (see `llm_translation/passthrough_client.py`), a frozen dataclass that holds the shared `ProxyClient` (as `.proxy`) and adds suite-specific routes. Cleanup runs through that same `ProxyClient`, so whatever keys or customers your test creates get torn down by the `resources` fixture -Request and response bodies are typed pydantic models in `models.py`; only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass +Request and response bodies are typed pydantic models in `models.py`. Only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass Mark live tests with `@pytest.mark.e2e` (on the class or the module). Coverage of the harness itself lives outside the suite in `tests/e2e_harness/` (see its `AGENTS.md`) and runs without a proxy, so nothing under `tests/e2e/` is a markerless test. Add `@pytest.mark.quiet_stack` to a test that measures the proxy itself (RSS, latency): the shared stack lock in `stack_lock.py` then runs it while no other test on the host is hitting the stack, marked or not, so the reading depends only on the test's own traffic. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache @@ -104,13 +104,13 @@ Mark live tests with `@pytest.mark.e2e` (on the class or the module). Coverage o `E2E_FIXTURE_MODE` scopes the proxy's provider-bound traffic: `live` (the default, and what an unset variable means: nothing changes), `record` (the proxy's provider calls are forwarded to the real provider through a local edge server and written to a fixture bundle), or `replay` (the edge answers those calls from the bundle, so the run makes zero provider calls and spends nothing). Test-to-proxy traffic always goes over the wire in every mode: record and replay both need the live proxy and database, because the point is that key auth, routing, cost calculation, and spend-log writes execute for real while only the provider is swapped out. Breaking any of those in the proxy turns a replay run red -The seam is `provider_edge.py`: `start_provider_edge` boots an in-process HTTP server (one shared instance per pytest process, `e2e_config.provider_edge_base` is the accessor) that mounts each supported provider under a path prefix (`EDGE_MOUNTS`: `/openai` -> `https://api.openai.com`, `/anthropic` -> `https://api.anthropic.com`). A test participates by registering its deployment with `api_base=provider_edge_base("openai")` plus the provider's path suffix; `quota_management/spend_tracking/test_provider_edge_spend_e2e.py` is the reference. In live mode the accessor returns None and the deployment defaults to the real provider, so an edge-wired test runs in all three modes unchanged. Non-wired tests hit their providers live in every mode. The edge binds `E2E_PROVIDER_EDGE_BIND_HOST` (default 127.0.0.1) and advertises `E2E_PROVIDER_EDGE_ADVERTISE_HOST` in the api_base it hands out, for proxies running in containers +The seam is `provider_edge.py`: `start_provider_edge` boots an in-process HTTP server (one shared instance per pytest process, `e2e_config.provider_edge_base` is the accessor) that mounts each supported provider under a path prefix (`EDGE_MOUNTS`: `/openai` -> `https://api.openai.com`, `/anthropic` -> `https://api.anthropic.com`). A test participates by registering its deployment with `api_base=provider_edge_base("openai")` plus the provider's path suffix. `quota_management/spend_tracking/test_provider_edge_spend_e2e.py` is the reference. In live mode the accessor returns None and the deployment defaults to the real provider, so an edge-wired test runs in all three modes unchanged. Non-wired tests hit their providers live in every mode. The edge binds `E2E_PROVIDER_EDGE_BIND_HOST` (default 127.0.0.1) and advertises `E2E_PROVIDER_EDGE_ADVERTISE_HOST` in the api_base it hands out, for proxies running in containers A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per provider call in call order (`0000-post-openai-v1-chat-completions.json`). Request headers are never stored (provider credentials never touch disk), non-JSON request bodies store a canonicalized sha256 digest instead of the bytes, `multipart/form-data` bodies store their ordinary fields plus a JSON list of the uploaded parts' `[field, filename, content-type]` triples and a digest of their content, so the per-request random boundary and the envelope never reach the key, and responses store status, filtered headers, and the verbatim body base64-encoded, which is part of why bundles are gitignored. Responses come in two shapes told apart by a `kind` tag: an ordinary one holding a single base64 body, and, for a response the provider streamed (`content-type: text/event-stream`), one holding its transfer chunks in order plus why the stream ended early if it did, so replay reproduces the split points the provider chose instead of one coalesced body. `fixture_bundle.py` owns the format, and `BUNDLE_FORMAT_VERSION` is checked on load, so a bundle recorded under older rules is refused by name rather than partially read. Record serves the proxy the same filtered stored response replay will serve later, chunk for chunk on a stream, so the two modes are byte-identical from the proxy's side of the socket. Replay does not reproduce the provider's inter-chunk timing (chunks go out as fast as the socket takes them), so a test that judges streaming on the clock, such as the `stream_event_arrivals` lead between the first content delta and `message_stop`, gates that assertion on `provider_paces_stream()` and proves only the event grammar in replay Multipart identity is the fiddly corner, and the rules exist because each one had a collision behind it. A part counts as an upload when it carries a filename or declares its own content type, and everything else is an ordinary field. Field names get a `name[n]` suffix on repeats, with a literal `[` doubled first, so a form that repeats `purpose` never keys the same as one that literally sends `purpose[1]`. A field whose name reads as a credential is stored as ``, which stays key-preserving because the key is recomputed from the stored request rather than saved alongside it, so the live request carrying the real value still matches its redacted fixture. A field value that is not UTF-8 is stored as a base64 sha256 digest, base64 and not hex because the canonicalizer rewrites any 64-character hex run to `` and would fold every binary value onto one key. The uploaded parts contribute a JSON list rather than a `field:filename` string, so a separator inside a filename cannot impersonate a field boundary, and their byte length is stored for a reader's benefit but deliberately left out of the key, since the canonicalizer absorbs timestamp and id drift inside a file that changes its length -Replay matches calls per test by canonical key: `fixture_canonical.py` canonicalizes the recorded request (volatile headers and credential fields out, unique markers, generated ids, uuids, and timestamps replaced with fixed placeholders, object keys sorted) and the key is the method, edge path, and a content hash, so identity survives re-records and machine changes while any real content drift comes back as an HTTP 599 naming the computed key, the closest recorded key with its file, and a content diff, and never falls through to a live call. Matching is order-independent across distinct keys (concurrent calls may interleave) and FIFO within one key (a retry loop replays its responses in recorded order); a passed test must also consume its whole recording, or teardown fails it naming a leftover key. Either way the fix is always to re-record with `E2E_FIXTURE_MODE=record`. Every rewrite rule lives in `fixture_canonical.py`, so a new volatile header, credential field name, or generated-id shape is one edit there. Record starts fresh every time: it wipes the previous bundle (refusing to wipe a directory that is not a bundle) and never reads it. A replay bundle whose manifest is older than seven days hard-fails at collection time naming the bundle's age, so replay can never certify against fixtures that have drifted more than a week from the live providers +Replay matches calls per test by canonical key: `fixture_canonical.py` canonicalizes the recorded request (volatile headers and credential fields out, unique markers, generated ids, uuids, and timestamps replaced with fixed placeholders, object keys sorted) and the key is the method, edge path, and a content hash, so identity survives re-records and machine changes while any real content drift comes back as an HTTP 599 naming the computed key, the closest recorded key with its file, and a content diff, and never falls through to a live call. Matching is order-independent across distinct keys (concurrent calls may interleave) and FIFO within one key (a retry loop replays its responses in recorded order). A passed test must also consume its whole recording, or teardown fails it naming a leftover key. Either way the fix is always to re-record with `E2E_FIXTURE_MODE=record`. Every rewrite rule lives in `fixture_canonical.py`, so a new volatile header, credential field name, or generated-id shape is one edit there. Record starts fresh every time: it wipes the previous bundle (refusing to wipe a directory that is not a bundle) and never reads it. A replay bundle whose manifest is older than seven days hard-fails at collection time naming the bundle's age, so replay can never certify against fixtures that have drifted more than a week from the live providers A replayed response carries the recorded provider response id, and `LiteLLM_SpendLogs.request_id` (the table's primary key) is that id, so a replay against a database that still holds the record run's rows silently dedupes its spend inserts and any spend assertion goes red with zero matching rows and nothing in the proxy log. Run both modes with `E2E_RESET_SPEND_LOGS=1` (plus `DATABASE_URL` in the runner env) so each session truncates the table after itself, or replay against a fresh database, which is the CI shape @@ -125,7 +125,7 @@ E2E_FIXTURE_MODE=replay E2E_FIXTURE_DIR=/tmp/e2e-fixtures E2E_RESET_SPEND_LOGS=1 Point the proxy at bogus provider credentials for the replay run and it still has to pass: that is the whole proof that nothing left the process. Bundles are never committed. `tests/e2e/.fixtures` is gitignored because a bundle holds verbatim provider response bodies and hard-fails after seven days. CI records and replays this lane on a schedule in `.github/workflows/e2e_record_replay.yml`, publishing the bundle as a private `e2e-fixtures-bundle` artifact instead of committing it, selecting the tests with the `@pytest.mark.replayable` marker, and proving the bogus-credentials replay hermetic by counting provider egress with `.github/scripts/e2e_egress_sentinel.py` -Current limits: Bedrock cannot be mounted in record or replay (SigV4 signs the Host header, so a rewritten api_base fails signature verification); a test that needs to observe the Converse body registers its own `LiveEdge` with `provider_edge_bedrock.bedrock_signer` re-signing the forwarded request, and carries the `provider_edge_host` opt-in marker because the gateway must reach the pytest host, which the Buildkite ephemeral stack cannot, so those tests are deselected unless `E2E_PROVIDER_EDGE_HOST_REACHABLE` is set, as on a local run whose gateway can reach the pytest host. Deployments baked into the proxy's config file cannot be edge-wired (only `/model/new` registrations can carry the edge api_base), and a file upload routed by `custom_llm_provider` through the proxy's `files_settings` block never passes a deployment at all, so the batches `model_param` and `provider_fallback` scenarios keep uploading live in every mode +Current limits: Bedrock cannot be mounted in record or replay (SigV4 signs the Host header, so a rewritten api_base fails signature verification). A test that needs to observe the Converse body registers its own `LiveEdge` with `provider_edge_bedrock.bedrock_signer` re-signing the forwarded request, and carries the `provider_edge_host` opt-in marker because the gateway must reach the pytest host, which the Buildkite ephemeral stack cannot, so those tests are deselected unless `E2E_PROVIDER_EDGE_HOST_REACHABLE` is set, as on a local run whose gateway can reach the pytest host. Deployments baked into the proxy's config file cannot be edge-wired (only `/model/new` registrations can carry the edge api_base), and a file upload routed by `custom_llm_provider` through the proxy's `files_settings` block never passes a deployment at all, so the batches `model_param` and `provider_fallback` scenarios keep uploading live in every mode ## Typing @@ -133,7 +133,7 @@ The harness is fully typed with no error budget: `make lint-e2e-basedpyright` mu ## Typed test metadata -Separate from the coverage registry and additive to it: `@meta(Subject(...))` from `e2e_metadata.py` says what a test DRIVES, as closed enums rather than a string id. `@pytest.mark.covers("cell.id")` is untouched and keeps working exactly as before; the two markers coexist on the same test, and `@meta` always goes BELOW `@covers` so `Item.location` still anchors at the first decorator and every `source` deep link stays put +Separate from the coverage registry and additive to it: `@meta(Subject(...))` from `e2e_metadata.py` says what a test DRIVES, as closed enums rather than a string id. `@pytest.mark.covers("cell.id")` is untouched and keeps working exactly as before. The two markers coexist on the same test, and `@meta` always goes BELOW `@covers` so `Item.location` still anchors at the first decorator and every `source` deep link stays put ```python @pytest.mark.covers("quota_management.budget.key.blocks_over_limit") @@ -150,7 +150,7 @@ def test_bare_key_blocks_over_its_own_budget(...) -> None: ... `route` is the endpoint the test is checking: `TEAM_MANAGEMENT` for a `/team/update` test, `SPEND_REPORTING` for a `/spend/logs` test, `MESSAGES` for a test of spend on `/v1/messages`. A test whose chat call only triggers the behavior under test, like the budget block above, leaves it unset, since its steps already name the call -Every pytest test under `tests/e2e/` declares a `Subject` with at least its `domain`, except the `load/` suite, which is kept out of the default collection. The harness's own tests in `tests/e2e_harness/` carry none, since they drive nothing. The fields themselves stay optional, since a test that makes no LLM call has no provider, model or mode to name. Every field is a closed enum, so a typo is a basedpyright error at the call site rather than a property that silently never appears. `providers`, `models` and `capabilities` are tuples even with one member, because one test node routinely drives several: the claude_code matrix runs haiku, sonnet and opus in a single body, and a spend test calls two providers on one key. Declare every provider and every model the test drives, fallbacks included. The three are independent sets with no positional pairing between them (one provider x three models is the common case), and each is deduped and sorted at declaration so the committed run artifacts diff cleanly. `models=("gpt-5.5")` is a str and not a tuple, so anything but a tuple raises a `TypeError` where the decorator runs and shows up as a collection error naming the file. `Subject` is serialized with `dataclasses.asdict`, so a new scalar field needs no serializer edit; empty fields emit no `` at all. A declared model names the constant the test drives (`CHEAP_ANTHROPIC_MODEL`, the file's own `BACKEND`), never a copy of its value, so the property cannot claim one model while an env override runs another. `e2e_metadata` and its call sites never import litellm, only the stdlib, pytest and pydantic: `Provider` mirrors litellm's `LlmProviders` values instead of importing them, because tests/e2e is shipped to the runner image on its own and a `from litellm...` at module scope would make the litellm package a hard dependency of COLLECTING the suite. `TestProviderMirrorsLitellm` in `tests/code_coverage_tests/test_e2e_metadata.py` fails on drift wherever litellm is importable and skips where it is not, so adding a provider is one line in `e2e_metadata` +Every pytest test under `tests/e2e/` declares a `Subject` with at least its `domain`, except the `load/` suite, which is kept out of the default collection. The harness's own tests in `tests/e2e_harness/` carry none, since they drive nothing. The fields themselves stay optional, since a test that makes no LLM call has no provider, model or mode to name. Every field is a closed enum, so a typo is a basedpyright error at the call site rather than a property that silently never appears. `providers`, `models` and `capabilities` are tuples even with one member, because one test node routinely drives several: the claude_code matrix runs haiku, sonnet and opus in a single body, and a spend test calls two providers on one key. Declare every provider and every model the test drives, fallbacks included. The three are independent sets with no positional pairing between them (one provider x three models is the common case), and each is deduped and sorted at declaration so the committed run artifacts diff cleanly. `models=("gpt-5.5")` is a str and not a tuple, so anything but a tuple raises a `TypeError` where the decorator runs and shows up as a collection error naming the file. `Subject` is serialized with `dataclasses.asdict`, so a new scalar field needs no serializer edit. Empty fields emit no `` at all. A declared model names the constant the test drives (`CHEAP_ANTHROPIC_MODEL`, the file's own `BACKEND`), never a copy of its value, so the property cannot claim one model while an env override runs another. `e2e_metadata` and its call sites never import litellm, only the stdlib, pytest and pydantic: `Provider` mirrors litellm's `LlmProviders` values instead of importing them, because tests/e2e is shipped to the runner image on its own and a `from litellm...` at module scope would make the litellm package a hard dependency of COLLECTING the suite. `TestProviderMirrorsLitellm` in `tests/code_coverage_tests/test_e2e_metadata.py` fails on drift wherever litellm is importable and skips where it is not, so adding a provider is one line in `e2e_metadata` Declared fields ride out as JUnit `` entries behind the fixed prefix, the same way steps do: each scalar under its field name, and each plural value as a repeated property under its SINGULAR name (`provider`, `model`, `capability`). The results JSON downstream regroups them under the plural key, so `providers`, `models` and `capabilities` are arrays there, `[]` when empty @@ -195,13 +195,13 @@ Each step is its own `` in the JUnit XML (`junit_propertie ## Coverage registry -The set of tests we want is a registry checked into this repo, one row per behavior; that file is the definition of done and the denominator. Each e2e test declares what it covers with `@pytest.mark.covers("...")`, and a small collector diffs the registry against the tests and ships coverage to the existing Grafana. No Allure, no new dependencies +The set of tests we want is a registry checked into this repo, one row per behavior. That file is the definition of done and the denominator. Each e2e test declares what it covers with `@pytest.mark.covers("...")`, and a small collector diffs the registry against the tests and ships coverage to the existing Grafana. No Allure, no new dependencies -Coverage is organized as module > feature > test. Dashboard modules are `Core LLMs`, `Non-Core LLMs`, `MCPs`, `Management/UI`, `Reliability & Performance`, `Quota Management`, `Logging & Guardrails`, and `Other`. The Loki stdout formatter maps those display modules to log-safe labels (`core_llms`, `non_core_llms`, `mcp`, `management_ui`, `reliability_performance`, `quota_management`, `logging_guardrails`, and `other`) without changing JSON or Prometheus labels. A feature is either an endpoint (`/chat/completions`) or a behavior (fallbacks, rate limits; config-driven, with no route of its own). A cell reads like `llm.chat_completions.bedrock_converse.tool_use.stream.works` +Coverage is organized as module > feature > test. Dashboard modules are `Core LLMs`, `Non-Core LLMs`, `MCPs`, `Management/UI`, `Reliability & Performance`, `Quota Management`, `Logging & Guardrails`, and `Other`. The Loki stdout formatter maps those display modules to log-safe labels (`core_llms`, `non_core_llms`, `mcp`, `management_ui`, `reliability_performance`, `quota_management`, `logging_guardrails`, and `other`) without changing JSON or Prometheus labels. A feature is either an endpoint (`/chat/completions`) or a behavior (fallbacks, rate limits) that is config-driven, with no route of its own. A cell reads like `llm.chat_completions.bedrock_converse.tool_use.stream.works` The metric is coverage: the share of registry rows that have a passing covering test, reported to Grafana per module so a gap surfaces as an uncovered row rather than a silent absence -Tests do not declare a dashboard module directly. They only declare the registry cell id with `@pytest.mark.covers("...")`; the registry row decides the module, tier, endpoint, and dashboard rollup. Run `python -m coverage_registry.collector --strict` when you want CI to reject unknown marker ids. Add `--fail-on-collection-errors` when the job should also fail on pytest collection errors. +Tests do not declare a dashboard module directly. They only declare the registry cell id with `@pytest.mark.covers("...")`. The registry row decides the module, tier, endpoint, and dashboard rollup. Run `python -m coverage_registry.collector --strict` when you want CI to reject unknown marker ids. Add `--fail-on-collection-errors` when the job should also fail on pytest collection errors. Skipping a test gives its cell back to the gap list: the collector counts a cell as covered only when a test pytest would actually run declares it, and prints the cells left claimed only by skipped tests. So a `@pytest.mark.skip` on a red cell is honest bookkeeping, not a way to keep the number up. @@ -216,7 +216,7 @@ llm..... | realtime route : openai | azure_openai | anthropic | bedrock_converse | bedrock_invoke | vertex | azure_foundry | cohere | together_ai | ollama | ollama_chat - (vocab varies per endpoint; messages is anthropic-format only) + (vocab varies per endpoint, messages is anthropic-format only) capability : basic | tool_use | prompt_cache_5m | vision | thinking | structured_output | service_tier | mid_conversation_system streaming : stream | nonstream (omit where n/a) @@ -247,21 +247,21 @@ mcp... e.g. mcp.call_tool.oauth.succeeds ``` -Reliability & Performance - behavior features (no route; endpoint is exercised_on) +Reliability & Performance - behavior features (no route, endpoint is exercised_on) ``` reliability... behavior : fallback | retry | cooldown | timeout | routing | cache | circuit_breaker | perf variant : 5xx | context_window | content_policy | 429 | timeout simple_shuffle | usage_based | latency_based | cost_based | least_busy - latency | throughput | session_anomaly | memory | idle_memory (perf only; SLO/threshold assertion, not binary) + latency | throughput | session_anomaly | memory | idle_memory (perf only, SLO/threshold assertion, not binary) assertion : routes_to_fallback | succeeds_within_retries | picks_under_tpm | returns_cached | trips_then_recovers | under_slo e.g. reliability.fallback.context_window.routes_to_fallback exercised_on=[chat_completions] reliability.cooldown.429.trips_then_recovers exercised_on=[chat_completions, messages] ``` -Quota Management - behavior features (entity- or config-driven caps and their accounting; endpoint is exercised_on) +Quota Management - behavior features (entity- or config-driven caps and their accounting, endpoint is exercised_on) ``` quota_management... @@ -286,7 +286,7 @@ quota_management... quota_management.budget.key.blocks_over_limit exercised_on=[chat_completions] ``` -Logging & Guardrails - behavior features (config-driven; endpoint is exercised_on) +Logging & Guardrails - behavior features (config-driven, endpoint is exercised_on) ``` logging... @@ -307,13 +307,13 @@ Other - holding pen (endpoint or behavior) ``` other... area : auth | lifecycle | config | ... - rule : audited periodically; a cluster here promotes to a new component + rule : audited periodically, and a cluster here promotes to a new component e.g. other.auth.jwt.valid_token_allows other.lifecycle.readiness.reports_db ``` ## Hard Rules -- no unit tests of a product feature under `tests/e2e`, and no mock tests or monkeypatching of code anywhere in it: a product feature is proven end to end against a live proxy, never with a unit test. if a contributor asks you to write an end to end test, do NOT stage a unit test with it; if you find a product gap, call it out in the PR description. the harness's own plumbing is tested outside the suite, in `tests/e2e_harness/` (mirroring this folder's layout), because the Buildkite e2e run copies `tests/e2e/` into the runner image and runs every test in it, so a harness test in here would count as a product test in the nightly numbers. those tests run without a proxy and take their inputs as arguments or env vars (setting an env var through pytest's `monkeypatch` fixture is fine, patching a function, class, or module is not), and no coverage-registry or compat-matrix cell rests on them. judge a change inside one of them by that standard, not as a misplaced product test +- no unit tests of a product feature under `tests/e2e`, and no mock tests or monkeypatching of code anywhere in it: a product feature is proven end to end against a live proxy, never with a unit test. if a contributor asks you to write an end to end test, do NOT stage a unit test with it. If you find a product gap, call it out in the PR description. the harness's own plumbing is tested outside the suite, in `tests/e2e_harness/` (mirroring this folder's layout), because the Buildkite e2e run copies `tests/e2e/` into the runner image and runs every test in it, so a harness test in here would count as a product test in the nightly numbers. those tests run without a proxy and take their inputs as arguments or env vars (setting an env var through pytest's `monkeypatch` fixture is fine, patching a function, class, or module is not), and no coverage-registry or compat-matrix cell rests on them. judge a change inside one of them by that standard, not as a misplaced product test - use model management endpoints to create new models for a test. this could be in a conftest / inline for each test. ask the user what they want. @@ -321,7 +321,7 @@ other... - when it comes to typing an input schema for an api endpoint, have it type X = A | B | C ... where X = exhaustive union of all supported input schemas and A, B, C typically are composed by a base type. types are only pretty for a api request / response body. make sure to compose types instead of repeating the same base attributes over and over again. -- spin up a local proxy by running the litellm proxy locally (`litellm --config .yml --port 4000`; see CONTRIBUTING.md), make sure all tests pass. if a test fails due to an internally found issue, let users know to create a linear ticket for it. +- spin up a local proxy by running the litellm proxy locally (`litellm --config .yml --port 4000`, see CONTRIBUTING.md), make sure all tests pass. if a test fails due to an internally found issue, let users know to create a linear ticket for it. - do not use xfail markers, tests should be written in a form that the end user expects it to pass diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 0009abe00df..b03d23fc473 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -195,7 +195,7 @@ If a step is missing, the test is not done. That is the whole pattern ## Style: lay the pattern down in a class -Keep the cases for one feature inside a class so the file reads as a spec for how that feature behaves in production. The class name says what is under test; each method is one behavior. Think of it as documenting the contract, with the rough intent being +Keep the cases for one feature inside a class so the file reads as a spec for how that feature behaves in production. The class name says what is under test. Each method is one behavior. Think of it as documenting the contract, with the rough intent being ```python # pseudo-code to convey intent @@ -211,9 +211,9 @@ class TestPromptCompression: assert response.cost == compressed_value # the cost was actually reduced ``` -That snippet only conveys intent. What you actually write uses the real harness: the `client` fixture for your suite, the `scoped_key` fixture for an auto-deleted key, typed pydantic bodies from `models.py`, and `unwrap(...)` on the tagged-union result. `tests/e2e/llm_translation/test_custom_pricing_e2e.py` is the reference to copy from; it creates a scoped key, drives a real gemini call, polls `/spend/logs` to a deadline for the cost-breakdown row, then asserts the input and output costs match the configured custom rates and that a sibling deployment kept its own price. Read it before writing yours +That snippet only conveys intent. What you actually write uses the real harness: the `client` fixture for your suite, the `scoped_key` fixture for an auto-deleted key, typed pydantic bodies from `models.py`, and `unwrap(...)` on the tagged-union result. `tests/e2e/llm_translation/test_custom_pricing_e2e.py` is the reference to copy from. It creates a scoped key, drives a real gemini call, polls `/spend/logs` to a deadline for the cost-breakdown row, then asserts the input and output costs match the configured custom rates and that a sibling deployment kept its own price. Read it before writing yours -## Use the shared transport; never touch requests directly +## Use the shared transport, never touch requests directly Every HTTP call goes through the shared transport, never through `requests.*` in a test. `e2e_http.py` is the only module permitted to call `requests.*`, and that is enforced in CI by `tests/code_coverage_tests/check_e2e_no_raw_requests.py`. A test that imports requests will fail the check @@ -227,7 +227,7 @@ The shape is layered so tests stay declarative Each suite provides its own `client` fixture (see `llm_translation/passthrough_client.py`), a frozen dataclass that holds the shared `ProxyClient` (as `.proxy`) and adds suite-specific routes. Cleanup runs through that same `ProxyClient`, so whatever keys or customers your test creates get torn down by the `resources` fixture -Request and response bodies are typed pydantic models in `models.py`; only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass +Request and response bodies are typed pydantic models in `models.py`. Only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure coverage of the harness itself lives in `tests/e2e_harness/` and runs without a proxy (`LITELLM_MASTER_KEY=sk-harness uv run pytest tests/e2e_harness`). A test that needs proxy configuration the default stack does not carry goes behind an opt-in marker (`managed_files`, `prompt_caching_stack`, `weekly`), each deselected unless its env var is set; `OPT_IN_MARKERS` in `conftest.py` maps marker to env var, and the coverage collector counts such a cell only where the env var is set. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache diff --git a/tests/e2e_harness/AGENTS.md b/tests/e2e_harness/AGENTS.md index 212e46d754c..a3e1794ce4a 100644 --- a/tests/e2e_harness/AGENTS.md +++ b/tests/e2e_harness/AGENTS.md @@ -10,8 +10,8 @@ Run them from the repo root. `e2e_config` reads `LITELLM_MASTER_KEY` at import a LITELLM_MASTER_KEY=sk-harness uv run pytest tests/e2e_harness ``` -`pytest.ini` here puts `tests/e2e` and the suite folders whose modules are under test on the path, so imports look exactly as they do inside the suite (`from e2e_http import ...`, `from batch_cleanup import ...`). `claude_code/test_request_determinism.py` drives the real `claude` CLI; deselect it with `-m "not cli_determinism"` when the CLI is not installed +`pytest.ini` here puts `tests/e2e` and the suite folders whose modules are under test on the path, so imports look exactly as they do inside the suite (`from e2e_http import ...`, `from batch_cleanup import ...`). `claude_code/test_request_determinism.py` drives the real `claude` CLI. Deselect it with `-m "not cli_determinism"` when the CLI is not installed -Rules: no `e2e` marker and no `@meta`, since nothing here drives the proxy; `@pytest.mark.covers` only where the test proves the collector or the JUnit properties read it; inputs via arguments or env vars (setting an env var through pytest's `monkeypatch` fixture is fine, patching a function, class or module is not); and the same typing bar as the suite, `make lint-e2e-basedpyright` covers this folder and allows zero errors. The raw HTTP client ban (`tests/code_coverage_tests/check_e2e_no_raw_requests.py`) applies here too +Rules: no `e2e` marker and no `@meta`, since nothing here drives the proxy. `@pytest.mark.covers` only where the test proves the collector or the JUnit properties read it. Inputs via arguments or env vars (setting an env var through pytest's `monkeypatch` fixture is fine, patching a function, class or module is not). And the same typing bar as the suite, `make lint-e2e-basedpyright` covers this folder and allows zero errors. The raw HTTP client ban (`tests/code_coverage_tests/check_e2e_no_raw_requests.py`) applies here too CI: the `python` job in `.github/workflows/test-linting.yml` runs this folder whenever anything under `tests/e2e/` (except `ui/`) or `tests/e2e_harness/` changes, with the `claude` CLI installed. The CircleCI `provider_replay_harness` job also runs the provider-edge and fixture tests at the root of this folder next to `tests/code_coverage_tests/test_provider_replay_harness.py`, which imports helpers from `test_provider_edge.py` diff --git a/tests/integration/AGENTS.md b/tests/integration/AGENTS.md index 0499bfc97c8..8df9c737ba0 100644 --- a/tests/integration/AGENTS.md +++ b/tests/integration/AGENTS.md @@ -16,12 +16,12 @@ assert float(rows[0]["spend"]) == pytest.approx(10 * 0.001 + 5 * 0.0001 + 7 * 0. ``` `sleep(3)` fails on a slow runner and taxes every fast one. Assert the outbound body in the upstream -handler; a leaked field is invisible from the response. `monkeypatch.setenv` is fine; patching our own +handler. A leaked field is invisible from the response. `monkeypatch.setenv` is fine. Patching our own function in a full stack is not ## Where it goes By the domain a user would name: `pricing`, `spend`, `routing`, `mcp`. A file only needs to live in a -directory that a `GROUPS` entry in `run.py` selects; there is no manifest and no `covers` marker on new +directory that a `GROUPS` entry in `run.py` selects. There is no manifest and no `covers` marker on new tests. A product bug the test exposes is `pytest.skip("BUG: ")` at the top of the body, not a fix in the test and not a deletion. Needs no proxy, DB or Redis: `tests/unit` diff --git a/tests/rust-python-harness/AGENTS.md b/tests/rust-python-harness/AGENTS.md index 6b0609a1c9a..3064960d66e 100644 --- a/tests/rust-python-harness/AGENTS.md +++ b/tests/rust-python-harness/AGENTS.md @@ -45,21 +45,21 @@ tests/rust-python-harness/ └── suite_runner.py ``` -- A strategy is a folder under `strategies/` with a one-line `AGENTS.md` and an `__init__.py` exporting exactly one `STRATEGY: StrategyDefinition`; its id must equal the folder name +- A strategy is a folder under `strategies/` with a one-line `AGENTS.md` and an `__init__.py` exporting exactly one `STRATEGY: StrategyDefinition`. Its id must equal the folder name - `shared/reporting/strategy.py` is the contract: runnable module/suite specs, not-implemented/skipped specs, the runner protocol, and `StrategyDefinition` -- Every `STRATEGY` explicitly classifies every SDK function; surface-aware strategies declare their surfaces and classify the complete surface-by-function matrix -- Run locally only; no CI integration -- `python -m tests.rust-python-harness run |all` runs the selected strategy; `--function` is common, while each strategy exposes only its supported options +- Every `STRATEGY` explicitly classifies every SDK function. Surface-aware strategies declare their surfaces and classify the complete surface-by-function matrix +- Run locally only. No CI integration +- `python -m tests.rust-python-harness run |all` runs the selected strategy. `--function` is common, while each strategy exposes only its supported options - Examples: `run e2e_parity --surface sdk --function ocr`, `run unit_tests_parity --function ocr --pytest-arg=-x`, or `run all --function ocr` -- `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them; `cli/__init__.py` builds the Click command tree; `cli/commands.py` runs selected cases +- `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them. `cli/__init__.py` builds the Click command tree. `cli/commands.py` runs selected cases - `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses -- `trace_parity/` profiles the Python call stack and prints every collected Python call under `litellm/`; it never collects Rust spans and never rebuilds the native extension +- `trace_parity/` profiles the Python call stack and prints every collected Python call under `litellm/`. It never collects Rust spans and never rebuilds the native extension - E2E and trace strategies load their registered module cases and run surface-specific execution from their folders - `shared/unit_runners/contracts.py` owns the typed per-function unit contracts consumed by `unit_tests_parity` and `unit_tests_rust` -- `unit_tests_parity/runner.py` runs each contract's `unit_parity_scope` with `LITELLM_RUST=0` and `LITELLM_RUST=1` in separate processes and requires matching outcomes, including failures; exclusions require a reason in the contract -- `unit_tests_rust/runner.py` runs each contract's focused Cargo test suite; native Rust unit tests stay beside their implementation +- `unit_tests_parity/runner.py` runs each contract's `unit_parity_scope` with `LITELLM_RUST=0` and `LITELLM_RUST=1` in separate processes and requires matching outcomes, including failures. Exclusions require a reason in the contract +- `unit_tests_rust/runner.py` runs each contract's focused Cargo test suite. Native Rust unit tests stay beside their implementation - `shared/unit_runners/suite_runner.py` runs typed suites registered in code with nodeids of the form `suite:::` -- Every strategy declares its report sections and presentation in its own `reporting.py`; shared reporting code only provides reusable models and cell-formatting primitives +- Every strategy declares its report sections and presentation in its own `reporting.py`. Shared reporting code only provides reusable models and cell-formatting primitives - `shared/` contains reusable parity, tracing, reporting primitives, and unit-runner machinery - Keep fixtures with their owning API and existing Python tests in their current locations - Each strategy folder carries an `AGENTS.md` one-liner stating what it should be doing diff --git a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md index 030caa557c8..742f8443d00 100644 --- a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md +++ b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md @@ -1 +1 @@ -Prints every collected Python call under litellm/ from live traces against replayed HTTP responses. API-key and Vertex credentials scenarios exercise separate authentication paths; credentials scenarios replay the token exchange locally. +Prints every collected Python call under litellm/ from live traces against replayed HTTP responses. API-key and Vertex credentials scenarios exercise separate authentication paths. Credentials scenarios replay the token exchange locally. diff --git a/tests/unit/AGENTS.md b/tests/unit/AGENTS.md index fc9798ad30e..80a3a87f959 100644 --- a/tests/unit/AGENTS.md +++ b/tests/unit/AGENTS.md @@ -23,7 +23,7 @@ with patch.object(streamer, "_group_by_date") as mock_group, patch.object(stream assert mock_send.call_count == 2 ``` -Green if `send_batched` drops every row. pydantic doubles in 12 of 203 files, fastapi 11 of 594; +Green if `send_batched` drops every row. pydantic doubles in 12 of 203 files, fastapi 11 of 594. `tests/test_litellm` 59 percent. Exception: the count is the behaviour (`test_dual_cache_async_batch_get_cache_coalesces_concurrent_redis_reads`, fifty readers, `call_count == 1`) diff --git a/ui/litellm-dashboard/AGENTS.md b/ui/litellm-dashboard/AGENTS.md index bfba65b2860..b6fe5219022 100644 --- a/ui/litellm-dashboard/AGENTS.md +++ b/ui/litellm-dashboard/AGENTS.md @@ -4,9 +4,9 @@ Never put LiteLLM tokens or API keys in `localStorage`. `localStorage` survives When you fix lint violations that are grandfathered in `eslint-suppressions.json`, run `eslint . --prune-suppressions` and commit the updated baseline so the gate ratchets down instead of leaving a stale suppression -`src/lib/http/schema.d.ts` is generated from the proxy's OpenAPI spec; never hand-edit it. After changing a backend route or response model that the dashboard consumes, run `npm run gen:api` and commit the result (CI `lint / ui-api-types` enforces this) +`src/lib/http/schema.d.ts` is generated from the proxy's OpenAPI spec. Never hand-edit it. After changing a backend route or response model that the dashboard consumes, run `npm run gen:api` and commit the result (CI `lint / ui-api-types` enforces this) -Tests come in three tiers, named by the standard definitions. `Foo.test.tsx` is a unit test: one module, collaborators replaced by doubles, no multi-component tree, and it should run in milliseconds. `Foo.integration.test.tsx` renders a real component tree with real children and only stubs the network boundary; it costs seconds per case, so it earns its place by proving wiring that a unit test cannot reach. Browser-level tests live in `tests/e2e/ui/` as Playwright specs against a live proxy +Tests come in three tiers, named by the standard definitions. `Foo.test.tsx` is a unit test: one module, collaborators replaced by doubles, no multi-component tree, and it should run in milliseconds. `Foo.integration.test.tsx` renders a real component tree with real children and only stubs the network boundary. It costs seconds per case, so it earns its place by proving wiring that a unit test cannot reach. Browser-level tests live in `tests/e2e/ui/` as Playwright specs against a live proxy When a component holds logic worth asserting, extract the logic and unit-test it there rather than driving it through a render. `CreateMCPServer` is the worked example: its payload building lives in `createServerPayload.ts` with 46 unit tests that run in single-digit milliseconds, while `CreateMCPServer.integration.test.tsx` keeps only the cases that prove a form field reaches the right payload key. A test that renders a whole modal to assert the shape of one object belongs in the first category, not the second diff --git a/ui/litellm-dashboard/src/components/chat/AGENTS.md b/ui/litellm-dashboard/src/components/chat/AGENTS.md index 4062b67a9ec..60a040e01f7 100644 --- a/ui/litellm-dashboard/src/components/chat/AGENTS.md +++ b/ui/litellm-dashboard/src/components/chat/AGENTS.md @@ -13,7 +13,7 @@ This directory has drifted from its own spec before — a hand-rolled ` + )} + + ), +})); vi.mock("@/components/model_info_view", () => ({ default: ({ modelId }: { modelId: string }) =>
model:{modelId}
, })); @@ -97,6 +107,27 @@ describe("ModelsAndEndpointsPage", () => { expect(screen.queryByTestId("panel-all-models")).not.toBeInTheDocument(); }); + it("opens the Add Model tab from the decision models banner", async () => { + const user = userEvent.setup(); + renderPage(); + await user.click(screen.getByRole("button", { name: "Add a decision model" })); + expect(screen.getByTestId("panel-add")).toBeInTheDocument(); + expect(screen.queryByTestId("panel-all-models")).not.toBeInTheDocument(); + }); + + it("does not offer the banner's Add a decision model to a session that cannot add models", () => { + mockUseAuthorized.mockReturnValue(NON_ADMIN); + renderPage(); + expect(screen.getByRole("region", { name: "Decision models" })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Add a decision model" })).not.toBeInTheDocument(); + }); + + it("does not show a view-only session the decision models banner, whose Playground it cannot open", () => { + mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN); + renderPage(); + expect(screen.queryByRole("region", { name: "Decision models" })).not.toBeInTheDocument(); + }); + it("renders the model detail overlay from the ?model drill-in and hides the tabs", () => { detailState.modelId = "abc-123"; renderPage(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx index cd2922dff4b..1eae68b9d43 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx @@ -9,7 +9,7 @@ import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings" import { all_admin_roles, internalUserRoles } from "@/utils/roles"; import { autoRouterCreationScope, canCreateModels } from "@/utils/modelPermissions"; import BetaBadge from "@/components/BetaBadge"; -import CostOptimizationFeedbackBanner from "@/components/molecules/cost_optimization_feedback_banner"; +import DecisionModelsBanner from "@/components/molecules/DecisionModelsBanner"; import ModelInfoView from "@/components/model_info_view"; import TeamInfoView from "@/components/team/TeamInfo"; import { useModelDetailRouting } from "@/app/(dashboard)/models-and-endpoints/detailNavigation"; @@ -178,7 +178,9 @@ export default function ModelsAndEndpointsPage() { - + {!isViewOnly && ( + setActiveKey("add") : undefined} /> + )} {modelId ? ( ({ default_model_placeholder: "claude-3-opus", credential_fields: [{ key: "api_key", label: "API Key", field_type: "password", required: false }], }, + { + provider: "TypeSafe", + provider_display_name: "TypeSafe", + litellm_provider: "typesafe", + default_model_placeholder: "jev-latest", + credential_fields: [{ key: "api_key", label: "API Key", field_type: "password", required: false }], + }, ], isLoading: false, error: null, @@ -538,3 +545,61 @@ describe("AddModelPanel behaviours the removed Advanced Settings form instance n }); }); }); + +describe("AddModelPanel decision models", () => { + const DECISION_COST_MAP = { + "typesafe/jev-latest": { litellm_provider: "typesafe", mode: "evaluation" }, + "gpt-4o-2024-08-06": { litellm_provider: "openai", mode: "chat" }, + "gpt-6-luna": { litellm_provider: "openai", mode: "chat", supported_endpoints: ["/v1/decisions"] }, + "claude-sonnet-4-5": { litellm_provider: "anthropic", mode: "chat" }, + }; + + beforeEach(() => { + vi.clearAllMocks(); + mockUseModelCostMap.mockReturnValue({ data: DECISION_COST_MAP }); + mockPtuEnabled.mockReturnValue(false); + mockAuthorized.mockReturnValue(PROXY_ADMIN); + }); + + const decisionNotice = () => screen.queryByRole("note", { name: "Decision model notice" }); + + it("finds a provider by the name of one of its decision models", async () => { + const { user } = await setup(); + await user.type(screen.getByRole("combobox", { name: /provider/i }), "jev"); + + expect(await screen.findByRole("option", { name: /TypeSafe/ })).toBeInTheDocument(); + expect(screen.queryByRole("option", { name: /Anthropic/ })).not.toBeInTheDocument(); + }); + + it("points a decision-only provider at the decision routes, the docs, and the playground", async () => { + const { user } = await setup(); + await user.type(screen.getByRole("combobox", { name: /provider/i }), "TypeSafe"); + await user.click(await screen.findByRole("option", { name: /TypeSafe/ })); + + const note = await screen.findByRole("note", { name: "Decision model notice" }); + expect(note).toHaveTextContent("/v1/decisions"); + expect(within(note).getByRole("link", { name: "How to call decision models" })).toHaveAttribute( + "href", + "https://docs.litellm.ai/docs/decisions", + ); + expect(within(note).getByRole("link", { name: "test it in the System One playground" })).toHaveAttribute( + "href", + "/ui/playground?tab=system-one", + ); + }); + + it("shows the notice on a chat provider only once one of its decision models is picked", async () => { + const { user } = await setup(); + await user.click(screen.getByRole("combobox", { name: /provider/i })); + await user.click(await screen.findByRole("option", { name: /OpenAI/ })); + await user.click(await screen.findByPlaceholderText("Select models")); + await user.click(await screen.findByText("gpt-4o-2024-08-06")); + expect(decisionNotice()).not.toBeInTheDocument(); + + await user.click(await screen.findByText("gpt-6-luna")); + + const note = await screen.findByRole("note", { name: "Decision model notice" }); + expect(note).toHaveTextContent("/v1/decisions"); + expect(note).not.toHaveTextContent("/chat/completions"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx index 9fbd138bda4..4703bda8d20 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx @@ -15,6 +15,7 @@ import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap" import { useCredentials } from "@/app/(dashboard)/hooks/credentials/useCredentials"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { buildDecisionCatalog } from "@/lib/decisionModels"; const INITIAL_VALUES: MountedFormValues = { litellm_credential_name: null }; @@ -32,6 +33,7 @@ export default function AddModelPanel() { () => (typeof pickedProvider === "string" ? getProviderModels(pickedProvider, modelCostMapData) : []), [pickedProvider, modelCostMapData], ); + const decisionCatalog = useMemo(() => buildDecisionCatalog(modelCostMapData), [modelCostMapData]); const [showAdvancedSettings, setShowAdvancedSettings] = useState(false); const refresh = () => queryClient.invalidateQueries({ queryKey: ["models", "list"] }); @@ -66,6 +68,7 @@ export default function AddModelPanel() { setShowAdvancedSettings={setShowAdvancedSettings} teams={teams ?? null} credentials={credentialsResponse?.credentials || []} + decisionCatalog={decisionCatalog} /> ); } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.integration.test.tsx index 0812e40d3e4..367b779ecf7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.integration.test.tsx @@ -13,6 +13,17 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import SystemOneUI from "./SystemOneUI"; import type { SystemOneResponse } from "./lib/schemas"; +interface ModelGroupInfo { + model_group: string; + mode: string; +} + +const DECISION_AND_CHAT_MODELS: ModelGroupInfo[] = [ + { model_group: "gpt-5.5", mode: "chat" }, + { model_group: "jev-latest", mode: "evaluation" }, + { model_group: "pplx-decider", mode: "evaluation" }, +]; + const responseBody: SystemOneResponse = { model: "jev-1.13.0", answers: { @@ -32,6 +43,21 @@ const render = (ui: ReactElement) => , ); +const modelGroupInfoResponse = (groups: ModelGroupInfo[]) => + ({ ok: true, status: 200, text: async () => JSON.stringify({ data: groups }) }) as Response; + +const EXPIRED_KEY_RESPONSE = { + ok: false, + status: 401, + text: async () => + JSON.stringify({ error: { message: "Authentication Error - Expired Key.", type: "expired_key", code: "401" } }), +} as Response; + +const isModelLookup = (input: RequestInfo | URL) => String(input).endsWith("/model_group/info"); + +const lookupAuthHeaders = (calls: Parameters[]) => + calls.map(([, init]) => Object.values((init?.headers ?? {}) as Record)); + const createResponse = (body: SystemOneResponse, status = 200, errorText = "") => ({ ok: status >= 200 && status < 300, @@ -41,17 +67,23 @@ const createResponse = (body: SystemOneResponse, status = 200, errorText = "") = describe("SystemOneUI integration", () => { const mockFetch = vi.fn(); + const mockModelLookup = vi.fn(); beforeEach(() => { sessionStorage.clear(); mockFetch.mockReset(); - vi.stubGlobal("fetch", mockFetch); + mockModelLookup.mockReset(); + vi.stubGlobal("fetch", (input: RequestInfo | URL, init?: RequestInit) => + isModelLookup(input) ? mockModelLookup(input, init) : mockFetch(input, init), + ); mockFetch.mockResolvedValue(createResponse(responseBody)); + mockModelLookup.mockResolvedValue(modelGroupInfoResponse([])); }); afterEach(() => { vi.unstubAllGlobals(); sessionStorage.clear(); + document.cookie = "token=; path=/; max-age=0"; }); it("restores the example after the request is edited", async () => { @@ -80,6 +112,10 @@ describe("SystemOneUI integration", () => { expect(screen.getByRole("note", { name: "Decision endpoint notice" })).toHaveTextContent( "omit model to use the proxy's configured default.", ); + expect(screen.getByRole("link", { name: "How to call /v1/decisions and /v1/systemone" })).toHaveAttribute( + "href", + "https://docs.litellm.ai/docs/decisions", + ); expect(screen.getByRole("link", { name: "Give us feedback on what you want for decision models" })).toHaveAttribute( "href", "https://github.com/BerriAI/litellm/discussions/44231", @@ -87,6 +123,106 @@ describe("SystemOneUI integration", () => { expect(screen.queryByRole("alert")).not.toBeInTheDocument(); }); + it("opens on /v1/systemone with the first decision model the key can call", async () => { + const user = userEvent.setup(); + mockModelLookup.mockResolvedValue(modelGroupInfoResponse(DECISION_AND_CHAT_MODELS)); + render(); + + expect(await screen.findByRole("combobox", { name: "Decision model" })).toHaveValue("jev-latest"); + expect(lookupAuthHeaders(mockModelLookup.mock.calls)).toEqual([expect.arrayContaining(["Bearer session-key"])]); + await user.click(screen.getByRole("button", { name: "Send" })); + + expect(await screen.findByText("Selected choice")).toBeInTheDocument(); + expect(mockFetch.mock.calls[0]?.[0]).toMatch(/\/v1\/systemone$/); + expect(mockFetch.mock.calls[0]?.[0]).not.toMatch(/typesafe/); + expect(JSON.parse(mockFetch.mock.calls[0]?.[1]?.body as string)).toMatchObject({ + model: "jev-latest", + questions: { has_repro_steps: { type: "noul" } }, + }); + }); + + it("offers only decision models and writes the picked one into the request", async () => { + const user = userEvent.setup(); + mockModelLookup.mockResolvedValue(modelGroupInfoResponse(DECISION_AND_CHAT_MODELS)); + render(); + const picker = await screen.findByRole("combobox", { name: "Decision model" }); + + await user.clear(picker); + await user.type(picker, "p"); + expect(await screen.findByRole("option", { name: "pplx-decider" })).toBeInTheDocument(); + expect(screen.queryByRole("option", { name: "gpt-5.5" })).not.toBeInTheDocument(); + await user.click(screen.getByRole("option", { name: "pplx-decider" })); + + const editor = screen.getByRole("textbox", { name: "System One JSON payload" }); + expect(JSON.parse((editor as HTMLTextAreaElement).value)).toMatchObject({ + model: "pplx-decider", + questions: { has_repro_steps: { type: "noul" } }, + }); + await user.click(screen.getByRole("button", { name: "Send" })); + expect(await screen.findByText("Selected choice")).toBeInTheDocument(); + expect(JSON.parse(mockFetch.mock.calls[0]?.[1]?.body as string).model).toBe("pplx-decider"); + }); + + it("links to Add Model when the key cannot call any decision model", async () => { + const user = userEvent.setup(); + render(); + screen.getByRole("combobox", { name: "Decision endpoint" }).focus(); + await user.keyboard("{ArrowDown}"); + await user.click(await screen.findByRole("option", { name: "System One · /v1/systemone" })); + + expect(await screen.findByRole("link", { name: "Add a decision model" })).toHaveAttribute( + "href", + "/ui/models-and-endpoints", + ); + expect(screen.getByRole("combobox", { name: "Decision model" })).toHaveAttribute( + "placeholder", + "No decision models yet", + ); + }); + + it("looks up a typed virtual key's decision models once, after typing settles", async () => { + const user = userEvent.setup(); + mockModelLookup.mockResolvedValue(modelGroupInfoResponse(DECISION_AND_CHAT_MODELS)); + render(); + + await user.type(screen.getByLabelText("Virtual Key", { exact: true }), "sk-typed-key"); + + expect(await screen.findByRole("combobox", { name: "Decision model" })).toHaveValue("jev-latest"); + expect(lookupAuthHeaders(mockModelLookup.mock.calls)).toEqual([expect.arrayContaining(["Bearer sk-typed-key"])]); + }); + + it("keeps the admin signed in when a typed virtual key has expired", async () => { + document.cookie = "token=ui-session; path=/"; + mockModelLookup.mockResolvedValue(EXPIRED_KEY_RESPONSE); + render(); + + fireEvent.change(screen.getByLabelText("Virtual Key", { exact: true }), { target: { value: "sk-expired" } }); + + await waitFor(() => expect(mockModelLookup).toHaveBeenCalled()); + await act(async () => { + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(document.cookie).toContain("token=ui-session"); + expect(screen.getByRole("combobox", { name: "Decision endpoint" })).toHaveTextContent("/typesafe/v1/systemone"); + }); + + it("keeps an edited TypeSafe request in place when decision models load afterwards", async () => { + const models = Promise.withResolvers(); + mockModelLookup.mockReturnValue(models.promise); + render(); + const editor = screen.getByRole("textbox", { name: "System One JSON payload" }); + const draft = JSON.stringify({ state: "edited", questions: { q: { type: "noul", instructions: "Yes?" } } }); + fireEvent.change(editor, { target: { value: draft } }); + + await act(async () => { + models.resolve(modelGroupInfoResponse(DECISION_AND_CHAT_MODELS)); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + + expect(editor).toHaveValue(draft); + expect(screen.queryByRole("combobox", { name: "Decision model" })).not.toBeInTheDocument(); + }); + it("shows invalid JSON and disables Send", async () => { render(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.tsx index e092efce161..3e61030da7b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.tsx @@ -4,24 +4,27 @@ import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { SearchSelect } from "@/components/shared/SearchSelect"; +import { DECISIONS_DOCS_URL } from "@/lib/decisionModels"; +import { uiHref } from "@/utils/uiHref"; import { useMutation } from "@tanstack/react-query"; import { Code, Info, LoaderCircle, RotateCcw, Send } from "lucide-react"; import { useEffect, useMemo, useRef, useState } from "react"; import { makeSystemOneRequest } from "../../llm_calls/system_one"; -import { DECISIONS_EXAMPLE, SYSTEM_ONE_EXAMPLE } from "./lib/example"; +import { PLACEHOLDER_DECISION_MODEL, SYSTEM_ONE_EXAMPLE, decisionsExample } from "./lib/example"; +import { payloadModel, withPayloadModel } from "./lib/payloadModel"; import type { DecisionEndpoint, PlaygroundRequest } from "./lib/schemas"; import JsonEditor from "./JsonEditor"; import QuestionBreakdown from "./QuestionBreakdown"; import ResponseView from "./ResponseView"; import { validateSystemOnePayload } from "./lib/validatePayload"; +import { useDecisionModels, type ApiKeySource } from "./useDecisionModels"; interface SystemOneUIProps { accessToken: string | null; disabledPersonalKeyCreation?: boolean; } -type ApiKeySource = "session" | "custom"; - interface SystemOneSendVariables { payload: PlaygroundRequest; endpoint: DecisionEndpoint; @@ -30,7 +33,6 @@ interface SystemOneSendVariables { } const EXAMPLE_PAYLOAD = JSON.stringify(SYSTEM_ONE_EXAMPLE, null, 2); -const DECISIONS_EXAMPLE_PAYLOAD = JSON.stringify(DECISIONS_EXAMPLE, null, 2); const DECISION_MODELS_DISCUSSION_URL = "https://github.com/BerriAI/litellm/discussions/44231"; function getCustomProxyBaseUrl(): string | undefined { @@ -40,16 +42,23 @@ function getCustomProxyBaseUrl(): string | undefined { export default function SystemOneUI({ accessToken, disabledPersonalKeyCreation = false }: SystemOneUIProps) { const [apiKeySource, setApiKeySource] = useState(disabledPersonalKeyCreation ? "custom" : "session"); const [customApiKey, setCustomApiKey] = useState(""); - const [endpoint, setEndpoint] = useState("/typesafe/v1/systemone"); - const [payloads, setPayloads] = useState>({ - "/v1/systemone": DECISIONS_EXAMPLE_PAYLOAD, - "/typesafe/v1/systemone": EXAMPLE_PAYLOAD, - }); - const rawPayload = payloads[endpoint]; - const examplePayload = endpoint === "/v1/systemone" ? DECISIONS_EXAMPLE_PAYLOAD : EXAMPLE_PAYLOAD; + const effectiveApiKey = apiKeySource === "session" ? accessToken || "" : customApiKey.trim(); + const { decisionModels, isLoaded: decisionModelsLoaded } = useDecisionModels( + apiKeySource, + effectiveApiKey, + getCustomProxyBaseUrl(), + ); + const [chosenEndpoint, setChosenEndpoint] = useState(null); + const endpoint: DecisionEndpoint = + chosenEndpoint ?? (decisionModels.length > 0 ? "/v1/systemone" : "/typesafe/v1/systemone"); + const [drafts, setDrafts] = useState>>({}); + const examplePayload = + endpoint === "/v1/systemone" + ? JSON.stringify(decisionsExample(decisionModels[0] ?? PLACEHOLDER_DECISION_MODEL), null, 2) + : EXAMPLE_PAYLOAD; + const rawPayload = drafts[endpoint] ?? examplePayload; const activeController = useRef(null); const validation = useMemo(() => validateSystemOnePayload(rawPayload, endpoint), [rawPayload, endpoint]); - const effectiveApiKey = apiKeySource === "session" ? accessToken || "" : customApiKey.trim(); const hasSyntaxError = validation.issues.some((issue) => issue.path === "syntax"); const systemOne = useMutation({ @@ -76,7 +85,21 @@ export default function SystemOneUI({ accessToken, disabledPersonalKeyCreation = function handlePayloadChange(value: string) { if (value !== rawPayload) { clearRequestState(); - setPayloads((current) => ({ ...current, [endpoint]: value })); + setChosenEndpoint(endpoint); + setDrafts((current) => ({ ...current, [endpoint]: value })); + } + } + + function handleResetExample() { + clearRequestState(); + setChosenEndpoint(endpoint); + setDrafts((current) => Object.fromEntries(Object.entries(current).filter(([key]) => key !== endpoint))); + } + + function handleModelPick(model: string | null) { + const next = model === null ? undefined : withPayloadModel(rawPayload, model); + if (next !== undefined) { + handlePayloadChange(next); } } @@ -91,6 +114,7 @@ export default function SystemOneUI({ accessToken, disabledPersonalKeyCreation = return; } clearRequestState(); + setChosenEndpoint(endpoint); const controller = new AbortController(); activeController.current = controller; const variables = { payload: validation.payload, apiKey: effectiveApiKey, signal: controller.signal, endpoint }; @@ -107,7 +131,7 @@ export default function SystemOneUI({ accessToken, disabledPersonalKeyCreation = onValueChange={(value) => { if (value === "/v1/systemone" || value === "/typesafe/v1/systemone") { clearRequestState(); - setEndpoint(value); + setChosenEndpoint(value); } }} > @@ -121,6 +145,28 @@ export default function SystemOneUI({ accessToken, disabledPersonalKeyCreation = TypeSafe · /typesafe/v1/systemone + {endpoint === "/v1/systemone" && ( + <> + Model +
+ ({ label: model, value: model }))} + value={payloadModel(rawPayload) ?? null} + onValueChange={handleModelPick} + placeholder={decisionModels.length > 0 ? "Pick a decision model" : "No decision models yet"} + emptyText="No decision models found" + disabled={hasSyntaxError} + allowClear={false} + /> +
+ {decisionModelsLoaded && decisionModels.length === 0 && ( + + Add a decision model + + )} + + )}
@@ -156,11 +202,7 @@ export default function SystemOneUI({ accessToken, disabledPersonalKeyCreation = )}
- @@ -186,8 +228,12 @@ export default function SystemOneUI({ accessToken, disabledPersonalKeyCreation = {endpoint === "/v1/systemone" - ? "Sends choice, noul, and score questions through /v1/systemone. Replace the example model with a decision model configured on your proxy, or omit model to use the proxy's configured default." + ? "Sends choice, noul, and score questions through /v1/systemone to a decision model on your proxy. Pick one under Model, or omit model to use the proxy's configured default." : "Sends requests through /typesafe/v1/systemone and requires TYPESAFE_API_KEY on the proxy."}{" "} + + How to call /v1/decisions and /v1/systemone + + {" · "} Give us feedback on what you want for decision models diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/example.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/example.ts index 2632b2e1a9b..796275212a1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/example.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/example.ts @@ -37,7 +37,6 @@ export const SYSTEM_ONE_EXAMPLE = { }, } satisfies SystemOneRequest; -export const DECISIONS_EXAMPLE: DecisionRequest = { - ...SYSTEM_ONE_EXAMPLE, - model: "your-decision-model", -}; +export const PLACEHOLDER_DECISION_MODEL = "your-decision-model"; + +export const decisionsExample = (model: string): DecisionRequest => ({ ...SYSTEM_ONE_EXAMPLE, model }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/payloadModel.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/payloadModel.test.ts new file mode 100644 index 00000000000..6e3ee62b217 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/payloadModel.test.ts @@ -0,0 +1,41 @@ +import { describe, expect, it } from "vitest"; +import { payloadModel, withPayloadModel } from "./payloadModel"; + +describe("payloadModel", () => { + it("reads the model from a JSON object payload", () => { + expect(payloadModel('{"model": "jev-latest", "state": "hi"}')).toBe("jev-latest"); + }); + + it("is undefined when the model is missing, empty, or not a string", () => { + expect(payloadModel('{"state": "hi"}')).toBeUndefined(); + expect(payloadModel('{"model": ""}')).toBeUndefined(); + expect(payloadModel('{"model": 3}')).toBeUndefined(); + }); + + it("is undefined when the payload is not a JSON object", () => { + expect(payloadModel("{not json")).toBeUndefined(); + expect(payloadModel('["jev-latest"]')).toBeUndefined(); + expect(payloadModel('"jev-latest"')).toBeUndefined(); + }); +}); + +describe("withPayloadModel", () => { + it("replaces the model and keeps every other field", () => { + const next = withPayloadModel( + '{"model": "old", "state": "hi", "questions": {"q": {"type": "noul"}}}', + "jev-latest", + ); + expect(next).toBeDefined(); + expect(JSON.parse(next ?? "")).toEqual({ model: "jev-latest", state: "hi", questions: { q: { type: "noul" } } }); + }); + + it("adds the model as the first field when the payload has none", () => { + const next = withPayloadModel('{"state": "hi"}', "jev-latest"); + expect(Object.keys(JSON.parse(next ?? "{}"))).toEqual(["model", "state"]); + }); + + it("leaves a payload that is not a JSON object alone", () => { + expect(withPayloadModel("{not json", "jev-latest")).toBeUndefined(); + expect(withPayloadModel("[1, 2]", "jev-latest")).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/payloadModel.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/payloadModel.ts new file mode 100644 index 00000000000..d347c7fb3c8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/payloadModel.ts @@ -0,0 +1,27 @@ +import { z } from "zod"; +import { parseJson } from "./validatePayload"; + +const jsonObject = z.record(z.string(), z.unknown()); + +const parseObject = (raw: string): Record | undefined => { + const json = parseJson(raw); + if (!json.ok) { + return undefined; + } + const parsed = jsonObject.safeParse(json.value); + return parsed.success ? parsed.data : undefined; +}; + +export const payloadModel = (raw: string): string | undefined => { + const model = parseObject(raw)?.model; + return typeof model === "string" && model !== "" ? model : undefined; +}; + +export const withPayloadModel = (raw: string, model: string): string | undefined => { + const payload = parseObject(raw); + if (payload === undefined) { + return undefined; + } + const rest = Object.fromEntries(Object.entries(payload).filter(([key]) => key !== "model")); + return JSON.stringify({ model, ...rest }, null, 2); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.ts index ac53f871333..9bed8e9924b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.ts @@ -24,7 +24,7 @@ const invalid = (path: string, message: string): SystemOnePayloadValidation => ( issues: [{ path, message, severity: "error" }], }); -function parseJson(raw: string): { ok: true; value: unknown } | { ok: false; message: string } { +export function parseJson(raw: string): { ok: true; value: unknown } | { ok: false; message: string } { try { return { ok: true, value: JSON.parse(raw) }; } catch (error: unknown) { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/useDecisionModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/useDecisionModels.ts new file mode 100644 index 00000000000..c694b6384e5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/useDecisionModels.ts @@ -0,0 +1,32 @@ +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; +import { useQuery } from "@tanstack/react-query"; +import { fetchDecisionModels } from "../../llm_calls/system_one"; + +export type ApiKeySource = "session" | "custom"; + +export interface DecisionModelsLookup { + decisionModels: string[]; + isLoaded: boolean; +} + +const TYPED_KEY_DEBOUNCE_WAIT_MS = 400; +const NO_DECISION_MODELS: string[] = []; + +export function useDecisionModels( + apiKeySource: ApiKeySource, + apiKey: string, + customBaseUrl: string | undefined, +): DecisionModelsLookup { + const typedApiKey = apiKeySource === "custom" ? apiKey : ""; + const [settledTypedApiKey] = useDebouncedValue(typedApiKey, { wait: TYPED_KEY_DEBOUNCE_WAIT_MS }); + const lookupApiKey = apiKeySource === "session" ? apiKey : settledTypedApiKey; + const queryOptions = { + queryKey: ["playground", "systemOne", "decisionModels", customBaseUrl ?? "", lookupApiKey], + queryFn: ({ signal }: { signal: AbortSignal }) => fetchDecisionModels(lookupApiKey, customBaseUrl, signal), + enabled: lookupApiKey !== "", + retry: false, + staleTime: 60 * 1000, + }; + const query = useQuery(queryOptions); + return { decisionModels: query.data ?? NO_DECISION_MODELS, isLoaded: query.isSuccess }; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/system_one.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/system_one.test.ts index c039b018b8f..e097f2d3d26 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/system_one.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/system_one.test.ts @@ -1,5 +1,5 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { makeSystemOneRequest } from "./system_one"; +import { fetchDecisionModels, makeSystemOneRequest } from "./system_one"; import type { SystemOneRequest, SystemOneResponse } from "../components/systemOneUI/lib/schemas"; vi.mock("@/components/networking", () => ({ @@ -106,3 +106,46 @@ describe("makeSystemOneRequest", () => { ); }); }); + +describe("fetchDecisionModels", () => { + const mockFetch = vi.fn<(input: RequestInfo | URL, init?: RequestInit) => Promise>(); + const modelGroupInfo = (data: unknown) => ({ ok: true, text: async () => JSON.stringify({ data }) }) as Response; + + beforeEach(() => { + vi.stubGlobal("fetch", mockFetch); + }); + + afterEach(() => { + vi.clearAllMocks(); + vi.unstubAllGlobals(); + }); + + it("lists the key's decision models from the proxy that answers the request, sorted by name", async () => { + mockFetch.mockResolvedValueOnce( + modelGroupInfo([ + { model_group: "pplx-decider", mode: "evaluation" }, + { model_group: "gpt-5.5", mode: "chat" }, + { model_group: "jev-latest", mode: "evaluation" }, + { model_group: "unknown-mode", mode: null }, + ]), + ); + + await expect(fetchDecisionModels("virtual-key", "https://custom.example.com/")).resolves.toEqual([ + "jev-latest", + "pplx-decider", + ]); + expect(mockFetch).toHaveBeenCalledWith( + "https://custom.example.com/model_group/info", + expect.objectContaining({ + method: "GET", + headers: expect.objectContaining({ Authorization: "Bearer virtual-key" }), + }), + ); + }); + + it("rejects a response that is not a model group list", async () => { + mockFetch.mockResolvedValueOnce(modelGroupInfo([{ mode: "evaluation" }])); + + await expect(fetchDecisionModels("virtual-key")).rejects.toThrow(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/system_one.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/system_one.ts index b71e1268ac2..6c29d9110ac 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/system_one.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/system_one.ts @@ -1,6 +1,8 @@ +import { z } from "zod"; import { getGlobalLitellmHeaderName, getProxyBaseUrl } from "@/components/networking"; import { withRequiredHeaders } from "@/components/llm_calls/request_headers"; import { createApiClient } from "@/lib/http/client"; +import { isDecisionMode } from "@/lib/decisionModels"; import { systemOneResponseSchema, type DecisionEndpoint, @@ -13,24 +15,54 @@ export interface SystemOneResult { latencyMs: number; } -export async function makeSystemOneRequest( - payload: PlaygroundRequest, - accessToken: string, - customBaseUrl?: string, - { signal, endpoint = "/typesafe/v1/systemone" }: { signal?: AbortSignal; endpoint?: DecisionEndpoint } = {}, -): Promise { +const modelGroupInfoSchema = z.object({ + data: z.array(z.object({ model_group: z.string(), mode: z.string().nullish() })), +}); + +function proxyClient(customBaseUrl?: string) { const proxyBaseUrl = customBaseUrl || getProxyBaseUrl(); const normalizedBaseUrl = proxyBaseUrl.endsWith("/") ? proxyBaseUrl.slice(0, -1) : proxyBaseUrl; - const startedAt = performance.now(); - const headers = withRequiredHeaders( + return createApiClient({ getBaseUrl: () => normalizedBaseUrl }); +} + +function authHeaders(accessToken: string): Record { + return withRequiredHeaders( {}, { "Content-Type": "application/json", [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, }, ); - const client = createApiClient({ getBaseUrl: () => normalizedBaseUrl }); - const body = await client.post(endpoint, { body: payload, headers, signal }); +} + +export async function fetchDecisionModels( + accessToken: string, + customBaseUrl?: string, + signal?: AbortSignal, +): Promise { + const body = await proxyClient(customBaseUrl).get("/model_group/info", { + headers: authHeaders(accessToken), + signal, + }); + return modelGroupInfoSchema + .parse(body) + .data.filter((group) => isDecisionMode(group.mode)) + .map((group) => group.model_group) + .sort((a, b) => a.localeCompare(b)); +} + +export async function makeSystemOneRequest( + payload: PlaygroundRequest, + accessToken: string, + customBaseUrl?: string, + { signal, endpoint = "/typesafe/v1/systemone" }: { signal?: AbortSignal; endpoint?: DecisionEndpoint } = {}, +): Promise { + const startedAt = performance.now(); + const body = await proxyClient(customBaseUrl).post(endpoint, { + body: payload, + headers: authHeaders(accessToken), + signal, + }); const parsed = systemOneResponseSchema.safeParse(body); if (!parsed.success) { throw new Error("System One response has an invalid shape."); diff --git a/ui/litellm-dashboard/src/app/model_hub_table/page.tsx b/ui/litellm-dashboard/src/app/model_hub_table/page.tsx index c399178e6aa..46c7394cff7 100644 --- a/ui/litellm-dashboard/src/app/model_hub_table/page.tsx +++ b/ui/litellm-dashboard/src/app/model_hub_table/page.tsx @@ -15,7 +15,15 @@ function PublicModelHubTableContent() { setAccessToken(key); }, [key]); - return ; + return ( + + ); } export default function PublicModelHubTable() { diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx index 1f052b8932a..cc3ac1d2591 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx @@ -17,6 +17,7 @@ vi.mock("@/components/networking", () => ({ modelHubPublicModelsCall: vi.fn(), modelHubCall: vi.fn(), getConfigFieldSetting: vi.fn(), + getGlobalLitellmHeaderName: vi.fn(() => "X-Litellm-Key"), getProxyBaseUrl: vi.fn(() => "http://localhost:4000"), getAgentsList: vi.fn(), fetchMCPServers: vi.fn(), @@ -115,7 +116,15 @@ describe("ModelHubTable", () => { it(description, async () => { setupAuthRedirectTest(requireAuth, tokenValue, isTokenValid); - renderWithProviders(); + renderWithProviders( + , + ); await waitFor(() => { if (shouldRedirect) { @@ -148,7 +157,13 @@ describe("ModelHubTable", () => { }); renderWithProviders( - , + , ); await waitFor(() => { @@ -165,7 +180,9 @@ describe("ModelHubTable", () => { isLoading: false, }); - renderWithProviders(); + renderWithProviders( + , + ); expect(await screen.findByText("No models yet")).toBeInTheDocument(); expect(networking.modelHubCall).not.toHaveBeenCalled(); @@ -191,7 +208,15 @@ describe("ModelHubTable", () => { isLoading: false, }); - renderWithProviders(); + renderWithProviders( + , + ); await waitFor(() => { expect(getUiConfigMock).toHaveBeenCalled(); @@ -204,6 +229,66 @@ describe("ModelHubTable", () => { expect(getUiConfigCallOrder).toBeLessThan(modelHubPublicModelsCallOrder); }); + describe("model details usage example", () => { + const openDetails = async ( + model: { model_group: string; providers: string[]; mode: string }, + canOpenPlayground = true, + ) => { + vi.mocked(networking.modelHubCall).mockResolvedValue({ data: [model] }); + vi.mocked(networking.getConfigFieldSetting).mockResolvedValue({ field_value: false }); + vi.mocked(networking.getAgentsList).mockResolvedValue({ agents: [] }); + vi.mocked(networking.fetchMCPServers).mockResolvedValue([]); + vi.mocked(networking.getUiSettings).mockResolvedValue({ values: {} }); + mockUseUISettings.mockReturnValue({ data: { values: {} }, isLoading: false }); + const user = userEvent.setup(); + renderWithProviders( + , + ); + await user.click(await screen.findByRole("button", { name: model.model_group })); + return screen.findByRole("dialog"); + }; + + it("shows a decision model how to call it on /v1/systemone, with the docs and the playground", async () => { + const dialog = await openDetails({ model_group: "jev-latest", providers: ["typesafe"], mode: "evaluation" }); + + expect(dialog).toHaveTextContent('"http://localhost:4000/v1/systemone"'); + expect(dialog).toHaveTextContent('headers={"X-Litellm-Key": "Bearer your_api_key"}'); + expect(dialog).not.toHaveTextContent("chat.completions"); + expect(screen.getByRole("link", { name: "How to call decision models" })).toHaveAttribute( + "href", + "https://docs.litellm.ai/docs/decisions", + ); + expect(screen.getByRole("link", { name: "Try it in the Playground" })).toHaveAttribute( + "href", + "/ui/playground?tab=system-one", + ); + }); + + it("does not link a view-only session to the Playground it cannot open", async () => { + const dialog = await openDetails( + { model_group: "jev-latest", providers: ["typesafe"], mode: "evaluation" }, + false, + ); + + expect(dialog).toHaveTextContent('"http://localhost:4000/v1/systemone"'); + expect(screen.getByRole("link", { name: "How to call decision models" })).toBeInTheDocument(); + expect(screen.queryByRole("link", { name: "Try it in the Playground" })).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" }); + + expect(dialog).toHaveTextContent("client.chat.completions.create("); + expect(screen.queryByRole("link", { name: "How to call decision models" })).not.toBeInTheDocument(); + }); + }); + describe("hub tabs", () => { const renderHub = async (agents: object[] = [], mcpServers: Promise = Promise.resolve([])) => { vi.mocked(networking.modelHubCall).mockResolvedValue({ @@ -217,7 +302,13 @@ describe("ModelHubTable", () => { const user = userEvent.setup(); renderWithProviders( - , + , ); return { user, search: await screen.findByPlaceholderText("Search model names...") }; }; diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index c4d776ea2fb..e5dcb8ebc73 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -4,6 +4,7 @@ import MakeMCPPublicForm from "@/components/AIHub/forms/MakeMCPPublicForm"; import MakeModelPublicForm from "@/components/AIHub/forms/MakeModelPublicForm"; import { getMCPHubTableColumns, MCPServerData } from "@/components/AIHub/MCPHubTableColumns"; import { getModelHubTableColumns, ModelHubData } from "@/components/AIHub/ModelHubTableColumns"; +import { modelUsageExample } from "@/components/AIHub/modelUsageExample"; import UsefulLinksManagement from "@/components/AIHub/UsefulLinksManagement"; import { getClaudeCodePluginsList } from "@/components/networking"; import { Plugin } from "@/components/claude_code_plugins/types"; @@ -15,6 +16,7 @@ import { fetchMCPServers, getAgentsList, getConfigFieldSetting, + getGlobalLitellmHeaderName, getProxyBaseUrl, getUiConfig, modelHubCall, @@ -41,12 +43,15 @@ import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings" import { checkTokenValidity } from "@/utils/jwtUtils"; import { getCookie } from "@/utils/cookieUtils"; import { getLoginUrl } from "@/utils/returnUrlUtils"; +import { uiHref } from "@/utils/uiHref"; +import { DECISIONS_DOCS_URL, SYSTEM_ONE_PLAYGROUND_ROUTE, isDecisionMode } from "@/lib/decisionModels"; interface ModelHubTableProps { accessToken: string | null; publicPage: boolean; premiumUser: boolean; userRole: string | null; + canOpenPlayground: boolean; } function isMCPHubVisibilityDisabled(isLoading: boolean, servers: readonly MCPServerData[] | null): boolean { @@ -65,7 +70,13 @@ function HubEmptyState({ title, body }: { title: string; body: string }) { ); } -const ModelHubTable: React.FC = ({ accessToken, publicPage, premiumUser, userRole }) => { +const ModelHubTable: React.FC = ({ + accessToken, + publicPage, + premiumUser, + userRole, + canOpenPlayground, +}) => { const syntaxTheme = useSyntaxTheme(prism); // Admin Viewer follows the read-parity rule: see the AI Hub catalog, but // cannot toggle public visibility (write). @@ -763,28 +774,32 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage,
)} - {/* Usage Example */}

Usage Example

+ {isDecisionMode(selectedModel.mode) && ( +

+ Decision models answer on /v1/decisions and /v1/systemone, not{" "} + /chat/completions.{" "} + + How to call decision models + + {canOpenPlayground && ( + <> + {" · "} + + Try it in the Playground + + + )} +

+ )} - {`import openai - -client = openai.OpenAI( - api_key="your_api_key", - base_url="${getProxyBaseUrl()}" # Your LiteLLM Proxy URL -) - -response = client.chat.completions.create( - model="${selectedModel.model_group}", - messages=[ - { - "role": "user", - "content": "Hello, how are you?" - } - ] -) - -print(response.choices[0].message.content)`} + {modelUsageExample( + selectedModel.mode, + getProxyBaseUrl(), + selectedModel.model_group, + getGlobalLitellmHeaderName(), + )}
diff --git a/ui/litellm-dashboard/src/components/AIHub/modelUsageExample.test.ts b/ui/litellm-dashboard/src/components/AIHub/modelUsageExample.test.ts new file mode 100644 index 00000000000..25d1b83ebe9 --- /dev/null +++ b/ui/litellm-dashboard/src/components/AIHub/modelUsageExample.test.ts @@ -0,0 +1,40 @@ +import { describe, expect, it } from "vitest"; +import { modelUsageExample } from "./modelUsageExample"; + +const BASE_URL = "https://proxy.example.com"; + +describe("modelUsageExample", () => { + it("posts a decision model to /v1/systemone on the proxy", () => { + const example = modelUsageExample("evaluation", BASE_URL, "jev-latest", "Authorization"); + expect(example).toContain('"https://proxy.example.com/v1/systemone"'); + expect(example).toContain('"model": "jev-latest"'); + expect(example).toContain('"type": "noul"'); + expect(example).not.toContain("chat.completions"); + }); + + it("sends the key in the header the proxy reads keys from", () => { + expect(modelUsageExample("evaluation", BASE_URL, "jev-latest", "Authorization")).toContain( + 'headers={"Authorization": "Bearer your_api_key"}', + ); + expect(modelUsageExample("evaluation", BASE_URL, "jev-latest", "X-Litellm-Key")).toContain( + 'headers={"X-Litellm-Key": "Bearer your_api_key"}', + ); + }); + + it("keeps the chat completions example for every other mode", () => { + for (const mode of ["chat", null, undefined]) { + const example = modelUsageExample(mode, BASE_URL, "gpt-5.5", "Authorization"); + expect(example).toContain("client.chat.completions.create("); + expect(example).toContain('base_url="https://proxy.example.com"'); + expect(example).toContain('model="gpt-5.5"'); + expect(example).not.toContain("/v1/systemone"); + } + }); + + it("escapes quotes in a model name so the snippet stays valid Python", () => { + expect(modelUsageExample("evaluation", BASE_URL, 'team"model', "Authorization")).toContain( + '"model": "team\\"model"', + ); + expect(modelUsageExample("chat", BASE_URL, 'team"model', "Authorization")).toContain('model="team\\"model"'); + }); +}); diff --git a/ui/litellm-dashboard/src/components/AIHub/modelUsageExample.ts b/ui/litellm-dashboard/src/components/AIHub/modelUsageExample.ts new file mode 100644 index 00000000000..2426ec09942 --- /dev/null +++ b/ui/litellm-dashboard/src/components/AIHub/modelUsageExample.ts @@ -0,0 +1,47 @@ +import { isDecisionMode } from "@/lib/decisionModels"; + +const chatCompletionsExample = (baseUrl: string, model: string): string => `import openai + +client = openai.OpenAI( + api_key="your_api_key", + base_url=${JSON.stringify(baseUrl)} # Your LiteLLM Proxy URL +) + +response = client.chat.completions.create( + model=${JSON.stringify(model)}, + messages=[ + { + "role": "user", + "content": "Hello, how are you?" + } + ] +) + +print(response.choices[0].message.content)`; + +const systemOneExample = (baseUrl: string, model: string, keyHeaderName: string): string => `import requests + +response = requests.post( + ${JSON.stringify(`${baseUrl}/v1/systemone`)}, # Your LiteLLM Proxy URL + headers={${JSON.stringify(keyHeaderName)}: "Bearer your_api_key"}, + json={ + "model": ${JSON.stringify(model)}, + "state": "I was charged twice for my subscription this month.", + "questions": { + "wants_refund": { + "type": "noul", + "instructions": "Is the customer asking for money back?" + } + } + } +) + +print(response.json()["answers"])`; + +export const modelUsageExample = ( + mode: string | null | undefined, + baseUrl: string, + model: string, + keyHeaderName: string, +): string => + isDecisionMode(mode) ? systemOneExample(baseUrl, model, keyHeaderName) : chatCompletionsExample(baseUrl, model); diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx index 75297a508d2..067adcaeb56 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx @@ -15,6 +15,7 @@ import { Providers } from "../provider_info_helpers"; import { projectMountedValues, useMountRegistry, type MountedFormValues } from "../common_components/MountedFormField"; import { useForm } from "react-hook-form"; import AddModelForm from "./AddModelForm"; +import { EMPTY_DECISION_CATALOG } from "@/lib/decisionModels"; import { handleAddModelSubmit } from "./handle_add_model_submit"; import { toast } from "@/lib/toast"; @@ -245,6 +246,7 @@ const createTestProps = (userRole = "proxy_admin", userId = "user-1", isTeamAdmi showAdvancedSettings: false, teams, credentials, + decisionCatalog: EMPTY_DECISION_CATALOG, userRole, userId, }; diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index fbdab07a545..e8b7ca62c3d 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -38,6 +38,8 @@ import ConnectionErrorDisplay from "./model_connection_test"; import ProviderSpecificFields from "./provider_specific_fields"; import { authTypesFor } from "./provider_auth_types"; import { TEST_MODES } from "./add_model_modes"; +import DecisionModelNote from "./DecisionModelNote"; +import { decisionModelsSublabel, type DecisionCatalog } from "@/lib/decisionModels"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { credentialsKeys } from "@/app/(dashboard)/hooks/credentials/useCredentials"; import { extractProxyErrorMessage } from "@/lib/http/client"; @@ -59,6 +61,7 @@ interface AddModelFormProps { setShowAdvancedSettings: (show: boolean) => void; teams: Team[] | null; credentials: CredentialItem[]; + decisionCatalog: DecisionCatalog; } const connectionTestModelName = (values: MountedFormValues): string | undefined => { @@ -82,6 +85,7 @@ const AddModelForm: React.FC = ({ setShowAdvancedSettings, teams, credentials, + decisionCatalog, }) => { const [testMode, setTestMode] = useState("chat"); const [isResultModalVisible, setIsResultModalVisible] = useState(false); @@ -99,6 +103,7 @@ const AddModelForm: React.FC = ({ const guardrailsList = guardrailsData?.guardrails.map((g) => g.guardrail_name); const { data: tagsList } = useTags(); const selectedCredentialName = useWatch({ control: form.control, name: "litellm_credential_name" }); + const pickedProvider = useWatch({ control: form.control, name: "custom_llm_provider" }); const queryClient = useQueryClient(); const [isCredentialModalOpen, setIsCredentialModalOpen] = useState(false); const [credentialModalAuthTypeId, setCredentialModalAuthTypeId] = useState(); @@ -168,11 +173,16 @@ const AddModelForm: React.FC = ({ sortedProviderMetadata.map((providerInfo) => ({ label: providerInfo.provider_display_name, value: providerInfo.provider, + sublabel: decisionModelsSublabel(decisionCatalog, providerInfo.litellm_provider), icon: , })), - [sortedProviderMetadata], + [sortedProviderMetadata, decisionCatalog], ); + const pickedLitellmProvider = sortedProviderMetadata.find( + (providerInfo) => providerInfo.provider === pickedProvider, + )?.litellm_provider; + const credentialSelectOptions: SearchSelectOption[] = useMemo(() => credentialOptions(credentials), [credentials]); const applyProviderSelection = (provider: string | null) => { @@ -274,6 +284,7 @@ const AddModelForm: React.FC = ({ providerModels={providerModels} getPlaceholder={getPlaceholder} /> + {/* Conditionally Render "Public Model Name" */} diff --git a/ui/litellm-dashboard/src/components/add_model/DecisionModelNote.tsx b/ui/litellm-dashboard/src/components/add_model/DecisionModelNote.tsx new file mode 100644 index 00000000000..b199facd484 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/DecisionModelNote.tsx @@ -0,0 +1,54 @@ +import { Info } from "lucide-react"; +import { useFormContext, useWatch } from "react-hook-form"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { + DECISIONS_DOCS_URL, + SYSTEM_ONE_PLAYGROUND_ROUTE, + isDecisionSelection, + type DecisionCatalog, +} from "@/lib/decisionModels"; +import { uiHref } from "@/utils/uiHref"; +import type { MountedFormValues } from "../common_components/MountedFormField"; + +interface DecisionModelNoteProps { + catalog: DecisionCatalog; + litellmProvider: string | undefined; +} + +const selectedModelNames = (value: unknown): string[] => { + if (Array.isArray(value)) { + return value.filter((model): model is string => typeof model === "string"); + } + return typeof value === "string" && value !== "" ? [value] : []; +}; + +export default function DecisionModelNote({ catalog, litellmProvider }: DecisionModelNoteProps) { + const { control } = useFormContext(); + const model = useWatch({ control, name: "model" }); + + if (!isDecisionSelection(catalog, litellmProvider, selectedModelNames(model))) { + return null; + } + + return ( + + + Decision model + +

+ For decision requests, call /v1/decisions (OpenAI format) or /v1/systemone (System + One format).{" "} + + How to call decision models + +

+

+ After you add it,{" "} + + test it in the System One playground + +

+ + + ); +} diff --git a/ui/litellm-dashboard/src/components/molecules/DecisionModelsBanner.integration.test.tsx b/ui/litellm-dashboard/src/components/molecules/DecisionModelsBanner.integration.test.tsx new file mode 100644 index 00000000000..b7d81a09f4e --- /dev/null +++ b/ui/litellm-dashboard/src/components/molecules/DecisionModelsBanner.integration.test.tsx @@ -0,0 +1,71 @@ +import { fireEvent, render, screen } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import DecisionModelsBanner from "./DecisionModelsBanner"; + +const STORAGE_KEY = "hideDecisionModelsBanner"; + +describe("DecisionModelsBanner", () => { + beforeEach(() => { + localStorage.removeItem(STORAGE_KEY); + localStorage.removeItem("hideCostOptimizationFeedbackBanner"); + }); + + it("links directly to the System One playground", () => { + render(); + const link = screen.getByRole("link", { name: "Try decision models" }); + expect(link).toHaveAttribute("href", "/ui/playground?tab=system-one"); + }); + + it("links to the docs on how to call decision models", () => { + render(); + expect(screen.getByRole("link", { name: "How to call them" })).toHaveAttribute( + "href", + "https://docs.litellm.ai/docs/decisions", + ); + }); + + it("offers Add a decision model only when the page can add models, and calls back on click", () => { + const onAddModel = vi.fn(); + const { rerender } = render(); + expect(screen.queryByRole("button", { name: "Add a decision model" })).not.toBeInTheDocument(); + + rerender(); + fireEvent.click(screen.getByRole("button", { name: "Add a decision model" })); + + expect(onAddModel).toHaveBeenCalledTimes(1); + }); + + it("points only users who can add models at Add Model", () => { + const { rerender } = render(); + expect(screen.getByRole("paragraph")).not.toHaveTextContent("Add Model"); + expect(screen.getByRole("paragraph")).toHaveTextContent("Call them at /v1/decisions or /v1/systemone"); + + rerender(); + + expect(screen.getByRole("paragraph")).toHaveTextContent( + "Search decision in Add Model to find them. Call them at /v1/decisions", + ); + }); + + it("shows the announcement even if the old feedback banner was dismissed", () => { + localStorage.setItem("hideCostOptimizationFeedbackBanner", "true"); + render(); + expect(screen.getByText("Decision models are now supported")).toBeInTheDocument(); + }); + + it("hides itself and persists the dismissal when the dismiss button is clicked", () => { + render(); + expect(screen.getByText("Decision models are now supported")).toBeInTheDocument(); + + fireEvent.click(screen.getByLabelText("Dismiss banner")); + + expect(screen.queryByText("Decision models are now supported")).not.toBeInTheDocument(); + expect(localStorage.getItem(STORAGE_KEY)).toBe("true"); + }); + + it("stays dismissed on remount once persisted", () => { + localStorage.setItem(STORAGE_KEY, "true"); + render(); + expect(screen.queryByText("Decision models are now supported")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/molecules/DecisionModelsBanner.tsx b/ui/litellm-dashboard/src/components/molecules/DecisionModelsBanner.tsx new file mode 100644 index 00000000000..5dbf5191198 --- /dev/null +++ b/ui/litellm-dashboard/src/components/molecules/DecisionModelsBanner.tsx @@ -0,0 +1,78 @@ +import React, { useState } from "react"; +import { ArrowRight, Sparkles, X } from "lucide-react"; + +import { DECISIONS_DOCS_URL, SYSTEM_ONE_PLAYGROUND_ROUTE } from "@/lib/decisionModels"; +import { uiHref } from "@/utils/uiHref"; + +import { Button } from "@/components/ui/button"; + +const STORAGE_KEY = "hideDecisionModelsBanner"; + +interface DecisionModelsBannerProps { + onAddModel?: () => void; +} + +const DecisionModelsBanner: React.FC = ({ onAddModel }) => { + const [dismissed, setDismissed] = useState(() => { + if (typeof window !== "undefined") { + return localStorage.getItem(STORAGE_KEY) === "true"; + } + return false; + }); + + if (dismissed) { + return null; + } + + return ( +
+
+ +
+
+

Decision models are now supported

+

+ {onAddModel && ( + <> + Search decision in Add Model to find them.{" "} + + )} + Call them at /v1/decisions or /v1/systemone, or try them in the System One + playground.{" "} + + How to call them + +

+
+ {onAddModel && ( + + )} + + +
+ ); +}; + +export default DecisionModelsBanner; diff --git a/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx b/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx deleted file mode 100644 index 3fc1d41dfdd..00000000000 --- a/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx +++ /dev/null @@ -1,33 +0,0 @@ -import { fireEvent, render, screen } from "@testing-library/react"; -import { beforeEach, describe, expect, it } from "vitest"; -import CostOptimizationFeedbackBanner from "./cost_optimization_feedback_banner"; - -const STORAGE_KEY = "hideCostOptimizationFeedbackBanner"; - -describe("CostOptimizationFeedbackBanner", () => { - beforeEach(() => { - localStorage.removeItem(STORAGE_KEY); - }); - - it("renders with a link to the feedback discussion", () => { - render(); - const link = screen.getByText("Share Feedback").closest("a"); - expect(link).toHaveAttribute("href", "https://github.com/BerriAI/litellm/discussions/32172"); - }); - - it("hides itself and persists the dismissal when the dismiss button is clicked", () => { - render(); - expect(screen.getByText("Help shape cost optimization")).toBeInTheDocument(); - - fireEvent.click(screen.getByLabelText("Dismiss banner")); - - expect(screen.queryByText("Help shape cost optimization")).not.toBeInTheDocument(); - expect(localStorage.getItem(STORAGE_KEY)).toBe("true"); - }); - - it("stays dismissed on remount once persisted", () => { - localStorage.setItem(STORAGE_KEY, "true"); - render(); - expect(screen.queryByText("Help shape cost optimization")).not.toBeInTheDocument(); - }); -}); diff --git a/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.tsx b/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.tsx deleted file mode 100644 index 176789db964..00000000000 --- a/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.tsx +++ /dev/null @@ -1,58 +0,0 @@ -import React, { useState } from "react"; -import { ExternalLink, MessageSquare, X } from "lucide-react"; - -import { Button } from "@/components/ui/button"; - -const STORAGE_KEY = "hideCostOptimizationFeedbackBanner"; -const DISCUSSION_URL = "https://github.com/BerriAI/litellm/discussions/32172"; - -const CostOptimizationFeedbackBanner: React.FC = () => { - const [dismissed, setDismissed] = useState(() => { - if (typeof window !== "undefined") { - return localStorage.getItem(STORAGE_KEY) === "true"; - } - return false; - }); - - if (dismissed) { - return null; - } - - return ( -
-
- -
-
-

Help shape cost optimization

-

- We're collecting suggestions for cost optimization improvements across routing, budgets, and more. Let us - know what you'd like to see. -

-
- - -
- ); -}; - -export default CostOptimizationFeedbackBanner; diff --git a/ui/litellm-dashboard/src/lib/decisionModels.test.ts b/ui/litellm-dashboard/src/lib/decisionModels.test.ts new file mode 100644 index 00000000000..c0f3842993c --- /dev/null +++ b/ui/litellm-dashboard/src/lib/decisionModels.test.ts @@ -0,0 +1,108 @@ +import { describe, expect, it } from "vitest"; +import { + EMPTY_DECISION_CATALOG, + buildDecisionCatalog, + decisionModelsSublabel, + isDecisionMode, + isDecisionSelection, +} from "./decisionModels"; + +const COST_MAP = { + "typesafe/jev-latest": { litellm_provider: "typesafe", mode: "evaluation" }, + "typesafe/jev-preview": { litellm_provider: "typesafe", mode: "evaluation" }, + "gpt-6-luna": { + litellm_provider: "openai", + mode: "chat", + supported_endpoints: ["/v1/chat/completions", "/v1/decisions"], + }, + "gpt-5.5": { litellm_provider: "openai", mode: "chat", supported_endpoints: ["/v1/chat/completions"] }, + "perplexity/pplx-decider-v1-27b": { litellm_provider: "perplexity", mode: "evaluation" }, + "perplexity/sonar-pro": { litellm_provider: "perplexity", mode: "chat" }, + "systemone-only": { litellm_provider: "somevendor", supported_endpoints: ["/v1/systemone"] }, + sample_spec: { max_tokens: "set to max_output_tokens if provider specifies it" }, +}; + +describe("isDecisionMode", () => { + it("is true only for the evaluation mode", () => { + expect(isDecisionMode("evaluation")).toBe(true); + expect(isDecisionMode("chat")).toBe(false); + expect(isDecisionMode(null)).toBe(false); + expect(isDecisionMode(undefined)).toBe(false); + }); +}); + +describe("buildDecisionCatalog", () => { + const catalog = buildDecisionCatalog(COST_MAP); + + it("collects evaluation models and models that list a decision endpoint", () => { + expect([...catalog.models].sort()).toEqual( + [ + "gpt-6-luna", + "perplexity/pplx-decider-v1-27b", + "systemone-only", + "typesafe/jev-latest", + "typesafe/jev-preview", + ].sort(), + ); + }); + + it("strips the provider prefix from the names shown per provider", () => { + expect(catalog.providers.get("typesafe")?.names).toEqual(["jev-latest", "jev-preview"]); + expect(catalog.providers.get("perplexity")?.names).toEqual(["pplx-decider-v1-27b"]); + expect(catalog.providers.get("openai")?.names).toEqual(["gpt-6-luna"]); + }); + + it("marks a provider decision-only when every one of its models is a decision model", () => { + expect(catalog.providers.get("typesafe")?.decisionOnly).toBe(true); + expect(catalog.providers.get("somevendor")?.decisionOnly).toBe(true); + expect(catalog.providers.get("perplexity")?.decisionOnly).toBe(false); + expect(catalog.providers.get("openai")?.decisionOnly).toBe(false); + }); + + it("leaves providers without decision models out", () => { + expect(catalog.providers.has("anthropic")).toBe(false); + }); + + it("returns an empty catalog for a missing or malformed cost map", () => { + expect(buildDecisionCatalog(undefined)).toEqual(EMPTY_DECISION_CATALOG); + expect(buildDecisionCatalog("not a map")).toEqual(EMPTY_DECISION_CATALOG); + expect(buildDecisionCatalog({ broken: { mode: "evaluation" } })).toEqual(EMPTY_DECISION_CATALOG); + }); +}); + +describe("decisionModelsSublabel", () => { + const catalog = buildDecisionCatalog(COST_MAP); + + it("lists the provider's decision models", () => { + expect(decisionModelsSublabel(catalog, "typesafe")).toBe("Decision models: jev-latest, jev-preview"); + }); + + it("is undefined for a provider with no decision models", () => { + expect(decisionModelsSublabel(catalog, "anthropic")).toBeUndefined(); + }); +}); + +describe("isDecisionSelection", () => { + const catalog = buildDecisionCatalog(COST_MAP); + + it("is true for a decision-only provider before any model is picked", () => { + expect(isDecisionSelection(catalog, "typesafe", [])).toBe(true); + }); + + it("waits for a decision model on a provider that also serves chat models", () => { + expect(isDecisionSelection(catalog, "perplexity", [])).toBe(false); + expect(isDecisionSelection(catalog, "perplexity", ["perplexity/sonar-pro"])).toBe(false); + expect(isDecisionSelection(catalog, "perplexity", ["perplexity/sonar-pro", "perplexity/pplx-decider-v1-27b"])).toBe( + true, + ); + }); + + it("is true for a chat-mode model that lists a decision endpoint", () => { + expect(isDecisionSelection(catalog, "openai", ["gpt-6-luna"])).toBe(true); + expect(isDecisionSelection(catalog, "openai", ["gpt-5.5"])).toBe(false); + }); + + it("is false with no provider and no models", () => { + expect(isDecisionSelection(catalog, undefined, [])).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/lib/decisionModels.ts b/ui/litellm-dashboard/src/lib/decisionModels.ts new file mode 100644 index 00000000000..78499952487 --- /dev/null +++ b/ui/litellm-dashboard/src/lib/decisionModels.ts @@ -0,0 +1,74 @@ +import { z } from "zod"; + +export const DECISIONS_DOCS_URL = "https://docs.litellm.ai/docs/decisions"; +export const SYSTEM_ONE_PLAYGROUND_ROUTE = "playground?tab=system-one"; + +const DECISION_ENDPOINTS: ReadonlySet = new Set(["/v1/decisions", "/v1/systemone"]); + +const costMapEntrySchema = z.object({ + litellm_provider: z.string(), + mode: z.string().nullish(), + supported_endpoints: z.array(z.string()).nullish(), +}); + +type CostMapEntry = z.infer; + +interface ProviderDecisionModels { + readonly names: readonly string[]; + readonly decisionOnly: boolean; +} + +export interface DecisionCatalog { + readonly models: ReadonlySet; + readonly providers: ReadonlyMap; +} + +export const EMPTY_DECISION_CATALOG: DecisionCatalog = { models: new Set(), providers: new Map() }; + +export const isDecisionMode = (mode: string | null | undefined): boolean => mode === "evaluation"; + +const isDecisionEntry = (entry: CostMapEntry): boolean => + isDecisionMode(entry.mode) || (entry.supported_endpoints ?? []).some((endpoint) => DECISION_ENDPOINTS.has(endpoint)); + +const withoutProviderPrefix = (key: string, provider: string): string => + key.startsWith(`${provider}/`) ? key.slice(provider.length + 1) : key; + +const parsedEntries = (costMap: unknown): ReadonlyArray => { + const record = z.record(z.string(), z.unknown()).safeParse(costMap); + if (!record.success) { + return []; + } + return Object.entries(record.data).flatMap(([key, value]) => { + const entry = costMapEntrySchema.safeParse(value); + return entry.success ? [[key, entry.data] as const] : []; + }); +}; + +export function buildDecisionCatalog(costMap: unknown): DecisionCatalog { + const entries = parsedEntries(costMap); + const decisionEntries = entries.filter(([, entry]) => isDecisionEntry(entry)); + const decisionProviders = new Set(decisionEntries.map(([, entry]) => entry.litellm_provider)); + const providers = new Map( + [...decisionProviders].map((provider) => { + const providerEntries = entries.filter(([, entry]) => entry.litellm_provider === provider); + const names = decisionEntries + .filter(([, entry]) => entry.litellm_provider === provider) + .map(([key]) => withoutProviderPrefix(key, provider)); + return [provider, { names, decisionOnly: providerEntries.every(([, entry]) => isDecisionEntry(entry)) }] as const; + }), + ); + return { models: new Set(decisionEntries.map(([key]) => key)), providers }; +} + +export const decisionModelsSublabel = (catalog: DecisionCatalog, litellmProvider: string): string | undefined => { + const names = catalog.providers.get(litellmProvider)?.names ?? []; + return names.length > 0 ? `Decision models: ${names.join(", ")}` : undefined; +}; + +export const isDecisionSelection = ( + catalog: DecisionCatalog, + litellmProvider: string | undefined, + selectedModels: readonly string[], +): boolean => + (litellmProvider !== undefined && catalog.providers.get(litellmProvider)?.decisionOnly === true) || + selectedModels.some((model) => catalog.models.has(model)); From fb547a095613f0b5e0cb5010d34b1a0a21ec520a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 10 Oct 2026 01:16:06 +0000 Subject: [PATCH 015/104] ci: finish the lint, unit and smoke check layout (#45706) * ci: delete legacy required-checks shim workflow * ci: fold semgrep into lint and helm into unit, drop the duplicate UI build semgrep runs the repo's own rules over the source, so it becomes lint / semgrep behind lint passed. The helm chart tests become unit / helm behind unit passed. UI Build Check ran the same docker build --target ui-builder as smoke / dashboard-build, so it is deleted * ci: merge the alphabetical llms shards into one llms-providers shard The two shards were split by letter range, which says nothing about what they run. Every provider without its own shard now runs in llms-providers, with the same 137 paths as before --------- Co-authored-by: yuneng --- .github/workflows/helm_unit_test.yml | 54 --- .github/workflows/required-checks-legacy.yml | 397 ------------------- .github/workflows/test-linting.yml | 25 +- .github/workflows/test-litellm-ui-build.yml | 40 -- .github/workflows/test-semgrep.yml | 37 -- .github/workflows/test-unit.yml | 57 ++- 6 files changed, 69 insertions(+), 541 deletions(-) delete mode 100644 .github/workflows/helm_unit_test.yml delete mode 100644 .github/workflows/required-checks-legacy.yml delete mode 100644 .github/workflows/test-litellm-ui-build.yml delete mode 100644 .github/workflows/test-semgrep.yml diff --git a/.github/workflows/helm_unit_test.yml b/.github/workflows/helm_unit_test.yml deleted file mode 100644 index f95848945a0..00000000000 --- a/.github/workflows/helm_unit_test.yml +++ /dev/null @@ -1,54 +0,0 @@ -name: Helm unit test - -on: - pull_request: - push: - branches: - - main - -permissions: - contents: read - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} - cancel-in-progress: ${{ github.event_name == 'pull_request' }} - -jobs: - unit-test: - runs-on: ubuntu-latest - steps: - - name: Checkout - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Set up Helm 3.11.1 - uses: azure/setup-helm@1a275c3b69536ee54be43f2070a358922e12c8d4 # v4.3.1 - with: - version: "3.11.1" - - - name: Download and verify Helm Unit Test Plugin - run: | - curl -fsSLo "$RUNNER_TEMP/helm-unittest.tgz" https://github.com/helm-unittest/helm-unittest/releases/download/v0.8.2/helm-unittest-linux-amd64-0.8.2.tgz - echo "56ab3091e6fa52a7c92ee951def9bed957f295d9ce98483aed404e748d7b3a94 $RUNNER_TEMP/helm-unittest.tgz" | sha256sum -c - - - - name: Install Helm Unit Test Plugin - run: | - PLUGIN_DIR="$(helm env HELM_PLUGINS)/helm-unittest" - mkdir -p "$PLUGIN_DIR" - tar -xzf "$RUNNER_TEMP/helm-unittest.tgz" -C "$PLUGIN_DIR" - helm plugin list - - - name: Run unit tests - run: | - for chart in helm/litellm-helm helm/litellm; do - declared="$(grep -h '^suite:' "$chart"/tests/*.yaml | wc -l | tr -d '[:space:]')" - output="$(mktemp)" - helm unittest -f 'tests/*.yaml' "$chart" | tee "$output" - executed="$(sed -n 's/^Test Suites:.*[[:space:]]\([0-9][0-9]*\) total$/\1/p' "$output")" - if [ "$declared" != "$executed" ]; then - echo "::error::$chart declares $declared test suites but helm-unittest ran $executed. Suites are being skipped silently, so their assertions never execute." - exit 1 - fi - echo "$chart: all $declared declared test suites ran" - done diff --git a/.github/workflows/required-checks-legacy.yml b/.github/workflows/required-checks-legacy.yml deleted file mode 100644 index 6bba690bd8a..00000000000 --- a/.github/workflows/required-checks-legacy.yml +++ /dev/null @@ -1,397 +0,0 @@ -name: Required checks (legacy) - -on: - pull_request: - branches: - - main - - "litellm_**" - -permissions: - contents: read - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} - cancel-in-progress: true - -jobs: - lint: - name: lint - permissions: - contents: read - pull-requests: read - actions: read - runs-on: ubuntu-latest - timeout-minutes: 15 - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - ref: ${{ github.event.pull_request.head.sha }} - fetch-depth: 1 - clean: true - persist-credentials: false - - - name: Detect relevant changes - id: changes - uses: ./.github/actions/detect-changes - - - name: Fetch gate base (merge-base with target branch) - if: steps.changes.outputs.decision != 'skip' - env: - GH_TOKEN: ${{ github.token }} - BASE_SHA: ${{ github.event.pull_request.base.sha }} - HEAD_SHA: ${{ github.event.pull_request.head.sha }} - run: | - retry() { "$@" || { sleep 15; "$@"; } || { sleep 30; "$@"; }; } - MERGE_BASE=$(retry gh api "repos/${{ github.repository }}/compare/${BASE_SHA}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha') - test -n "$MERGE_BASE" - retry git fetch --no-tags --depth=1 origin "$MERGE_BASE" - echo "GATE_BASE_SHA=$MERGE_BASE" >> "$GITHUB_ENV" - - - name: Set up Python - if: steps.changes.outputs.decision != 'skip' - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Set up uv - if: steps.changes.outputs.decision != 'skip' - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Cache uv dependencies - if: steps.changes.outputs.decision != 'skip' - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 - with: - path: | - ~/.cache/uv - .venv - key: ${{ runner.os }}-uv-lint-${{ hashFiles('uv.lock') }} - restore-keys: | - ${{ runner.os }}-uv-lint- - - - name: Clean Python cache - if: steps.changes.outputs.decision != 'skip' - run: | - find . -type d -name "__pycache__" -exec rm -rf {} + || true - find . -name "*.pyc" -delete || true - - - name: Check uv.lock is up to date - if: steps.changes.outputs.decision != 'skip' - run: | - uv lock --check || (echo "❌ uv.lock is out of sync with pyproject.toml. Run 'uv lock' locally and commit the result." && exit 1) - - - name: Cache the Rust build - if: steps.changes.outputs.decision != 'skip' - uses: ./.github/actions/cache-cargo-build - - - name: Install dependencies - if: steps.changes.outputs.decision != 'skip' - run: | - uv sync --frozen --group proxy-dev --group e2e-dev - - - name: Cache Prisma binaries - if: steps.changes.outputs.decision != 'skip' - uses: ./.github/actions/cache-prisma-binaries - - - name: Generate Prisma client - if: steps.changes.outputs.decision != 'skip' - run: | - uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - - - name: Check ruff format - if: steps.changes.outputs.decision != 'skip' - run: | - git diff --name-only --diff-filter=ACMR "$GATE_BASE_SHA" HEAD -- ':(glob)litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true - if [ ! -s "$RUNNER_TEMP/ruff_format_files.txt" ]; then - echo "No changed litellm Python files to check with ruff format." - exit 0 - fi - xargs uv run --no-sync ruff format --check --exclude '/enterprise/' < "$RUNNER_TEMP/ruff_format_files.txt" - - - name: Debug - Check file state - if: steps.changes.outputs.decision != 'skip' - run: | - echo "Current branch:" - git branch --show-current - echo "Last 3 commits:" - git log --oneline -3 - echo "File content around line 43:" - head -50 litellm/litellm_core_utils/custom_logger_registry.py | tail -10 - - - name: Check MCP operation boundary - if: steps.changes.outputs.decision != 'skip' - run: uv run --no-sync python scripts/check_mcp_operation_boundary.py - - - name: Run Ruff linting - if: steps.changes.outputs.decision != 'skip' - run: | - cd litellm - uv run --no-sync ruff check . - cd .. - - - name: Run Ruff linting (test tree) - if: steps.changes.outputs.decision != 'skip' - run: | - uv run --no-sync ruff check --config ruff-tests.toml tests - - - name: Check strict ruff rules (delta vs merge-base counts) - if: steps.changes.outputs.decision != 'skip' - env: - GH_TOKEN: ${{ github.token }} - run: | - uv run --no-sync python scripts/ruff_strict_gate.py --base "$GATE_BASE_SHA" - - - name: Check type discipline (mutable collections / casts / type guards / kwargs / unexplained suppressions, delta vs merge-base counts) - if: steps.changes.outputs.decision != 'skip' - env: - GH_TOKEN: ${{ github.token }} - run: | - uv run --no-sync python scripts/type_discipline_gate.py --base "$GATE_BASE_SHA" - - - name: Check test quality (zero-assert / mock-echo tests, sys.path.insert, raw env writes, litellm global mutation, credential-gated skips, conftest snapshot inventory, delta vs merge-base counts) - if: steps.changes.outputs.decision != 'skip' - env: - GH_TOKEN: ${{ github.token }} - run: | - uv run --no-sync python scripts/test_quality_gate.py --base "$GATE_BASE_SHA" - - - name: Print OpenAI version - if: steps.changes.outputs.decision != 'skip' - run: | - uv run --no-sync python -c "import openai; print(f'OpenAI version: {openai.__version__}')" - - - name: Check basedpyright (delta vs merge-base counts) - if: steps.changes.outputs.decision != 'skip' - env: - GH_TOKEN: ${{ github.token }} - run: | - uv run --no-sync python scripts/type_check_gate.py --base "$GATE_BASE_SHA" - - - name: Check tests/e2e basedpyright (zero errors) - if: steps.changes.outputs.decision != 'skip' - run: | - if git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- ':(glob)tests/e2e/**/*.py' ':(glob)tests/e2e_harness/**/*.py' pyrightconfig.json | grep -q .; then - uv run --no-sync basedpyright tests/e2e tests/e2e_harness - else - echo "No changed tests/e2e Python files; skipping." - fi - - - name: Run the e2e harness tests - if: steps.changes.outputs.decision != 'skip' - env: - LITELLM_MASTER_KEY: sk-e2e-harness-tests-reach-no-proxy - run: | - if ! git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- tests/e2e tests/e2e_harness ':(exclude)tests/e2e/ui' pyproject.toml uv.lock .github/workflows/test-linting.yml | grep -q .; then - echo "No changed e2e harness files; skipping." - exit 0 - fi - retry() { "$@" || { sleep 15; "$@"; } || { sleep 30; "$@"; }; } - CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/claude_code/pr_gate_version_resolver.py)" - tests/e2e/claude_code/cron_vm/install_claude_code.sh "$CLAUDE_VERSION" "$RUNNER_TEMP/claude-cli" - PATH="$RUNNER_TEMP/claude-cli:$PATH" uv run --no-sync pytest -q tests/e2e_harness - - - name: Check for circular imports - if: steps.changes.outputs.decision != 'skip' - run: | - cd litellm - uv run --no-sync python ../tests/documentation_tests/test_circular_imports.py - cd .. - - - name: Check import safety - if: steps.changes.outputs.decision != 'skip' - run: | - uv run --no-sync python -c "from litellm import *" || (echo '🚨 import failed, this means you introduced unprotected imports! 🚨'; exit 1) - - frontend-lint: - name: frontend-lint - permissions: - contents: read - runs-on: ubuntu-latest - timeout-minutes: 8 - defaults: - run: - working-directory: ui/litellm-dashboard - - steps: - - name: Checkout repository - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - fetch-depth: 1 - persist-credentials: false - - - name: Collect changed files - id: changed - env: - GH_TOKEN: ${{ github.token }} - BASE_SHA: ${{ github.event.pull_request.base.sha }} - HEAD_SHA: ${{ github.event.pull_request.head.sha }} - run: | - merge_base=$(gh api "repos/${{ github.repository }}/compare/${BASE_SHA}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha') - test -n "$merge_base" - git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA" - : > "$RUNNER_TEMP/prettier_files.txt" - : > "$RUNNER_TEMP/eslint_files.txt" - while IFS= read -r f; do - [ -f "$f" ] || continue - case "$f" in - *.js | *.jsx | *.ts | *.tsx | *.mjs | *.cjs) - printf '%s\n' "$f" >> "$RUNNER_TEMP/prettier_files.txt" - printf '%s\n' "$f" >> "$RUNNER_TEMP/eslint_files.txt" ;; - *.json | *.css | *.scss | *.md | *.mdx | *.yml | *.yaml | *.html) - printf '%s\n' "$f" >> "$RUNNER_TEMP/prettier_files.txt" ;; - esac - done < <(git diff --name-only --diff-filter=ACMR --relative "$merge_base" "$HEAD_SHA" -- .) - if [ -s "$RUNNER_TEMP/prettier_files.txt" ] || [ -s "$RUNNER_TEMP/eslint_files.txt" ]; then - echo "has_files=true" >> "$GITHUB_OUTPUT" - else - echo "has_files=false" >> "$GITHUB_OUTPUT" - echo "No lintable UI files changed in this PR; nothing to check." - fi - - - name: Setup Node.js - if: steps.changed.outputs.has_files == 'true' - uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0 - with: - node-version-file: ui/litellm-dashboard/.nvmrc - cache: "npm" - cache-dependency-path: ui/litellm-dashboard/package-lock.json - - - name: Install dependencies - if: steps.changed.outputs.has_files == 'true' - run: npm ci - - - name: Lint changed files (prettier + eslint) - if: steps.changed.outputs.has_files == 'true' - run: | - prettier_files=() - eslint_files=() - while IFS= read -r f; do prettier_files+=("$f"); done < "$RUNNER_TEMP/prettier_files.txt" - while IFS= read -r f; do eslint_files+=("$f"); done < "$RUNNER_TEMP/eslint_files.txt" - status=0 - if [ ${#prettier_files[@]} -gt 0 ]; then - echo "::group::Prettier (${#prettier_files[@]} files)" - npx prettier --check "${prettier_files[@]}" || { status=1; echo "::error::Unformatted files. Fix with: npm run format"; } - echo "::endgroup::" - fi - if [ ${#eslint_files[@]} -gt 0 ]; then - echo "::group::ESLint (${#eslint_files[@]} files)" - npx eslint --no-warn-ignored --pass-on-unpruned-suppressions "${eslint_files[@]}" || status=1 - echo "::endgroup::" - fi - exit $status - - - name: Check lint budgets - if: ${{ !cancelled() && steps.changed.outputs.has_files == 'true' }} - run: | - npx eslint . -f json -o "$RUNNER_TEMP/lint-report.json" || true - node scripts/check-lint-budgets.mjs "$RUNNER_TEMP/lint-report.json" eslint-budgets.json - - - name: Check for dead code (knip) - if: ${{ !cancelled() && steps.changed.outputs.has_files == 'true' }} - run: npm run knip:ci - - assert-ci-coverage: - name: assert-ci-coverage - permissions: - contents: read - runs-on: ubuntu-latest - timeout-minutes: 5 - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Assert every test file and Dockerfile is invoked by a job - run: | - python -m pip install "pyyaml==6.0.3" - python .github/scripts/assert_ci_coverage.py - - - name: Assert no -k expression deselects a file from every job that globs it - run: python .github/scripts/assert_ci_coverage.py --slices - - - name: Assert .github/workflows/ holds only workflows, correctly named - run: python .github/scripts/assert_workflow_dir_hygiene.py - - dashboard-build: - name: Dashboard build - runs-on: ubuntu-24.04 - timeout-minutes: 30 - steps: - - name: Checkout - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Build the dashboard stage - run: docker build --target ui-builder -f Dockerfile . - - core-checks: - name: Core checks (Python ${{ matrix.python-version }}) - runs-on: ubuntu-24.04 - timeout-minutes: 30 - strategy: - fail-fast: false - matrix: - python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] - env: - LITELLM_LOCAL_MODEL_COST_MAP: "True" - steps: - - name: Checkout - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: ${{ matrix.python-version }} - - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Install dependencies - run: .github/scripts/uv_sync_with_retries.sh --frozen --extra proxy --extra cli --group dev --group proxy-dev --python ${{ matrix.python-version }} - - - name: Create the loopback-only network namespace - run: | - sudo ip netns add smoke - sudo ip netns exec smoke ip link set lo up - cat > "${RUNNER_TEMP}/in-netns" <<'WRAP' - #!/usr/bin/env bash - set -euo pipefail - exec sudo --preserve-env=LITELLM_LOCAL_MODEL_COST_MAP ip netns exec smoke setpriv --reuid "$(id -u)" --regid "$(id -g)" --init-groups -- env HOME="${HOME}" PATH="${PATH}" "$@" - WRAP - chmod +x "${RUNNER_TEMP}/in-netns" - echo "IN_NETNS=${RUNNER_TEMP}/in-netns" >> "${GITHUB_ENV}" - - - name: Verify namespace isolation - run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py verify-isolation - - - name: Verify interpreter version - run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py interpreter --expect ${{ matrix.python-version }} - - - name: Import and CLI checks - run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py cli - - - name: Proxy startup check - run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py proxy-startup --diagnostics-dir "${RUNNER_TEMP}/smoke-diagnostics" - - - name: Upload smoke diagnostics - if: always() - uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 - with: - name: merge-smoke-diagnostics-py${{ matrix.python-version }} - path: ${{ runner.temp }}/smoke-diagnostics - if-no-files-found: ignore - - - name: Remove the network namespace - if: always() - run: sudo ip netns delete smoke diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index f00e9f28516..e85aaacd3f8 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -645,9 +645,32 @@ jobs: - name: Assert .github/workflows/ holds only workflows, correctly named run: python .github/scripts/assert_workflow_dir_hygiene.py + semgrep: + name: semgrep + permissions: + contents: read + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Run Semgrep (custom rules) + run: uv tool run --from 'semgrep==1.157.0' semgrep scan --config .semgrep/rules . --error lint-passed: name: lint passed - needs: [python, secret-scan, ui, ui-api-types, code-quality, ci-coverage] + needs: [python, secret-scan, ui, ui-api-types, code-quality, ci-coverage, semgrep] if: always() runs-on: ubuntu-latest timeout-minutes: 2 diff --git a/.github/workflows/test-litellm-ui-build.yml b/.github/workflows/test-litellm-ui-build.yml deleted file mode 100644 index eace78fc2cb..00000000000 --- a/.github/workflows/test-litellm-ui-build.yml +++ /dev/null @@ -1,40 +0,0 @@ -name: UI Build Check -permissions: - contents: read - pull-requests: read - -on: - pull_request: - branches: - - main - - "litellm_**" - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} - cancel-in-progress: ${{ github.event_name == 'pull_request' }} - -jobs: - build-ui: - runs-on: ubuntu-latest - timeout-minutes: 10 - - steps: - - name: Checkout repository - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Detect relevant changes - id: changes - uses: ./.github/actions/detect-changes - with: - category: ui - - # Built through the image stage rather than the checkout, because the - # stage copies ui/litellm-dashboard/ alone: an import reaching above the - # dashboard root resolves in a checkout and fails in every image we ship. - # Dockerfile, docker/Dockerfile.non_root and ui/Dockerfile share this - # stage verbatim, so building one covers all three. - - name: Build the dashboard as the shipped images build it - if: steps.changes.outputs.decision != 'skip' - run: docker build --target ui-builder -f Dockerfile . diff --git a/.github/workflows/test-semgrep.yml b/.github/workflows/test-semgrep.yml deleted file mode 100644 index d375824e3c9..00000000000 --- a/.github/workflows/test-semgrep.yml +++ /dev/null @@ -1,37 +0,0 @@ -name: Semgrep - -on: - pull_request: - branches: - - main - - "litellm_**" - -permissions: - contents: read - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true - -jobs: - semgrep: - runs-on: ubuntu-latest - timeout-minutes: 10 - - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Run Semgrep (custom rules) - run: uv tool run --from 'semgrep==1.157.0' semgrep scan --config .semgrep/rules . --error diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 7ea3ed691a9..1cf79049a66 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -214,7 +214,7 @@ jobs: job-timeout-minutes: 60 dist: loadscope - - shard: llms-a-to-g + - shard: llms-providers test-path: |- tests/unit/llms/test_cache_control_and_reasoning.py tests/unit/llms/test_custom_llm.py @@ -272,16 +272,6 @@ jobs: tests/unit/llms/github_copilot tests/unit/llms/gradient_ai tests/unit/llms/groq - --ignore=tests/unit/llms/base_llm/batches/base_batches_config_test.py - python-version: "3.12" - workers: 4 - reruns: 0 - timeout-minutes: 20 - job-timeout-minutes: 60 - dist: loadscope - - - shard: llms-h-to-z - test-path: |- tests/unit/llms/heroku tests/unit/llms/hosted_vllm tests/unit/llms/huggingface @@ -362,6 +352,7 @@ jobs: tests/unit/llms/xinference tests/unit/llms/you_com tests/unit/llms/zai + --ignore=tests/unit/llms/base_llm/batches/base_batches_config_test.py python-version: "3.12" workers: 4 reruns: 0 @@ -1235,6 +1226,48 @@ jobs: uv run --no-sync python ./tests/documentation_tests/test_router_settings.py uv run --no-sync python ./tests/documentation_tests/test_api_docs.py uv run --no-sync python ./tests/documentation_tests/test_circular_imports.py + helm: + name: helm + permissions: + contents: read + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - name: Checkout + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Helm 3.11.1 + uses: azure/setup-helm@1a275c3b69536ee54be43f2070a358922e12c8d4 # v4.3.1 + with: + version: "3.11.1" + + - name: Download and verify Helm Unit Test Plugin + run: | + curl -fsSLo "$RUNNER_TEMP/helm-unittest.tgz" https://github.com/helm-unittest/helm-unittest/releases/download/v0.8.2/helm-unittest-linux-amd64-0.8.2.tgz + echo "56ab3091e6fa52a7c92ee951def9bed957f295d9ce98483aed404e748d7b3a94 $RUNNER_TEMP/helm-unittest.tgz" | sha256sum -c - + + - name: Install Helm Unit Test Plugin + run: | + PLUGIN_DIR="$(helm env HELM_PLUGINS)/helm-unittest" + mkdir -p "$PLUGIN_DIR" + tar -xzf "$RUNNER_TEMP/helm-unittest.tgz" -C "$PLUGIN_DIR" + helm plugin list + + - name: Run unit tests + run: | + for chart in helm/litellm-helm helm/litellm; do + declared="$(grep -h '^suite:' "$chart"/tests/*.yaml | wc -l | tr -d '[:space:]')" + output="$(mktemp)" + helm unittest -f 'tests/*.yaml' "$chart" | tee "$output" + executed="$(sed -n 's/^Test Suites:.*[[:space:]]\([0-9][0-9]*\) total$/\1/p' "$output")" + if [ "$declared" != "$executed" ]; then + echo "::error::$chart declares $declared test suites but helm-unittest ran $executed. Suites are being skipped silently, so their assertions never execute." + exit 1 + fi + echo "$chart: all $declared declared test suites ran" + done coverage: name: coverage needs: unit @@ -1283,7 +1316,7 @@ jobs: unit-passed: name: unit passed permissions: {} - needs: [rust-bridge, assert-shard-coverage, unit, ui-unit, docs] + needs: [rust-bridge, assert-shard-coverage, unit, ui-unit, docs, helm] if: always() runs-on: ubuntu-latest timeout-minutes: 2 From fa7ad80a622afb8fe64dce9365abd4358ae834b1 Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Fri, 9 Oct 2026 18:51:36 -0700 Subject: [PATCH 016/104] fix(databricks): keep #/$defs refs in json_schema response_format (#45659) * fix(databricks): keep #/$defs refs in json_schema response_format for non-Claude models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(databricks): type the json_schema response_format override Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(databricks): integration cells for json_schema $defs refs across endpoints Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(databricks): split stacked loop in streaming json_schema cell Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(databricks): mark BaseConfig dict contract on json_schema override Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(databricks): cover json_schema refs on every endpoint, edge refs, sad paths and an outage burst --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../llms/databricks/chat/transformation.py | 17 +- .../providers/test_databricks_chat_wire.py | 747 +++++++++++++++++- .../test_databricks_json_schema_refs_sdk.py | 186 +++++ .../test_databricks_chat_transformation.py | 86 ++ 4 files changed, 1032 insertions(+), 4 deletions(-) create mode 100644 tests/integration/sdk/test_databricks_json_schema_refs_sdk.py diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 89b614d6258..50d3b82a989 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -7,7 +7,7 @@ from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequenc from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload import httpx -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( @@ -21,6 +21,9 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_name_from_message, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator +from litellm.llms.base_llm.base_utils import ( + type_to_response_format_param, # pyright: ignore[reportUnknownVariableType] # base_utils helper returns an untyped dict +) from litellm.types.llms.anthropic import AllAnthropicToolsValues from litellm.types.llms.databricks import ( AllDatabricksContentValues, @@ -59,6 +62,10 @@ from ...anthropic.chat.transformation import ( from ...openai_like.chat.transformation import OpenAILikeChatConfig from ..common_utils import DatabricksBase, DatabricksException +_RESPONSE_FORMAT_ADAPTER: ( # mutable-ok: mirrors the dict return contract of get_json_schema_from_pydantic_object + Final[TypeAdapter[dict[str, object] | None]] +) = TypeAdapter(dict[str, object] | None) + def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool: """Databricks rejects assistant messages with neither content nor tool calls, e.g. a replayed @@ -191,6 +198,14 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): def get_config(cls, *, model: str | None = None): return super().get_config() + def get_json_schema_from_pydantic_object( + self, + response_format: ( # mutable-ok: matches BaseConfig override signature + type[BaseModel] | dict[str, object] | None + ), + ) -> dict[str, object] | None: # mutable-ok: BaseConfig contract returns a dict + return _RESPONSE_FORMAT_ADAPTER.validate_python(type_to_response_format_param(response_format=response_format)) + def get_required_params(self) -> list[ProviderField]: """For a given provider, return it's required fields with a description""" return [ diff --git a/tests/integration/providers/test_databricks_chat_wire.py b/tests/integration/providers/test_databricks_chat_wire.py index 614382a77f0..f6192d32c31 100644 --- a/tests/integration/providers/test_databricks_chat_wire.py +++ b/tests/integration/providers/test_databricks_chat_wire.py @@ -1,7 +1,9 @@ +import itertools import json import uuid from collections.abc import Mapping -from typing import Final +from concurrent.futures import ThreadPoolExecutor +from typing import Final, Literal, TypeAlias import pytest from integration._support.client import Gateway, eventually @@ -12,6 +14,10 @@ from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter _BACKEND: Final = "databricks-glm-5-2" _API_KEY: Final = "synthetic-databricks-key" _PROMPT: Final = "Summarise the cached briefing in one sentence." +_JSON_SCHEMA_BACKEND: Final = "databricks-qwen35-122b-a10b" +_CLAUDE_BACKEND: Final = "databricks-claude-haiku-4-5" +_JSON_SCHEMA_PROMPT: Final = "Return the requested JSON." +_JSON_CONTENT: Final = '{"p":{"name":"Ada"}}' _PROVIDER_USAGE: Final[Mapping[str, JsonValue]] = { "prompt_tokens": 12011, "completion_tokens": 8, @@ -20,6 +26,67 @@ _PROVIDER_USAGE: Final[Mapping[str, JsonValue]] = { "cache_creation_input_tokens": 0, } _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_PERSON_SCHEMA: Final[Mapping[str, JsonValue]] = { + "$defs": { + "Person": { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + } + }, + "type": "object", + "properties": {"p": {"$ref": "#/$defs/Person"}}, + "required": ["p"], +} +_JSON_SCHEMA_RESPONSE_FORMAT: Final[Mapping[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "P", "strict": True, "schema": _PERSON_SCHEMA}, +} +_MESSAGES_PERSON_SCHEMA: Final[Mapping[str, JsonValue]] = { + "$defs": { + "Person": { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + "additionalProperties": False, + } + }, + "type": "object", + "properties": {"p": {"$ref": "#/$defs/Person"}}, + "required": ["p"], + "additionalProperties": False, +} +_MESSAGES_RESPONSE_FORMAT: Final[Mapping[str, JsonValue]] = { + "type": "json_schema", + "json_schema": { + "name": "structured_output", + "schema": _MESSAGES_PERSON_SCHEMA, + "strict": True, + }, +} +_JSON_TOOL_CALL_MESSAGE: Final[dict[str, JsonValue]] = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "json-tool-call", + "type": "function", + "function": {"name": "json_tool_call", "arguments": _JSON_CONTENT}, + } + ], +} + + +def _chat_completion_response(model: str, message: dict[str, JsonValue], finish_reason: str) -> bytes: + response: Final = { + "id": "databricks-json-schema-response", + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + return json.dumps(response).encode() class _PromptTokensDetails(BaseModel): @@ -52,12 +119,17 @@ class _Chunk(BaseModel): usage: _Usage | None = None -def _frame(identity: str, choices: list[Mapping[str, object]], usage: Mapping[str, JsonValue] | None = None) -> bytes: +def _frame( + identity: str, + choices: list[Mapping[str, object]], + usage: Mapping[str, JsonValue] | None = None, + model: str = _BACKEND, +) -> bytes: value: Final = { "id": identity, "object": "chat.completion.chunk", "created": 1, - "model": _BACKEND, + "model": model, "choices": choices, **({} if usage is None else {"usage": usage}), } @@ -127,3 +199,672 @@ def test_databricks_stream_final_usage_chunk_reaches_client_and_spend_log(gatewa seconds=70, ) assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"], rows[0]["total_tokens"]) == (12011, 8, 12019) + + +def test_databricks_json_schema_refs_are_preserved_in_chat_completions(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/chat/completions" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _JSON_SCHEMA_BACKEND + assert body["messages"] == [{"role": "user", "content": _JSON_SCHEMA_PROMPT}] + assert body["stream"] is False + assert body["response_format"] == _JSON_SCHEMA_RESPONSE_FORMAT + return Reply( + body=_chat_completion_response( + _JSON_SCHEMA_BACKEND, + {"role": "assistant", "content": _JSON_CONTENT}, + "stop", + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}], + "response_format": _JSON_SCHEMA_RESPONSE_FORMAT, + }, + ) + assert response.status_code == 200, response.text + body: Final = _JSON_OBJECT.validate_json(response.content) + assert body["choices"] == [ + { + "index": 0, + "message": {"role": "assistant", "content": _JSON_CONTENT}, + "finish_reason": "stop", + } + ] + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] + + +def test_databricks_streaming_json_schema_refs_are_preserved(gateway: Gateway) -> None: + identity: Final = f"databricks-json-schema-stream-{uuid.uuid4().hex}" + frames: Final = ( + _frame( + identity, + [{"index": 0, "delta": {"role": "assistant", "content": '{"p":'}, "finish_reason": None}], + model=_JSON_SCHEMA_BACKEND, + ), + _frame( + identity, + [{"index": 0, "delta": {"content": '{"name":"Ada"}}'}, "finish_reason": None}], + model=_JSON_SCHEMA_BACKEND, + ), + _frame(identity, [{"index": 0, "delta": {}, "finish_reason": "stop"}], model=_JSON_SCHEMA_BACKEND), + b"data: [DONE]\n\n", + ) + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/chat/completions" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _JSON_SCHEMA_BACKEND + assert body["messages"] == [{"role": "user", "content": _JSON_SCHEMA_PROMPT}] + assert body["stream"] is True + assert body["response_format"] == _JSON_SCHEMA_RESPONSE_FORMAT + return Reply(content_type="text/event-stream", chunks=frames) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + with gateway.client.stream( + "POST", + "/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}], + "response_format": _JSON_SCHEMA_RESPONSE_FORMAT, + "stream": True, + }, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + assert response.status_code == 200, response.read() + lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: ")) + assert lines[-1] == "data: [DONE]", lines + chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1]) + choices: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks)) + assert "".join(choice.delta.content or "" for choice in choices) == _JSON_CONTENT + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] + + +def test_databricks_claude_json_schema_refs_are_preserved_in_tool_parameters(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/chat/completions" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _CLAUDE_BACKEND + assert body["messages"] == [{"role": "user", "content": _JSON_SCHEMA_PROMPT}] + assert "response_format" not in body + assert body["tools"] == [ + { + "type": "function", + "function": {"name": "json_tool_call", "parameters": _PERSON_SCHEMA}, + } + ] + assert body["tool_choice"] == {"type": "function", "function": {"name": "json_tool_call"}} + return Reply(body=_chat_completion_response(_CLAUDE_BACKEND, _JSON_TOOL_CALL_MESSAGE, "tool_calls")) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{_CLAUDE_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}], + "response_format": _JSON_SCHEMA_RESPONSE_FORMAT, + }, + ) + assert response.status_code == 200, response.text + body: Final = _JSON_OBJECT.validate_json(response.content) + assert body["choices"] == [ + { + "index": 0, + "message": {"role": "assistant", "content": _JSON_CONTENT}, + "finish_reason": "stop", + } + ] + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] + + +def test_databricks_responses_json_schema_refs_are_preserved_in_chat_bridge(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/chat/completions" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _JSON_SCHEMA_BACKEND + assert body["messages"] == [{"role": "user", "content": _JSON_SCHEMA_PROMPT}] + assert body["response_format"] == { + "type": "json_schema", + "json_schema": {"name": "P", "schema": _PERSON_SCHEMA, "strict": True}, + } + return Reply( + body=_chat_completion_response( + _JSON_SCHEMA_BACKEND, + {"role": "assistant", "content": _JSON_CONTENT}, + "stop", + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/responses", + { + "model": model, + "input": _JSON_SCHEMA_PROMPT, + "text": { + "format": { + "type": "json_schema", + "name": "P", + "strict": True, + "schema": _PERSON_SCHEMA, + } + }, + }, + ) + assert response.status_code == 200, response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] + + +def test_databricks_messages_json_schema_refs_are_preserved_for_qwen(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/chat/completions" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _JSON_SCHEMA_BACKEND + assert body["response_format"] == _MESSAGES_RESPONSE_FORMAT + return Reply( + body=_chat_completion_response( + _JSON_SCHEMA_BACKEND, + {"role": "assistant", "content": _JSON_CONTENT}, + "stop", + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}], + "output_format": {"type": "json_schema", "schema": _PERSON_SCHEMA}, + }, + ) + assert response.status_code == 200, response.text + body: Final = _JSON_OBJECT.validate_json(response.content) + assert body["content"] == [{"type": "text", "text": _JSON_CONTENT}] + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] + + +def test_databricks_messages_claude_json_schema_refs_are_preserved_in_tool_parameters(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/chat/completions" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _CLAUDE_BACKEND + assert "response_format" not in body + assert body["tools"] == [ + { + "type": "function", + "function": {"name": "json_tool_call", "parameters": _MESSAGES_PERSON_SCHEMA}, + } + ] + assert body["tool_choice"] == {"type": "function", "function": {"name": "json_tool_call"}} + return Reply(body=_chat_completion_response(_CLAUDE_BACKEND, _JSON_TOOL_CALL_MESSAGE, "tool_calls")) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{_CLAUDE_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}], + "output_format": {"type": "json_schema", "schema": _PERSON_SCHEMA}, + }, + ) + assert response.status_code == 200, response.text + body: Final = _JSON_OBJECT.validate_json(response.content) + assert body["content"] == [{"type": "text", "text": _JSON_CONTENT}] + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] + + +_Endpoint: TypeAlias = Literal["chat", "responses", "messages"] +_PATHS: Final[Mapping[_Endpoint, str]] = { + "chat": "/v1/chat/completions", + "responses": "/v1/responses", + "messages": "/v1/messages", +} +_OUTAGE: Final = "upstream-outage" +_OUTAGE_MESSAGE: Final = "The service is temporarily unavailable" +_EXPECTED_JSON: Final[JsonValue] = _JSON_OBJECT.validate_json(_JSON_CONTENT) +_CONTENT_DELTAS: Final[tuple[Mapping[str, JsonValue], ...]] = ( + {"role": "assistant", "content": '{"p":'}, + {"content": '{"name":"Ada"}}'}, +) +_TOOL_CALL_DELTAS: Final[tuple[Mapping[str, JsonValue], ...]] = ( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "index": 0, + "id": "json-tool-call", + "type": "function", + "function": {"name": "json_tool_call", "arguments": ""}, + } + ], + }, + {"tool_calls": [{"index": 0, "function": {"arguments": _JSON_CONTENT}}]}, +) + + +class _ChatMessage(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + content: str | None = None + + +class _ChatChoice(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + message: _ChatMessage + + +class _ChatCompletion(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + choices: tuple[_ChatChoice, ...] + + +class _OutputPart(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + type: str + text: str | None = None + + +class _OutputItem(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + type: str + content: tuple[_OutputPart, ...] = () + name: str | None = None + arguments: str | None = None + + +class _ResponsesBody(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + output: tuple[_OutputItem, ...] + + +class _ResponsesEvent(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + type: str + response: _ResponsesBody | None = None + + +class _MessagesBlock(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + type: str + text: str | None = None + + +class _MessagesBody(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + content: tuple[_MessagesBlock, ...] + + +class _MessagesDelta(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + type: str | None = None + text: str | None = None + partial_json: str | None = None + + +class _MessagesEvent(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + type: str + delta: _MessagesDelta | None = None + + +class _FunctionDelta(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + arguments: str | None = None + + +class _ToolCallDelta(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + function: _FunctionDelta | None = None + + +class _StreamDelta(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + content: str | None = None + tool_calls: tuple[_ToolCallDelta, ...] | None = None + + +class _StreamChoice(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + delta: _StreamDelta + + +class _StreamChunk(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + choices: tuple[_StreamChoice, ...] + + +def _json_schema_format(schema: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {"type": "json_schema", "json_schema": {"name": "P", "strict": True, "schema": dict(schema)}} + + +def _client_body( + endpoint: _Endpoint, + model: str, + stream: bool, + prompt: str = _JSON_SCHEMA_PROMPT, +) -> dict[str, JsonValue]: + match endpoint: + case "chat": + return { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "response_format": _json_schema_format(_PERSON_SCHEMA), + "stream": stream, + } + case "responses": + return { + "model": model, + "input": prompt, + "text": { + "format": {"type": "json_schema", "name": "P", "strict": True, "schema": dict(_PERSON_SCHEMA)} + }, + "stream": stream, + } + case "messages": + return { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": prompt}], + "output_format": {"type": "json_schema", "schema": dict(_PERSON_SCHEMA)}, + "stream": stream, + } + + +def _expected_outbound(endpoint: _Endpoint, backend: str) -> dict[str, JsonValue]: + schema: Final = dict(_MESSAGES_PERSON_SCHEMA if endpoint == "messages" else _PERSON_SCHEMA) + if "claude" in backend: + return _JSON_OBJECT.validate_python( + { + "tools": [{"type": "function", "function": {"name": "json_tool_call", "parameters": schema}}], + "tool_choice": {"type": "function", "function": {"name": "json_tool_call"}}, + } + ) + name: Final = "structured_output" if endpoint == "messages" else "P" + return _JSON_OBJECT.validate_python( + {"response_format": {"type": "json_schema", "json_schema": {"name": name, "schema": schema, "strict": True}}} + ) + + +def _assert_outbound(request: Request, endpoint: _Endpoint, backend: str, stream: bool) -> None: + assert (request.method, request.target) == ("POST", "/chat/completions") + body: Final = _JSON_OBJECT.validate_json(request.body) + expected: Final = _expected_outbound(endpoint, backend) + assert {key: body.get(key) for key in expected} == expected + assert (body["model"], body.get("stream") is True) == (backend, stream) + assert "claude" not in backend or "response_format" not in body + + +def _upstream_reply(backend: str, stream: bool) -> Reply: + claude: Final = "claude" in backend + finish_reason: Final = "tool_calls" if claude else "stop" + if not stream: + message: Final[dict[str, JsonValue]] = ( + _JSON_TOOL_CALL_MESSAGE if claude else {"role": "assistant", "content": _JSON_CONTENT} + ) + return Reply(body=_chat_completion_response(backend, message, finish_reason)) + identity: Final = f"databricks-json-schema-{uuid.uuid4().hex}" + deltas: Final = _TOOL_CALL_DELTAS if claude else _CONTENT_DELTAS + frames: Final = ( + *(_frame(identity, [{"index": 0, "delta": delta, "finish_reason": None}], model=backend) for delta in deltas), + _frame(identity, [{"index": 0, "delta": {}, "finish_reason": finish_reason}], model=backend), + b"data: [DONE]\n\n", + ) + return Reply(content_type="text/event-stream", chunks=frames) + + +def _outage_or_reply(request: Request) -> Reply: + if _OUTAGE.encode() in request.body: + error: Final = {"error_code": "TEMPORARILY_UNAVAILABLE", "message": _OUTAGE_MESSAGE} + return Reply(status=503, body=json.dumps(error).encode()) + body: Final = _JSON_OBJECT.validate_json(request.body) + return _upstream_reply(str(body["model"]), body.get("stream") is True) + + +def _item_text(item: _OutputItem) -> str: + match item.type: + case "message": + return "".join(part.text or "" for part in item.content) + case "function_call" if item.name == "json_tool_call": + return item.arguments or "" + case _: + return "" + + +def _responses_text(body: _ResponsesBody) -> str: + return "".join(_item_text(item) for item in body.output) + + +def _chat_delta_text(delta: _StreamDelta) -> str: + calls: Final = delta.tool_calls or () + return (delta.content or "") + "".join(call.function.arguments or "" for call in calls if call.function is not None) + + +def _messages_delta_text(delta: _MessagesDelta) -> str: + match delta.type: + case "text_delta": + return delta.text or "" + case "input_json_delta": + return delta.partial_json or "" + case _: + return "" + + +def _final_text(endpoint: _Endpoint, content: bytes) -> str: + match endpoint: + case "chat": + return _ChatCompletion.model_validate_json(content).choices[0].message.content or "" + case "responses": + return _responses_text(_ResponsesBody.model_validate_json(content)) + case "messages": + return "".join(block.text or "" for block in _MessagesBody.model_validate_json(content).content) + + +def _streamed_text(endpoint: _Endpoint, data: tuple[str, ...]) -> str: + match endpoint: + case "chat": + assert data[-1] == "[DONE]", data + chunks: Final = tuple(_StreamChunk.model_validate_json(line) for line in data[:-1]) + choices: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks)) + return "".join(_chat_delta_text(choice.delta) for choice in choices) + case "responses": + events: Final = tuple(_ResponsesEvent.model_validate_json(line) for line in data if line != "[DONE]") + completed: Final = tuple(event.response for event in events if event.type == "response.completed") + assert len(completed) == 1 and completed[0] is not None, data + return _responses_text(completed[0]) + case "messages": + message_events: Final = tuple(_MessagesEvent.model_validate_json(line) for line in data) + assert message_events[-1].type == "message_stop", data + deltas: Final = tuple(event.delta for event in message_events if event.delta is not None) + return "".join(_messages_delta_text(delta) for delta in deltas) + + +def _outcome(status: int, text: str) -> tuple[int, JsonValue]: + if status == 200: + return status, _JSON_OBJECT.validate_json(text) + return status, _OUTAGE_MESSAGE if _OUTAGE_MESSAGE in text else text + + +def _call(gateway: Gateway, endpoint: _Endpoint, body: Mapping[str, JsonValue]) -> tuple[int, str]: + if body["stream"] is not True: + response: Final = gateway.request("POST", _PATHS[endpoint], body) + return response.status_code, _final_text(endpoint, response.content) if response.is_success else response.text + with gateway.client.stream( + "POST", _PATHS[endpoint], json=dict(body), headers={"Authorization": f"Bearer {gateway.key}"} + ) as streamed: + if not streamed.is_success: + return streamed.status_code, streamed.read().decode() + data: Final = tuple(line.removeprefix("data: ") for line in streamed.iter_lines() if line.startswith("data: ")) + return 200, _streamed_text(endpoint, data) + + +@pytest.mark.parametrize( + ("endpoint", "backend", "stream"), + [ + pytest.param("chat", _CLAUDE_BACKEND, True, id="chat-stream-claude"), + pytest.param("responses", _CLAUDE_BACKEND, False, id="responses-claude"), + pytest.param("responses", _JSON_SCHEMA_BACKEND, True, id="responses-stream-qwen"), + pytest.param("responses", _CLAUDE_BACKEND, True, id="responses-stream-claude"), + pytest.param("messages", _JSON_SCHEMA_BACKEND, True, id="messages-stream-qwen"), + pytest.param("messages", _CLAUDE_BACKEND, True, id="messages-stream-claude"), + ], +) +def test_databricks_json_schema_refs_are_preserved_on_every_endpoint( + gateway: Gateway, endpoint: _Endpoint, backend: str, stream: bool +) -> None: + with wire_server(lambda _request: _upstream_reply(backend, stream)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{backend}", api_base=wire.url, api_key=_API_KEY) + status, text = _call(gateway, endpoint, _client_body(endpoint, model, stream)) + received: Final = wire.drain() + assert len(received) == 1, received + _assert_outbound(received[0], endpoint, backend, stream) + assert _outcome(status, text) == (200, _EXPECTED_JSON), text + + +@pytest.mark.parametrize( + "response_format", + [ + pytest.param( + _json_schema_format( + { + "definitions": { + "Person": {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]} + }, + "type": "object", + "properties": {"p": {"$ref": "#/definitions/Person"}}, + "required": ["p"], + } + ), + id="definitions-ref", + ), + pytest.param( + _json_schema_format( + { + "type": "object", + "properties": { + "name": {"type": "string"}, + "children": {"type": "array", "items": {"$ref": "#"}}, + }, + "required": ["name", "children"], + } + ), + id="recursive-root-ref", + ), + pytest.param( + _json_schema_format({"type": "object", "properties": {"$ref": {"type": "string"}}, "required": ["$ref"]}), + id="property-named-ref", + ), + pytest.param( + _json_schema_format({"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}), + id="flat-schema", + ), + pytest.param({"type": "json_object"}, id="json-object"), + ], +) +def test_databricks_response_format_reaches_upstream_as_sent( + gateway: Gateway, response_format: dict[str, JsonValue] +) -> None: + with ( + wire_server(lambda _request: _upstream_reply(_JSON_SCHEMA_BACKEND, False)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}], + "response_format": response_format, + }, + ) + received: Final = wire.drain() + assert response.status_code == 200, response.text + assert _final_text("chat", response.content) == _JSON_CONTENT + assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")] + assert _JSON_OBJECT.validate_json(received[0].body)["response_format"] == response_format + + +def test_databricks_json_schema_without_schema_returns_the_upstream_error(gateway: Gateway) -> None: + response_format: Final[dict[str, JsonValue]] = {"type": "json_schema", "json_schema": {"name": "P", "strict": True}} + upstream_error: Final = {"error_code": "INVALID_PARAMETER_VALUE", "message": "json_schema.schema is required"} + + def respond(_request: Request) -> Reply: + return Reply(status=400, body=json.dumps(upstream_error).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}], + "response_format": response_format, + }, + ) + received: Final = wire.drain() + assert (response.status_code, "json_schema.schema is required" in response.text) == (400, True), response.text + assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")] + assert _JSON_OBJECT.validate_json(received[0].body)["response_format"] == response_format + + +def _assert_burst_outbound(request: Request) -> None: + body: Final = _JSON_OBJECT.validate_json(request.body) + endpoint: Final[_Endpoint] = "messages" if body["response_format"] == _MESSAGES_RESPONSE_FORMAT else "chat" + _assert_outbound(request, endpoint, _JSON_SCHEMA_BACKEND, body.get("stream") is True) + + +def _burst_request(index: int) -> tuple[_Endpoint, bool, bool]: + endpoints: Final[tuple[_Endpoint, ...]] = ("chat", "responses", "messages") + return endpoints[index % 3], index % 2 == 1, index % 5 == 2 + + +def test_databricks_json_schema_burst_through_an_upstream_outage_keeps_every_ref(gateway: Gateway) -> None: + plan: Final = tuple(_burst_request(index) for index in range(24)) + recovery: Final[tuple[_Endpoint, ...]] = ("chat", "responses", "messages") + with wire_server(_outage_or_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + endpoints: Final = tuple(endpoint for endpoint, _stream, _down in plan) + bodies: Final = tuple( + _client_body(endpoint, model, stream, f"{_JSON_SCHEMA_PROMPT} {index} {_OUTAGE if down else ''}") + for index, (endpoint, stream, down) in enumerate(plan) + ) + with ThreadPoolExecutor(max_workers=12) as pool: + results: Final = tuple(pool.map(_call, itertools.repeat(gateway), endpoints, bodies)) + recovered: Final = tuple( + _call(gateway, endpoint, _client_body(endpoint, model, False)) for endpoint in recovery + ) + received: Final = wire.drain() + assert len(received) == len(plan) + len(recovery), received + for request in received: + _assert_burst_outbound(request) + expected: Final = tuple( + (503, _OUTAGE_MESSAGE) if down else (200, _EXPECTED_JSON) for _endpoint, _stream, down in plan + ) + assert tuple(_outcome(*result) for result in results) == expected + assert tuple(_outcome(*result) for result in recovered) == ((200, _EXPECTED_JSON),) * len(recovery) diff --git a/tests/integration/sdk/test_databricks_json_schema_refs_sdk.py b/tests/integration/sdk/test_databricks_json_schema_refs_sdk.py new file mode 100644 index 00000000000..009679fa4b5 --- /dev/null +++ b/tests/integration/sdk/test_databricks_json_schema_refs_sdk.py @@ -0,0 +1,186 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from typing import Final + +import pytest +from integration._support.wire import Reply, Request, wire_server +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +import litellm + +_QWEN_BACKEND: Final = "databricks-qwen35-122b-a10b" +_CLAUDE_BACKEND: Final = "databricks-claude-haiku-4-5" +_API_KEY: Final = "synthetic-databricks-key" +_ORDER_JSON: Final = '{"buyer":{"name":"Ada"},"seller":{"name":"Bo"}}' +_MESSAGES: Final[tuple[Mapping[str, str], ...]] = ({"role": "user", "content": "Return the order as JSON."},) + + +class Person(BaseModel): + model_config = ConfigDict(frozen=True) + name: str + + +class Order(BaseModel): + model_config = ConfigDict(frozen=True) + buyer: Person + seller: Person + + +class _StrictObject(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + additional_properties: bool | None = Field(default=None, alias="additionalProperties") + + +class _OrderSchema(_StrictObject): + defs: Mapping[str, _StrictObject] = Field(alias="$defs") + + +class _JsonSchema(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + name: str + strict: bool | None = None + json_schema: dict[str, JsonValue] = Field(alias="schema") + + +class _ResponseFormat(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + type: str + json_schema: _JsonSchema + + +class _Function(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + name: str + parameters: dict[str, JsonValue] + + +class _Tool(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + function: _Function + + +class _Outbound(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + model: str + response_format: _ResponseFormat | None = None + tools: tuple[_Tool, ...] = () + + +class _CallerMessage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + content: str + + +class _CallerChoice(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + message: _CallerMessage + + +class _CallerCompletion(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + choices: tuple[_CallerChoice, ...] + + +def _completion(backend: str) -> Reply: + claude: Final = "claude" in backend + tool_call: Final[dict[str, JsonValue]] = { + "id": "json-tool-call", + "type": "function", + "function": {"name": "json_tool_call", "arguments": _ORDER_JSON}, + } + message: Final[dict[str, JsonValue]] = ( + {"role": "assistant", "content": None, "tool_calls": [tool_call]} + if claude + else {"role": "assistant", "content": _ORDER_JSON} + ) + body: Final[dict[str, JsonValue]] = { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": backend, + "choices": [{"index": 0, "message": message, "finish_reason": "tool_calls" if claude else "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 7, "total_tokens": 16}, + } + return Reply(body=json.dumps(body).encode()) + + +def _refs(node: JsonValue) -> Iterator[str]: + if isinstance(node, list): + for item in node: + yield from _refs(item) + return + if not isinstance(node, dict): + return + for key, value in node.items(): + if key == "$ref" and isinstance(value, str): + yield value + else: + yield from _refs(value) + + +def _sent_schema(backend: str, received: tuple[Request, ...]) -> dict[str, JsonValue]: + assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")] + outbound: Final = _Outbound.model_validate_json(received[0].body) + assert outbound.model == backend + if "claude" in backend: + assert outbound.response_format is None + assert [tool.function.name for tool in outbound.tools] == ["json_tool_call"] + return outbound.tools[0].function.parameters + assert outbound.response_format is not None + assert (outbound.response_format.type, outbound.response_format.json_schema.name) == ("json_schema", "Order") + assert outbound.response_format.json_schema.strict is True + return outbound.response_format.json_schema.json_schema + + +def _assert_strict_order_schema(schema: dict[str, JsonValue]) -> None: + assert set(_refs(schema)) == {"#/$defs/Person"}, schema + parsed: Final = _OrderSchema.model_validate(schema) + assert (parsed.additional_properties, parsed.defs["Person"].additional_properties) == (False, False), schema + + +def _caller_order(response: litellm.ModelResponse) -> Order: + return Order.model_validate_json( + _CallerCompletion.model_validate_json(response.model_dump_json()).choices[0].message.content + ) + + +_BACKENDS: Final = ( + pytest.param(_QWEN_BACKEND, id="qwen"), + pytest.param(_CLAUDE_BACKEND, id="claude"), +) +_EXPECTED_ORDER: Final = Order(buyer=Person(name="Ada"), seller=Person(name="Bo")) + + +@pytest.mark.parametrize("backend", _BACKENDS) +def test_databricks_pydantic_response_format_sends_strict_defs_refs(backend: str) -> None: + with wire_server(lambda _request: _completion(backend)) as wire: + response: Final = litellm.completion( + model=f"databricks/{backend}", + api_base=wire.url, + api_key=_API_KEY, + messages=[dict(message) for message in _MESSAGES], + response_format=Order, + num_retries=0, + ) + received: Final = wire.drain() + _assert_strict_order_schema(_sent_schema(backend, received)) + assert isinstance(response, litellm.ModelResponse) + assert _caller_order(response) == _EXPECTED_ORDER + + +@pytest.mark.parametrize("backend", _BACKENDS) +async def test_databricks_pydantic_response_format_sends_strict_defs_refs_async(backend: str) -> None: + with wire_server(lambda _request: _completion(backend)) as wire: + response: Final = await litellm.acompletion( + model=f"databricks/{backend}", + api_base=wire.url, + api_key=_API_KEY, + messages=[dict(message) for message in _MESSAGES], + response_format=Order, + num_retries=0, + ) + received: Final = wire.drain() + _assert_strict_order_schema(_sent_schema(backend, received)) + assert isinstance(response, litellm.ModelResponse) + assert _caller_order(response) == _EXPECTED_ORDER diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index c8e92145ab3..a366019d310 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -1,3 +1,4 @@ +import copy import json from collections.abc import Iterator from typing import Final @@ -7,6 +8,7 @@ import httpx import pytest import respx from fastapi.testclient import TestClient +from pydantic import BaseModel, ConfigDict import litellm from litellm.caching.llm_caching_handler import LLMClientCache @@ -32,6 +34,25 @@ DATABRICKS_API_BASE: Final = "https://my.workspace.cloud.databricks.com/serving- DATABRICKS_API_KEY: Final = "dapimykey" DATABRICKS_CHAT_COMPLETIONS_URL: Final = f"{DATABRICKS_API_BASE}/chat/completions" DATABRICKS_EMBEDDINGS_URL: Final = f"{DATABRICKS_API_BASE}/embeddings" +JSON_SCHEMA_RESPONSE_FORMAT: Final = { + "type": "json_schema", + "json_schema": { + "name": "P", + "strict": True, + "schema": { + "$defs": { + "Person": { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + } + }, + "type": "object", + "properties": {"p": {"$ref": "#/$defs/Person"}}, + "required": ["p"], + }, + }, +} @pytest.fixture() @@ -2088,3 +2109,68 @@ def mock_chat_streaming_response_chunks() -> List[str]: } ), ] + + +@pytest.mark.parametrize( + "model", + ( + "databricks-meta-llama-3-3-70b-instruct", + "databricks-qwen35-122b-a10b", + "databricks-gpt-oss-120b", + ), +) +def test_databricks_non_claude_json_schema_preserves_refs(model: str) -> None: + optional_params: Final = litellm.utils.get_optional_params( + model=model, + custom_llm_provider="databricks", + response_format=copy.deepcopy(JSON_SCHEMA_RESPONSE_FORMAT), + ) + + assert optional_params["response_format"] == JSON_SCHEMA_RESPONSE_FORMAT + + +def test_databricks_claude_json_schema_preserves_refs_in_tool_parameters() -> None: + optional_params: Final = litellm.utils.get_optional_params( + model="databricks-claude-haiku-4-5", + custom_llm_provider="databricks", + response_format=copy.deepcopy(JSON_SCHEMA_RESPONSE_FORMAT), + ) + + assert "response_format" not in optional_params + assert optional_params["tools"] == [ + { + "type": "function", + "function": { + "name": "json_tool_call", + "parameters": JSON_SCHEMA_RESPONSE_FORMAT["json_schema"]["schema"], + }, + } + ] + + +def test_databricks_pydantic_json_schema_preserves_nested_refs() -> None: + class Person(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str + + class Order(BaseModel): + model_config = ConfigDict(frozen=True) + + person: Person + + optional_params: Final = litellm.utils.get_optional_params( + model="databricks-meta-llama-3-3-70b-instruct", + custom_llm_provider="databricks", + response_format=Order, + ) + schema: Final = optional_params["response_format"]["json_schema"]["schema"] + + assert schema["properties"]["person"] == {"$ref": "#/$defs/Person"} + assert schema["$defs"]["Person"] == { + "additionalProperties": False, + "properties": {"name": {"title": "Name", "type": "string"}}, + "required": ["name"], + "title": "Person", + "type": "object", + } From d6db8e8744e36e970989aad2bb66b1b518355175 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 18:55:22 -0700 Subject: [PATCH 017/104] fix(gemini): map OpenAI voice names to Gemini prebuilt TTS voices (#45303) * fix(gemini): map OpenAI voice names to Gemini prebuilt TTS voices * test(gemini): type the voice mapping test helpers * fix(gemini): leave nova out of the voice mapping since Gemini accepts it as given * test(gemini): add integration cells for the OpenAI voice mapping on Gemini and Vertex TTS --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../vertex_and_google_ai_studio_gemini.py | 28 +- .../providers/_gemini_tts_voice_support.py | 451 ++++++++++++ ...st_gemini_tts_voice_mapping_owned_proxy.py | 354 ++++++++++ .../test_gemini_tts_voice_mapping_wire.py | 652 ++++++++++++++++++ ...test_vertex_and_google_ai_studio_gemini.py | 46 ++ 5 files changed, 1529 insertions(+), 2 deletions(-) create mode 100644 tests/integration/providers/_gemini_tts_voice_support.py create mode 100644 tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py create mode 100644 tests/integration/providers/test_gemini_tts_voice_mapping_wire.py diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 6bdce393b82..fb846d6b196 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -7,6 +7,7 @@ import time from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from functools import partial +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args import httpx @@ -115,6 +116,25 @@ else: SUPPORTED_REASONING_EFFORTS: Final = ("minimal", "low", "medium", "high", "none", "disable") +# Gemini prebuilt voices per https://ai.google.dev/gemini-api/docs/speech-generation (2026-10-01), by nearest character; +# nova is left out since Gemini 3.x TTS models accept it as given (live check 2026-10-08) +OPENAI_TO_GEMINI_TTS_VOICES: Final[Mapping[str, str]] = MappingProxyType( + { + "alloy": "Kore", + "ash": "Iapetus", + "ballad": "Algieba", + "cedar": "Achird", + "coral": "Sulafat", + "echo": "Charon", + "fable": "Umbriel", + "marin": "Despina", + "onyx": "Orus", + "sage": "Vindemiatrix", + "shimmer": "Achernar", + "verse": "Puck", + } +) + def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsError: return UnsupportedParamsError( @@ -1071,7 +1091,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): speechConfig = { voiceConfig: { prebuiltVoiceConfig: { - voiceName: "alloy", + voiceName: "Kore", } }, languageCode: "en-US", @@ -1096,7 +1116,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): speech_config: Final[SpeechConfig] = {} if "voice" in value: - prebuilt_voice_config: Final[PrebuiltVoiceConfig] = {"voiceName": value["voice"]} + voice: Final = value["voice"] + voice_name: Final = ( + OPENAI_TO_GEMINI_TTS_VOICES.get(voice.lower(), voice) if isinstance(voice, str) else voice + ) + prebuilt_voice_config: Final[PrebuiltVoiceConfig] = {"voiceName": voice_name} voice_config: Final[VoiceConfig] = {"prebuiltVoiceConfig": prebuilt_voice_config} speech_config["voiceConfig"] = voice_config diff --git a/tests/integration/providers/_gemini_tts_voice_support.py b/tests/integration/providers/_gemini_tts_voice_support.py new file mode 100644 index 00000000000..7c453a63ef5 --- /dev/null +++ b/tests/integration/providers/_gemini_tts_voice_support.py @@ -0,0 +1,451 @@ +from __future__ import annotations + +import asyncio +import base64 +import json +import uuid +from collections.abc import AsyncIterator, Callable, Generator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from hashlib import sha256 +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from queue import SimpleQueue +from threading import Thread +from typing import Final + +import anthropic +import httpx +import openai +import websockets +import yaml +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.upstream import GEMINI_LIVE_PATH, ScenarioHandle, delete_scenario, register_scenario +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request +from integration.cost_calculation.cost_tracking_case import RealtimeResponse +from pydantic import JsonValue +from websockets.asyncio.client import ClientConnection +from websockets.exceptions import ConnectionClosed + +BACKEND: Final = "gemini-3.8-flash-tts" +GEMINI_API_KEY: Final = "synthetic-gemini-key" +VERTEX_PROJECT: Final = "scripted-tts-project" +VERTEX_LOCATION: Final = "us-central1" +VERTEX_MODEL_PATH: Final = ( + f"/v1/projects/{VERTEX_PROJECT}/locations/{VERTEX_LOCATION}/publishers/google/models/{BACKEND}" +) +GEMINI_PREBUILT_VOICES: Final = frozenset( + { + "Zephyr", + "Puck", + "Charon", + "Kore", + "Fenrir", + "Leda", + "Orus", + "Aoede", + "Callirrhoe", + "Autonoe", + "Enceladus", + "Iapetus", + "Umbriel", + "Algieba", + "Despina", + "Erinome", + "Algenib", + "Rasalgethi", + "Laomedeia", + "Achernar", + "Alnilam", + "Schedar", + "Gacrux", + "Pulcherrima", + "Achird", + "Zubenelgenubi", + "Vindemiatrix", + "Sadachbia", + "Sadaltager", + "Sulafat", + } +) +ACCEPTED_VOICES: Final = GEMINI_PREBUILT_VOICES | {"nova"} +REJECTED_PREFIX: Final = "No matching speaker voice found for name: " +NO_VOICE: Final = "" +PCM: Final = bytes(range(256)) * 4 +PCM_B64: Final = base64.b64encode(PCM).decode() +AUDIO_MIME: Final = "audio/L16;codec=pcm;rate=24000" +PROMPT_TOKENS: Final = 5 +AUDIO_TOKENS: Final = 50 +SPEND_BY_ID_SQL: Final = 'SELECT call_type, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id = %s' +SPEND_BY_KEY_SQL: Final = 'SELECT call_type, request_id FROM "LiteLLM_SpendLogs" WHERE api_key = %s' + + +def marker() -> str: + return f"tts-{uuid.uuid4().hex[:12]}" + + +def audio_param(voice: JsonValue) -> dict[str, JsonValue]: + return {"voice": voice, "format": "pcm16"} + + +def chat_body(model: str, voice: JsonValue, text: str, *, stream: bool = False) -> dict[str, JsonValue]: + return { + "model": model, + "messages": [{"role": "user", "content": text}], + "modalities": ["audio"], + "audio": audio_param(voice), + **({"stream": True} if stream else {}), + } + + +def received_voice(request: Request) -> JsonValue: + return voice_in(JSON_OBJECT.validate_json(request.body)) + + +def text_in(body: Mapping[str, JsonValue]) -> str: + contents: Final = body["contents"] + assert isinstance(contents, list) and contents, body + parts: Final = object_value(contents[0])["parts"] + assert isinstance(parts, list) and parts, body + return string_value(object_value(parts[0])["text"]) + + +def received_text(request: Request) -> str: + return text_in(JSON_OBJECT.validate_json(request.body)) + + +def _candidate(text: str) -> dict[str, JsonValue]: + return { + "content": { + "role": "model", + "parts": [{"text": text}, {"inlineData": {"mimeType": AUDIO_MIME, "data": PCM_B64}}], + }, + "finishReason": "STOP", + } + + +def _usage() -> dict[str, JsonValue]: + return { + "promptTokenCount": PROMPT_TOKENS, + "candidatesTokenCount": AUDIO_TOKENS, + "totalTokenCount": PROMPT_TOKENS + AUDIO_TOKENS, + "candidatesTokensDetails": [{"modality": "AUDIO", "tokenCount": AUDIO_TOKENS}], + } + + +def _rejection(voice: JsonValue) -> Reply: + error: Final = {"code": 400, "message": f"{REJECTED_PREFIX}{voice} and language: ", "status": "INVALID_ARGUMENT"} + return Reply(status=400, body=json.dumps({"error": error}).encode()) + + +def _accepted(voice: JsonValue) -> bool: + return voice == NO_VOICE or (isinstance(voice, str) and voice in ACCEPTED_VOICES) + + +def gemini_peer(request: Request) -> Reply: + assert request.method == "POST", request.method + voice: Final = received_voice(request) + if not _accepted(voice): + return _rejection(voice) + text: Final = received_text(request) + if "streamGenerateContent" in request.target: + first: Final = json.dumps({"candidates": [{"content": {"role": "model", "parts": [{"text": text}]}}]}) + last: Final = json.dumps({"candidates": [_candidate("")], "usageMetadata": _usage(), "modelVersion": BACKEND}) + return Reply( + content_type="text/event-stream", chunks=(f"data: {first}\n\n".encode(), f"data: {last}\n\n".encode()) + ) + body: Final = {"candidates": [_candidate(text)], "usageMetadata": _usage(), "modelVersion": BACKEND} + return Reply(body=json.dumps(body).encode()) + + +def gemini_deployment(scenario: Scenario, wire_url: str, **extra: JsonValue) -> str: + return scenario.model(model=f"gemini/{BACKEND}", api_base=wire_url, api_key=GEMINI_API_KEY, **extra) + + +def vertex_deployment( + scenario: Scenario, wire_url: str, token_url: str, *, api_base: str | None = None, **extra: JsonValue +) -> str: + return scenario.model( + model=f"vertex_ai/{BACKEND}", + api_base=api_base or f"{wire_url}{VERTEX_MODEL_PATH}", + api_key=None, + vertex_project=VERTEX_PROJECT, + vertex_location=VERTEX_LOCATION, + vertex_credentials=service_account_json(VERTEX_PROJECT, token_url.rstrip("/")), + **extra, + ) + + +def openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=60), + ) + + +def async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False, timeout=60), + ) + + +def anthropic_client(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic( + base_url=str(gateway.client.base_url), + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=60), + ) + + +def spend_row(response_id: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows(SPEND_BY_ID_SQL, (response_id,)), lambda values: len(values) == 1, seconds=70 + ) + return rows[0] + + +def spend_rows_for_key(key: str, count: int) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows(SPEND_BY_KEY_SQL, (sha256(key.encode()).hexdigest(),)), + lambda rows: len(rows) == count, + seconds=70, + ) + + +def error_message(response: httpx.Response) -> str: + error: Final = JSON_OBJECT.validate_json(response.content)["error"] + return string_value(object_value(error)["message"]) if isinstance(error, dict) else string_value(error) + + +def audio_data(payload: Mapping[str, JsonValue]) -> str: + choices: Final = payload["choices"] + assert isinstance(choices, list) and len(choices) == 1, payload + message: Final = object_value(object_value(choices[0])["message"]) + return string_value(object_value(message["audio"])["data"]) + + +LIVE_BACKEND: Final = "gemini-3.8-live" +LIVE_INPUT_TOKENS: Final = 7 +LIVE_OUTPUT_TOKENS: Final = 5 +SPEND_BY_CALL_SQL: Final = 'SELECT call_type, completion_tokens FROM "LiteLLM_SpendLogs" WHERE litellm_call_id = %s' + + +def voice_in(body: Mapping[str, JsonValue]) -> JsonValue: + generation: Final = body.get("generationConfig") + speech: Final = object_value(generation).get("speechConfig") if isinstance(generation, dict) else None + voice_config: Final = object_value(speech).get("voiceConfig") if isinstance(speech, dict) else None + if not isinstance(voice_config, dict): + return NO_VOICE + return object_value(voice_config["prebuiltVoiceConfig"])["voiceName"] + + +def spend_row_by_call(call_id: str) -> dict[str, JsonValue]: + assert call_id, "the response carried no x-litellm-call-id" + rows: Final = eventually( + lambda: read_rows(SPEND_BY_CALL_SQL, (call_id,)), lambda values: len(values) == 1, seconds=70 + ) + return rows[0] + + +def live_turn_script() -> RealtimeResponse: + return RealtimeResponse( + content_type="application/x-realtime", + events=( + {"serverContent": {"modelTurn": {"parts": [{"text": "scripted $REQUEST_ID"}]}}}, + { + "serverContent": {"turnComplete": True}, + "usageMetadata": { + "promptTokenCount": LIVE_INPUT_TOKENS, + "responseTokenCount": LIVE_OUTPUT_TOKENS, + "totalTokenCount": LIVE_INPUT_TOKENS + LIVE_OUTPUT_TOKENS, + }, + }, + ), + ) + + +def live_scenario(scenario: Scenario) -> ScenarioHandle: + handle: Final = register_scenario(f"tts-live-{uuid.uuid4().hex[:12]}", live_turn_script()) + scenario.cleanups.callback(delete_scenario, handle) + return handle + + +def live_deployment(scenario: Scenario, project: str, upstream_url: str) -> str: + return scenario.model( + model=f"vertex_ai/{LIVE_BACKEND}", + api_base=upstream_url.rstrip("/"), + api_key=None, + vertex_project=project, + vertex_location=VERTEX_LOCATION, + vertex_credentials=service_account_json(project, upstream_url.rstrip("/")), + model_info={"mode": "realtime"}, + ) + + +def _user_turn() -> str: + item: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "say the line"}]} + return json.dumps({"type": "conversation.item.create", "item": item}) + + +async def _live_frames(socket: ClientConnection, voice: str) -> AsyncIterator[dict[str, JsonValue]]: + try: + first: Final = JSON_OBJECT.validate_json(await socket.recv()) + yield first + if first.get("type") != "session.created": + return + await socket.send(json.dumps({"type": "session.update", "session": {"voice": voice}})) + await socket.send(_user_turn()) + async for message in socket: + event: Final = JSON_OBJECT.validate_json(message) + yield event + if event.get("type") == "response.done": + return + except ConnectionClosed as closed: + yield {"type": "closed", "code": 1006 if closed.rcvd is None else closed.rcvd.code} + + +async def _live_session(ws_base: str, model: str, key: str, voice: str) -> tuple[dict[str, JsonValue], ...]: + headers: Final = {"Authorization": f"Bearer {key}"} + async with websockets.connect(f"{ws_base}/v1/realtime?model={model}", additional_headers=headers) as socket: + return tuple([frame async for frame in _live_frames(socket, voice)]) + + +def live_turn(gateway: Gateway, model: str, key: str, voice: str) -> tuple[dict[str, JsonValue], ...]: + ws_base: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://") + return asyncio.run(asyncio.wait_for(_live_session(ws_base, model, key, voice), 60)) + + +def frame_types(frames: tuple[dict[str, JsonValue], ...]) -> tuple[str, ...]: + return tuple(string_value(frame["type"]) for frame in frames) + + +def live_setups(gateway: Gateway, project: str) -> tuple[dict[str, JsonValue], ...]: + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + observed: Final = tuple(map(object_value, upstream.get("/__observations").json()["requests"])) + frames: Final = tuple( + object_value(request["body"]) + for request in observed + if request["api_key"] == project and request["method"] == "WEBSOCKET_FRAME" + ) + assert all(request["path"] == GEMINI_LIVE_PATH for request in observed if request["api_key"] == project) + return tuple(object_value(frame["setup"]) for frame in frames if "setup" in frame) + + +HEALTH_DEFAULT: Final = "tts-health-default-voice" +HEALTH_NOVA: Final = "tts-health-nova" +HEALTH_FABLE: Final = "tts-health-fable" + + +def _health_model(name: str, **model_info: JsonValue) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": {"model": f"gemini/{BACKEND}", "api_key": GEMINI_API_KEY}, + "model_info": {"mode": "audio_speech", **model_info}, + } + + +def owned_config(directory: Path, model_list: Sequence[JsonValue]) -> Path: + config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + settings: Final = object_value(config["litellm_settings"]) + cache_params: Final = {**object_value(settings["cache_params"]), "namespace": f"tts-{uuid.uuid4().hex}"} + merged: Final = { + **config, + "model_list": list(model_list), + "litellm_settings": {**settings, "cache_params": cache_params}, + "router_settings": {**object_value(config["router_settings"]), "num_retries": 0}, + } + path: Final = directory / f"gemini-tts-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(merged)) + return path + + +def health_config(directory: Path) -> Path: + return owned_config( + directory, + ( + _health_model(HEALTH_DEFAULT), + _health_model(HEALTH_NOVA, health_check_voice="nova"), + _health_model(HEALTH_FABLE, health_check_voice="fable"), + ), + ) + + +def speech_deployment(scenario: Scenario, wire_url: str) -> str: + return scenario.model( + model=f"gemini/{BACKEND}", api_base=wire_url, api_key=GEMINI_API_KEY, model_info={"mode": "audio_speech"} + ) + + +def wav_payload(body: bytes) -> bytes: + assert body[:4] == b"RIFF" and body[8:12] == b"WAVE", body[:12] + return body[44:] + + +@dataclass(frozen=True, slots=True) +class ChunkedPeer: + url: str + received: SimpleQueue[Request] + + def drain(self) -> tuple[Request, ...]: + return tuple(self.received.get_nowait() for _ in range(self.received.qsize())) + + +def _read_chunked(handler: BaseHTTPRequestHandler) -> bytes: + def chunks() -> Generator[bytes, None, None]: + while True: + size: Final = int(handler.rfile.readline().split(b";", 1)[0].strip(), 16) + if size == 0: + handler.rfile.readline() + return + yield handler.rfile.read(size) + handler.rfile.readline() + + return b"".join(chunks()) + + +@contextmanager +def chunked_peer(respond: Callable[[Request], Reply]) -> Generator[ChunkedPeer, None, None]: + """Owned TCP peer like ``wire_server`` that also reads a ``transfer-encoding: chunked`` upload body, + the shape the Vertex files route streams a batch file to GCS in.""" + received: Final[SimpleQueue[Request]] = SimpleQueue() + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + chunked: Final = self.headers.get("transfer-encoding", "").lower() == "chunked" + body: Final = ( + _read_chunked(self) if chunked else self.rfile.read(int(self.headers.get("content-length", "0"))) + ) + request: Final = Request( + self.command, self.path, {name.lower(): value for name, value in self.headers.items()}, body + ) + received.put(request) + reply: Final = respond(request) + self.send_response(reply.status) + self.send_header("content-type", reply.content_type) + self.send_header("content-length", str(len(reply.body))) + self.send_header("connection", "close") + self.end_headers() + self.wfile.write(reply.body) + + def log_message(self, format: str, *args: object) -> None: + return + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}, daemon=True) + thread.start() + try: + yield ChunkedPeer(f"http://127.0.0.1:{server.server_address[1]}", received) + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) diff --git a/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py b/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py new file mode 100644 index 00000000000..3608e743506 --- /dev/null +++ b/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py @@ -0,0 +1,354 @@ +"""The speech route, the audio_speech health check and deferred Vertex Live setups map OpenAI voices. + +The speech bridge calls the Gemini chat path without the deployment's ``api_base`` (tracked separately), +so these rows boot their own proxy with ``GEMINI_API_BASE`` pointed at the scripted peer, and with +``LITELLM_GEMINI_LIVE_DEFER_SETUP`` so a client ``session.update`` voice builds the Vertex Live setup. +""" + +from __future__ import annotations + +import asyncio +import os +import re +import signal +import threading +from collections.abc import Callable, Iterator +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +from integration._support.client import ( + JSON_OBJECT, + Gateway, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._gemini_tts_voice_support import ( + BACKEND, + GEMINI_API_KEY, + HEALTH_DEFAULT, + HEALTH_FABLE, + HEALTH_NOVA, + NO_VOICE, + PCM, + REJECTED_PREFIX, + async_openai_client, + frame_types, + gemini_peer, + health_config, + live_deployment, + live_scenario, + live_setups, + live_turn, + marker, + openai_client, + received_text, + received_voice, + speech_deployment, + spend_row_by_call, + voice_in, + wav_payload, +) +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +BURST: Final = 12 +VOICES: Final = ("alloy", "fable", "Kore") +MAPPED: Final = MappingProxyType({"alloy": "Kore", "fable": "Umbriel", "Kore": "Kore"}) + + +def _overrides(wire_url: str) -> dict[str, str]: + return {"GEMINI_API_BASE": wire_url, "LITELLM_GEMINI_LIVE_DEFER_SETUP": "true"} + + +@pytest.fixture(scope="module") +def shared_wire() -> Iterator[Wire]: + with wire_server(gemini_peer) as wire: + yield wire + + +@pytest.fixture(scope="module") +def owned(shared_wire: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + directory: Final = tmp_path_factory.mktemp("gemini-tts") + with gateway_from_environment() as gateway: + config: Final = health_config(directory) + with owned_proxy_process(gateway, directory, _overrides(shared_wire.url), config=config, workers=2) as proxy: + yield proxy + + +def _speech_body(model: str, text: str, voice: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "input": text, "voice": voice, **extra} + + +def _received(wire: Wire, text: str) -> Request: + matching: Final = tuple(request for request in wire.drain() if received_text(request) == text) + assert len(matching) == 1, [request.target for request in matching] + return matching[0] + + +def _assert_speech_logged(call_id: str) -> None: + row: Final = spend_row_by_call(call_id) + assert row["call_type"] == "aspeech", row + + +def test_speech_route_maps_alloy_and_answers_wav(owned: OwnedProxy, shared_wire: Wire) -> None: + with owned.gateway.scenario() as scenario: + model: Final = speech_deployment(scenario, shared_wire.url) + text: Final = marker() + response: Final = owned.gateway.request("POST", "/v1/audio/speech", _speech_body(model, text, "alloy")) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("audio/wav"), response.headers + assert wav_payload(response.content) == PCM + assert received_voice(_received(shared_wire, text)) == "Kore" + _assert_speech_logged(response.headers.get("x-litellm-call-id", "")) + + +def test_speech_sdk_maps_alloy_and_answers_wav(owned: OwnedProxy, shared_wire: Wire) -> None: + with owned.gateway.scenario() as scenario: + model: Final = speech_deployment(scenario, shared_wire.url) + text: Final = marker() + raw: Final = openai_client(owned.gateway).audio.speech.with_raw_response.create( + model=model, input=text, voice="alloy" + ) + assert wav_payload(raw.content) == PCM + assert received_voice(_received(shared_wire, text)) == "Kore" + _assert_speech_logged(raw.headers.get("x-litellm-call-id", "")) + + +async def _async_speech(gateway: Gateway, model: str, text: str) -> tuple[bytes, str, str]: + client: Final = async_openai_client(gateway) + try: + raw: Final = await client.audio.speech.with_raw_response.create( + model=model, input=text, voice="Alloy", response_format="pcm" + ) + return raw.content, raw.headers.get("content-type", ""), raw.headers.get("x-litellm-call-id", "") + finally: + await client.close() + + +def test_async_speech_sdk_maps_a_capitalised_alloy_and_answers_raw_pcm(owned: OwnedProxy, shared_wire: Wire) -> None: + with owned.gateway.scenario() as scenario: + model: Final = speech_deployment(scenario, shared_wire.url) + text: Final = marker() + content, content_type, call_id = asyncio.run(_async_speech(owned.gateway, model, text)) + assert content == PCM, content[:16] + assert content_type.startswith("audio/pcm"), content_type + assert received_voice(_received(shared_wire, text)) == "Kore" + _assert_speech_logged(call_id) + + +@pytest.mark.parametrize("voice", ["nova", "Kore"]) +def test_speech_route_forwards_gemini_voice_names_verbatim(owned: OwnedProxy, shared_wire: Wire, voice: str) -> None: + with owned.gateway.scenario() as scenario: + model: Final = speech_deployment(scenario, shared_wire.url) + text: Final = marker() + response: Final = owned.gateway.request("POST", "/v1/audio/speech", _speech_body(model, text, voice)) + assert response.status_code == 200, response.text + assert wav_payload(response.content) == PCM + assert received_voice(_received(shared_wire, text)) == voice + _assert_speech_logged(response.headers.get("x-litellm-call-id", "")) + + +def _health(gateway: Gateway, model: str) -> dict[str, JsonValue]: + response: Final = gateway.request("GET", "/health", params={"model": model}) + return JSON_OBJECT.validate_json(response.content) + + +def _counts(report: dict[str, JsonValue]) -> tuple[JsonValue, JsonValue]: + return report["healthy_count"], report["unhealthy_count"] + + +def _only_voice(wire: Wire) -> JsonValue: + requests: Final = wire.drain() + assert len(requests) == 1, [request.target for request in requests] + return received_voice(requests[0]) + + +@pytest.mark.parametrize( + ("model", "mapped"), [(HEALTH_DEFAULT, "Kore"), (HEALTH_NOVA, "nova"), (HEALTH_FABLE, "Umbriel")] +) +def test_audio_speech_health_check_reaches_the_peer_with_a_gemini_voice( + owned: OwnedProxy, shared_wire: Wire, model: str, mapped: str +) -> None: + shared_wire.drain() + report: Final = _health(owned.gateway, model) + assert _counts(report) == (1, 0), report + assert _only_voice(shared_wire) == mapped + + +def test_unknown_health_check_voice_reports_the_vendor_rejection(owned: OwnedProxy, shared_wire: Wire) -> None: + with owned.gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"gemini/{BACKEND}", + api_base=shared_wire.url, + api_key=GEMINI_API_KEY, + model_info={"mode": "audio_speech", "health_check_voice": "Custom-Voice"}, + ) + shared_wire.drain() + report: Final = _health(owned.gateway, model) + assert _counts(report) == (0, 1), report + unhealthy: Final = report["unhealthy_endpoints"] + assert isinstance(unhealthy, list) and len(unhealthy) == 1, report + assert f"{REJECTED_PREFIX}Custom-Voice" in string_value(object_value(unhealthy[0])["error"]), report + assert _only_voice(shared_wire) == "Custom-Voice" + + +@pytest.mark.parametrize( + ("voice", "expected"), [("fable", "Umbriel"), ("alloy", NO_VOICE), ("Kore", "Kore")], ids=["fable", "alloy", "Kore"] +) +def test_deferred_vertex_live_setup_carries_the_session_voice(owned: OwnedProxy, voice: str, expected: str) -> None: + with owned.gateway.scenario() as scenario: + handle: Final = live_scenario(scenario) + key: Final = scenario.key() + model: Final = live_deployment(scenario, handle.scenario_id, owned.gateway.upstream_url) + frames: Final = live_turn(owned.gateway, model, key, voice) + types: Final = frame_types(frames) + assert types[0] == "session.created" and types[-1] == "response.done", frames + setups: Final = live_setups(owned.gateway, handle.scenario_id) + assert len(setups) == 1, setups + assert voice_in(setups[0]) == expected, setups[0] + + +@dataclass(frozen=True, slots=True) +class Sent: + voice: str + text: str + status: int + body: bytes + call_id: str + + +async def _fire(url: str, key: str, model: str, *, tolerate_transport_errors: bool = False) -> tuple[Sent, ...]: + async def one(client: httpx.AsyncClient, index: int) -> Sent: + voice: Final = VOICES[index % len(VOICES)] + text: Final = marker() + try: + response: Final = await client.post( + "/v1/audio/speech", json=_speech_body(model, text, voice), headers={"Authorization": f"Bearer {key}"} + ) + except httpx.TransportError as error: + if not tolerate_transport_errors: + raise + return Sent(voice, text, 0, repr(error).encode(), "") + return Sent(voice, text, response.status_code, response.content, response.headers.get("x-litellm-call-id", "")) + + async with httpx.AsyncClient(base_url=url, timeout=60, trust_env=False) as client: + return tuple(await asyncio.gather(*(one(client, index) for index in range(BURST)))) + + +def _held_peer(release: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert release.wait(timeout=120), "Held peer was never released" + return gemini_peer(request) + + return respond + + +def _held_connections(pid: int, wire_url: str) -> int: + port: Final = urlsplit(wire_url).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +def _by_text(received: tuple[Request, ...]) -> dict[str, tuple[Request, ...]]: + texts: Final = {received_text(request) for request in received} + return {text: tuple(request for request in received if received_text(request) == text) for text in texts} + + +def _assert_served(items: tuple[Sent, ...], received: tuple[Request, ...]) -> None: + by_text: Final = _by_text(received) + for item in items: + assert item.status == 200, (item.voice, item.body[:200]) + assert wav_payload(item.body) == PCM, item.text + assert len(by_text.get(item.text, ())) == 1, item.text + assert received_voice(by_text[item.text][0]) == MAPPED[item.voice], item.text + _assert_speech_logged(item.call_id) + + +def _worker_pids(proxy: OwnedProxy) -> tuple[int, ...]: + return eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(proxy.log.read_text())), + lambda pids: len(pids) == 2, + seconds=graceful_stop_seconds(), + ) + + +def test_worker_sigkill_mid_speech_burst_leaves_the_sibling_serving(gateway: Gateway, tmp_path: Path) -> None: + release: Final = threading.Event() + with wire_server(_held_peer(release)) as wire: + config: Final = health_config(tmp_path) + with owned_proxy_process(gateway, tmp_path, _overrides(wire.url), config=config, workers=2) as proxy: + workers: Final = _worker_pids(proxy) + url: Final = str(proxy.gateway.client.base_url) + with proxy.gateway.scenario() as scenario: + model: Final = speech_deployment(scenario, wire.url) + loop: Final = asyncio.new_event_loop() + burst: Final = loop.run_in_executor( + None, lambda: asyncio.run(_fire(url, proxy.gateway.key, model, tolerate_transport_errors=True)) + ) + try: + eventually(lambda: wire.received.qsize(), lambda size: size >= BURST, 90) + held_by: Final = MappingProxyType({pid: _held_connections(pid, wire.url) for pid in workers}) + victim: Final = max(workers, key=held_by.__getitem__) + os.kill(victim, signal.SIGKILL) + finally: + release.set() + served: Final = loop.run_until_complete(burst) + loop.close() + during: Final = wire.drain() + after: Final = asyncio.run(_fire(url, proxy.gateway.key, model)) + after_received: Final = wire.drain() + assert sum(held_by.values()) == BURST, held_by + assert held_by[victim] > 0, held_by + completed: Final = tuple(item for item in served if item.status == 200) + assert len(completed) == BURST - held_by[victim], (len(completed), held_by) + assert all(len(requests) == 1 for requests in _by_text(during).values()) + _assert_served(completed, during) + _assert_served(after, after_received) + + +def test_proxy_restart_mid_speech_burst_serves_the_next_burst_after_the_reboot( + gateway: Gateway, tmp_path: Path +) -> None: + release: Final = threading.Event() + with wire_server(_held_peer(release)) as wire, gateway.scenario() as scenario: + config: Final = health_config(tmp_path) + model: Final = speech_deployment(scenario, wire.url) + with owned_proxy_process(gateway, tmp_path, _overrides(wire.url), config=config, workers=2) as first: + url: Final = str(first.gateway.client.base_url) + loop: Final = asyncio.new_event_loop() + burst: Final = loop.run_in_executor( + None, lambda: asyncio.run(_fire(url, first.gateway.key, model, tolerate_transport_errors=True)) + ) + try: + eventually(lambda: wire.received.qsize(), lambda size: size >= BURST, 90) + first.process.terminate() + finally: + release.set() + served: Final = loop.run_until_complete(burst) + loop.close() + during: Final = wire.drain() + with owned_proxy_process(gateway, tmp_path, _overrides(wire.url), config=config, workers=2) as second: + after: Final = asyncio.run(_fire(str(second.gateway.client.base_url), second.gateway.key, model)) + after_received: Final = wire.drain() + completed: Final = tuple(item for item in served if item.status == 200) + assert all(len(requests) == 1 for requests in _by_text(during).values()) + _assert_served(completed, during) + for item in served: + if item.status != 200: + assert item.status == 0 or b'"error"' in item.body or item.body == b"", (item.status, item.body[:200]) + _assert_served(after, after_received) diff --git a/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py b/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py new file mode 100644 index 00000000000..1241c2a7f49 --- /dev/null +++ b/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py @@ -0,0 +1,652 @@ +"""Gemini and Vertex TTS deployments map OpenAI voice names onto Gemini prebuilt voices. + +The scripted peer on 127.0.0.1 answers an audio part when the ``voiceName`` it receives is one of +Gemini's prebuilt voices and the vendor's ``No matching speaker voice found`` 400 otherwise, so a row +is green only when the voice the peer received is the mapped one. Every row runs through the rig proxy. +""" + +from __future__ import annotations + +import asyncio +import base64 +import dataclasses +import json +import socket +import threading +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._gemini_tts_voice_support import ( + AUDIO_TOKENS, + BACKEND, + GEMINI_API_KEY, + NO_VOICE, + PCM_B64, + REJECTED_PREFIX, + VERTEX_MODEL_PATH, + anthropic_client, + async_openai_client, + audio_data, + audio_param, + chat_body, + chunked_peer, + error_message, + frame_types, + gemini_deployment, + gemini_peer, + live_deployment, + live_scenario, + live_setups, + live_turn, + marker, + openai_client, + received_text, + received_voice, + spend_row_by_call, + text_in, + vertex_deployment, + voice_in, +) +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(180) + +BUCKET: Final = "integration-tts-bucket" +FIVE_KB: Final = "v" * 5120 +BURST: Final = 24 +SLOW_BURST: Final = 12 +OUTAGE_AFTER: Final = 6 +SLOW_SECONDS: Final = 1.0 + + +@dataclass(frozen=True, slots=True) +class Answer: + status: int + text: str + response_id: str + audio: str | None + content: str + call_id: str + + +ChatClient = Callable[..., Answer] + + +def _failed(status: int, text: str) -> Answer: + return Answer(status, text, "", None, "", "") + + +def _delta_text(chunk: Mapping[str, JsonValue]) -> str: + choices: Final = chunk["choices"] + if not isinstance(choices, list) or not choices: + return "" + content: Final = object_value(object_value(choices[0])["delta"]).get("content") + return content if isinstance(content, str) else "" + + +def _message_text(body: Mapping[str, JsonValue]) -> str: + choices: Final = body["choices"] + assert isinstance(choices, list) and len(choices) == 1, body + content: Final = object_value(object_value(choices[0])["message"]).get("content") + return content if isinstance(content, str) else "" + + +def _sse_chunks(text: str) -> tuple[dict[str, JsonValue], ...]: + lines: Final = tuple(line[6:] for line in text.splitlines() if line.startswith("data: ")) + return tuple(JSON_OBJECT.validate_json(line) for line in lines if line != "[DONE]") + + +def _error_text(response: httpx.Response) -> str: + response.read() + return response.text + + +async def _async_error_text(response: httpx.Response) -> str: + await response.aread() + return response.text + + +def _httpx_chat(gateway: Gateway, model: str, voice: JsonValue, text: str, *, stream: bool) -> Answer: + response: Final = gateway.request("POST", "/v1/chat/completions", chat_body(model, voice, text, stream=stream)) + call_id: Final = response.headers.get("x-litellm-call-id", "") + if response.status_code != 200: + return _failed(response.status_code, response.text) + if stream: + chunks: Final = _sse_chunks(response.text) + content: Final = "".join(_delta_text(chunk) for chunk in chunks) + return Answer(200, response.text, "", None, content, call_id) + body: Final = JSON_OBJECT.validate_json(response.content) + return Answer(200, response.text, string_value(body["id"]), audio_data(body), _message_text(body), call_id) + + +def _sdk_chat(gateway: Gateway, model: str, voice: str, text: str, *, stream: bool) -> Answer: + client: Final = openai_client(gateway) + messages: Final = [{"role": "user", "content": text}] + try: + if stream: + raw: Final = client.chat.completions.with_raw_response.create( + model=model, + messages=messages, + modalities=["audio"], + audio={"voice": voice, "format": "pcm16"}, + stream=True, + ) + chunks: Final = list(raw.parse()) + content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) + return Answer(200, "", "", None, content, raw.headers.get("x-litellm-call-id", "")) + completed: Final = client.chat.completions.with_raw_response.create( + model=model, messages=messages, modalities=["audio"], audio={"voice": voice, "format": "pcm16"} + ) + completion: Final = completed.parse() + message: Final = completion.choices[0].message + audio: Final = message.audio.data if message.audio is not None else None + return Answer( + 200, + completion.model_dump_json(), + completion.id, + audio, + message.content or "", + completed.headers.get("x-litellm-call-id", ""), + ) + except openai.APIStatusError as error: + return _failed(error.status_code, _error_text(error.response)) + + +async def _async_sdk_chat(gateway: Gateway, model: str, voice: str, text: str, *, stream: bool) -> Answer: + client: Final = async_openai_client(gateway) + messages: Final = [{"role": "user", "content": text}] + try: + if stream: + raw: Final = await client.chat.completions.with_raw_response.create( + model=model, + messages=messages, + modalities=["audio"], + audio={"voice": voice, "format": "pcm16"}, + stream=True, + ) + chunks: Final = [chunk async for chunk in raw.parse()] + content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) + return Answer(200, "", "", None, content, raw.headers.get("x-litellm-call-id", "")) + completed: Final = await client.chat.completions.with_raw_response.create( + model=model, messages=messages, modalities=["audio"], audio={"voice": voice, "format": "pcm16"} + ) + completion: Final = completed.parse() + message: Final = completion.choices[0].message + audio: Final = message.audio.data if message.audio is not None else None + return Answer( + 200, + completion.model_dump_json(), + completion.id, + audio, + message.content or "", + completed.headers.get("x-litellm-call-id", ""), + ) + except openai.APIStatusError as error: + return _failed(error.status_code, await _async_error_text(error.response)) + finally: + await client.close() + + +def _sdk_async_chat(gateway: Gateway, model: str, voice: str, text: str, *, stream: bool) -> Answer: + return asyncio.run(_async_sdk_chat(gateway, model, voice, text, stream=stream)) + + +CLIENTS: Final[Mapping[str, ChatClient]] = {"httpx": _httpx_chat, "sdk": _sdk_chat, "sdk-async": _sdk_async_chat} + + +def _deployment(scenario: Scenario, provider: str, wire_url: str, upstream_url: str, **extra: JsonValue) -> str: + if provider == "gemini": + return gemini_deployment(scenario, wire_url, **extra) + return vertex_deployment(scenario, wire_url, upstream_url, **extra) + + +def _received(wire: Wire, text: str) -> Request: + matching: Final = tuple(request for request in wire.drain() if received_text(request) == text) + assert len(matching) == 1, [request.target for request in matching] + return matching[0] + + +def _assert_target(request: Request, provider: str, *, stream: bool) -> None: + suffix: Final = ":streamGenerateContent?alt=sse" if stream else ":generateContent" + if provider == "gemini": + assert request.target == f"/models/{BACKEND}{suffix}", request.target + assert request.headers["x-goog-api-key"] == GEMINI_API_KEY, request.headers + return + assert request.target == f"{VERTEX_MODEL_PATH}{suffix}", request.target + assert request.headers["authorization"] == "Bearer scripted-token", request.headers + + +def _assert_logged(answer: Answer) -> None: + row: Final = spend_row_by_call(answer.call_id) + assert (row["call_type"], row["completion_tokens"]) == ("acompletion", AUDIO_TOKENS), row + + +MAPPED: Final = ( + pytest.param("gemini", "httpx", False, "alloy", "Kore", id="gemini-httpx-alloy"), + pytest.param("vertex", "httpx", False, "alloy", "Kore", id="vertex-httpx-alloy"), + pytest.param("gemini", "sdk", False, "ALLOY", "Kore", id="gemini-sdk-ALLOY"), + pytest.param("gemini", "sdk-async", False, "fable", "Umbriel", id="gemini-sdk-async-fable"), + pytest.param("vertex", "sdk", False, "onyx", "Orus", id="vertex-sdk-onyx"), + pytest.param("gemini", "httpx", True, "echo", "Charon", id="gemini-httpx-stream-echo"), + pytest.param("vertex", "sdk", True, "onyx", "Orus", id="vertex-sdk-stream-onyx"), + pytest.param("vertex", "sdk-async", True, "alloy", "Kore", id="vertex-sdk-async-stream-alloy"), +) + + +@pytest.mark.parametrize(("provider", "client", "stream", "voice", "mapped"), MAPPED) +def test_openai_voice_reaches_the_peer_as_the_gemini_voice( + gateway: Gateway, provider: str, client: str, stream: bool, voice: str, mapped: str +) -> None: + with wire_server(gemini_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, provider, wire.url, gateway.upstream_url) + text: Final = marker() + answer: Final = CLIENTS[client](gateway, model, voice, text, stream=stream) + assert answer.status == 200, answer.text + request: Final = _received(wire, text) + _assert_target(request, provider, stream=stream) + assert received_voice(request) == mapped, request.body + assert answer.content == text, answer.text + assert answer.audio == (None if stream else PCM_B64), answer.text + _assert_logged(answer) + + +@pytest.mark.parametrize(("provider", "voice"), [("gemini", "Kore"), ("vertex", "nova")]) +def test_gemini_voice_names_pass_through_verbatim(gateway: Gateway, provider: str, voice: str) -> None: + with wire_server(gemini_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, provider, wire.url, gateway.upstream_url) + text: Final = marker() + answer: Final = _httpx_chat(gateway, model, voice, text, stream=False) + assert answer.status == 200, answer.text + assert received_voice(_received(wire, text)) == voice + assert answer.audio == PCM_B64, answer.text + _assert_logged(answer) + + +def _raw_chat(gateway: Gateway, body: str) -> httpx.Response: + headers: Final = {"Authorization": f"Bearer {gateway.key}", "content-type": "application/json"} + return gateway.client.post("/v1/chat/completions", content=body.encode(), headers=headers) + + +def _raw_body(model: str, text: str, audio_literal: str) -> str: + return ( + f'{{"model": {json.dumps(model)}, "messages": [{{"role": "user", "content": {json.dumps(text)}}}], ' + f'"modalities": ["audio"], "audio": {audio_literal}}}' + ) + + +def test_vendor_rejection_of_an_unknown_voice_reaches_the_caller_and_the_proxy_keeps_serving(gateway: Gateway) -> None: + with wire_server(gemini_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_deployment(scenario, wire.url) + text: Final = marker() + refused: Final = gateway.request("POST", "/v1/chat/completions", chat_body(model, "Custom-Voice", text)) + assert refused.status_code == 400, refused.text + assert f"{REJECTED_PREFIX}Custom-Voice" in error_message(refused), refused.text + assert received_voice(_received(wire, text)) == "Custom-Voice" + follow_up: Final = marker() + served: Final = _httpx_chat(gateway, model, "Kore", follow_up, stream=False) + assert served.status == 200, served.text + assert received_voice(_received(wire, follow_up)) == "Kore" + + +MALFORMED: Final = ( + pytest.param('{"voice": 7, "format": "pcm16"}', "7", id="int"), + pytest.param('{"voice": null, "format": "pcm16"}', "None", id="null"), + pytest.param('{"voice": "", "format": "pcm16"}', "", id="empty"), + pytest.param(f'{{"voice": "{FIVE_KB}", "format": "pcm16"}}', FIVE_KB, id="5kb"), + pytest.param('{"voice": ["alloy"], "format": "pcm16"}', "['alloy']", id="list"), +) + + +@pytest.mark.parametrize(("audio_literal", "rejected_as"), MALFORMED) +def test_malformed_voice_values_are_refused_without_taking_the_proxy_down( + gateway: Gateway, audio_literal: str, rejected_as: str +) -> None: + with wire_server(gemini_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_deployment(scenario, wire.url) + text: Final = marker() + refused: Final = _raw_chat(gateway, _raw_body(model, text, audio_literal)) + assert refused.status_code == 400, refused.text + assert f"{REJECTED_PREFIX}{rejected_as} and language: " in error_message(refused), refused.text + received: Final = _received(wire, text) + assert received_voice(received) == json.loads(audio_literal)["voice"], received.body + follow_up: Final = marker() + served: Final = _httpx_chat(gateway, model, "alloy", follow_up, stream=False) + assert served.status == 200, served.text + assert received_voice(_received(wire, follow_up)) == "Kore" + + +def test_duplicated_voice_key_keeps_the_last_value(gateway: Gateway) -> None: + with wire_server(gemini_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_deployment(scenario, wire.url) + text: Final = marker() + response: Final = _raw_chat( + gateway, _raw_body(model, text, '{"voice": "alloy", "voice": "fable", "format": "pcm16"}') + ) + assert response.status_code == 200, response.text + assert received_voice(_received(wire, text)) == "Umbriel" + assert audio_data(JSON_OBJECT.validate_json(response.content)) == PCM_B64 + + +def test_audio_without_a_voice_sends_no_voice_config(gateway: Gateway) -> None: + with wire_server(gemini_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_deployment(scenario, wire.url) + text: Final = marker() + body: Final = {**chat_body(model, "alloy", text), "audio": {"format": "pcm16"}} + response: Final = gateway.request("POST", "/v1/chat/completions", body) + assert response.status_code == 200, response.text + assert received_voice(_received(wire, text)) == NO_VOICE + assert audio_data(JSON_OBJECT.validate_json(response.content)) == PCM_B64 + + +def _messages_bridge(gateway: Gateway, model: str, text: str, *, stream: bool) -> Answer: + client: Final = anthropic_client(gateway) + try: + if stream: + raw: Final = client.messages.with_raw_response.create( + model=model, max_tokens=64, messages=[{"role": "user", "content": text}], stream=True + ) + collected: Final = "".join( + event.delta.text + for event in raw.parse() + if event.type == "content_block_delta" and event.delta.type == "text_delta" + ) + return Answer(200, "", "", None, collected, raw.headers.get("x-litellm-call-id", "")) + completed: Final = client.messages.with_raw_response.create( + model=model, max_tokens=64, messages=[{"role": "user", "content": text}] + ) + message: Final = completed.parse() + content: Final = "".join(block.text for block in message.content if block.type == "text") + return Answer( + 200, message.model_dump_json(), message.id, None, content, completed.headers.get("x-litellm-call-id", "") + ) + except anthropic.APIStatusError as error: + return _failed(error.status_code, _error_text(error.response)) + + +def _responses_bridge(gateway: Gateway, model: str, text: str, *, stream: bool) -> Answer: + client: Final = openai_client(gateway) + try: + if stream: + raw: Final = client.responses.with_raw_response.create(model=model, input=text, stream=True) + events: Final = list(raw.parse()) + collected: Final = "".join(event.delta for event in events if event.type == "response.output_text.delta") + assert events[-1].type == "response.completed", [event.type for event in events] + return Answer(200, "", "", None, collected, raw.headers.get("x-litellm-call-id", "")) + completed: Final = client.responses.with_raw_response.create(model=model, input=text) + response: Final = completed.parse() + return Answer( + 200, + response.model_dump_json(), + response.id, + None, + response.output_text, + completed.headers.get("x-litellm-call-id", ""), + ) + except openai.APIStatusError as error: + return _failed(error.status_code, _error_text(error.response)) + + +def _completions_bridge(gateway: Gateway, model: str, text: str, *, stream: bool) -> Answer: + response: Final = gateway.request("POST", "/v1/completions", {"model": model, "prompt": text, "max_tokens": 64}) + if response.status_code != 200: + return _failed(response.status_code, response.text) + body: Final = JSON_OBJECT.validate_json(response.content) + choices: Final = body["choices"] + assert isinstance(choices, list) and len(choices) == 1, response.text + return Answer( + 200, + response.text, + string_value(body["id"]), + None, + string_value(object_value(choices[0])["text"]), + response.headers.get("x-litellm-call-id", ""), + ) + + +BRIDGES: Final[Mapping[str, Callable[..., Answer]]] = { + "messages": _messages_bridge, + "responses": _responses_bridge, + "completions": _completions_bridge, +} + + +@pytest.mark.parametrize( + ("bridge", "stream"), + [("messages", False), ("messages", True), ("responses", False), ("responses", True), ("completions", False)], + ids=["messages", "messages-stream", "responses", "responses-stream", "completions"], +) +def test_bridged_endpoints_carry_the_deployment_voice(gateway: Gateway, bridge: str, stream: bool) -> None: + with wire_server(gemini_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_deployment(scenario, wire.url, audio=audio_param("alloy"), modalities=["audio"]) + text: Final = marker() + answer: Final = BRIDGES[bridge](gateway, model, text, stream=stream) + assert answer.status == 200, answer.text + request: Final = _received(wire, text) + _assert_target(request, "gemini", stream=stream) + assert received_voice(request) == "Kore", request.body + assert answer.content == text, answer.text + row: Final = spend_row_by_call(answer.call_id) + assert row["completion_tokens"] == AUDIO_TOKENS, row + + +def _gcs_peer(request: Request) -> Reply: + assert request.method == "POST", request.method + assert request.target.startswith(f"/upload/storage/v1/b/{BUCKET}/o?uploadType=media&name="), request.target + assert request.headers["authorization"] == "Bearer scripted-token", request.headers + name: Final = request.target.rsplit("name=", 1)[1] + body: Final = { + "kind": "storage#object", + "id": f"{BUCKET}/{name}/1", + "name": name, + "bucket": BUCKET, + "size": str(len(request.body)), + "timeCreated": "2026-10-08T00:00:00.000Z", + "contentType": "application/json", + } + return Reply(body=json.dumps(body).encode()) + + +def _batch_line(custom_id: str, text: str, voice: str) -> str: + body: Final = { + "model": BACKEND, + "messages": [{"role": "user", "content": text}], + "modalities": ["audio"], + "audio": audio_param(voice), + } + return json.dumps({"custom_id": custom_id, "method": "POST", "url": "/v1/chat/completions", "body": body}) + + +def _uploaded_rows(request: Request) -> tuple[dict[str, JsonValue], ...]: + lines: Final = tuple(line for line in request.body.decode().splitlines() if line.strip()) + return tuple(object_value(JSON_OBJECT.validate_json(line)["request"]) for line in lines) + + +def _decoded_file_id(file_id: str) -> tuple[str, str]: + assert file_id.startswith("file-"), file_id + encoded: Final = file_id[len("file-") :] + decoded: Final = base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)).decode() + assert decoded.startswith("litellm:"), decoded + raw, model = decoded[len("litellm:") :].rsplit(";model,", 1) + return raw, model + + +def test_vertex_batch_file_rows_carry_the_mapped_voice(gateway: Gateway) -> None: + with chunked_peer(_gcs_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_deployment( + scenario, wire.url, gateway.upstream_url, api_base=wire.url, gcs_bucket_name=BUCKET + ) + first: Final = marker() + second: Final = marker() + content: Final = f"{_batch_line('row-1', first, 'alloy')}\n{_batch_line('row-2', second, 'Kore')}\n".encode() + response: Final = gateway.request_multipart( + "/v1/files", {"purpose": "batch", "model": model}, {"file": ("rows.jsonl", content, "application/jsonl")} + ) + assert response.status_code == 200, response.text + uploaded: Final = JSON_OBJECT.validate_json(response.content) + requests: Final = wire.drain() + assert len(requests) == 1, [request.target for request in requests] + assert (uploaded["object"], uploaded["purpose"], uploaded["bytes"]) == ("file", "batch", len(requests[0].body)) + raw, file_model = _decoded_file_id(string_value(uploaded["id"])) + assert raw.startswith(f"gs://{BUCKET}/") and file_model == model, response.text + rows: Final = _uploaded_rows(requests[0]) + assert tuple(voice_in(row) for row in rows) == ("Kore", "Kore"), requests[0].body + assert tuple(text_in(row) for row in rows) == (first, second) + + +def test_vertex_live_default_mode_sends_one_setup_without_the_client_voice(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = live_scenario(scenario) + key: Final = scenario.key() + model: Final = live_deployment(scenario, handle.scenario_id, gateway.upstream_url) + frames: Final = live_turn(gateway, model, key, "fable") + types: Final = frame_types(frames) + assert types[0] == "session.created" and types[-1] == "response.done", frames + setups: Final = live_setups(gateway, handle.scenario_id) + assert len(setups) == 1, setups + assert voice_in(setups[0]) == NO_VOICE, setups[0] + + +def _free_port() -> int: + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + return int(reserve.getsockname()[1]) + + +@dataclass(frozen=True, slots=True) +class Fired: + provider: str + stream: bool + voice: str + text: str + status: int + text_body: str + call_id: str + + +VOICES: Final = ("alloy", "fable", "Kore") + + +async def _fire( + gateway: Gateway, models: Mapping[str, str], count: int, *, tolerate_transport_errors: bool = False +) -> tuple[Fired, ...]: + async def one(client: httpx.AsyncClient, index: int) -> Fired: + provider: Final = "gemini" if index % 2 == 0 else "vertex" + stream: Final = index % 4 >= 2 + voice: Final = VOICES[index % len(VOICES)] + text: Final = marker() + body: Final = chat_body(models[provider], voice, text, stream=stream) + try: + response: Final = await client.post( + "/v1/chat/completions", json=body, headers={"Authorization": f"Bearer {gateway.key}"} + ) + except httpx.TransportError as error: + if not tolerate_transport_errors: + raise + return Fired(provider, stream, voice, text, 0, repr(error), "") + call_id: Final = response.headers.get("x-litellm-call-id", "") + if response.status_code != 200: + return Fired(provider, stream, voice, text, response.status_code, response.text, call_id) + content: Final = ( + "".join(_delta_text(chunk) for chunk in _sse_chunks(response.text)) + if stream + else _message_text(JSON_OBJECT.validate_json(response.content)) + ) + assert content == text, response.text + return Fired(provider, stream, voice, text, 200, response.text, call_id) + + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client: + return tuple(await asyncio.gather(*(one(client, index) for index in range(count)))) + + +def _by_text(received: tuple[Request, ...]) -> dict[str, tuple[Request, ...]]: + texts: Final = {received_text(request) for request in received} + return {text: tuple(request for request in received if received_text(request) == text) for text in texts} + + +def _expected_voice(voice: str) -> str: + return {"alloy": "Kore", "fable": "Umbriel"}.get(voice, voice) + + +def _assert_served(items: tuple[Fired, ...], received: tuple[Request, ...]) -> None: + by_text: Final = _by_text(received) + for item in items: + assert item.status == 200, (item.provider, item.voice, item.text_body) + assert len(by_text.get(item.text, ())) == 1, item.text + assert received_voice(by_text[item.text][0]) == _expected_voice(item.voice), item.text + row: Final = spend_row_by_call(item.call_id) + assert row["call_type"] == "acompletion", row + + +def test_peer_outage_mid_burst_ends_each_request_once_and_recovers_on_the_same_port(gateway: Gateway) -> None: + port: Final = _free_port() + with gateway.scenario() as scenario: + url: Final = f"http://127.0.0.1:{port}" + models: Final = { + "gemini": gemini_deployment(scenario, url), + "vertex": vertex_deployment(scenario, url, gateway.upstream_url), + } + with wire_server(gemini_peer, port=port) as wire: + healthy: Final = asyncio.run(_fire(gateway, models, BURST)) + _assert_served(healthy, wire.drain()) + racing: Final = _start_in_thread(lambda: asyncio.run(_fire(gateway, models, BURST))) + eventually(lambda: wire.received.qsize(), lambda size: size >= OUTAGE_AFTER, 30) + down: Final = asyncio.run(_fire(gateway, models, BURST)) + raced: Final = racing() + raced_received: Final = wire.drain() + with wire_server(gemini_peer, port=port) as restarted: + recovered: Final = asyncio.run(_fire(gateway, models, BURST)) + recovered_received: Final = restarted.drain() + assert all(len(requests) == 1 for requests in _by_text(raced_received).values()) + _assert_served(tuple(item for item in raced if item.status == 200), raced_received) + lost: Final = tuple(item for item in raced if item.status != 200) + for item in (*lost, *down): + assert item.status >= 500, (item.status, item.text_body) + assert '"error"' in item.text_body, item.text_body + assert all(item.status != 200 for item in down), [item.status for item in down] + _assert_served(recovered, recovered_received) + assert {received_text(request) for request in recovered_received} == {item.text for item in recovered} + + +def _start_in_thread(work: Callable[[], tuple[Fired, ...]]) -> Callable[[], tuple[Fired, ...]]: + results: Final[list[tuple[Fired, ...]]] = [] + thread: Final = threading.Thread(target=lambda: results.append(work())) + thread.start() + + def finish() -> tuple[Fired, ...]: + thread.join(timeout=120) + assert not thread.is_alive(), "the racing burst never finished" + return results[0] + + return finish + + +def _slow_peer(request: Request) -> Reply: + reply: Final = gemini_peer(request) + if reply.chunks is not None: + return dataclasses.replace(reply, pause_between_chunks=SLOW_SECONDS) + time.sleep(SLOW_SECONDS) + return reply + + +def test_slow_peer_under_a_burst_answers_every_request_once_while_liveliness_stays_up(gateway: Gateway) -> None: + with wire_server(_slow_peer) as wire, gateway.scenario() as scenario: + models: Final = { + "gemini": gemini_deployment(scenario, wire.url), + "vertex": vertex_deployment(scenario, wire.url, gateway.upstream_url), + } + burst: Final = _start_in_thread(lambda: asyncio.run(_fire(gateway, models, SLOW_BURST))) + eventually(lambda: wire.received.qsize(), lambda size: size >= SLOW_BURST, 30) + liveliness: Final = gateway.request("GET", "/health/liveliness") + assert liveliness.status_code == 200, liveliness.text + served: Final = burst() + received: Final = wire.drain() + assert len(received) == SLOW_BURST, [received_text(request) for request in received] + _assert_served(served, received) diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 65ac3bc544c..7b2a639776a 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -8899,3 +8899,49 @@ def load_vertex_ai_credentials(): # Export the temporary file as GOOGLE_APPLICATION_CREDENTIALS os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name) + + +# Gemini prebuilt TTS voices per https://ai.google.dev/gemini-api/docs/speech-generation (page dated 2026-10-01, read 2026-10-07) +GEMINI_TTS_PREBUILT_VOICES: Final[frozenset[str]] = frozenset( + { + "Zephyr", "Puck", "Charon", "Kore", "Fenrir", "Leda", "Orus", "Aoede", "Callirrhoe", "Autonoe", + "Enceladus", "Iapetus", "Umbriel", "Algieba", "Despina", "Erinome", "Algenib", "Rasalgethi", + "Laomedeia", "Achernar", "Alnilam", "Schedar", "Gacrux", "Pulcherrima", "Achird", "Zubenelgenubi", + "Vindemiatrix", "Sadachbia", "Sadaltager", "Sulafat", + } +) +OPENAI_TTS_VOICES: Final[tuple[str, ...]] = ( + "alloy", "ash", "ballad", "cedar", "coral", "echo", "fable", "marin", "onyx", "sage", "shimmer", "verse", +) + + +def _mapped_voice_name(audio: dict[str, str]) -> str: + return VertexGeminiConfig()._map_audio_params(audio)["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] + + +@pytest.mark.parametrize("openai_voice", OPENAI_TTS_VOICES) +def test_map_audio_params_maps_openai_voice_names_to_gemini_prebuilt_voices(openai_voice: str) -> None: + """Regression for the audio_speech health check default voice (alloy) 400ing on Gemini TTS models.""" + mapped: Final = _mapped_voice_name({"voice": openai_voice, "format": "pcm16"}) + assert mapped in GEMINI_TTS_PREBUILT_VOICES + assert _mapped_voice_name({"voice": openai_voice.capitalize()}) == mapped + + +def test_map_audio_params_maps_the_health_check_default_voice_to_kore() -> None: + assert _mapped_voice_name({"voice": "alloy", "format": "pcm16"}) == "Kore" + + +# nova is accepted by Gemini 3.x TTS models as given (live check 2026-10-08), so it must not be remapped +@pytest.mark.parametrize("gemini_voice", ("Kore", "Puck", "Sulafat", "nova", "custom-voice")) +def test_map_audio_params_passes_non_openai_voice_names_through(gemini_voice: str) -> None: + assert _mapped_voice_name({"voice": gemini_voice, "format": "pcm16"}) == gemini_voice + + +def test_map_openai_params_audio_voice_reaches_speech_config_mapped() -> None: + optional_params: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"modalities": ["audio"], "audio": {"voice": "alloy", "format": "pcm16"}}, + optional_params={}, + model="gemini-3.8-flash-tts", + drop_params=False, + ) + assert optional_params["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" From 546402c98c7ecea5f32d263fadd3c95bd9b934ff Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 19:14:48 -0700 Subject: [PATCH 018/104] fix(decisions)!: refuse safety_identifier on providers that cannot take it unless drop_params drops it (#44955) * feat(decisions): add the OpenAI Decisions spec types and the System One translation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(decisions): share one DecisionsModel config and require model and usage on responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(decisions): rename the shared pydantic parent to DecisionsObjectBase Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(decisions): dispatch /v1/decisions through provider configs and the shared HTTP handler Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(decisions): map OpenRouter connection failures to APIConnectionError and drop explanatory docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(decisions): switch /v1/decisions to the OpenAI Decisions schema Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(decisions): build the usage the OpenAI response schema now requires Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(lens): ask signal questions through the OpenAI Decisions schema Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(decisions): refuse safety_identifier for System One providers unless drop_params * feat(decisions): refuse safety_identifier on providers that cannot take it unless drop_params drops it * test(decisions): cover hosted_vllm in the safety_identifier and per-provider wire tests * fix(decisions): check a System One request's safety_identifier against the provider too * fix(decisions): drop a non-string safety_identifier under drop_params A malformed safety_identifier now follows the drop_params convention in both body shapes: it answers 400 without drop_params and is dropped before the provider call with it. A string identifier on OpenAI stays on the wire either way. * fix(decisions): let /v1/decisions drop a non-string safety_identifier under drop_params The route checked the whole body before routing, so a non-string safety_identifier answered 400 even when the deployment or the body set drop_params. The route now leaves that field to the Decisions call, which knows every drop_params source * refactor(decisions): return the unsupported safety_identifier as a value and raise it in _prepare_call --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/decisions/main.py | 66 +- .../llms/base_llm/decisions/transformation.py | 1 + .../llms/openai/decisions/transformation.py | 1 + .../proxy/decisions_endpoints/endpoints.py | 7 +- .../test_databricks_decisions_wire.py | 1 - .../test_decisions_openai_format_wire.py | 586 ++++++++++++++++++ .../translation/decisions/bases/cloudflare.py | 97 +++ .../translation/decisions/bases/openrouter.py | 94 +++ .../translation/decisions/bases/perplexity.py | 95 +++ .../decisions/bases/strands_decider.py | 34 + .../translation/decisions/bases/typesafe.py | 92 +++ .../basic/test_decisions_basic_cloudflare.py | 4 +- .../basic/test_decisions_basic_openrouter.py | 9 +- .../basic/test_decisions_basic_perplexity.py | 9 +- .../test_decisions_basic_strands_decider.py | 11 +- .../basic/test_decisions_basic_typesafe.py | 4 +- tests/unit/decisions/test_main.py | 135 ++++ .../decisions_endpoints/test_endpoints.py | 62 ++ 18 files changed, 1289 insertions(+), 19 deletions(-) create mode 100644 tests/integration/providers/test_decisions_openai_format_wire.py diff --git a/litellm/decisions/main.py b/litellm/decisions/main.py index 4a4af2b580e..032cad00454 100644 --- a/litellm/decisions/main.py +++ b/litellm/decisions/main.py @@ -1,13 +1,13 @@ from collections.abc import Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Final, TypeAlias import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError from typing_extensions import assert_never import litellm -from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER +from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER, normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.decisions.transformation import ( @@ -40,6 +40,9 @@ DecisionsRequestFormat: TypeAlias = DecisionsRequestBody | OpenAIDecisionRequest _SYSTEMONE_REQUEST_ADAPTER: Final[TypeAdapter[DecisionsRequestBody]] = TypeAdapter(DecisionsRequestBody) _OPENAI_REQUEST_ADAPTER: Final[TypeAdapter[OpenAIDecisionRequestBody]] = TypeAdapter(OpenAIDecisionRequestBody) +_SAFETY_IDENTIFIER_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter( + str | None, config=ConfigDict(title="safety_identifier") +) _HANDLER: Final = BaseLLMHTTPHandler() @@ -96,16 +99,46 @@ def _validate_request( ) -def _ir_request(request: DecisionsRequestFormat) -> DecisionsIRRequest: +def _ir_request(request: DecisionsRequestFormat, safety_identifier: str | None) -> DecisionsIRRequest: match request: case DecisionsRequestBody(): - return systemone_request_to_ir(request) + return replace(systemone_request_to_ir(request), safety_identifier=safety_identifier) case OpenAIDecisionRequestBody(): return openai_request_to_ir(request) case _: assert_never(request) +def _drops_params(kwargs: Mapping[str, object]) -> bool: + return litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True + + +def _request_safety_identifier(safety_identifier: object, kwargs: Mapping[str, object]) -> str | None: + if isinstance(safety_identifier, str) or not _drops_params(kwargs): + return _SAFETY_IDENTIFIER_ADAPTER.validate_python(safety_identifier) + return None + + +@dataclass(frozen=True, slots=True) +class _UnsupportedSafetyIdentifier: + pass + + +def _provider_ir_request( + request: DecisionsRequestFormat, + *, + safety_identifier: str | None, + provider_config: BaseDecisionsConfig, + kwargs: Mapping[str, object], +) -> DecisionsIRRequest | _UnsupportedSafetyIdentifier: + ir_request: Final = _ir_request(request, safety_identifier) + if ir_request.safety_identifier is None or provider_config.supports_safety_identifier: + return ir_request + if _drops_params(kwargs): + return replace(ir_request, safety_identifier=None) + return _UnsupportedSafetyIdentifier() + + def _prepare_call( *, model: str, @@ -144,8 +177,12 @@ def _prepare_call( except ValueError as error: raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error try: + request_safety_identifier: Final = _request_safety_identifier(safety_identifier, kwargs) request: Final = _validate_request( - state=state, questions=questions, decision_input=decision_input, safety_identifier=safety_identifier + state=state, + questions=questions, + decision_input=decision_input, + safety_identifier=request_safety_identifier, ) except ValidationError as error: raise litellm.BadRequestError( @@ -169,7 +206,22 @@ def _prepare_call( llm_provider=provider, ) - ir_request: Final = _ir_request(request) + ir_request: Final = _provider_ir_request( + request, + safety_identifier=request_safety_identifier, + provider_config=provider_config, + kwargs=kwargs, + ) + if isinstance(ir_request, _UnsupportedSafetyIdentifier): + raise litellm.UnsupportedParamsError( + message=( + f"{provider} does not support parameters: ['safety_identifier'], for model={model}. " + "To drop these, set `litellm.drop_params=True` or for proxy:\n\n" + "`litellm_settings:\n drop_params: true`\n" + ), + model=model, + llm_provider=provider, + ) body: Final = provider_config.transform_decisions_request( model=canonical_model, request=ir_request, custom_llm_provider=provider ) diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index c57c643983a..baeeb3f0449 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -301,6 +301,7 @@ class BaseDecisionsConfig(ABC): api_key_env: tuple[str, ...] = () api_base_env: tuple[str, ...] = () api_key_required: bool = True + supports_safety_identifier: bool = False health_check_questions: Mapping[str, Mapping[str, object]] = MappingProxyType( {"reachable": MappingProxyType({"type": "noul", "instructions": "Is the service reachable?"})} ) diff --git a/litellm/llms/openai/decisions/transformation.py b/litellm/llms/openai/decisions/transformation.py index 0afbe1a434e..317e5abeeb5 100644 --- a/litellm/llms/openai/decisions/transformation.py +++ b/litellm/llms/openai/decisions/transformation.py @@ -292,6 +292,7 @@ def ir_to_openai_response( class OpenAIDecisionsConfig(BaseDecisionsConfig): path = "/v1/decisions" + supports_safety_identifier = True def get_default_api_base(self) -> str | None: return "https://api.openai.com" diff --git a/litellm/proxy/decisions_endpoints/endpoints.py b/litellm/proxy/decisions_endpoints/endpoints.py index f2891c184bf..4f4209208ac 100644 --- a/litellm/proxy/decisions_endpoints/endpoints.py +++ b/litellm/proxy/decisions_endpoints/endpoints.py @@ -1,4 +1,5 @@ from collections.abc import Mapping +from types import MappingProxyType from typing import Annotated, Final from fastapi import APIRouter, Depends, Request, Response @@ -38,6 +39,10 @@ async def _invalid_request( ) +def _fields_checked_before_routing(data: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({key: value for key, value in data.items() if key != "safety_identifier"}) + + async def _request_data(request: Request, user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]: body: Final = await request.body() try: @@ -75,7 +80,7 @@ async def _process_decisions( data: Final = await _request_data(request, user_api_key_dict) try: - body_adapter.validate_python(data) + body_adapter.validate_python(_fields_checked_before_routing(data)) except ValidationError as error: raise await _invalid_request(raw_data=data, error=error, user_api_key_dict=user_api_key_dict) general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(proxy_general_settings) diff --git a/tests/integration/providers/test_databricks_decisions_wire.py b/tests/integration/providers/test_databricks_decisions_wire.py index 5c3b7d1b1a1..93dc2ea04fa 100644 --- a/tests/integration/providers/test_databricks_decisions_wire.py +++ b/tests/integration/providers/test_databricks_decisions_wire.py @@ -208,7 +208,6 @@ def test_an_openai_format_request_at_v1_decisions_reaches_the_serving_endpoint_a "model": deployment, "input": _OPENAI_INPUT, "questions": [_OPENAI_TIER_QUESTION], - "safety_identifier": "end-user-1", }, ) assert response.status_code == 200, response.text diff --git a/tests/integration/providers/test_decisions_openai_format_wire.py b/tests/integration/providers/test_decisions_openai_format_wire.py new file mode 100644 index 00000000000..f0e7163e264 --- /dev/null +++ b/tests/integration/providers/test_decisions_openai_format_wire.py @@ -0,0 +1,586 @@ +import math +import uuid +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy_process +from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import JsonResponse +from pydantic import JsonValue, TypeAdapter + +import litellm + +_API_KEY: Final = "synthetic-decisions-key" +_SAFETY_IDENTIFIER: Final = "end-user-7" +_DROPPING_DEPLOYMENT: Final = "decisions-under-litellm-settings-drop-params" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_INPUT: Final = "Ticket (billing): The export job hangs at 99%" +_FOLLOW_UP: Final = "Customer: still stuck after retrying" +_QUESTIONS: Final[list[JsonValue]] = [ + {"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}, + { + "type": "choice", + "name": "severity", + "instructions": "How severe is it?", + "choices": [{"value": "low", "description": "cosmetic"}, {"value": "high", "description": "blocks users"}], + }, + { + "type": "score", + "name": "confidence", + "instructions": "How sure are you?", + "levels": [{"label": "unsure"}, {"label": "sure"}], + }, +] +_SYSTEM_ONE_QUESTIONS: Final[dict[str, JsonValue]] = { + "defect": {"type": "noul", "instructions": "Is this a defect?"}, + "severity": { + "type": "choice", + "instructions": "How severe is it?", + "criteria": {"low": "cosmetic", "high": "blocks users"}, + }, + "confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]}, +} +_SYSTEM_ONE_ANSWERS: Final[dict[str, JsonValue]] = { + "defect": {"type": "noul", "noul": 0.93}, + "severity": {"type": "choice", "choice": "high", "confidence": 0.8, "probabilities": {"low": 0.2, "high": 0.8}}, + "confidence": { + "type": "score", + "score": 1.0, + "confidence": 0.7, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.3, "1": 0.7}, + }, +} +_ANSWERS: Final[list[JsonValue]] = [ + {"type": "predicate", "name": "defect", "probability": 0.93}, + { + "type": "choice", + "name": "severity", + "choice": "high", + "probabilities": [{"value": "low", "probability": 0.2}, {"value": "high", "probability": 0.8}], + "confidence": 0.8, + }, + { + "type": "score", + "name": "confidence", + "score": 1.0, + "probabilities": [ + {"value": 0, "label": "unsure", "probability": 0.3}, + {"value": 1, "label": "sure", "probability": 0.7}, + ], + "confidence": 0.7, + }, +] +_INPUT_TOKENS: Final = 367 +_OUTPUT_TOKENS: Final = 3 +_SYSTEM_ONE_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": _INPUT_TOKENS, "output_tokens": _OUTPUT_TOKENS} +_USAGE: Final[dict[str, JsonValue]] = { + "input_tokens": _INPUT_TOKENS, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": _OUTPUT_TOKENS, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": _INPUT_TOKENS + _OUTPUT_TOKENS, +} +_SDK_QUESTIONS: Final = TypeAdapter(list[dict[str, object]]).validate_python(_QUESTIONS) +_SPEND_QUERY: Final = ( + "SELECT spend, status, call_type, model_group, custom_llm_provider, api_base, prompt_tokens, completion_tokens, " + 'request_tags FROM "LiteLLM_SpendLogs" WHERE request_id = %s' +) + + +@dataclass(frozen=True, slots=True) +class _Provider: + name: str + model: str + path: str + body_model: str + api_key: str | None + wraps_result: bool + cost_map_key: str | None + speaks_openai: bool = False + + def upstream_body(self) -> dict[str, JsonValue]: + if self.speaks_openai: + return {"model": self.body_model, "input": _INPUT, "questions": _QUESTIONS} + return {"model": self.body_model, "state": _INPUT, "questions": _SYSTEM_ONE_QUESTIONS} + + def upstream_reply(self) -> dict[str, JsonValue]: + if self.speaks_openai: + return {"model": self.body_model, "answers": _ANSWERS, "usage": _USAGE} + answer: Final[dict[str, JsonValue]] = { + "model": self.body_model, + "answers": _SYSTEM_ONE_ANSWERS, + "usage": _SYSTEM_ONE_USAGE, + } + return {"result": answer, "success": True} if self.wraps_result else answer + + def litellm_response(self) -> dict[str, JsonValue]: + return {"model": self.body_model, "answers": _ANSWERS, "usage": _USAGE} + + +_PROVIDERS: Final = ( + _Provider( + "perplexity", + "perplexity/pplx-decider-v1-27b", + "/v1/decisions", + "pplx-decider-v1-27b", + _API_KEY, + False, + "perplexity/pplx-decider-v1-27b", + ), + _Provider("typesafe", "typesafe/jev-1.13.0", "/v1/systemone", "jev-1.13.0", _API_KEY, False, "typesafe/jev-1.13.0"), + _Provider( + "openrouter", + "openrouter/typesafe/jev-1.13", + "/alpha/decisions", + "typesafe/jev-1.13", + _API_KEY, + False, + "openrouter/typesafe/jev-1.13", + ), + _Provider( + "strands_decider", "strands_decider/systemone-decider", "/v1/systemone", "systemone-decider", None, False, None + ), + _Provider( + "cloudflare", + "cloudflare/clef", + "/ai/run/@cf/cloudflare/clef", + "clef", + _API_KEY, + True, + "cloudflare/@cf/cloudflare/clef", + ), + _Provider("hosted_vllm", "hosted_vllm/Qwen/Qwen3-0.6B", "/v1/systemone", "Qwen/Qwen3-0.6B", None, False, None), + _Provider( + "databricks", + "databricks/databricks-openjev-qwen35-4b", + "/databricks-openjev-qwen35-4b/invocations", + "databricks-openjev-qwen35-4b", + _API_KEY, + False, + None, + ), + _Provider( + "azure_ai", "azure_ai/decision-1", "/providers/microsoft/v1/systemone", "decision-1", _API_KEY, False, None + ), + _Provider("openai", "openai/gpt-6-luna", "/v1/decisions", "gpt-6-luna", _API_KEY, False, "gpt-6-luna", True), +) +_SYSTEM_ONE_PROVIDERS: Final = tuple(provider for provider in _PROVIDERS if not provider.speaks_openai) +_PERPLEXITY: Final = _PROVIDERS[0] +_TYPESAFE: Final = _PROVIDERS[1] +_OPENAI: Final = next(provider for provider in _PROVIDERS if provider.speaks_openai) +_PREDICATE: Final[dict[str, JsonValue]] = {"type": "predicate", "name": "q", "instructions": "Is it?"} +_INVALID_BODIES: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = ( + ("missing questions", {"input": _INPUT}), + ("missing input", {"questions": _QUESTIONS}), + ("numeric input", {"input": 5, "questions": _QUESTIONS}), + ("assistant message input", {"input": [{"role": "assistant", "content": "hi"}], "questions": _QUESTIONS}), + ("empty questions", {"input": _INPUT, "questions": []}), + ("questions as a map", {"input": _INPUT, "questions": {"q": _PREDICATE}}), + ("predicate without instructions", {"input": _INPUT, "questions": [{"type": "predicate", "name": "q"}]}), + ( + "choice with one choice", + {"input": _INPUT, "questions": [{**_PREDICATE, "type": "choice", "choices": [{"value": "only"}]}]}, + ), + ( + "score with one level", + {"input": _INPUT, "questions": [{**_PREDICATE, "type": "score", "levels": [{"label": "only"}]}]}, + ), + ("unknown question type", {"input": _INPUT, "questions": [{**_PREDICATE, "type": "ranking"}]}), + ("numeric safety_identifier", {"input": _INPUT, "questions": _QUESTIONS, "safety_identifier": 7}), + ("list safety_identifier", {"input": _INPUT, "questions": _QUESTIONS, "safety_identifier": [_SAFETY_IDENTIFIER]}), +) +_IMAGE_INPUT: Final[list[JsonValue]] = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": _INPUT}, + {"type": "input_image", "image_url": "data:image/png;base64,iVBORw0KGgo=", "detail": "auto"}, + ], + } +] +_OPENAI_IMAGE_MESSAGES: Final[list[JsonValue]] = [ + { + "role": "user", + "type": "message", + "content": [ + {"type": "input_text", "text": _INPUT}, + {"type": "input_image", "image_url": "data:image/png;base64,iVBORw0KGgo=", "detail": "auto"}, + ], + } +] +_MESSAGE_LIST_INPUT: Final[list[JsonValue]] = [ + {"role": "user", "content": _INPUT}, + {"role": "user", "content": [{"type": "input_text", "text": _FOLLOW_UP}]}, +] + + +def _response_cost(response: httpx.Response) -> float: + return float(response.headers["x-litellm-response-cost"]) if "x-litellm-response-cost" in response.headers else 0.0 + + +def _provider_id(provider: _Provider) -> str: + return provider.name + + +def _number(value: JsonValue) -> float: + assert isinstance(value, (int, float)) and not isinstance(value, bool), value + return float(value) + + +def _expected_spend(cost_map_key: str | None) -> float: + if cost_map_key is None: + return 0.0 + cost_map: Final = _JSON_OBJECT.validate_json(Path("model_prices_and_context_window.json").read_bytes()) + prices: Final = object_value(cost_map[cost_map_key]) + return _INPUT_TOKENS * _number(prices["input_cost_per_token"]) + _OUTPUT_TOKENS * _number( + prices["output_cost_per_token"] + ) + + +def _register(scenario: Scenario, body: dict[str, JsonValue], *, status: int = 200) -> ScenarioHandle: + handle: Final = register_scenario( + f"decisions-{uuid.uuid4().hex[:12]}", JsonResponse(content_type="application/json", body=body, status=status) + ) + scenario.cleanups.callback(delete_scenario, handle) + return handle + + +def _deployment(scenario: Scenario, handle: ScenarioHandle, provider: _Provider, *, drop_params: bool = False) -> str: + if drop_params: + return scenario.model( + model=provider.model, api_base=handle.api_base(), api_key=provider.api_key, drop_params=True + ) + return scenario.model(model=provider.model, api_base=handle.api_base(), api_key=provider.api_key) + + +def _decide(gateway: Gateway, model: str, *, route: str = "/v1/decisions", **extra: JsonValue) -> httpx.Response: + return gateway.request("POST", route, {"model": model, "input": _INPUT, "questions": _QUESTIONS, **extra}) + + +def _decide_in_system_one_format(gateway: Gateway, model: str, **extra: JsonValue) -> httpx.Response: + return gateway.request( + "POST", "/v1/systemone", {"model": model, "state": _INPUT, "questions": _SYSTEM_ONE_QUESTIONS, **extra} + ) + + +def _observed_requests(gateway: Gateway) -> tuple[dict[str, JsonValue], ...]: + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + observations: Final = _JSON_OBJECT.validate_json(upstream.get("/__observations").content)["requests"] + assert isinstance(observations, list), observations + return tuple(map(object_value, observations)) + + +def _calls_to(requests: Sequence[dict[str, JsonValue]], handle: ScenarioHandle) -> list[dict[str, JsonValue]]: + return [request for request in requests if string_value(request["path"]).startswith(f"/{handle.scenario_id}/")] + + +def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str, JsonValue]]: + return _calls_to(_observed_requests(gateway), handle) + + +def _spend_row(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70) + return rows[0] + + +def _assert_refused_as_invalid(gateway: Gateway, model: str, label: str, body: dict[str, JsonValue]) -> None: + response: Final = gateway.request("POST", "/v1/decisions", {"model": model, **body}) + assert response.status_code == 400, (label, response.text) + assert "Invalid Decisions request" in response.text, (label, response.text) + + +def _assert_refused_for_safety_identifier(response: httpx.Response) -> None: + assert response.status_code == 400, response.text + assert "safety_identifier" in response.text and "drop_params" in response.text, response.text + + +def _drop_params_config(directory: Path, api_base: str) -> Path: + base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **base, + "litellm_settings": {**object_value(base["litellm_settings"]), "drop_params": True}, + "model_list": [ + { + "model_name": _DROPPING_DEPLOYMENT, + "litellm_params": {"model": _PERPLEXITY.model, "api_base": api_base, "api_key": _API_KEY}, + } + ], + } + path: Final = directory / "decisions-drop-params.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.mark.parametrize("provider", _PROVIDERS, ids=_provider_id) +def test_each_provider_gets_its_own_path_key_and_body_and_is_billed_from_the_cost_map( + gateway: Gateway, provider: _Provider +) -> None: + expected_spend: Final = _expected_spend(provider.cost_map_key) + with gateway.scenario() as scenario: + handle: Final = _register(scenario, provider.upstream_reply()) + model: Final = _deployment(scenario, handle, provider) + response: Final = _decide(gateway, model) + assert response.status_code == 200, response.text + assert response.json() == provider.litellm_response() + assert response.headers["x-litellm-model-group"] == model + assert math.isclose(_response_cost(response), expected_spend, rel_tol=1e-9) + (call,) = _upstream_calls(gateway, handle) + assert call["path"] == f"/{handle.scenario_id}{provider.path}" + assert call["authorization"] == (f"Bearer {provider.api_key}" if provider.api_key else "") + assert call["body"] == provider.upstream_body() + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert ( + row["status"], + row["call_type"], + row["custom_llm_provider"], + row["model_group"], + row["api_base"], + row["prompt_tokens"], + row["completion_tokens"], + ) == ( + "success", + "adecisions", + provider.name, + model, + f"{handle.api_base()}{provider.path}", + _INPUT_TOKENS, + _OUTPUT_TOKENS, + ) + assert math.isclose(_number(row["spend"]), expected_spend, rel_tol=1e-9), row + + +def test_the_unversioned_alias_serves_the_same_request(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + model: Final = _deployment(scenario, handle, _PERPLEXITY) + response: Final = _decide(gateway, model, route="/decisions") + assert response.status_code == 200, response.text + assert response.json() == _PERPLEXITY.litellm_response() + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == _PERPLEXITY.upstream_body() + + +def test_repeated_identical_requests_each_reach_the_upstream_and_are_each_billed(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + model: Final = _deployment(scenario, handle, _PERPLEXITY) + responses: Final = tuple(_decide(gateway, model) for _ in range(2)) + assert [response.status_code for response in responses] == [200, 200], [r.text for r in responses] + call_ids: Final = tuple(response.headers["x-litellm-call-id"] for response in responses) + assert len(set(call_ids)) == 2, call_ids + assert len(_upstream_calls(gateway, handle)) == 2 + for call_id in call_ids: + assert _spend_row(call_id)["status"] == "success" + + +async def test_sdk_sync_and_async_clients_send_the_same_request(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _TYPESAFE.upstream_reply()) + synchronous: Final = litellm.decisions( + model=_TYPESAFE.model, input=_INPUT, questions=_SDK_QUESTIONS, api_base=handle.api_base(), api_key=_API_KEY + ) + asynchronous: Final = await litellm.adecisions( + model=_TYPESAFE.model, input=_INPUT, questions=_SDK_QUESTIONS, api_base=handle.api_base(), api_key=_API_KEY + ) + for response in (synchronous, asynchronous): + assert response.model_dump(mode="json") == _TYPESAFE.litellm_response() + calls: Final = _upstream_calls(gateway, handle) + assert len(calls) == 2, calls + for call in calls: + assert call["path"] == f"/{handle.scenario_id}{_TYPESAFE.path}" + assert call["authorization"] == f"Bearer {_API_KEY}" + assert call["body"] == _TYPESAFE.upstream_body() + + +def test_gateway_only_fields_stay_at_the_gateway_and_tags_reach_the_spend_log(gateway: Gateway) -> None: + tag: Final = f"decisions-openai-format-{uuid.uuid4().hex[:8]}" + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + model: Final = _deployment(scenario, handle, _PERPLEXITY) + response: Final = _decide( + gateway, model, user="auditor", num_retries=0, temperature=0.2, metadata={"tags": [tag]} + ) + assert response.status_code == 200, response.text + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == _PERPLEXITY.upstream_body() + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + tags: Final = row["request_tags"] + assert isinstance(tags, list) and tag in tags, row + + +def test_invalid_bodies_are_refused_at_the_gateway_without_an_upstream_call(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + model: Final = _deployment(scenario, handle, _PERPLEXITY) + for label, body in _INVALID_BODIES: + _assert_refused_as_invalid(gateway, model, label, body) + assert _upstream_calls(gateway, handle) == [] + + +@pytest.mark.parametrize("provider", _SYSTEM_ONE_PROVIDERS, ids=_provider_id) +def test_system_one_providers_refuse_images_at_the_gateway_without_an_upstream_call( + gateway: Gateway, provider: _Provider +) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, provider.upstream_reply()) + model: Final = _deployment(scenario, handle, provider) + response: Final = _decide(gateway, model, input=_IMAGE_INPUT) + assert response.status_code == 400, response.text + assert "input_image" in response.text + assert _upstream_calls(gateway, handle) == [] + + +def test_openai_forwards_image_input_in_its_own_message_shape(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _OPENAI.upstream_reply()) + model: Final = _deployment(scenario, handle, _OPENAI) + response: Final = _decide(gateway, model, input=_IMAGE_INPUT) + assert response.status_code == 200, response.text + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == {"model": _OPENAI.body_model, "input": _OPENAI_IMAGE_MESSAGES, "questions": _QUESTIONS} + + +def test_a_message_list_input_reaches_a_system_one_provider_as_its_flattened_text(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + model: Final = _deployment(scenario, handle, _PERPLEXITY) + response: Final = _decide(gateway, model, input=_MESSAGE_LIST_INPUT) + assert response.status_code == 200, response.text + assert response.json() == _PERPLEXITY.litellm_response() + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == { + "model": _PERPLEXITY.body_model, + "state": f"{_INPUT}\n\n{_FOLLOW_UP}", + "questions": _SYSTEM_ONE_QUESTIONS, + } + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert (row["status"], row["call_type"], row["prompt_tokens"], row["completion_tokens"]) == ( + "success", + "adecisions", + _INPUT_TOKENS, + _OUTPUT_TOKENS, + ) + + +@pytest.mark.parametrize("provider", _SYSTEM_ONE_PROVIDERS, ids=_provider_id) +def test_safety_identifier_is_refused_by_system_one_providers_unless_the_deployment_drops_params( + gateway: Gateway, provider: _Provider +) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, provider.upstream_reply()) + strict: Final = _deployment(scenario, handle, provider) + refused: Final = _decide(gateway, strict, safety_identifier=_SAFETY_IDENTIFIER) + _assert_refused_for_safety_identifier(refused) + assert _upstream_calls(gateway, handle) == [] + + dropping: Final = _deployment(scenario, handle, provider, drop_params=True) + accepted: Final = _decide(gateway, dropping, safety_identifier=_SAFETY_IDENTIFIER) + assert accepted.status_code == 200, accepted.text + assert accepted.json() == provider.litellm_response() + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == provider.upstream_body() + row: Final = _spend_row(accepted.headers["x-litellm-call-id"]) + assert (row["status"], row["call_type"], row["model_group"]) == ("success", "adecisions", dropping) + + +@pytest.mark.parametrize("value", ("", "x" * 5000), ids=("empty", "5kb")) +def test_every_string_safety_identifier_is_refused_or_dropped_like_the_usual_one(gateway: Gateway, value: str) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + strict: Final = _deployment(scenario, handle, _PERPLEXITY) + _assert_refused_for_safety_identifier(_decide(gateway, strict, safety_identifier=value)) + dropping: Final = _deployment(scenario, handle, _PERPLEXITY, drop_params=True) + accepted: Final = _decide(gateway, dropping, safety_identifier=value) + assert accepted.status_code == 200, accepted.text + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == _PERPLEXITY.upstream_body() + + +def test_a_request_body_drop_params_drops_the_safety_identifier_like_chat(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + strict: Final = _deployment(scenario, handle, _PERPLEXITY) + response: Final = _decide(gateway, strict, safety_identifier=_SAFETY_IDENTIFIER, drop_params=True) + assert response.status_code == 200, response.text + assert response.json() == _PERPLEXITY.litellm_response() + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == _PERPLEXITY.upstream_body() + + +def test_openai_keeps_the_safety_identifier_on_the_wire_without_drop_params(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _OPENAI.upstream_reply()) + model: Final = _deployment(scenario, handle, _OPENAI) + response: Final = _decide(gateway, model, safety_identifier=_SAFETY_IDENTIFIER) + assert response.status_code == 200, response.text + assert response.json() == _OPENAI.litellm_response() + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == {**_OPENAI.upstream_body(), "safety_identifier": _SAFETY_IDENTIFIER} + + +def test_a_system_one_format_safety_identifier_is_refused_unless_the_deployment_drops_params(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + strict: Final = _deployment(scenario, handle, _PERPLEXITY) + refused: Final = _decide_in_system_one_format(gateway, strict, safety_identifier=_SAFETY_IDENTIFIER) + _assert_refused_for_safety_identifier(refused) + assert _upstream_calls(gateway, handle) == [] + + dropping: Final = _deployment(scenario, handle, _PERPLEXITY, drop_params=True) + accepted: Final = _decide_in_system_one_format(gateway, dropping, safety_identifier=_SAFETY_IDENTIFIER) + assert accepted.status_code == 200, accepted.text + assert accepted.json()["answers"] == _SYSTEM_ONE_ANSWERS + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == _PERPLEXITY.upstream_body() + + +def test_openai_keeps_a_system_one_format_safety_identifier_on_the_wire(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _OPENAI.upstream_reply()) + model: Final = _deployment(scenario, handle, _OPENAI) + response: Final = _decide_in_system_one_format(gateway, model, safety_identifier=_SAFETY_IDENTIFIER) + assert response.status_code == 200, response.text + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == {**_OPENAI.upstream_body(), "safety_identifier": _SAFETY_IDENTIFIER} + + +@pytest.mark.parametrize("decide", (_decide, _decide_in_system_one_format), ids=("openai_format", "system_one_format")) +@pytest.mark.parametrize("value", (7, [_SAFETY_IDENTIFIER]), ids=("numeric", "list")) +def test_a_non_string_safety_identifier_is_refused_as_invalid_unless_the_deployment_drops_params( + gateway: Gateway, value: JsonValue, decide: Callable[..., httpx.Response] +) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _OPENAI.upstream_reply()) + strict: Final = _deployment(scenario, handle, _OPENAI) + refused: Final = decide(gateway, strict, safety_identifier=value) + assert refused.status_code == 400, refused.text + assert "Invalid Decisions request" in refused.text, refused.text + assert _upstream_calls(gateway, handle) == [] + + dropping: Final = _deployment(scenario, handle, _OPENAI, drop_params=True) + accepted: Final = decide(gateway, dropping, safety_identifier=value) + assert accepted.status_code == 200, accepted.text + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == _OPENAI.upstream_body() + + +def test_litellm_settings_drop_params_drops_the_safety_identifier_for_a_strict_deployment( + gateway: Gateway, tmp_path: Path +) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + config: Final = _drop_params_config(tmp_path, handle.api_base()) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as owned: + response: Final = _decide(owned.gateway, _DROPPING_DEPLOYMENT, safety_identifier=_SAFETY_IDENTIFIER) + assert response.status_code == 200, response.text + assert response.json() == _PERPLEXITY.litellm_response() + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == _PERPLEXITY.upstream_body() diff --git a/tests/integration/translation/decisions/bases/cloudflare.py b/tests/integration/translation/decisions/bases/cloudflare.py index ae86a2ccd87..3efc96eeb87 100644 --- a/tests/integration/translation/decisions/bases/cloudflare.py +++ b/tests/integration/translation/decisions/bases/cloudflare.py @@ -6,6 +6,103 @@ from integration.translation.case import TranslationTestCase """ CLEF_TEST_CASE: Final = TranslationTestCase( scenario="basic", + litellm_endpoint="/v1/decisions", + litellm_request={ + "model": "cloudflare/@cf/cloudflare/clef", + "input": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": [ + {"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}, + { + "type": "choice", + "name": "severity", + "instructions": "How severe is it?", + "choices": [ + {"value": "low", "description": "cosmetic"}, + {"value": "high", "description": "blocks users"}, + ], + }, + { + "type": "score", + "name": "confidence", + "instructions": "How sure are you?", + "levels": [{"label": "unsure"}, {"label": "sure"}], + }, + ], + "cache": {"no-cache": True}, + }, + expected_provider_endpoint="/ai/run/@cf/cloudflare/clef", + expected_provider_headers={"authorization": "Bearer synthetic-cloudflare-key", "content-type": "application/json"}, + expected_provider_request={ + "model": "clef", + "state": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": { + "defect": {"type": "noul", "instructions": "Is this a defect?"}, + "severity": { + "type": "choice", + "instructions": "How severe is it?", + "criteria": {"low": "cosmetic", "high": "blocks users"}, + }, + "confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]}, + }, + }, + mock_provider_response={ + "result": { + "model": "clef", + "answers": { + "defect": {"type": "noul", "noul": 0.9345}, + "severity": { + "type": "choice", + "choice": "high", + "confidence": 0.8067, + "probabilities": {"low": 0.0509, "high": 0.9491}, + }, + "confidence": { + "type": "score", + "score": 0.9036, + "confidence": 0.6515, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.0964, "1": 0.9036}, + }, + }, + "usage": {"input_tokens": 290, "output_tokens": 0}, + }, + "success": True, + "errors": [], + "messages": [], + }, + expected_litellm_response={ + "model": "clef", + "answers": [ + {"type": "predicate", "name": "defect", "probability": 0.9345}, + { + "type": "choice", + "name": "severity", + "choice": "high", + "probabilities": [{"value": "low", "probability": 0.0509}, {"value": "high", "probability": 0.9491}], + "confidence": 0.8067, + }, + { + "type": "score", + "name": "confidence", + "score": 0.9036, + "probabilities": [ + {"value": 0, "label": "unsure", "probability": 0.0964}, + {"value": 1, "label": "sure", "probability": 0.9036}, + ], + "confidence": 0.6515, + }, + ], + "usage": { + "input_tokens": 290, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 290, + }, + }, +) +CLEF_SYSTEMONE_TEST_CASE: Final = TranslationTestCase( + scenario="systemone", litellm_endpoint="/v1/systemone", litellm_request={ "model": "cloudflare/@cf/cloudflare/clef", diff --git a/tests/integration/translation/decisions/bases/openrouter.py b/tests/integration/translation/decisions/bases/openrouter.py index b65acff6b06..ed39bc1a55b 100644 --- a/tests/integration/translation/decisions/bases/openrouter.py +++ b/tests/integration/translation/decisions/bases/openrouter.py @@ -6,6 +6,100 @@ from integration.translation.case import TranslationTestCase """ TYPESAFE_JEV_1_13_TEST_CASE: Final = TranslationTestCase( scenario="basic", + litellm_endpoint="/v1/decisions", + litellm_request={ + "model": "openrouter/typesafe/jev-1.13", + "input": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": [ + {"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}, + { + "type": "choice", + "name": "severity", + "instructions": "How severe is it?", + "choices": [ + {"value": "low", "description": "cosmetic"}, + {"value": "high", "description": "blocks users"}, + ], + }, + { + "type": "score", + "name": "confidence", + "instructions": "How sure are you?", + "levels": [{"label": "unsure"}, {"label": "sure"}], + }, + ], + "cache": {"no-cache": True}, + }, + expected_provider_endpoint="/alpha/decisions", + expected_provider_headers={"authorization": "Bearer synthetic-openrouter-key", "content-type": "application/json"}, + expected_provider_request={ + "model": "typesafe/jev-1.13", + "state": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": { + "defect": {"type": "noul", "instructions": "Is this a defect?"}, + "severity": { + "type": "choice", + "instructions": "How severe is it?", + "criteria": {"low": "cosmetic", "high": "blocks users"}, + }, + "confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]}, + }, + }, + mock_provider_response={ + "model": "typesafe/jev-1.13-20260917", + "answers": { + "defect": {"type": "noul", "noul": 0.81}, + "severity": { + "type": "choice", + "choice": "high", + "confidence": 0.99, + "probabilities": {"low": 0.01, "high": 0.99}, + }, + "confidence": { + "type": "score", + "score": 0.5, + "confidence": 0, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.5, "1": 0.5}, + }, + }, + "usage": {"input_tokens": 377, "output_tokens": 62, "cost": 1.5834e-05}, + "id": "gen-dec-1791323839-GA15kY0nt34oiJ7srfki", + "provider": "TypeSafe", + }, + expected_litellm_response={ + "model": "typesafe/jev-1.13-20260917", + "answers": [ + {"type": "predicate", "name": "defect", "probability": 0.81}, + { + "type": "choice", + "name": "severity", + "choice": "high", + "probabilities": [{"value": "low", "probability": 0.01}, {"value": "high", "probability": 0.99}], + "confidence": 0.99, + }, + { + "type": "score", + "name": "confidence", + "score": 0.5, + "probabilities": [ + {"value": 0, "label": "unsure", "probability": 0.5}, + {"value": 1, "label": "sure", "probability": 0.5}, + ], + "confidence": 0, + }, + ], + "usage": { + "input_tokens": 377, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 62, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 439, + }, + }, +) +TYPESAFE_JEV_1_13_SYSTEMONE_TEST_CASE: Final = TranslationTestCase( + scenario="systemone", litellm_endpoint="/v1/systemone", litellm_request={ "model": "openrouter/typesafe/jev-1.13", diff --git a/tests/integration/translation/decisions/bases/perplexity.py b/tests/integration/translation/decisions/bases/perplexity.py index 02b1de51f24..9e29ceab8cc 100644 --- a/tests/integration/translation/decisions/bases/perplexity.py +++ b/tests/integration/translation/decisions/bases/perplexity.py @@ -6,6 +6,101 @@ from integration.translation.case import TranslationTestCase """ PPLX_DECIDER_V1_27B_TEST_CASE: Final = TranslationTestCase( scenario="basic", + litellm_endpoint="/v1/decisions", + litellm_request={ + "model": "perplexity/pplx-decider-v1-27b", + "input": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": [ + {"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}, + { + "type": "choice", + "name": "severity", + "instructions": "How severe is it?", + "choices": [ + {"value": "low", "description": "cosmetic"}, + {"value": "high", "description": "blocks users"}, + ], + }, + { + "type": "score", + "name": "confidence", + "instructions": "How sure are you?", + "levels": [{"label": "unsure"}, {"label": "sure"}], + }, + ], + "cache": {"no-cache": True}, + }, + expected_provider_endpoint="/v1/decisions", + expected_provider_headers={"authorization": "Bearer synthetic-perplexity-key", "content-type": "application/json"}, + expected_provider_request={ + "model": "pplx-decider-v1-27b", + "state": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": { + "defect": {"type": "noul", "instructions": "Is this a defect?"}, + "severity": { + "type": "choice", + "instructions": "How severe is it?", + "criteria": {"low": "cosmetic", "high": "blocks users"}, + }, + "confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]}, + }, + }, + mock_provider_response={ + "model": "pplx-decider-v1-27b", + "answers": { + "defect": {"type": "noul", "noul": 0.9989100737587077}, + "severity": { + "type": "choice", + "choice": "high", + "confidence": 0.9964631215356778, + "probabilities": {"low": 0.0017684392321610232, "high": 0.9982315607678389}, + }, + "confidence": { + "type": "score", + "score": 0.07367392327139817, + "confidence": 0.8526521534572037, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.9263260767286018, "1": 0.07367392327139817}, + }, + }, + "usage": {"input_tokens": 318, "output_tokens": 3}, + }, + expected_litellm_response={ + "model": "pplx-decider-v1-27b", + "answers": [ + {"type": "predicate", "name": "defect", "probability": 0.9989100737587077}, + { + "type": "choice", + "name": "severity", + "choice": "high", + "probabilities": [ + {"value": "low", "probability": 0.0017684392321610232}, + {"value": "high", "probability": 0.9982315607678389}, + ], + "confidence": 0.9964631215356778, + }, + { + "type": "score", + "name": "confidence", + "score": 0.07367392327139817, + "probabilities": [ + {"value": 0, "label": "unsure", "probability": 0.9263260767286018}, + {"value": 1, "label": "sure", "probability": 0.07367392327139817}, + ], + "confidence": 0.8526521534572037, + }, + ], + "usage": { + "input_tokens": 318, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 3, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 321, + }, + }, +) +PPLX_DECIDER_V1_27B_SYSTEMONE_TEST_CASE: Final = TranslationTestCase( + scenario="systemone", litellm_endpoint="/v1/systemone", litellm_request={ "model": "perplexity/pplx-decider-v1-27b", diff --git a/tests/integration/translation/decisions/bases/strands_decider.py b/tests/integration/translation/decisions/bases/strands_decider.py index 8a1a51f8c3e..2950b090737 100644 --- a/tests/integration/translation/decisions/bases/strands_decider.py +++ b/tests/integration/translation/decisions/bases/strands_decider.py @@ -6,6 +6,40 @@ from integration.translation.case import TranslationTestCase """ STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE: Final = TranslationTestCase( scenario="basic", + litellm_endpoint="/v1/decisions", + litellm_request={ + "model": "strands_decider/strands-decider-2B-hobson-v19", + "input": "Help! My payouts have been failing for 3 days!", + "questions": [{"type": "predicate", "name": "is_urgent", "instructions": "Does this convey urgency?"}], + "cache": {"no-cache": True}, + }, + expected_provider_endpoint="/v1/systemone", + expected_provider_headers={"content-type": "application/json"}, + expected_provider_request={ + "model": "strands-decider-2B-hobson-v19", + "state": "Help! My payouts have been failing for 3 days!", + "questions": {"is_urgent": {"type": "noul", "instructions": "Does this convey urgency?"}}, + }, + mock_provider_response={ + "model": "strands-decider-2B-hobson-v19", + "answers": {"is_urgent": {"type": "noul", "noul": 0.8277}}, + "usage": {"input_tokens": 86, "output_tokens": 1}, + "latency_ms": 140.03, + }, + expected_litellm_response={ + "model": "strands-decider-2B-hobson-v19", + "answers": [{"type": "predicate", "name": "is_urgent", "probability": 0.8277}], + "usage": { + "input_tokens": 86, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 1, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 87, + }, + }, +) +STRANDS_DECIDER_2B_HOBSON_V19_SYSTEMONE_TEST_CASE: Final = TranslationTestCase( + scenario="systemone", litellm_endpoint="/v1/systemone", litellm_request={ "model": "strands_decider/strands-decider-2B-hobson-v19", diff --git a/tests/integration/translation/decisions/bases/typesafe.py b/tests/integration/translation/decisions/bases/typesafe.py index 5da5b05cb09..7d47ce6b082 100644 --- a/tests/integration/translation/decisions/bases/typesafe.py +++ b/tests/integration/translation/decisions/bases/typesafe.py @@ -6,6 +6,98 @@ from integration.translation.case import TranslationTestCase """ JEV_1_13_0_TEST_CASE: Final = TranslationTestCase( scenario="basic", + litellm_endpoint="/v1/decisions", + litellm_request={ + "model": "typesafe/jev-1.13.0", + "input": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": [ + {"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}, + { + "type": "choice", + "name": "severity", + "instructions": "How severe is it?", + "choices": [ + {"value": "low", "description": "cosmetic"}, + {"value": "high", "description": "blocks users"}, + ], + }, + { + "type": "score", + "name": "confidence", + "instructions": "How sure are you?", + "levels": [{"label": "unsure"}, {"label": "sure"}], + }, + ], + "cache": {"no-cache": True}, + }, + expected_provider_endpoint="/v1/systemone", + expected_provider_headers={"authorization": "Bearer synthetic-typesafe-key", "content-type": "application/json"}, + expected_provider_request={ + "model": "jev-1.13.0", + "state": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": { + "defect": {"type": "noul", "instructions": "Is this a defect?"}, + "severity": { + "type": "choice", + "instructions": "How severe is it?", + "criteria": {"low": "cosmetic", "high": "blocks users"}, + }, + "confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]}, + }, + }, + mock_provider_response={ + "model": "jev-1.13.0", + "answers": { + "defect": {"type": "noul", "noul": 0.78}, + "severity": { + "type": "choice", + "choice": "high", + "confidence": 0.99, + "probabilities": {"low": 0.01, "high": 0.99}, + }, + "confidence": { + "type": "score", + "score": 0.51, + "confidence": 0.03, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.49, "1": 0.51}, + }, + }, + "usage": {"input_tokens": 377, "output_tokens": 62}, + }, + expected_litellm_response={ + "model": "jev-1.13.0", + "answers": [ + {"type": "predicate", "name": "defect", "probability": 0.78}, + { + "type": "choice", + "name": "severity", + "choice": "high", + "probabilities": [{"value": "low", "probability": 0.01}, {"value": "high", "probability": 0.99}], + "confidence": 0.99, + }, + { + "type": "score", + "name": "confidence", + "score": 0.51, + "probabilities": [ + {"value": 0, "label": "unsure", "probability": 0.49}, + {"value": 1, "label": "sure", "probability": 0.51}, + ], + "confidence": 0.03, + }, + ], + "usage": { + "input_tokens": 377, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 62, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 439, + }, + }, +) +JEV_1_13_0_SYSTEMONE_TEST_CASE: Final = TranslationTestCase( + scenario="systemone", litellm_endpoint="/v1/systemone", litellm_request={ "model": "typesafe/jev-1.13.0", diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_cloudflare.py b/tests/integration/translation/decisions/basic/test_decisions_basic_cloudflare.py index 79fc0574b3b..21e0538897d 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_cloudflare.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_cloudflare.py @@ -2,10 +2,10 @@ import pytest from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration.translation.case import TranslationTestCase -from integration.translation.decisions.bases.cloudflare import CLEF_TEST_CASE +from integration.translation.decisions.bases.cloudflare import CLEF_TEST_CASE, CLEF_SYSTEMONE_TEST_CASE from integration.translation.runner import assert_translation -@pytest.mark.parametrize("case", [CLEF_TEST_CASE], ids=lambda case: case.id) +@pytest.mark.parametrize("case", [CLEF_TEST_CASE, CLEF_SYSTEMONE_TEST_CASE], ids=lambda case: case.id) def test_decisions_basic_cloudflare(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_openrouter.py b/tests/integration/translation/decisions/basic/test_decisions_basic_openrouter.py index a74cb23f6b4..973526d7272 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_openrouter.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_openrouter.py @@ -2,10 +2,15 @@ import pytest from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration.translation.case import TranslationTestCase -from integration.translation.decisions.bases.openrouter import TYPESAFE_JEV_1_13_TEST_CASE +from integration.translation.decisions.bases.openrouter import ( + TYPESAFE_JEV_1_13_TEST_CASE, + TYPESAFE_JEV_1_13_SYSTEMONE_TEST_CASE, +) from integration.translation.runner import assert_translation -@pytest.mark.parametrize("case", [TYPESAFE_JEV_1_13_TEST_CASE], ids=lambda case: case.id) +@pytest.mark.parametrize( + "case", [TYPESAFE_JEV_1_13_TEST_CASE, TYPESAFE_JEV_1_13_SYSTEMONE_TEST_CASE], ids=lambda case: case.id +) def test_decisions_basic_openrouter(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_perplexity.py b/tests/integration/translation/decisions/basic/test_decisions_basic_perplexity.py index e341d24ba53..d9662ec4b9f 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_perplexity.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_perplexity.py @@ -2,10 +2,15 @@ import pytest from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration.translation.case import TranslationTestCase -from integration.translation.decisions.bases.perplexity import PPLX_DECIDER_V1_27B_TEST_CASE +from integration.translation.decisions.bases.perplexity import ( + PPLX_DECIDER_V1_27B_TEST_CASE, + PPLX_DECIDER_V1_27B_SYSTEMONE_TEST_CASE, +) from integration.translation.runner import assert_translation -@pytest.mark.parametrize("case", [PPLX_DECIDER_V1_27B_TEST_CASE], ids=lambda case: case.id) +@pytest.mark.parametrize( + "case", [PPLX_DECIDER_V1_27B_TEST_CASE, PPLX_DECIDER_V1_27B_SYSTEMONE_TEST_CASE], ids=lambda case: case.id +) def test_decisions_basic_perplexity(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_strands_decider.py b/tests/integration/translation/decisions/basic/test_decisions_basic_strands_decider.py index cae7ca70c78..2d2cd9caaf5 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_strands_decider.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_strands_decider.py @@ -2,10 +2,17 @@ import pytest from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration.translation.case import TranslationTestCase -from integration.translation.decisions.bases.strands_decider import STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE +from integration.translation.decisions.bases.strands_decider import ( + STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE, + STRANDS_DECIDER_2B_HOBSON_V19_SYSTEMONE_TEST_CASE, +) from integration.translation.runner import assert_translation -@pytest.mark.parametrize("case", [STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE], ids=lambda case: case.id) +@pytest.mark.parametrize( + "case", + [STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE, STRANDS_DECIDER_2B_HOBSON_V19_SYSTEMONE_TEST_CASE], + ids=lambda case: case.id, +) def test_decisions_basic_strands_decider(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_typesafe.py b/tests/integration/translation/decisions/basic/test_decisions_basic_typesafe.py index 302636e6fb4..c1111cc6ba9 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_typesafe.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_typesafe.py @@ -2,10 +2,10 @@ import pytest from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration.translation.case import TranslationTestCase -from integration.translation.decisions.bases.typesafe import JEV_1_13_0_TEST_CASE +from integration.translation.decisions.bases.typesafe import JEV_1_13_0_TEST_CASE, JEV_1_13_0_SYSTEMONE_TEST_CASE from integration.translation.runner import assert_translation -@pytest.mark.parametrize("case", [JEV_1_13_0_TEST_CASE], ids=lambda case: case.id) +@pytest.mark.parametrize("case", [JEV_1_13_0_TEST_CASE, JEV_1_13_0_SYSTEMONE_TEST_CASE], ids=lambda case: case.id) def test_decisions_basic_typesafe(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: assert_translation(case, gateway, provider) diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 8f05a492119..c1cac2d386f 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -1262,6 +1262,141 @@ async def test_openrouter_decisions_uses_provider_reported_cost_without_cost_map assert get_response_cost_from_hidden_params(response.hidden_params) == cost +_SAFETY_IDENTIFIER: Final = "end-user-7" +_PREDICATE_QUESTIONS: Final[tuple[Mapping[str, object], ...]] = ( + {"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}, +) +_PREDICATE_REQUESTS: Final[tuple[Mapping[str, object], ...]] = ( + MappingProxyType({"input": "review", "questions": _PREDICATE_QUESTIONS}), + MappingProxyType( + {"state": "review", "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}} + ), +) +_PREDICATE_REQUEST_IDS: Final = ("input_format", "state_format") + + +@pytest.mark.parametrize("request_kwargs", _PREDICATE_REQUESTS, ids=_PREDICATE_REQUEST_IDS) +def test_safety_identifier_is_refused_before_http_when_the_provider_cannot_take_it( + request_kwargs: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond(json=_RESPONSE) + + with pytest.raises(litellm.UnsupportedParamsError, match=r"safety_identifier.*drop_params") as caught: + litellm.decisions( + model="perplexity/pplx-decider-v1-27b", + safety_identifier=_SAFETY_IDENTIFIER, + api_key="caller-key", + **request_kwargs, + ) + + assert caught.value.status_code == 400 + assert not route.called + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope", ("call", "global")) +async def test_safety_identifier_is_dropped_from_the_wire_under_drop_params( + scope: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", scope == "global") + route: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond(json=_RESPONSE) + + response: Final = await litellm.adecisions( + model="perplexity/pplx-decider-v1-27b", + input="review", + questions=_PREDICATE_QUESTIONS, + safety_identifier=_SAFETY_IDENTIFIER, + api_key="caller-key", + **({"drop_params": True} if scope == "call" else {}), + ) + + assert route.called + assert json.loads(respx_mock.calls[0].request.content) == { + "model": "pplx-decider-v1-27b", + "state": "review", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert isinstance(response, OpenAIDecisionResponse) + assert [answer.model_dump(mode="json") for answer in response.answers] == [ + {"type": "predicate", "name": "is_defect", "probability": 0.9} + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_kwargs", _PREDICATE_REQUESTS, ids=_PREDICATE_REQUEST_IDS) +@pytest.mark.parametrize("drop_params", (False, True), ids=("strict", "drop_params")) +async def test_safety_identifier_reaches_openai_whether_or_not_params_are_dropped( + drop_params: bool, + request_kwargs: Mapping[str, object], + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", drop_params) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE) + + await litellm.adecisions( + model="openai/gpt-6-luna", + safety_identifier=_SAFETY_IDENTIFIER, + api_key="caller-key", + **request_kwargs, + ) + + assert route.called + assert json.loads(respx_mock.calls[0].request.content) == { + "model": "gpt-6-luna", + "input": "review", + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + "safety_identifier": _SAFETY_IDENTIFIER, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_kwargs", _PREDICATE_REQUESTS, ids=_PREDICATE_REQUEST_IDS) +async def test_non_string_safety_identifier_is_rejected_before_http( + request_kwargs: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE) + invalid_safety_identifier: Final[Mapping[str, object]] = MappingProxyType({"safety_identifier": 7}) + + with pytest.raises(litellm.BadRequestError, match=r"(?s)Invalid Decisions request.*safety_identifier"): + await litellm.adecisions( + model="openai/gpt-6-luna", api_key="caller-key", **request_kwargs, **invalid_safety_identifier + ) + + assert not route.called + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_kwargs", _PREDICATE_REQUESTS, ids=_PREDICATE_REQUEST_IDS) +async def test_non_string_safety_identifier_is_dropped_from_the_wire_under_drop_params( + request_kwargs: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE) + dropped_safety_identifier: Final[Mapping[str, object]] = MappingProxyType( + {"safety_identifier": 7, "drop_params": True} + ) + + await litellm.adecisions( + model="openai/gpt-6-luna", api_key="caller-key", **request_kwargs, **dropped_safety_identifier + ) + + assert route.called + assert json.loads(respx_mock.calls[0].request.content) == { + "model": "gpt-6-luna", + "input": "review", + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + } + + @pytest.mark.asyncio @pytest.mark.parametrize( "api_base", diff --git a/tests/unit/proxy/decisions_endpoints/test_endpoints.py b/tests/unit/proxy/decisions_endpoints/test_endpoints.py index 6ce15636f1e..ed78fc7063a 100644 --- a/tests/unit/proxy/decisions_endpoints/test_endpoints.py +++ b/tests/unit/proxy/decisions_endpoints/test_endpoints.py @@ -575,6 +575,68 @@ def test_openai_format_decisions_reach_an_openai_deployment_unchanged_including_ assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected_cost) +@pytest.mark.parametrize("safety_identifier", (7, ["end-user-1"]), ids=("numeric", "list")) +@pytest.mark.parametrize("deployment_drops_params", (False, True), ids=("strict", "drop_params")) +def test_a_non_string_safety_identifier_is_refused_unless_the_deployment_drops_params( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + safety_identifier: object, + deployment_drops_params: bool, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "decider", + "litellm_params": { + "model": "openai/gpt-6-luna", + "api_key": "k", + "drop_params": deployment_drops_params, + }, + } + ] + ), + ) + upstream: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond( + json={ + "model": "gpt-6-luna", + "answers": _OPENAI_FORMAT_ANSWERS[:1], + "usage": { + "input_tokens": _INPUT_TOKENS, + "output_tokens": _OUTPUT_TOKENS, + "total_tokens": _INPUT_TOKENS + _OUTPUT_TOKENS, + }, + } + ) + request_body: Final = { + "model": "decider", + "input": "The package arrived with a broken screen.", + "questions": _OPENAI_FORMAT_REQUEST["questions"][:1], + "safety_identifier": safety_identifier, + } + + response: Final = client.post("/v1/decisions", json=request_body) + + if not deployment_drops_params: + assert response.status_code == 400, response.text + assert "safety_identifier" in response.json()["error"]["message"] + assert not upstream.called + return + assert response.status_code == 200, response.text + assert json.loads(upstream.calls[0].request.content) == { + "model": "gpt-6-luna", + "input": "The package arrived with a broken screen.", + "questions": _OPENAI_FORMAT_REQUEST["questions"][:1], + } + + def _decisions_feature() -> LazyFeature: return next(feature for feature in LAZY_FEATURES if feature.name == "decisions") From 6830a5c33844afe78497b9f20dd53d5b9140b947 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 9 Oct 2026 19:36:49 -0700 Subject: [PATCH 019/104] fix(python-bridge): preserve callback keyword deletions in native requests (#45692) * fix(messages): preserve callback keyword deletions in native requests * fix(python-bridge): share hook-aware request view across native route hosts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/routes/chat_completions/host.rs | 10 +- .../src/routes/chat_completions/mod.rs | 2 +- .../python-bridge/src/routes/inference.rs | 32 ++++-- .../python-bridge/src/routes/messages/host.rs | 34 +++--- .../python-bridge/src/routes/messages/mod.rs | 2 +- .../crates/python-bridge/src/routes/mod.rs | 101 ++++++++++++++++++ .../python-bridge/src/routes/ocr/host.rs | 19 ++-- .../python-bridge/src/routes/ocr/mod.rs | 2 +- .../src/routes/responses/host.rs | 10 +- .../python-bridge/src/routes/responses/mod.rs | 2 +- .../messages/test_callbacks.py | 32 ++++++ tests/test_litellm_rust/ocr/test_callbacks.py | 19 +++- tests/test_litellm_rust/test_inference.py | 22 +++- 13 files changed, 242 insertions(+), 45 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs index 1547a4daa93..5ad9f8df8b8 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs @@ -12,7 +12,7 @@ use pyo3::{ pub(super) struct ChatCompletionsPythonHost(pub InferenceHost); pub(super) fn project( - host: &InferenceHost, + host: &mut InferenceHost, py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> PyResult { @@ -38,7 +38,7 @@ impl PythonBinding for ChatCompletionsPythonHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> Result> { - let call = project(&self.0, py, arguments).map_err(InvokeError::Python)?; + let call = project(&mut self.0, py, arguments).map_err(InvokeError::Python)?; if call .optional_params .get("stream") @@ -94,8 +94,10 @@ impl PythonHostCalls for ChatCompletionsPythonHost { } impl PythonOwned for ChatCompletionsPythonHost { - fn close(&mut self, _: Python<'_>) {} + fn close(&mut self, _: Python<'_>) { + self.0.close(); + } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.0.request) + self.0.traverse(visit) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs index 2b69d0e6321..00ca5e092a8 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs @@ -22,7 +22,7 @@ fn run_chat_completions( asynchronous: bool, ) -> PyResult> { let host = InferenceHost::new( - call.resolved()?.unbind(), + call.view()?, "litellm.rust_bridge.chat_completions.route_host", ); run_inference::( diff --git a/litellm-rust/crates/python-bridge/src/routes/inference.rs b/litellm-rust/crates/python-bridge/src/routes/inference.rs index 0333aaeb04a..bc442d1cf36 100644 --- a/litellm-rust/crates/python-bridge/src/routes/inference.rs +++ b/litellm-rust/crates/python-bridge/src/routes/inference.rs @@ -5,15 +5,20 @@ use litellm_cache_response::{CachePolicy, CacheScope, ScopedCache}; use litellm_callbacks_legacy_python::LoggingOperation; use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; use litellm_host::{call::HostedMachine, protocol::Protocol}; -use litellm_host_python::{PythonBinding, PythonHostCalls, from_py, present}; +use litellm_host_python::{PythonBinding, PythonHostCalls, from_py}; use litellm_http::transport::Error as TransportError; use litellm_inference::RouteError; use litellm_secrets::source::SecretSource; -use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use pyo3::{ + exceptions::PyValueError, + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::PyDict, +}; use serde::Serialize; use serde_json::{Map, Value}; -use super::NativeCall; +use super::{NativeCall, RequestView}; use crate::{ errors::{RustUpstreamError, route_error_to_pyerr}, marshal::{ @@ -23,7 +28,7 @@ use crate::{ }; pub(super) struct InferenceHost { - pub request: Py, + view: RequestView, module: &'static str, } @@ -34,16 +39,17 @@ pub(super) struct ProjectedCall { } impl InferenceHost { - pub fn new(request: Py, module: &'static str) -> Self { - Self { request, module } + pub fn new(view: RequestView, module: &'static str) -> Self { + Self { view, module } } pub fn project( - &self, + &mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>, input: &str, ) -> PyResult { + self.view.resolve(py, arguments)?; let argument = |name: &str| self.argument(py, arguments, name); let string = |name: &str| -> PyResult> { argument(name)?.map(|value| value.extract()).transpose() @@ -93,7 +99,15 @@ impl InferenceHost { arguments: &Bound<'py, PyDict>, name: &str, ) -> PyResult>> { - present(arguments, self.request.bind(py), name) + self.view.argument(py, arguments, name) + } + + pub fn close(&mut self) { + self.view.close(); + } + + pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.view.traverse(visit) } pub fn parameters( @@ -130,7 +144,7 @@ impl InferenceHost { let mapped = py .import(self.module)? .getattr("map_failure")? - .call1((native.value(py), self.request.bind(py)))?; + .call1((native.value(py), self.view.request(py)?))?; Ok(PyErr::from_value(mapped)) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index d34047c735d..adeba397b26 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -2,7 +2,7 @@ use crate::cache::{CacheCall, Cached, PythonCache, Selection}; use litellm_host_python::{PythonHostCalls, PythonOwned}; use bytes::Bytes; -use litellm_host_python::{InvokeError, PythonBinding, from_py, present, to_py}; +use litellm_host_python::{InvokeError, PythonBinding, from_py, to_py}; use litellm_http::transport::Error as TransportError; use litellm_inference_messages::{ Error, MessagesCall, MessagesSettings, MessagesShaping, litellm_params, messages_body, @@ -23,6 +23,7 @@ use crate::{ errors::{RustUpstreamError, route_error_to_pyerr}, marshal::{optional_timeout, project_optional_fields, public_response, python_timeout_seconds}, python_settings::missing_module, + routes::RequestView, }; const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.messages.route_host"; @@ -118,14 +119,14 @@ fn native_error(py: Python<'_>, error: Error) -> PyResult { /// The Python side of the Messages route: projects the prepared arguments and builds the /// public response, chunks and exceptions. pub(super) struct MessagesPythonHost { - request: Py, + view: RequestView, cache: PythonCache, } impl MessagesPythonHost { - pub(super) fn new(request: Py, asynchronous: bool) -> Self { + pub(super) fn new(view: RequestView, asynchronous: bool) -> Self { Self { - request, + view, cache: PythonCache::new(asynchronous), } } @@ -135,8 +136,7 @@ impl MessagesPythonHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> PyResult> { - let request = self.request.bind(py); - let argument = |name: &str| present(arguments, request, name); + let argument = |name: &str| self.view.argument(py, arguments, name); let string = |name: &str| -> PyResult> { argument(name)?.map(|value| value.extract()).transpose() }; @@ -182,9 +182,9 @@ impl MessagesPythonHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> PyResult>> { - let request = self.request.bind(py); let mapping = |name: &str| -> PyResult>> { - present(arguments, request, name)? + self.view + .argument(py, arguments, name)? .map(|value| from_py(&value)) .transpose() }; @@ -199,7 +199,8 @@ impl MessagesPythonHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> PyResult> { - present(arguments, self.request.bind(py), "provider_specific_header")? + self.view + .argument(py, arguments, "provider_specific_header")? .map(|value| from_py(&value)) .transpose() } @@ -226,9 +227,9 @@ impl MessagesPythonHost { } fn provider(&self, py: Python<'_>) -> String { - self.request - .bind(py) - .get_item("custom_llm_provider") + self.view + .request(py) + .and_then(|request| request.get_item("custom_llm_provider")) .ok() .flatten() .and_then(|value| value.extract::>().ok().flatten()) @@ -242,7 +243,7 @@ impl MessagesPythonHost { let mapped = py .import(ROUTE_HOST_MODULE) .and_then(|module| module.getattr("map_failure")) - .and_then(|map| map.call1((error.value(py), self.request.bind(py), self.provider(py)))) + .and_then(|map| map.call1((error.value(py), self.view.request(py)?, self.provider(py)))) .and_then(|mapped| { mapped .extract::>() @@ -264,6 +265,9 @@ impl PythonBinding for MessagesPythonHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> Result<(MessagesCall, Selection), InvokeError> { + self.view + .resolve(py, arguments) + .map_err(InvokeError::Python)?; let selection = crate::cache::configure(&mut self.cache, py, arguments, "anthropic_messages") .map_err(InvokeError::Python)?; @@ -341,16 +345,18 @@ impl PythonHostCalls> for MessagesPythonHost { impl PythonOwned for MessagesPythonHost { fn close(&mut self, _: Python<'_>) { + self.view.close(); self.cache.close(); } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.request)?; + self.view.traverse(visit)?; self.cache.traverse(visit) } } #[cfg(test)] mod tests { + use litellm_host_python::present; use rstest::rstest; use serde_json::json; diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index 5f69cf078b7..c788760b268 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -44,7 +44,7 @@ fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyR }, )) }, - MessagesPythonHost::new(call.resolved()?.unbind(), asynchronous), + MessagesPythonHost::new(call.view()?, asynchronous), hooks, asynchronous, ) diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 7427ec8da7d..416b94db55b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -14,7 +14,9 @@ use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol} use litellm_host_python::{ HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args, }; +use litellm_host_python::{missing_state, present}; use pyo3::{ + gc::{PyTraverseError, PyVisit}, prelude::*, types::{PyDict, PyMapping, PyTuple}, }; @@ -36,6 +38,69 @@ impl<'py> NativeCall<'py> { fn resolved(&self) -> PyResult> { effective_py_args(&self.base, &self.kwargs) } + + /// The view a route host reads its arguments from. + fn view(&self) -> PyResult { + RequestView::new(&self.base, &self.kwargs) + } +} + +/// The keyword view a route host reads: the signature base under the caller's keywords, +/// laid again once the hooks have rewritten them so a keyword a hook deleted stays deleted +/// instead of falling back to the caller's value. +pub(crate) struct RequestView(Option); + +struct RequestDicts { + base: Py, + request: Py, +} + +impl RequestView { + pub(crate) fn new(base: &Bound<'_, PyDict>, kwargs: &Bound<'_, PyDict>) -> PyResult { + Ok(Self(Some(RequestDicts { + base: base.clone().unbind(), + request: effective_py_args(base, kwargs)?.unbind(), + }))) + } + + fn dicts(&self) -> PyResult<&RequestDicts> { + self.0.as_ref().ok_or_else(missing_state) + } + + pub(crate) fn resolve( + &mut self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + ) -> PyResult<()> { + let dicts = self.0.as_mut().ok_or_else(missing_state)?; + dicts.request = effective_py_args(dicts.base.bind(py), arguments)?.unbind(); + Ok(()) + } + + pub(crate) fn request<'py>(&self, py: Python<'py>) -> PyResult<&Bound<'py, PyDict>> { + Ok(self.dicts()?.request.bind(py)) + } + + pub(crate) fn argument<'py>( + &self, + py: Python<'py>, + arguments: &Bound<'py, PyDict>, + name: &str, + ) -> PyResult>> { + present(arguments, self.request(py)?, name) + } + + pub(crate) fn close(&mut self) { + self.0 = None; + } + + pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + let Some(dicts) = &self.0 else { + return Ok(()); + }; + visit.call(&dicts.base)?; + visit.call(&dicts.request) + } } fn mapping_dict<'py>(value: &Bound<'py, PyAny>) -> PyResult> { @@ -130,6 +195,42 @@ mod tests { .unwrap() } + fn dict<'py>(py: Python<'py>, source: &str) -> Bound<'py, PyDict> { + py.eval(&std::ffi::CString::new(source).unwrap(), None, None) + .unwrap() + .cast_into::() + .unwrap() + } + + #[rstest] + #[case::deleted_keyword_stays_deleted("{'system': None}", "{'system': 'caller'}", "{}", None)] + #[case::deleted_keyword_without_default("{}", "{'system': 'caller'}", "{}", None)] + #[case::rewritten_keyword_wins( + "{'system': None}", + "{'system': 'caller'}", + "{'system': 'hook'}", + Some("hook") + )] + #[case::signature_default_survives("{'system': 'default'}", "{}", "{}", Some("default"))] + fn request_view_reads_the_hook_rewritten_keywords_over_the_base( + #[case] base: &str, + #[case] kwargs: &str, + #[case] prepared: &str, + #[case] expected: Option<&str>, + ) { + Python::initialize(); + Python::attach(|py| { + let mut view = super::RequestView::new(&dict(py, base), &dict(py, kwargs)).unwrap(); + let prepared = dict(py, prepared); + view.resolve(py, &prepared).unwrap(); + let value = view + .argument(py, &prepared, "system") + .unwrap() + .map(|value| value.extract::().unwrap()); + assert_eq!(value.as_deref(), expected); + }); + } + #[rstest] #[case::sync("transcription")] #[case::asynchronous("atranscription")] diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 5ae72acd104..dd3c648cc3e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -15,7 +15,7 @@ use super::{ errors::to_pyerr as ocr_error_to_pyerr, project::{OcrHostHandles, SETTINGS_ERROR_MARKER, project_request}, }; -use crate::marshal::public_response; +use crate::{marshal::public_response, routes::RequestView}; enum OcrHostData { Unprojected, @@ -27,14 +27,14 @@ enum OcrHostData { /// document as it goes), acquires Azure AD tokens, and builds the public response and /// exception. pub(super) struct OcrPythonHost { - request: Py, + view: RequestView, data: OcrHostData, } impl OcrPythonHost { - pub(super) fn new(request: Py) -> Self { + pub(super) fn new(view: RequestView) -> Self { Self { - request, + view, data: OcrHostData::Unprojected, } } @@ -58,7 +58,8 @@ impl OcrPythonHost { let OcrHostData::Unprojected = self.data else { return Err(missing_state()); }; - let (request, handles) = project_request(self.request.bind(py), arguments)?; + self.view.resolve(py, arguments)?; + let (request, handles) = project_request(self.view.request(py)?, arguments)?; let caller_token = handles.azure_ad_token_provider.is_some(); self.data = OcrHostData::Projected(Box::new(handles)); Ok(OcrCall { @@ -78,7 +79,7 @@ impl OcrPythonHost { let mapped = py .import("litellm.rust_bridge.ocr.route_host") .and_then(|module| module.getattr("map_failure")) - .and_then(|map| map.call1((error.value(py), self.request.bind(py), provider))) + .and_then(|map| map.call1((error.value(py), self.view.request(py)?, provider))) .and_then(|mapped| mapped.extract::>().map_err(PyErr::from)); match mapped { Ok(mapped) => PyErr::from_value(mapped.into_bound(py).into_any()), @@ -160,10 +161,11 @@ impl PythonHostCalls for OcrPythonHost { impl PythonOwned for OcrPythonHost { fn close(&mut self, _: Python<'_>) { + self.view.close(); self.data = OcrHostData::Released; } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.request)?; + self.view.traverse(visit)?; if let OcrHostData::Projected(handles) = &self.data && let Some(provider) = &handles.azure_ad_token_provider { @@ -218,7 +220,8 @@ del provider .unwrap() .cast_into::() .unwrap(); - let mut host = OcrPythonHost::new(PyDict::new(py).unbind()); + let empty = PyDict::new(py); + let mut host = OcrPythonHost::new(RequestView::new(&empty, &empty).unwrap()); assert!(host.decode_request(py, &kwargs).unwrap().caller_token); locals.del_item("kwargs").unwrap(); drop(kwargs); diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index fab9ecc13c1..ec7f2ce79d9 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -44,7 +44,7 @@ fn run_ocr(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult let route = litellm_inference_ocr::OcrRoute::new(client); Ok(route.machine(request, None)) }, - OcrPythonHost::new(call.resolved()?.unbind()), + OcrPythonHost::new(call.view()?), hooks, asynchronous, ) diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs index b22bbc6192b..8261e66406b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs @@ -12,7 +12,7 @@ use pyo3::{ pub(super) struct ResponsesPythonHost(pub InferenceHost); pub(super) fn project( - host: &InferenceHost, + host: &mut InferenceHost, py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> PyResult { @@ -59,7 +59,7 @@ impl PythonBinding for ResponsesPythonHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> Result> { - let call = project(&self.0, py, arguments).map_err(InvokeError::Python)?; + let call = project(&mut self.0, py, arguments).map_err(InvokeError::Python)?; if call .optional_params .get("stream") @@ -121,8 +121,10 @@ impl PythonHostCalls for ResponsesPythonHost { } impl PythonOwned for ResponsesPythonHost { - fn close(&mut self, _: Python<'_>) {} + fn close(&mut self, _: Python<'_>) { + self.0.close(); + } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.0.request) + self.0.traverse(visit) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs index 50e5790a532..1bf24813103 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs @@ -61,7 +61,7 @@ fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> Py "native Python responses streaming", )); } - let host = InferenceHost::new(resolved.unbind(), ROUTE_HOST_MODULE); + let host = InferenceHost::new(call.view()?, ROUTE_HOST_MODULE); run_inference::(py, call, asynchronous, ResponsesPythonHost(host)) } diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index dc66852d214..b3cd7bd01f9 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -239,3 +239,35 @@ async def test_native_messages_stream_success_log_carries_usage_rebuilt_from_the assert usage.completion_tokens == MESSAGES_EVENTS[4][1]["usage"]["output_tokens"] assert usage.prompt_tokens == MESSAGES_RESPONSE["usage"]["input_tokens"] assert success[0].response.choices[0].message.content == "Hello from native Messages" + + +@pytest.mark.asyncio +async def test_native_messages_does_not_restore_keywords_deleted_by_a_pre_call_hook( + messages_server: RecordingServer, +) -> None: + class DropKeywords(CustomLogger): + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: object + ) -> dict[str, object]: + return {name: value for name, value in kwargs.items() if name not in ("system", "extra_headers")} + + call_arguments: Final = arguments( + messages_server, + system="remove this instruction", + extra_headers={"x-deleted": "remove this header"}, + headers={"x-kept": "keep this header"}, + stream=False, + ) + with rebound(litellm, "callbacks", [DropKeywords()]): + await litellm.anthropic.messages.acreate(**call_arguments) + + assert_served_natively(messages_server) + sent: Final = messages_server.requests[0] + assert sent.body == { + "model": MESSAGES_MODEL.removeprefix("anthropic/"), + "messages": call_arguments["messages"], + "max_tokens": call_arguments["max_tokens"], + "stream": call_arguments["stream"], + } + assert sent.headers.get("x-deleted") is None + assert sent.headers.get("x-kept") == "keep this header" diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index d3b04c8bb6f..9e3945f11b1 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -16,7 +16,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.ocr.transformation import OCRResponse from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging -from tests.test_litellm_rust.support.isolation import isolated_callback_registries +from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import ( OCR_DOCUMENT, @@ -65,6 +65,23 @@ def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_ } +@pytest.mark.asyncio +async def test_native_ocr_does_not_restore_keywords_deleted_by_a_pre_call_hook(ocr_server: RecordingServer) -> None: + class DropKeywords(CustomLogger): + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: object + ) -> dict[str, object]: + return {name: value for name, value in kwargs.items() if name not in ("pages", "extra_headers")} + + with rebound(litellm, "callbacks", [DropKeywords()]): + await call_native_aocr(ocr_server, pages=[0], extra_headers={"x-deleted": "remove this header"}) + + assert len(ocr_server.requests) == 1 + sent: Final = ocr_server.requests[0] + assert sent.body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT} + assert sent.headers.get("x-deleted") is None + + @pytest.mark.parametrize("raise_after_edit", [False, True], ids=["callback-returns", "callback-raises"]) def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider( ocr_server: RecordingServer, raise_after_edit: bool diff --git a/tests/test_litellm_rust/test_inference.py b/tests/test_litellm_rust/test_inference.py index 8211b50cb8d..83a01e8ad4d 100644 --- a/tests/test_litellm_rust/test_inference.py +++ b/tests/test_litellm_rust/test_inference.py @@ -87,7 +87,6 @@ def native_call( "api_base": server.base_url, "custom_llm_provider": "openai", "extra_headers": None, - **response_kwargs, }, ) return (_native.aresponses if asynchronous else _native.responses)(response_request) @@ -147,6 +146,27 @@ async def test_native_inference_pre_call_edits_reach_the_provider( assert _OBJECT.validate_python(recording_server.requests[0].body)["temperature"] == 0.75 +@pytest.mark.asyncio +async def test_native_inference_does_not_restore_keywords_deleted_by_a_pre_call_hook( + route: Route, recording_server: RecordingServer +) -> None: + class DropKeywords(CustomLogger): + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object]: + return {name: value for name, value in kwargs.items() if name not in ("temperature", "extra_headers")} + + litellm.callbacks.append(DropKeywords()) + await execute( + route, True, recording_server, {"temperature": 0.25, "extra_headers": {"x-deleted": "remove this header"}} + ) + + assert len(recording_server.requests) == 1 + sent: Final = recording_server.requests[0] + assert "temperature" not in _OBJECT.validate_python(sent.body) + assert sent.headers.get("x-deleted") is None + + @pytest.mark.asyncio @pytest.mark.parametrize("from_credentials", (False, True)) async def test_native_resource_setup_uses_deployment_hook_arguments( From 674208bb7dac89b177df00293bc8cd03725352e8 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 10 Oct 2026 02:40:28 +0000 Subject: [PATCH 020/104] fix(model_prices): drop unverified output price for azure_ai/Microsoft-Decision-1 (#45717) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ci_cd/cost_map_guard.py | 56 +++++++++++++++-- ...odel_prices_and_context_window_backup.json | 2 +- model_prices_and_context_window.json | 2 +- tests/unit/test_cost_map_guard.py | 62 ++++++++++++++++++- 4 files changed, 114 insertions(+), 8 deletions(-) diff --git a/ci_cd/cost_map_guard.py b/ci_cd/cost_map_guard.py index 5842cf6f1ac..14eecf01461 100644 --- a/ci_cd/cost_map_guard.py +++ b/ci_cd/cost_map_guard.py @@ -14,8 +14,9 @@ import argparse import json import subprocess import sys -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass +from types import MappingProxyType from typing import Final from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors @@ -25,6 +26,7 @@ BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json" SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json" GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH) BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_" +EVALUATION_OUTPUT_PRICE_SOURCES: Final[Mapping[str, str]] = MappingProxyType({}) CostMap = dict[str, object] @@ -51,7 +53,31 @@ def _rendered_schema(cost_map: CostMap) -> str: return str(error) -def _file_failures(head: Snapshot, head_map: CostMap) -> tuple[str, ...]: +def _evaluation_output_price_failure( + key: str, entry: object, evaluation_output_price_sources: Mapping[str, str] +) -> str | None: + if not isinstance(entry, dict) or entry.get("mode") != "evaluation": + return None + output_cost: Final = entry.get("output_cost_per_token") + if ( + not isinstance(output_cost, (int, float)) + or isinstance(output_cost, bool) + or output_cost <= 0 + or key in evaluation_output_price_sources + ): + return None + return ( + f"{key}: mode evaluation rows bill input only; output_cost_per_token must be 0 unless the vendor " + "prices output tokens for this model (add the source to EVALUATION_OUTPUT_PRICE_SOURCES)" + ) + + +def _file_failures( + head: Snapshot, + head_map: CostMap, + *, + evaluation_output_price_sources: Mapping[str, str] = EVALUATION_OUTPUT_PRICE_SOURCES, +) -> tuple[str, ...]: schema_text: Final = _rendered_schema(head_map) if not schema_text.startswith("{"): return (schema_text,) @@ -75,6 +101,16 @@ def _file_failures(head: Snapshot, head_map: CostMap) -> tuple[str, ...]: f"{COST_MAP_PATH} does not validate against its schema: {error}" for error in validation_errors(head_map, json.loads(schema_text))[:20] ), + *( + failure + for key, entry in head_map.items() + if ( + failure := _evaluation_output_price_failure( + key, entry, evaluation_output_price_sources + ) + ) + is not None + ), ) @@ -121,13 +157,25 @@ def contract_for(bot: bool, changed_files: Sequence[str]) -> str: return "human PR, file checks only" if touches_cost_map(changed_files) else "human PR, cost map untouched" -def guard_failures(base: Snapshot, head: Snapshot, changed_files: Sequence[str], bot: bool) -> tuple[str, ...]: +def guard_failures( + base: Snapshot, + head: Snapshot, + changed_files: Sequence[str], + bot: bool, + *, + evaluation_output_price_sources: Mapping[str, str] = EVALUATION_OUTPUT_PRICE_SOURCES, +) -> tuple[str, ...]: if not bot and not touches_cost_map(changed_files): return () head_map: Final = _parse_object(head.cost_map, COST_MAP_PATH) if isinstance(head_map, str): return (head_map,) - return (*_file_failures(head, head_map), *(_bot_failures(base, head_map, changed_files) if bot else ())) + return ( + *_file_failures( + head, head_map, evaluation_output_price_sources=evaluation_output_price_sources + ), + *(_bot_failures(base, head_map, changed_files) if bot else ()), + ) def _git(*args: str) -> str | None: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6d49e7cd7fd..400f153c115 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -71768,7 +71768,7 @@ "input_cost_per_token": 4.2e-08, "litellm_provider": "azure_ai", "mode": "evaluation", - "output_cost_per_token": 4.2e-08, + "output_cost_per_token": 0.0, "source": "https://prices.azure.com/api/retail/prices?$filter=productName eq 'Microsoft Decision Models'" }, "gpt-rosalind-discovery": { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6d49e7cd7fd..400f153c115 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -71768,7 +71768,7 @@ "input_cost_per_token": 4.2e-08, "litellm_provider": "azure_ai", "mode": "evaluation", - "output_cost_per_token": 4.2e-08, + "output_cost_per_token": 0.0, "source": "https://prices.azure.com/api/retail/prices?$filter=productName eq 'Microsoft Decision Models'" }, "gpt-rosalind-discovery": { diff --git a/tests/unit/test_cost_map_guard.py b/tests/unit/test_cost_map_guard.py index 1595bd9864f..e127aa08d29 100644 --- a/tests/unit/test_cost_map_guard.py +++ b/tests/unit/test_cost_map_guard.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping import importlib.util import json import subprocess @@ -62,8 +63,19 @@ def _snapshot(cost_map: dict[str, object], backup: str | None = None, schema: st BASE: Final = _snapshot(BASE_MAP) -def _failures(head: object, changed_files: tuple[str, ...] = MAP_FILES, bot: bool = True) -> tuple[str, ...]: - return guard.guard_failures(BASE, head, changed_files, bot) +def _failures( + head: object, + changed_files: tuple[str, ...] = MAP_FILES, + bot: bool = True, + evaluation_output_price_sources: Mapping[str, str] = guard.EVALUATION_OUTPUT_PRICE_SOURCES, +) -> tuple[str, ...]: + return guard.guard_failures( + BASE, + head, + changed_files, + bot, + evaluation_output_price_sources=evaluation_output_price_sources, + ) def test_in_sync_files_pass_for_humans_and_bots() -> None: @@ -71,6 +83,52 @@ def test_in_sync_files_pass_for_humans_and_bots() -> None: assert _failures(BASE, bot=True) == () +EVALUATION_MODEL: Final = "azure_ai/Microsoft-Decision-1" +EVALUATION_OUTPUT_PRICE_FAILURE: Final = ( + f"{EVALUATION_MODEL}: mode evaluation rows bill input only; output_cost_per_token must be 0 unless the vendor " + "prices output tokens for this model (add the source to EVALUATION_OUTPUT_PRICE_SOURCES)" +) + + +@pytest.mark.parametrize("bot", (True, False)) +def test_evaluation_rows_with_output_price_fail_for_humans_and_bots(bot: bool) -> None: + head: Final = _snapshot( + {**BASE_MAP, EVALUATION_MODEL: _entry(mode="evaluation", output_cost_per_token=4.2e-08)} + ) + assert _failures(head, bot=bot) == (EVALUATION_OUTPUT_PRICE_FAILURE,) + + +def test_evaluation_rows_with_zero_or_missing_output_price_pass() -> None: + zero_output: Final = _snapshot( + {**BASE_MAP, EVALUATION_MODEL: _entry(mode="evaluation", output_cost_per_token=0.0)} + ) + missing_output_entry: Final = { + key: value + for key, value in _entry(mode="evaluation").items() + if key != "output_cost_per_token" + } + missing_output: Final = _snapshot({**BASE_MAP, EVALUATION_MODEL: missing_output_entry}) + + assert _failures(zero_output) == () + assert _failures(missing_output) == () + + +def test_chat_rows_with_output_price_pass() -> None: + head: Final = _snapshot( + {**BASE_MAP, "openrouter/chat-with-output": _entry(output_cost_per_token=4.2e-08)} + ) + assert _failures(head) == () + + +def test_evaluation_output_price_source_allowlist_passes() -> None: + head: Final = _snapshot( + {**BASE_MAP, EVALUATION_MODEL: _entry(mode="evaluation", output_cost_per_token=4.2e-08)} + ) + source: Final = 'https://vendor.example/pricing "Output tokens are charged per token."' + + assert _failures(head, evaluation_output_price_sources={EVALUATION_MODEL: source}) == () + + def test_bot_may_add_and_reprice_models() -> None: head = _snapshot({**BASE_MAP, "openrouter/a": _entry(9e-06, supports_vision=True), "openrouter/c": _entry()}) assert _failures(head) == () From 63d823d28aff0fef8b0630e59f52c4d3adcdf7d2 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 19:45:50 -0700 Subject: [PATCH 021/104] fix: support OpenAI SDK 3 while retaining SDK 2 compatibility (#44927) * fix: support OpenAI SDK 3.x while retaining 2.x compatibility * fix: preserve HTTPX compatibility across OpenAI SDK versions * refactor: trim OpenAI SDK compatibility patch * refactor: simplify SDK client defaults * test: cover SDK compatibility across provider routes * fix: use shared names for SDK transport helpers * fix: address SDK compatibility review feedback * test(openai): send requests through the SDK API factory clients * fix(exceptions): keep the provider response on PaymentRequiredError under httpx 2 * ci(base-sdk): count httpx2 as base-only when a base dependency brings it in OpenAI SDK 3 depends on httpx2, so the litellm-core install at highest resolution pulls it in through openai. The base-only check now resolves the declared base dependency closure of litellm and litellm-core and accepts extras-only modules owned by a distribution inside that closure * test: cover OpenAI SDK 3 client paths end to end and read SDK HTTP headers case-insensitively * test: read C02 spend rows once, wait out the worker healthcheck on respawn, and resolve module owners on Python 3.10 --------- Co-authored-by: Marcus Wood Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .github/workflows/test-unit.yml | 71 +- litellm/__init__.py | 5 +- litellm/exceptions.py | 62 +- .../litellm_core_utils/completion_timeout.py | 25 +- litellm/llms/azure/common_utils.py | 8 +- litellm/llms/openai/common_utils.py | 49 +- litellm/llms/openai/completion/handler.py | 8 +- litellm/llms/openai/fine_tuning/handler.py | 5 +- .../llms/openai/image_variations/handler.py | 24 +- litellm/llms/openai/openai.py | 14 +- litellm/main.py | 42 +- litellm/proxy/proxy_server.py | 2 +- litellm/types/llms/openai.py | 26 +- litellm/types/router.py | 13 +- packaging/litellm-core/pyproject.toml | 2 +- pyproject.toml | 2 +- .../base_sdk_tests/check_base_sdk_install.py | 41 +- tests/base_sdk_tests/check_sdk_http.py | 7 +- .../test_openai_sdk_client_paths.py | 1200 +++++++++++++++++ .../test_openai_sdk_client_paths_chaos.py | 424 ++++++ .../test_openai_sdk_client_sessions_wire.py | 488 +++++++ tests/typing/openai_sdk_compat.py | 66 + tests/typing/pyrightconfig.json | 8 + .../otel/test_otel_v2_sources_of_truth.py | 14 +- tests/unit/llms/openai/test_openai.py | 130 +- .../llms/openai/test_openai_common_utils.py | 126 +- .../openai/test_openai_workload_identity.py | 20 +- .../test_completion_timeout_resolution.py | 35 + .../test_exception_header_preservation.py | 33 +- ...penai_embedding_encoding_format_default.py | 42 +- tests/unit/types/test_router.py | 10 + uv.lock | 2 +- 32 files changed, 2851 insertions(+), 153 deletions(-) create mode 100644 tests/integration/compatibility/test_openai_sdk_client_paths.py create mode 100644 tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py create mode 100644 tests/integration/sdk/test_openai_sdk_client_sessions_wire.py create mode 100644 tests/typing/openai_sdk_compat.py create mode 100644 tests/typing/pyrightconfig.json diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 1cf79049a66..2b9849793e9 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -1268,6 +1268,75 @@ jobs: fi echo "$chart: all $declared declared test suites ran" done + openai-sdk-compat: + name: OpenAI ${{ matrix.openai-version }} compatibility + permissions: + contents: read + pull-requests: read + runs-on: ubuntu-latest + timeout-minutes: 20 + strategy: + fail-fast: false + matrix: + openai-version: ["2.20.0", "3.0.0", "3.25.0"] + env: + LITELLM_LOCAL_MODEL_COST_MAP: "True" + PYTHONPATH: ${{ github.workspace }} + UV_PYTHON: "3.10" + OPENAI_VERSION: ${{ matrix.openai-version }} + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Detect relevant changes + id: changes + timeout-minutes: 2 + uses: ./.github/actions/detect-changes + + - name: Set up Python + if: steps.changes.outputs.decision != 'skip' + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: ${{ env.UV_PYTHON }} + + - name: Set up uv + if: steps.changes.outputs.decision != 'skip' + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Install dependencies + if: steps.changes.outputs.decision != 'skip' + run: | + .github/scripts/uv_sync_with_retries.sh --frozen --group dev --group proxy-dev --extra proxy --no-install-workspace + uv pip install "openai==${OPENAI_VERSION:?}" --exclude-newer-package openai=2026-10-07T00:00:00Z + uv run --no-sync python -c 'import os, openai; assert openai.__version__ == os.environ["OPENAI_VERSION"], openai.__version__' + + - name: Run OpenAI SDK compatibility tests + if: steps.changes.outputs.decision != 'skip' + run: | + uv run --no-sync pytest \ + tests/unit/llms/openai/test_openai.py \ + tests/unit/llms/openai/test_openai_common_utils.py \ + tests/unit/llms/openai/test_openai_file_content_streaming.py \ + tests/unit/test_openai_embedding_encoding_format_default.py \ + tests/unit/llms/inception/test_inception_completion_transformation.py::test_inception_fim_targets_fim_endpoint \ + tests/unit/llms/inception/test_inception_completion_transformation.py::test_inception_fim_async \ + tests/unit/llms/azure/test_azure_exception_mapping.py \ + tests/unit/test_completion_timeout_resolution.py \ + tests/unit/types/test_router.py \ + tests/unit/test_exception_header_preservation.py \ + tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py::test_speech_binary_response_still_streaming_reports_the_bytes_downloaded_so_far \ + --tb=short -q + + - name: Run workload identity tests with SDK 3 + if: steps.changes.outputs.decision != 'skip' && startsWith(matrix.openai-version, '3.') + run: uv run --no-sync pytest tests/unit/llms/openai/test_openai_workload_identity.py --tb=short -q + + - name: Check consumer types with the installed SDK + if: steps.changes.outputs.decision != 'skip' + run: uv run --no-sync basedpyright --pythonpath .venv/bin/python --project tests/typing/pyrightconfig.json coverage: name: coverage needs: unit @@ -1316,7 +1385,7 @@ jobs: unit-passed: name: unit passed permissions: {} - needs: [rust-bridge, assert-shard-coverage, unit, ui-unit, docs, helm] + needs: [rust-bridge, assert-shard-coverage, unit, ui-unit, docs, helm, openai-sdk-compat] if: always() runs-on: ubuntu-latest timeout-minutes: 2 diff --git a/litellm/__init__.py b/litellm/__init__.py index ac57a9d0fee..4196912488f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -100,6 +100,7 @@ from litellm.constants import ( DEFAULT_ALLOWED_FAILS, ) import httpx +from openai import DefaultAsyncHttpxClient, DefaultHttpxClient # register_async_client_cleanup is lazy-loaded and called on first access @@ -420,8 +421,8 @@ error_logs: Dict = {} add_function_to_prompt: bool = ( False # if function calling not supported by api, append function call details to system prompt ) -client_session: Optional[httpx.Client] = None -aclient_session: Optional[httpx.AsyncClient] = None +client_session: Optional[Union[httpx.Client, DefaultHttpxClient]] = None +aclient_session: Optional[Union[httpx.AsyncClient, DefaultAsyncHttpxClient]] = None model_fallbacks: Optional[List] = None # Deprecated for 'litellm.fallbacks' model_cost_map_url: str = os.getenv( "LITELLM_MODEL_COST_MAP_URL", diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 439eadce13d..ac461a2472e 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -241,14 +241,7 @@ class BadRequestError(openai.BadRequestError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - # Use response if it's a valid httpx.Response with a request, otherwise use minimal error response - # Note: We check _request (not .request property) to avoid RuntimeError when _request is None - if ( - response is not None - and isinstance(response, httpx.Response) - and hasattr(response, "_request") - and getattr(response, "_request", None) is not None - ): + if response is not None and getattr(response, "_request", None) is not None: self.response = response else: self.response = _get_minimal_error_response() @@ -360,6 +353,8 @@ class UnprocessableEntityError(openai.UnprocessableEntityError): class Timeout(openai.APITimeoutError): + request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request + def __init__( self, message, @@ -458,6 +453,9 @@ class RateLimitError(openai.RateLimitError): :class:`RateLimitErrorCategory` for the available values. """ + response: httpx.Response # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX response + request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request + def __init__( self, message, @@ -515,7 +513,9 @@ class RateLimitError(openai.RateLimitError): ), ) super().__init__( - self.message, response=self.response, body=body + self.message, + response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime + body=body, ) # Call the base class constructor with the parameters it needs self.code = "429" self.type = "throttling_error" @@ -594,12 +594,7 @@ class PaymentRequiredError(BadRequestError): response: httpx.Response | None = None, litellm_debug_info: str | None = None, ) -> None: - response_is_valid: Final = ( - response is not None - and isinstance(response, httpx.Response) - and hasattr(response, "_request") - and getattr(response, "_request", None) is not None - ) + response_is_valid: Final = response is not None and getattr(response, "_request", None) is not None response_for_parent: Final = ( response if response_is_valid @@ -770,6 +765,9 @@ class ServiceUnavailableError(openai.APIStatusError): class BadGatewayError(openai.APIStatusError): + response: httpx.Response # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX response + request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request + def __init__( self, message, @@ -797,7 +795,9 @@ class BadGatewayError(openai.APIStatusError): ), ) super().__init__( - self.message, response=self.response, body=None + self.message, + response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime + body=None, ) # Call the base class constructor with the parameters it needs def __str__(self): @@ -818,6 +818,9 @@ class BadGatewayError(openai.APIStatusError): class InternalServerError(openai.InternalServerError): + response: httpx.Response # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX response + request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request + def __init__( self, message, @@ -846,7 +849,9 @@ class InternalServerError(openai.InternalServerError): ), ) super().__init__( - self.message, response=self.response, body=body + self.message, + response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime + body=body, ) # Call the base class constructor with the parameters it needs self.type = "internal_server_error" @@ -911,6 +916,8 @@ class APIError(openai.APIError): # raised if an invalid request (not get, delete, put, post) is made class APIConnectionError(openai.APIConnectionError): + request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request + def __init__( self, message, @@ -929,7 +936,10 @@ class APIConnectionError(openai.APIConnectionError): self.request = httpx.Request(method="POST", url="https://api.openai.com/v1") self.max_retries = max_retries self.num_retries = num_retries - super().__init__(message=self.message, request=self.request) + super().__init__( + message=self.message, + request=self.request, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime + ) def __str__(self): _message = self.message @@ -950,6 +960,9 @@ class APIConnectionError(openai.APIConnectionError): # raised if an invalid request (not get, delete, put, post) is made class APIResponseValidationError(openai.APIResponseValidationError): + response: httpx.Response # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX response + request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request + def __init__( self, message, @@ -967,7 +980,11 @@ class APIResponseValidationError(openai.APIResponseValidationError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - super().__init__(response=response, body=None, message=message) + super().__init__( + response=response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime + body=None, + message=message, + ) def __str__(self): _message = self.message @@ -1087,6 +1104,9 @@ class BudgetExceededError(Exception): ## DEPRECATED ## class InvalidRequestError(openai.BadRequestError): + response: httpx.Response # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX response + request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request + def __init__(self, message, model, llm_provider): self.status_code = 400 self.message = message @@ -1097,7 +1117,9 @@ class InvalidRequestError(openai.BadRequestError): request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object ) super().__init__( - message=self.message, response=self.response, body=None + message=self.message, + response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime + body=None, ) # Call the base class constructor with the parameters it needs diff --git a/litellm/litellm_core_utils/completion_timeout.py b/litellm/litellm_core_utils/completion_timeout.py index 163a4a6b9d6..e6f18257138 100644 --- a/litellm/litellm_core_utils/completion_timeout.py +++ b/litellm/litellm_core_utils/completion_timeout.py @@ -6,6 +6,7 @@ from collections.abc import Callable from typing import Final import httpx +from openai import Timeout as SDKTimeout from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS @@ -13,6 +14,12 @@ from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS class CompletionTimeout: """Resolves HTTP timeout for ``completion()`` from model vs global settings.""" + @staticmethod + def normalize(timeout: httpx.Timeout | SDKTimeout) -> httpx.Timeout: + if isinstance(timeout, httpx.Timeout): # pyright: ignore[reportUnnecessaryIsInstance] # SDK 2 aliases both classes. + return timeout + return httpx.Timeout(connect=timeout.connect, read=timeout.read, write=timeout.write, pool=timeout.pool) + @staticmethod def _fallback_when_no_explicit_timeout( global_timeout: float | str | None, @@ -31,7 +38,7 @@ class CompletionTimeout: @staticmethod def resolve( - model_timeout: float | str | httpx.Timeout | None, + model_timeout: float | str | httpx.Timeout | SDKTimeout | None, kwargs: dict, custom_llm_provider: str, *, @@ -49,7 +56,7 @@ class CompletionTimeout: Coerce :class:`httpx.Timeout` when the provider does not support it. """ - resolved: float | str | httpx.Timeout + resolved: float | str | httpx.Timeout | SDKTimeout if model_timeout is not None: resolved = model_timeout elif kwargs.get("timeout") is not None: @@ -59,12 +66,12 @@ class CompletionTimeout: else: resolved = CompletionTimeout._fallback_when_no_explicit_timeout(global_timeout) - if isinstance(resolved, httpx.Timeout) and not supports_httpx_timeout(custom_llm_provider): - read_timeout: Final = resolved.read - resolved = ( + if isinstance(resolved, (httpx.Timeout, SDKTimeout)): + normalized: Final = CompletionTimeout.normalize(resolved) + if supports_httpx_timeout(custom_llm_provider): + return normalized + read_timeout: Final = normalized.read + return ( float(read_timeout) if read_timeout is not None else COMPLETION_HTTP_FALLBACK_SECONDS ) # default 10 min timeout - elif not isinstance(resolved, httpx.Timeout): - resolved = float(resolved) - - return resolved + return float(resolved) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index c640405c269..3564d7a1364 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -16,7 +16,7 @@ from litellm._logging import verbose_logger from litellm.caching.caching import DualCache from litellm.constants import DEFAULT_MAX_RETRIES from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.llms.openai.common_utils import BaseOpenAILLM +from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIHTTPClient from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, ) @@ -733,7 +733,11 @@ class BaseAzureLLM(BaseOpenAILLM): azure_client_params: Final[_AzureGatewayClientParams] = { "api_version": api_version, "base_url": f"{api_base}", - "http_client": litellm.client_session, + "http_client": ( + litellm.aclient_session or OpenAIAsyncHTTPClient() + if acompletion + else litellm.client_session or OpenAIHTTPClient() + ), "max_retries": max_retries, "timeout": timeout, } diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index ffc5a5a71e0..839ca04e082 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -2,6 +2,7 @@ Common helpers / utils across al OpenAI endpoints """ +import asyncio import hashlib import inspect import json @@ -10,12 +11,13 @@ import ssl import time import uuid from collections.abc import AsyncIterator, Iterator, Mapping +from contextlib import suppress from typing import TYPE_CHECKING, Final, Literal, NamedTuple, Optional from urllib.parse import urlsplit import httpx import openai -from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI +from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, DefaultAsyncHttpxClient, DefaultHttpxClient, OpenAI from openai.types.chat import ChatCompletion, ChatCompletionChunk, ChatCompletionMessage from openai.types.chat.chat_completion import Choice from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice @@ -26,6 +28,7 @@ if TYPE_CHECKING: from aiohttp import ClientSession import litellm +from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.token_counter import token_counter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import ( @@ -48,6 +51,42 @@ _AZURE_OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(Azur _OPENAI_API_HOST: Final[str] = "api.openai.com" +_OPENAI_HTTPX_DEFAULT_TIMEOUT: Final = CompletionTimeout.normalize(openai.DEFAULT_TIMEOUT) +_OPENAI_HTTPX_CONNECTION_LIMITS: Final = httpx.Limits( + max_connections=openai.DEFAULT_CONNECTION_LIMITS.max_connections, + max_keepalive_connections=openai.DEFAULT_CONNECTION_LIMITS.max_keepalive_connections, + keepalive_expiry=openai.DEFAULT_CONNECTION_LIMITS.keepalive_expiry, +) + + +class OpenAIHTTPClient(httpx.Client): + def __init__(self) -> None: + super().__init__( + timeout=_OPENAI_HTTPX_DEFAULT_TIMEOUT, + limits=_OPENAI_HTTPX_CONNECTION_LIMITS, + follow_redirects=True, + ) + + def __del__(self) -> None: + with suppress(Exception): + if not self.is_closed: + self.close() + + +class OpenAIAsyncHTTPClient(httpx.AsyncClient): + def __init__(self) -> None: + super().__init__( + timeout=_OPENAI_HTTPX_DEFAULT_TIMEOUT, + limits=_OPENAI_HTTPX_CONNECTION_LIMITS, + follow_redirects=True, + ) + + def __del__(self) -> None: + with suppress(Exception): + if not self.is_closed: + asyncio.get_running_loop().create_task(self.aclose()) + + def is_openai_backed_api_base(api_base: str) -> bool: hostname: Final = urlsplit(api_base).hostname return hostname is not None and (hostname == _OPENAI_API_HOST or hostname.endswith(f".{_OPENAI_API_HOST}")) @@ -221,7 +260,9 @@ class BaseOpenAILLM: return _cached_client @staticmethod - def owns_wrapped_http_client(http_client: httpx.Client | httpx.AsyncClient | None) -> bool: + def owns_wrapped_http_client( + http_client: httpx.Client | httpx.AsyncClient | DefaultHttpxClient | DefaultAsyncHttpxClient | None, + ) -> bool: """Whether litellm may close an SDK client built around ``http_client``. ``_get_async_http_client`` / ``_get_sync_http_client`` hand back @@ -304,7 +345,7 @@ class BaseOpenAILLM: @staticmethod def _get_async_http_client( shared_session: Optional["ClientSession"] = None, - ) -> httpx.AsyncClient | None: + ) -> httpx.AsyncClient | DefaultAsyncHttpxClient | None: if litellm.aclient_session is not None: return litellm.aclient_session @@ -337,7 +378,7 @@ class BaseOpenAILLM: return cls._get_async_http_client(shared_session) @staticmethod - def _get_sync_http_client() -> httpx.Client | None: + def _get_sync_http_client() -> httpx.Client | DefaultHttpxClient | None: if litellm.client_session is not None: return litellm.client_session diff --git a/litellm/llms/openai/completion/handler.py b/litellm/llms/openai/completion/handler.py index 1bc752b6714..174aefd29e5 100644 --- a/litellm/llms/openai/completion/handler.py +++ b/litellm/llms/openai/completion/handler.py @@ -12,7 +12,7 @@ from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUser from litellm.types.utils import LlmProviders, ModelResponse, TextCompletionResponse from litellm.utils import ProviderConfigManager -from ..common_utils import BaseOpenAILLM, OpenAIError +from ..common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIError, OpenAIHTTPClient from .transformation import OpenAITextCompletionConfig @@ -129,7 +129,7 @@ class OpenAITextCompletion(BaseLLM): openai_client = OpenAI( api_key=api_key, base_url=api_base, - http_client=litellm.client_session, + http_client=litellm.client_session or OpenAIHTTPClient(), timeout=timeout, max_retries=max_retries, organization=organization, @@ -233,7 +233,7 @@ class OpenAITextCompletion(BaseLLM): openai_client = OpenAI( api_key=api_key, base_url=api_base, - http_client=litellm.client_session, + http_client=litellm.client_session or OpenAIHTTPClient(), timeout=timeout, max_retries=max_retries, organization=organization, @@ -290,7 +290,7 @@ class OpenAITextCompletion(BaseLLM): openai_client = AsyncOpenAI( api_key=api_key, base_url=api_base, - http_client=litellm.aclient_session, + http_client=litellm.aclient_session or OpenAIAsyncHTTPClient(), timeout=timeout, max_retries=max_retries, organization=organization, diff --git a/litellm/llms/openai/fine_tuning/handler.py b/litellm/llms/openai/fine_tuning/handler.py index c54263f4c35..851113bbd37 100644 --- a/litellm/llms/openai/fine_tuning/handler.py +++ b/litellm/llms/openai/fine_tuning/handler.py @@ -6,6 +6,7 @@ from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI from openai.types.fine_tuning import FineTuningJob from litellm._logging import verbose_logger +from litellm.llms.openai.common_utils import OpenAIAsyncHTTPClient, OpenAIHTTPClient from litellm.types.utils import LiteLLMFineTuningJob _AZURE_STATUS_MAP: Final[Mapping[object, str]] = { @@ -87,9 +88,9 @@ class OpenAIFineTuningAPI: elif v is not None: data[k] = v if _is_async is True: - openai_client = AsyncOpenAI(**data) + openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient()) else: - openai_client = OpenAI(**data) + openai_client = OpenAI(**data, http_client=OpenAIHTTPClient()) else: openai_client = client diff --git a/litellm/llms/openai/image_variations/handler.py b/litellm/llms/openai/image_variations/handler.py index bc02d274f24..664d16c48bf 100644 --- a/litellm/llms/openai/image_variations/handler.py +++ b/litellm/llms/openai/image_variations/handler.py @@ -14,7 +14,7 @@ from litellm.utils import ProviderConfigManager from ...base_llm.image_variations.transformation import BaseImageVariationConfig from ...custom_httpx.llm_http_handler import LiteLLMLoggingObj -from ..common_utils import OpenAIError +from ..common_utils import OpenAIAsyncHTTPClient, OpenAIError, OpenAIHTTPClient class OpenAIImageVariationsHandler: @@ -23,22 +23,14 @@ class OpenAIImageVariationsHandler: client: OpenAI | None, init_client_params: dict, ): - if client is None: - openai_client = OpenAI( - **init_client_params, - ) - else: - openai_client = client - return openai_client + if client is not None: + return client + return OpenAI(**init_client_params, http_client=litellm.client_session or OpenAIHTTPClient()) def get_async_client(self, client: AsyncOpenAI | None, init_client_params: dict) -> AsyncOpenAI: - if client is None: - openai_client = AsyncOpenAI( - **init_client_params, - ) - else: - openai_client = client - return openai_client + if client is not None: + return client + return AsyncOpenAI(**init_client_params, http_client=litellm.aclient_session or OpenAIAsyncHTTPClient()) async def async_image_variations( self, @@ -62,7 +54,6 @@ class OpenAIImageVariationsHandler: init_client_params: Final = { "api_key": api_key, "base_url": api_base, - "http_client": litellm.client_session, "timeout": timeout, "max_retries": max_retries, "organization": organization, @@ -179,7 +170,6 @@ class OpenAIImageVariationsHandler: init_client_params: Final = { "api_key": api_key, "base_url": api_base, - "http_client": litellm.client_session, "timeout": timeout, "max_retries": max_retries, "organization": organization, diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index a782bf35ac8..3a12326d433 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -53,7 +53,9 @@ from .chat.gpt_transformation import OpenAIGPTConfig, OpenAIUnknownModelConfig from .chat.o_series_transformation import OpenAIOSeriesConfig from .common_utils import ( BaseOpenAILLM, + OpenAIAsyncHTTPClient, OpenAIError, + OpenAIHTTPClient, build_output_token_limit_response, drop_params_from_unprocessable_entity_error, is_openai_backed_api_base, @@ -1749,9 +1751,9 @@ class OpenAIFilesAPI(BaseLLM): elif v is not None: data[k] = v if _is_async is True: - openai_client = AsyncOpenAI(**data) + openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient()) else: - openai_client = OpenAI(**data) + openai_client = OpenAI(**data, http_client=OpenAIHTTPClient()) else: openai_client = client @@ -2107,9 +2109,9 @@ class OpenAIBatchesAPI(BaseLLM): elif v is not None: data[k] = v if _is_async is True: - openai_client = AsyncOpenAI(**data) + openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient()) else: - openai_client = OpenAI(**data) + openai_client = OpenAI(**data, http_client=OpenAIHTTPClient()) else: openai_client = client @@ -2317,7 +2319,7 @@ class OpenAIAssistantsAPI(BaseLLM): data["base_url"] = v elif v is not None: data[k] = v - openai_client = OpenAI(**data) + openai_client = OpenAI(**data, http_client=OpenAIHTTPClient()) else: openai_client = client @@ -2342,7 +2344,7 @@ class OpenAIAssistantsAPI(BaseLLM): data["base_url"] = v elif v is not None: data[k] = v - openai_client = AsyncOpenAI(**data) + openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient()) else: openai_client = client diff --git a/litellm/main.py b/litellm/main.py index b80819354f0..4744a502194 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -432,7 +432,7 @@ async def acompletion( messages: list = [], functions: list | None = None, function_call: str | None = None, - timeout: float | None = None, + timeout: float | httpx.Timeout | openai.Timeout | None = None, temperature: float | None = None, top_p: float | None = None, n: int | None = None, @@ -5225,7 +5225,7 @@ def completion( model: str, # Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create messages: list = [], - timeout: float | str | httpx.Timeout | None = None, + timeout: float | str | httpx.Timeout | openai.Timeout | None = None, temperature: float | None = None, top_p: float | None = None, n: int | None = None, @@ -7880,6 +7880,8 @@ def adapter_completion(*, adapter_id: str, **kwargs) -> BaseModel | AdapterCompl def moderation(input: str, model: str | None = None, api_key: str | None = None, **kwargs) -> OpenAIModerationResponse: + from litellm.llms.openai.common_utils import OpenAIHTTPClient + # only supports open ai for now api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") @@ -7888,10 +7890,7 @@ def moderation(input: str, model: str | None = None, api_key: str | None = None, openai_client = kwargs.get("client", None) if openai_client is None: - if api_base is not None: - openai_client = openai.OpenAI(api_key=api_key, base_url=api_base) - else: - openai_client = openai.OpenAI(api_key=api_key) + openai_client = openai.OpenAI(api_key=api_key, base_url=api_base, http_client=OpenAIHTTPClient()) if model is not None: response = openai_client.moderations.create(input=input, model=model) @@ -8364,7 +8363,7 @@ def speech( project: str | None = None, max_retries: int | None = None, metadata: dict | None = None, - timeout: float | httpx.Timeout | None = None, + timeout: float | httpx.Timeout | openai.Timeout | None = None, response_format: str | None = None, speed: int | None = None, instructions: str | None = None, @@ -8394,9 +8393,12 @@ def speech( if instructions is not None: optional_params["instructions"] = instructions - if timeout is None: - timeout = litellm.request_timeout - + timeout_or_default: Final = litellm.request_timeout if timeout is None else timeout + http_timeout: Final = ( + CompletionTimeout.normalize(timeout_or_default) + if isinstance(timeout_or_default, openai.Timeout) + else timeout_or_default + ) if max_retries is None: max_retries = litellm.num_retries or openai.DEFAULT_MAX_RETRIES litellm_params_dict: Final = get_litellm_params(metadata=metadata, api_key=api_key or dynamic_api_key, **kwargs) @@ -8445,7 +8447,7 @@ def speech( custom_llm_provider=custom_llm_provider, litellm_params=litellm_params_dict, logging_obj=logging_obj, - timeout=timeout, + timeout=http_timeout, extra_headers=extra_headers, client=client, _is_async=aspeech or False, @@ -8501,7 +8503,7 @@ def speech( organization=organization, project=project, max_retries=max_retries, - timeout=timeout, + timeout=http_timeout, logging_obj=logging_obj, client=client, # pass AsyncOpenAI, OpenAI client aspeech=aspeech, @@ -8532,7 +8534,7 @@ def speech( optional_params=optional_params, litellm_params_dict=litellm_params_dict, logging_obj=logging_obj, - timeout=timeout, + timeout=http_timeout, extra_headers=extra_headers, base_llm_http_handler=base_llm_http_handler, aspeech=aspeech or False, @@ -8580,7 +8582,7 @@ def speech( azure_ad_token_provider=azure_ad_token_provider, organization=organization, max_retries=max_retries, - timeout=timeout, + timeout=http_timeout, logging_obj=logging_obj, client=client, # pass AsyncOpenAI, OpenAI client aspeech=aspeech, @@ -8625,7 +8627,7 @@ def speech( custom_llm_provider=custom_llm_provider, litellm_params=litellm_params_dict, logging_obj=logging_obj, - timeout=timeout, + timeout=http_timeout, extra_headers=extra_headers, client=client, _is_async=aspeech or False, @@ -8682,7 +8684,7 @@ def speech( optional_params=optional_params, litellm_params_dict=litellm_params_dict, logging_obj=logging_obj, - timeout=timeout, + timeout=http_timeout, extra_headers=headers, base_llm_http_handler=base_llm_http_handler, aspeech=aspeech or False, @@ -8728,7 +8730,7 @@ def speech( optional_params=optional_params, litellm_params_dict=litellm_params_dict, logging_obj=logging_obj, - timeout=timeout, + timeout=http_timeout, extra_headers=extra_headers, base_llm_http_handler=base_llm_http_handler, aspeech=aspeech or False, @@ -8769,7 +8771,7 @@ def speech( custom_llm_provider=custom_llm_provider, litellm_params=litellm_params_dict, logging_obj=logging_obj, - timeout=timeout, + timeout=http_timeout, extra_headers=extra_headers, client=client, _is_async=aspeech or False, @@ -8797,7 +8799,7 @@ def speech( custom_llm_provider=custom_llm_provider, litellm_params=litellm_params_dict, logging_obj=logging_obj, - timeout=timeout, + timeout=http_timeout, extra_headers=extra_headers, client=client, _is_async=aspeech or False, @@ -8821,7 +8823,7 @@ def speech( optional_params=optional_params, litellm_params_dict=litellm_params_dict, logging_obj=logging_obj, - timeout=timeout, + timeout=http_timeout, extra_headers=extra_headers, base_llm_http_handler=base_llm_http_handler, aspeech=aspeech or False, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 53f49f78c9f..a7198fd1a95 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12919,7 +12919,7 @@ async def audio_speech( requested_format: Final = data.get("response_format") upstream_content_type: Final = ( - response.response.headers.get("content-type") if isinstance(response, HttpxBinaryResponseContent) else None + response.content_type if isinstance(response, HttpxBinaryResponseContent) else None ) media_type: Final = resolve_speech_media_type( upstream_content_type=upstream_content_type, diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index cd77d6088f6..26f65fee4fe 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -2,13 +2,14 @@ import builtins from collections.abc import Iterable, Mapping from enum import Enum from os import PathLike -from typing import IO, Any, Final, Literal, Optional, TypeAlias, Union +from typing import IO, Any, Final, Generic, Literal, Optional, TypeAlias, Union import httpx from openai import Omit from openai._legacy_response import ( HttpxBinaryResponseContent as _HttpxBinaryResponseContent, ) +from openai._types import Response as SDKResponse from openai.lib.streaming._assistants import ( AssistantEventHandler, AssistantStreamManager, @@ -79,6 +80,7 @@ from typing_extensions import ( ReadOnly, Required, TypedDict, + TypeVar, override, ) @@ -118,8 +120,12 @@ class BinaryResponseSummary(TypedDict): num_bytes: ReadOnly[int] -class HttpxBinaryResponseContent(_HttpxBinaryResponseContent): +_ResponseT = TypeVar("_ResponseT", bound=httpx.Response | SDKResponse, default=httpx.Response) + + +class HttpxBinaryResponseContent(_HttpxBinaryResponseContent, Generic[_ResponseT]): _hidden_params: dict + response: _ResponseT # pyright: ignore[reportIncompatibleVariableOverride] # SDK accepts both backends at runtime @property def hidden_params(self) -> dict[str, builtins.object]: # mutable-ok: API requires mutation @@ -129,22 +135,28 @@ class HttpxBinaryResponseContent(_HttpxBinaryResponseContent): def hidden_params(self, hidden_params: dict[str, builtins.object]) -> None: # mutable-ok: API requires mutation self._hidden_params = hidden_params - def __init__(self, response: httpx.Response) -> None: - super().__init__(response) + def __init__(self, response: _ResponseT) -> None: + super().__init__(response) # pyright: ignore[reportArgumentType] # SDK accepts both backends at runtime self._hidden_params = {} + @property + def content_type(self) -> str | None: + headers: Final[Mapping[str, str]] = self.response.headers + return headers.get("content-type") + def logging_summary(self) -> BinaryResponseSummary: return { "object": "binary", - "content_type": self.response.headers.get("content-type"), + "content_type": self.content_type, "num_bytes": self._num_bytes(), } def _num_bytes(self) -> int: try: - return len(self.response.content) - except httpx.ResponseNotRead: + content: Final = self.response.content + except RuntimeError: return self.response.num_bytes_downloaded + return len(content) def set_response_cost(self, response_cost: float | None) -> None: if response_cost is None: diff --git a/litellm/types/router.py b/litellm/types/router.py index 53df3ce6773..cbe332dab12 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -21,11 +21,13 @@ from typing import ( from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import httpx +from openai import Timeout as SDKTimeout from pydantic import ConfigDict, Field, JsonValue, field_validator, model_validator from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_checkable from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.provider_affinity import validate_provider_affinity_header_name from litellm.types.llms.base import LiteLLMBaseModel @@ -481,7 +483,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): rpm: int | None = None itpm: int | None = None otpm: int | None = None - timeout: float | str | httpx.Timeout | None = None # if str, pass in as os.environ/ + timeout: float | str | httpx.Timeout | SDKTimeout | None = None # if str, pass in as os.environ/ stream_timeout: float | str | None = None # timeout when making stream=True calls, if str, pass in as os.environ/ max_retries: int | None = None drop_params: bool | str | None = None @@ -564,6 +566,13 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): valkey_text_field: str | None = None valkey_embedding_field: str | None = None + @field_validator("timeout") + @classmethod + def normalize_timeout( + cls, value: float | str | httpx.Timeout | SDKTimeout | None + ) -> float | str | httpx.Timeout | None: + return CompletionTimeout.normalize(value) if isinstance(value, (httpx.Timeout, SDKTimeout)) else value + @field_validator("provider_affinity_header") @classmethod def validate_provider_affinity_header(cls, value: str | None) -> str | None: @@ -677,7 +686,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): api_key: str | None api_base: str | None api_version: str | None - timeout: float | str | httpx.Timeout | None + timeout: float | str | httpx.Timeout | SDKTimeout | None # writable-ok: preserve timeout updates stream_timeout: float | str | None max_retries: int | None organization: list | str | None # for openai orgs diff --git a/packaging/litellm-core/pyproject.toml b/packaging/litellm-core/pyproject.toml index 9653e6f983a..cf87cc6b936 100644 --- a/packaging/litellm-core/pyproject.toml +++ b/packaging/litellm-core/pyproject.toml @@ -13,7 +13,7 @@ dependencies = [ "fastuuid>=0.14.0,<1.0", "filelock>=3.16.1,<4.0", "httpx[http2]>=0.28.0,<1.0", - "openai>=2.20.0,<3.0.0", + "openai>=2.20.0,<4.0.0", "python-dateutil>=2.8.2,<3.0", "python-dotenv>=1.0.0,<2.0", "pyyaml>=6.0.3,<7.0", diff --git a/pyproject.toml b/pyproject.toml index 924e3a939b2..133fc735df1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ dependencies = [ "fastuuid>=0.14.0,<1.0", "filelock>=3.16.1,<4.0", "httpx[http2]>=0.28.0,<1.0", - "openai>=2.20.0,<3.0.0", + "openai>=2.20.0,<4.0.0", "python-dotenv>=1.0.0,<2.0", "pyyaml>=6.0.3,<7.0", "packaging>=24.0", diff --git a/tests/base_sdk_tests/check_base_sdk_install.py b/tests/base_sdk_tests/check_base_sdk_install.py index e56d48f5e6e..3abea2a9707 100644 --- a/tests/base_sdk_tests/check_base_sdk_install.py +++ b/tests/base_sdk_tests/check_base_sdk_install.py @@ -8,11 +8,18 @@ the very class of undeclared-dependency bug this guards against. import argparse import importlib.util +import re import sys import traceback from collections.abc import Callable +from importlib.metadata import PackageNotFoundError, packages_distributions, requires +from itertools import chain +from typing import Final EXTRAS_ONLY_MODULES = ("fastapi", "uvicorn", "keyring", "mcp", "mcp_types", "httpx2", "httpcore2") +SDK_DISTRIBUTIONS: Final = frozenset({"litellm", "litellm-core"}) +REQUIREMENT_NAME: Final = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*") +EXTRA_MARKER: Final = re.compile(r";.*\bextra\s*==") def _require(condition: bool, message: str) -> None: @@ -20,13 +27,43 @@ def _require(condition: bool, message: str) -> None: raise AssertionError(message) +def _canonical(name: str) -> str: + return re.sub(r"[-_.]+", "-", name).lower() + + +def _base_requirements(distribution: str) -> frozenset[str]: + try: + declared: Final = requires(distribution) or () + except PackageNotFoundError: + return frozenset() + return frozenset( + _canonical(match.group(0)) + for requirement in declared + if not EXTRA_MARKER.search(requirement) and (match := REQUIREMENT_NAME.match(requirement)) + ) + + +def _base_closure(pending: frozenset[str], reached: frozenset[str] = frozenset()) -> frozenset[str]: + if not pending: + return reached + expanded: Final = reached | pending + return _base_closure(frozenset(chain.from_iterable(_base_requirements(name) for name in pending)) - expanded, expanded) + + def check_environment_is_base_only() -> str: - present = tuple(name for name in EXTRAS_ONLY_MODULES if importlib.util.find_spec(name) is not None) + from_base: Final = _base_closure(SDK_DISTRIBUTIONS) + owners: Final = packages_distributions() + installed: Final = tuple(name for name in EXTRAS_ONLY_MODULES if importlib.util.find_spec(name) is not None) + inherited: Final = tuple( + name for name in installed if from_base & {_canonical(owner) for owner in owners.get(name) or (name,)} + ) + present: Final = tuple(name for name in installed if name not in inherited) _require( not present, f"{', '.join(present)} installed, so this environment is not base-only and the run proves nothing", ) - return f"no extras-only packages present ({', '.join(EXTRAS_ONLY_MODULES)})" + from_dependencies: Final = f"; {', '.join(inherited)} come from base dependencies" if inherited else "" + return f"no extras-only packages present ({', '.join(EXTRAS_ONLY_MODULES)}){from_dependencies}" def check_import() -> str: diff --git a/tests/base_sdk_tests/check_sdk_http.py b/tests/base_sdk_tests/check_sdk_http.py index d773e252127..c7fb6a542ff 100644 --- a/tests/base_sdk_tests/check_sdk_http.py +++ b/tests/base_sdk_tests/check_sdk_http.py @@ -37,7 +37,8 @@ class RecordingUpstream(BaseHTTPRequestHandler): pass def do_POST(self) -> None: - REQUESTS.put((self.path, dict(self.headers), self.rfile.read(int(self.headers["Content-Length"])))) + headers: Final = {name.lower(): value for name, value in self.headers.items()} + REQUESTS.put((self.path, headers, self.rfile.read(int(self.headers["Content-Length"])))) status, body = RESPONSES.get(timeout=10) ARRIVED.set() if status == 0: @@ -78,7 +79,7 @@ def check_http(base: str) -> None: assert response.usage.total_tokens == 5 path, headers, body = REQUESTS.get(timeout=10) assert path == "/v1/chat/completions" - assert headers["Authorization"] == "Bearer test-key" + assert headers["authorization"] == "Bearer test-key" assert json.loads(body)["messages"] == MESSAGES enqueue_stream() assert "".join(part.choices[0].delta.content or "" for part in litellm.completion(**arguments, stream=True)) == "pong" @@ -156,7 +157,7 @@ def check_http(base: str) -> None: assert result.usage.total_tokens == 5 path, headers, body = REQUESTS.get(timeout=10) assert path.endswith("/converse") - assert {key.lower(): value for key, value in headers.items()}["authorization"].startswith( + assert headers["authorization"].startswith( "AWS4-HMAC-SHA256" if signed else "Bearer bearer-key" ) assert json.loads(body)["messages"][0]["content"][0]["text"] == "ping" diff --git a/tests/integration/compatibility/test_openai_sdk_client_paths.py b/tests/integration/compatibility/test_openai_sdk_client_paths.py new file mode 100644 index 00000000000..f8d7449038a --- /dev/null +++ b/tests/integration/compatibility/test_openai_sdk_client_paths.py @@ -0,0 +1,1200 @@ +from __future__ import annotations + +import asyncio +import base64 +import binascii +import json +import os +import threading +import uuid +from collections.abc import Awaitable, Callable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, TypeVar + +import anthropic +import httpx +import openai +import pytest +import yaml +from anthropic.types import RawContentBlockDeltaEvent, RawMessageStartEvent, RawMessageStreamEvent, TextDelta +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.openai_wire import chat_reply, openai_error, responses_reply +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.responses_vendor import response_identities +from integration._support.wire import Reply, Request, Wire, wire_server +from openai.types import Completion, ModerationCreateResponse +from openai.types.chat import ChatCompletionChunk +from openai.types.responses import ResponseStreamEvent +from pydantic import JsonValue + +from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS + +PROXY_WORKERS: Final = int(os.environ.get("INTEGRATION_PROXY_WORKERS", "1")) +WORKER_SYNC_SECONDS: Final = 0.0 if PROXY_WORKERS == 1 else PROXY_CONFIG_RELOAD_INTERVAL_SECONDS + 5.0 +ROOT: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) +TEXT: Final = "openai sdk client path" +PROVIDER_KEY: Final = "integration-provider-key" +PROVIDER_AUTHORIZATION: Final = f"Bearer {PROVIDER_KEY}" +UPSTREAM_MODEL: Final = "gpt-4o-mini" +DEPLOYMENT_MODEL: Final = f"openai/{UPSTREAM_MODEL}" +INSTRUCT_MODEL: Final = "gpt-3.5-turbo-instruct" +INSTRUCT_DEPLOYMENT_MODEL: Final = f"openai/{INSTRUCT_MODEL}" +NO_CACHE: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache": {"no-cache": True}}) +USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8} +) +FIVE_KB: Final = "x" * 5120 +REJECTED_CHAT_TIMEOUTS: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"list": [1, 2], "dict": {"read": 5}, "5 KB string": FIVE_KB} +) +ACCEPTED_CHAT_TIMEOUTS: Final[Mapping[str, JsonValue]] = MappingProxyType({"empty string": "", "int": 30}) +NO_EXTRA: Final[Mapping[str, JsonValue]] = MappingProxyType({}) +REJECTED_DEPLOYMENT_TIMEOUTS: Final[Mapping[str, JsonValue]] = MappingProxyType({"list": [1, 2], "dict": {"read": 5}}) +ACCEPTED_DEPLOYMENT_TIMEOUTS: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"empty string": "", "5 KB string": FIVE_KB} +) +MODERATION_CATEGORIES: Final = ("hate", "harassment", "self-harm", "sexual", "violence") +LISTED_JOB: Final = "ftjob-listed" +ASSISTANT: Final = "asst_integration" +CLOUDFLARE_GATEWAY: Final = "/gateway.ai.cloudflare.com/v1/account/gateway/azure-openai/resource" +AZURE_API_VERSION: Final = "2024-10-21" +STREAM_HEAD: Final = b"h" * 65536 +STREAM_TAIL: Final = b'{"custom_id": "tail", "response": {"status_code": 200}}\n' +SPEECH_AUDIO: Final = b"OggS" + bytes(range(60)) +SPEECH_UPSTREAM_MODEL: Final = "gpt-4o-mini-tts" +PAYMENT_REQUIRED_DEPLOYMENTS: Final[Mapping[str, str]] = MappingProxyType( + {"openai": "openai/gpt-4o-mini", "anthropic": "anthropic/claude-sonnet-4-5"} +) +_R: Final = TypeVar("_R") + + +def _identity(prefix: str) -> str: + return f"{prefix}-{uuid.uuid4().hex}" + + +def _path(request: Request) -> str: + return request.target.split("?", 1)[0] + + +def _body(request: Request) -> Mapping[str, JsonValue]: + return JSON_OBJECT.validate_json(request.body) + + +def _streams(request: Request) -> bool: + return _body(request).get("stream") is True + + +def _upstream_model(request: Request) -> str: + return string_value(_body(request).get("model") or UPSTREAM_MODEL) + + +def _json( + payload: Mapping[str, JsonValue], *, status: int = 200, headers: Mapping[str, str] = MappingProxyType({}) +) -> Reply: + return Reply(status=status, body=json.dumps(payload).encode(), headers=headers) + + +def _peer(route: Callable[[Request], Reply | None]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method == "GET" and _path(request).endswith("/models"): + return _json({"object": "list", "data": []}) + reply: Final = route(request) + if reply is None: + return _json({"error": {"message": f"unrouted {request.method} {request.target}"}}, status=404) + return reply + + return respond + + +def _is_model_discovery(request: Request) -> bool: + return request.method == "GET" and _path(request).endswith("/models") + + +def _served(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if not _is_model_discovery(request)) + + +def _routes(served: tuple[Request, ...]) -> tuple[tuple[str, str], ...]: + return tuple((request.method, _path(request)) for request in served) + + +def _is_chat(request: Request) -> bool: + return request.method == "POST" and _path(request).endswith("/chat/completions") + + +def _chat_peer(identity: str) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + if not _is_chat(request): + return None + return chat_reply(identity, _upstream_model(request), TEXT, stream=_streams(request)) + + return _peer(route) + + +def _fresh_chat_peer(served: SimpleQueue[str]) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + if not _is_chat(request): + return None + identity: Final = _identity("chatcmpl") + served.put(identity) + return chat_reply(identity, _upstream_model(request), TEXT, stream=False) + + return _peer(route) + + +def _held_chat_peer(gate: threading.Event, identity: str) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + if not _is_chat(request): + return None + assert gate.wait(timeout=10), "the timeout cell never released its peer" + return chat_reply(identity, _upstream_model(request), TEXT, stream=False) + + return _peer(route) + + +def _status_peer(reply: Reply) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + return reply if request.method == "POST" else None + + return _peer(route) + + +def _completion_frame(identity: str, model: str, text: str, finish_reason: str | None) -> bytes: + frame: Final = { + "id": identity, + "object": "text_completion", + "created": 1, + "model": model, + "choices": [{"text": text, "index": 0, "finish_reason": finish_reason, "logprobs": None}], + } + return b"data: " + json.dumps(frame).encode() + b"\n\n" + + +def _completion_reply(identity: str, model: str, *, stream: bool) -> Reply: + if stream: + return Reply( + content_type="text/event-stream", + chunks=( + _completion_frame(identity, model, TEXT, None), + _completion_frame(identity, model, "", "stop") + b"data: [DONE]\n\n", + ), + ) + return _json( + { + "id": identity, + "object": "text_completion", + "created": 1, + "model": model, + "choices": [{"text": TEXT, "index": 0, "finish_reason": "stop", "logprobs": None}], + "usage": dict(USAGE), + } + ) + + +def _completion_peer(identity: str) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + if request.method != "POST" or _path(request) != "/v1/completions": + return None + return _completion_reply(identity, _upstream_model(request), stream=_streams(request)) + + return _peer(route) + + +def _responses_peer(identity: str) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + if request.method != "POST" or _path(request) != "/v1/responses": + return None + return responses_reply(identity, _upstream_model(request), TEXT, stream=_streams(request)) + + return _peer(route) + + +def _moderation(identity: str, model: str) -> Mapping[str, JsonValue]: + return { + "id": identity, + "model": model, + "results": [ + { + "flagged": False, + "categories": {category: False for category in MODERATION_CATEGORIES}, + "category_scores": {category: 0.0 for category in MODERATION_CATEGORIES}, + "category_applied_input_types": {category: ["text"] for category in MODERATION_CATEGORIES}, + } + ], + } + + +def _moderation_peer(identity: str) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + if request.method != "POST" or _path(request) != "/v1/moderations": + return None + return _json(_moderation(identity, _upstream_model(request))) + + return _peer(route) + + +def _batch_line(custom_id: str) -> bytes: + line: Final = { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": UPSTREAM_MODEL, "messages": [_user_message()]}, + } + return json.dumps(line).encode() + b"\n" + + +def _file_object(file_id: str, size: int) -> Mapping[str, JsonValue]: + return { + "id": file_id, + "object": "file", + "purpose": "batch", + "bytes": size, + "created_at": 1, + "filename": "batch.jsonl", + "status": "processed", + } + + +def _files_peer(file_id: str, content: bytes, chunks: tuple[bytes, ...] | None = None) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + path: Final = _path(request) + if request.method == "POST" and path == "/v1/files": + return _json(_file_object(file_id, len(content))) + if request.method == "GET" and path == f"/v1/files/{file_id}": + return _json(_file_object(file_id, len(content))) + if request.method == "GET" and path == f"/v1/files/{file_id}/content": + return Reply(content_type="application/octet-stream", body=content, chunks=chunks) + return None + + return _peer(route) + + +def _batch_object(batch_id: str, input_file_id: str, status: str) -> Mapping[str, JsonValue]: + return { + "id": batch_id, + "object": "batch", + "endpoint": "/v1/chat/completions", + "errors": None, + "input_file_id": input_file_id, + "completion_window": "24h", + "status": status, + "output_file_id": None, + "error_file_id": None, + "created_at": 1, + "in_progress_at": None, + "expires_at": None, + "finalizing_at": None, + "completed_at": None, + "failed_at": None, + "expired_at": None, + "cancelling_at": None, + "cancelled_at": None, + "request_counts": {"total": 1, "completed": 0, "failed": 0}, + "metadata": None, + } + + +def _batches_peer(file_id: str, batch_id: str) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + path: Final = _path(request) + if request.method == "POST" and path == "/v1/files": + return _json(_file_object(file_id, 1)) + if request.method == "POST" and path == "/v1/batches": + return _json(_batch_object(batch_id, string_value(_body(request)["input_file_id"]), "validating")) + if request.method == "GET" and path == f"/v1/batches/{batch_id}": + return _json(_batch_object(batch_id, file_id, "in_progress")) + if request.method == "GET" and path == "/v1/batches": + return _json( + { + "object": "list", + "data": [_batch_object(batch_id, file_id, "in_progress")], + "first_id": batch_id, + "last_id": batch_id, + "has_more": False, + } + ) + if request.method == "POST" and path == f"/v1/batches/{batch_id}/cancel": + return _json(_batch_object(batch_id, file_id, "cancelling")) + return None + + return _peer(route) + + +def _fine_tuning_job(job_id: str, status: str) -> Mapping[str, JsonValue]: + return { + "id": job_id, + "object": "fine_tuning.job", + "model": UPSTREAM_MODEL, + "created_at": 1, + "fine_tuned_model": None, + "finished_at": None, + "hyperparameters": {"n_epochs": 1, "batch_size": 1, "learning_rate_multiplier": 1.0}, + "organization_id": "org-integration", + "result_files": [], + "status": status, + "trained_tokens": None, + "training_file": f"file-{job_id.removeprefix('ftjob-')}", + "validation_file": None, + "seed": 1, + "error": None, + } + + +def _assistant(identity: str) -> Mapping[str, JsonValue]: + return { + "id": identity, + "object": "assistant", + "created_at": 1, + "name": "integration", + "description": None, + "model": UPSTREAM_MODEL, + "instructions": None, + "tools": [], + "metadata": {}, + "temperature": 1.0, + "top_p": 1.0, + "response_format": "auto", + "tool_resources": None, + } + + +def _configured_peer() -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + path: Final = _path(request) + if request.method == "POST" and path == "/v1/fine_tuning/jobs": + training_file: Final = string_value(_body(request)["training_file"]) + return _json(_fine_tuning_job(f"ftjob-{training_file.removeprefix('file-')}", "queued")) + if request.method == "GET" and path == "/v1/fine_tuning/jobs": + return _json({"object": "list", "data": [_fine_tuning_job(LISTED_JOB, "running")], "has_more": False}) + if request.method == "POST" and path.startswith("/v1/fine_tuning/jobs/") and path.endswith("/cancel"): + return _json( + _fine_tuning_job(path.removeprefix("/v1/fine_tuning/jobs/").removesuffix("/cancel"), "cancelled") + ) + if request.method == "GET" and path.startswith("/v1/fine_tuning/jobs/"): + return _json(_fine_tuning_job(path.removeprefix("/v1/fine_tuning/jobs/"), "running")) + if request.method == "GET" and path == "/v1/assistants": + return _json( + { + "object": "list", + "data": [_assistant(ASSISTANT)], + "first_id": ASSISTANT, + "last_id": ASSISTANT, + "has_more": False, + } + ) + return None + + return _peer(route) + + +def _config_with_provider_settings(directory: Path, api_base: str) -> Path: + shared: Final = object_value(yaml.safe_load((ROOT / "tests/integration/proxy_config.yaml").read_text())) + path: Final = directory / "proxy_config_with_provider_settings.yaml" + path.write_text( + yaml.safe_dump( + { + **shared, + "finetune_settings": [{"custom_llm_provider": "openai", "api_base": api_base, "api_key": PROVIDER_KEY}], + "assistant_settings": { + "custom_llm_provider": "openai", + "litellm_params": {"api_base": api_base, "api_key": PROVIDER_KEY}, + }, + } + ) + ) + return path + + +@dataclass(frozen=True, slots=True) +class ConfiguredProxy: + gateway: Gateway + wire: Wire + + +@pytest.fixture(scope="module") +def configured_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[ConfiguredProxy]: + with gateway_from_environment() as rig_gateway, wire_server(_configured_peer()) as wire: + directory: Final = tmp_path_factory.mktemp("openai_sdk_client_paths") + config: Final = _config_with_provider_settings(directory, f"{wire.url}/v1") + with owned_proxy(rig_gateway, directory, {}, config=config, workers=PROXY_WORKERS) as owned: + yield ConfiguredProxy(owned, wire) + + +def _v1(gateway: Gateway) -> str: + return f"{str(gateway.client.base_url).rstrip('/')}/v1" + + +def _origin(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +@contextmanager +def _openai(gateway: Gateway) -> Iterator[openai.OpenAI]: + with ( + httpx.Client(timeout=15, trust_env=False) as transport, + openai.OpenAI(api_key=gateway.key, base_url=_v1(gateway), max_retries=0, http_client=transport) as client, + ): + yield client + + +def _async_openai(gateway: Gateway, call: Callable[[openai.AsyncOpenAI], Awaitable[_R]]) -> _R: + async def run() -> _R: + async with ( + httpx.AsyncClient(timeout=15, trust_env=False) as transport, + openai.AsyncOpenAI( + api_key=gateway.key, base_url=_v1(gateway), max_retries=0, http_client=transport + ) as client, + ): + return await call(client) + + return asyncio.run(run()) + + +@contextmanager +def _anthropic(gateway: Gateway) -> Iterator[anthropic.Anthropic]: + with ( + httpx.Client(timeout=15, trust_env=False) as transport, + anthropic.Anthropic( + api_key=gateway.key, base_url=_origin(gateway), max_retries=0, http_client=transport + ) as client, + ): + yield client + + +def _async_anthropic(gateway: Gateway, call: Callable[[anthropic.AsyncAnthropic], Awaitable[_R]]) -> _R: + async def run() -> _R: + async with ( + httpx.AsyncClient(timeout=15, trust_env=False) as transport, + anthropic.AsyncAnthropic( + api_key=gateway.key, base_url=_origin(gateway), max_retries=0, http_client=transport + ) as client, + ): + return await call(client) + + return asyncio.run(run()) + + +def _user_message() -> Mapping[str, JsonValue]: + return {"role": "user", "content": TEXT} + + +def _chat_body(model: str, extra: Mapping[str, JsonValue] = NO_EXTRA) -> Mapping[str, JsonValue]: + return {"model": model, "messages": [_user_message()], **NO_CACHE, **extra} + + +def _duplicated_timeout_body(model: str) -> bytes: + fields: Final = json.dumps(_chat_body(model)).removesuffix("}") + return f'{fields}, "timeout": 30, "timeout": 30}}'.encode() + + +def _chat(gateway: Gateway, model: str, extra: Mapping[str, JsonValue] = NO_EXTRA) -> httpx.Response: + return gateway.request("POST", "/v1/chat/completions", _chat_body(model, extra)) + + +def _response_id(response: httpx.Response) -> str: + assert response.status_code == 200, response.text + return string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + +def _error_message(response: httpx.Response) -> str: + body: Final = JSON_OBJECT.validate_json(response.content) + if "error" in body: + return string_value(object_value(body["error"])["message"]) + detail: Final = body["detail"] + if isinstance(detail, str): + return detail + assert isinstance(detail, list) and detail, response.text + return string_value(object_value(detail[0])["msg"]) + + +def _assert_rejected(response: httpx.Response, shape: str) -> None: + assert response.status_code >= 400, (shape, response.status_code, response.text) + assert _error_message(response), (shape, response.text) + + +def _assert_served(response: httpx.Response, identity: str, shape: str) -> None: + assert response.status_code == 200, (shape, response.status_code, response.text) + assert _response_id(response) == identity, (shape, response.text) + + +def _moderation_through_the_sdk(gateway: Gateway, model: str) -> ModerationCreateResponse | openai.APIStatusError: + with _openai(gateway) as client: + try: + return client.moderations.create(model=model, input=TEXT) + except openai.APIStatusError as error: + return error + + +def _moderation_after_worker_sync(gateway: Gateway, model: str) -> ModerationCreateResponse | openai.APIStatusError: + return eventually( + lambda: _moderation_through_the_sdk(gateway, model), + lambda outcome: not isinstance(outcome, openai.InternalServerError), + seconds=WORKER_SYNC_SECONDS + 10, + ) + + +def _speech_peer(content_type: str) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + if request.method != "POST" or not _path(request).endswith("/audio/speech"): + return None + return Reply(body=SPEECH_AUDIO, content_type=content_type) + + return _peer(route) + + +def _speech_through_the_sdk(gateway: Gateway, model: str) -> httpx.Response | openai.APIStatusError: + with _openai(gateway) as client: + try: + return client.audio.speech.create(model=model, voice="alloy", input=TEXT, response_format="mp3").response + except openai.APIStatusError as error: + return error + + +def _speech_after_worker_sync(gateway: Gateway, model: str) -> httpx.Response | openai.APIStatusError: + return eventually( + lambda: _speech_through_the_sdk(gateway, model), + lambda outcome: not isinstance(outcome, openai.InternalServerError), + seconds=WORKER_SYNC_SECONDS + 10, + ) + + +def _attempt(call: Callable[[], _R]) -> _R | openai.APIStatusError: + try: + return call() + except openai.APIStatusError as error: + return error + + +def _unknown_to_the_serving_worker(outcome: object) -> bool: + return isinstance(outcome, openai.BadRequestError) and "Invalid model name" in outcome.response.text + + +def _outcome_after_worker_sync(call: Callable[[], _R]) -> _R | openai.APIStatusError: + return eventually( + lambda: _attempt(call), + lambda outcome: not _unknown_to_the_serving_worker(outcome), + seconds=WORKER_SYNC_SECONDS + 10, + ) + + +def _after_worker_sync(call: Callable[[], _R]) -> _R: + outcome: Final = _outcome_after_worker_sync(call) + assert not isinstance(outcome, openai.APIStatusError), outcome.response.text + return outcome + + +def _async_after_worker_sync(gateway: Gateway, call: Callable[[openai.AsyncOpenAI], Awaitable[_R]]) -> _R: + return _after_worker_sync(lambda: _async_openai(gateway, call)) + + +def _streamed_file_content(gateway: Gateway, file_id: str, model: str) -> tuple[int, tuple[bytes, ...]]: + with gateway.client.stream( + "GET", + f"/v1/files/{file_id}/content", + params={"model": model}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + return response.status_code, tuple(response.iter_bytes()) + + +def _upstream_id(managed: str, prefix: str) -> str: + encoded: Final = managed.removeprefix(prefix) + try: + decoded: Final = base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)).decode() + except (binascii.Error, UnicodeDecodeError): + return managed + return decoded.removeprefix("litellm:").split(";", 1)[0] if decoded.startswith("litellm:") else managed + + +def _spend_row(identity: str) -> Mapping[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, call_type, model FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (identity,) + ), + lambda rows: len(rows) == 1, + seconds=70, + ) + return rows[0] + + +def _spend_row_keyed_by_any(identities: frozenset[str]) -> Mapping[str, JsonValue]: + keys: Final = tuple(sorted(identities)) + placeholders: Final = ", ".join("%s" for _ in keys) + rows: Final = eventually( + lambda: read_rows( + f'SELECT request_id, call_type, model FROM "LiteLLM_SpendLogs" WHERE request_id IN ({placeholders})', keys + ), + lambda rows: len(rows) == 1, + seconds=70, + ) + return rows[0] + + +def _spend_request_ids(identities: tuple[str, str, str]) -> frozenset[str]: + rows: Final = eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id IN (%s, %s, %s)', identities), + lambda rows: len(rows) == 3, + seconds=70, + ) + return frozenset(string_value(row["request_id"]) for row in rows) + + +def _anthropic_text(event: RawMessageStreamEvent) -> str: + if isinstance(event, RawContentBlockDeltaEvent) and isinstance(event.delta, TextDelta): + return event.delta.text + return "" + + +def _chunk_text(chunk: ChatCompletionChunk) -> str: + return chunk.choices[0].delta.content or "" if chunk.choices else "" + + +def _chunk_finished(chunk: ChatCompletionChunk) -> bool: + return bool(chunk.choices) and chunk.choices[0].finish_reason == "stop" + + +def _new_model(gateway: Gateway, api_base: str, timeout: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/model/new", + { + "model_name": f"integration-{uuid.uuid4().hex}", + "litellm_params": { + "model": f"openai/{UPSTREAM_MODEL}", + "api_key": PROVIDER_KEY, + "api_base": api_base, + "timeout": timeout, + }, + "model_info": {}, + }, + ) + + +def _register_accepted_model(scenario: Scenario, response: httpx.Response, shape: str) -> None: + assert response.status_code == 200, (shape, response.status_code, response.text) + identity: Final = string_value(object_value(JSON_OBJECT.validate_json(response.content)["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + + +def test_r01_chat_completions_non_stream_through_the_sync_openai_sdk(gateway: Gateway) -> None: + identity: Final = _identity("chatcmpl") + with gateway.scenario() as scenario, wire_server(_chat_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + with _openai(gateway) as client: + response: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": TEXT}], extra_body=dict(NO_CACHE) + ) + assert response.id == identity, response + assert response.choices[0].message.content == TEXT, response + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/chat/completions"),), served + assert served[0].headers["authorization"] == PROVIDER_AUTHORIZATION, served[0].headers + upstream_body: Final = _body(served[0]) + assert upstream_body["model"] == UPSTREAM_MODEL, upstream_body + assert upstream_body["messages"] == [_user_message()], upstream_body + assert _spend_row(identity)["model"] == DEPLOYMENT_MODEL + + +def test_r02_chat_completions_stream_through_the_async_openai_sdk(gateway: Gateway) -> None: + identity: Final = _identity("chatcmpl") + with gateway.scenario() as scenario, wire_server(_chat_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + + async def stream(client: openai.AsyncOpenAI) -> tuple[ChatCompletionChunk, ...]: + chunks: Final = await client.chat.completions.create( + model=model, messages=[{"role": "user", "content": TEXT}], stream=True, extra_body=dict(NO_CACHE) + ) + return tuple([chunk async for chunk in chunks]) + + received: Final = _async_openai(gateway, stream) + assert {chunk.id for chunk in received} == {identity}, received + assert "".join(_chunk_text(chunk) for chunk in received) == TEXT, received + assert any(_chunk_finished(chunk) for chunk in received), received + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/chat/completions"),), served + assert _body(served[0])["stream"] is True, served[0].body + assert _spend_row(identity)["model"] == DEPLOYMENT_MODEL + + +def test_r03_chat_completions_through_raw_httpx_with_a_body_timeout(gateway: Gateway) -> None: + identity: Final = _identity("chatcmpl") + with gateway.scenario() as scenario, wire_server(_chat_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + assert _response_id(_chat(gateway, model, {"timeout": 30})) == identity + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/chat/completions"),), served + assert "timeout" not in _body(served[0]), served[0].body + assert _spend_row(identity)["model"] == DEPLOYMENT_MODEL + + +def test_r04_messages_non_stream_through_the_sync_anthropic_sdk(gateway: Gateway) -> None: + identity: Final = _identity("resp") + with gateway.scenario() as scenario, wire_server(_responses_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + with _anthropic(gateway) as client: + message: Final = client.messages.create( + model=model, max_tokens=64, messages=[{"role": "user", "content": TEXT}] + ) + assert identity in response_identities(message.id), message.id + first_block: Final = message.content[0] + assert first_block.type == "text" and first_block.text == TEXT, message + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/responses"),), served + assert TEXT in served[0].body.decode(), served[0].body + assert _spend_row(message.id)["call_type"] == "anthropic_messages" + + +def test_r05_messages_stream_through_the_async_anthropic_sdk(gateway: Gateway) -> None: + identity: Final = _identity("resp") + with gateway.scenario() as scenario, wire_server(_responses_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + + async def stream(client: anthropic.AsyncAnthropic) -> tuple[RawMessageStreamEvent, ...]: + events: Final = await client.messages.create( + model=model, max_tokens=64, messages=[{"role": "user", "content": TEXT}], stream=True + ) + return tuple([event async for event in events]) + + received: Final = _async_anthropic(gateway, stream) + first: Final = received[0] + assert isinstance(first, RawMessageStartEvent), received + assert received[-1].type == "message_stop", received + assert "".join(_anthropic_text(event) for event in received) == TEXT, received + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/responses"),), served + assert _body(served[0])["stream"] is True, served[0].body + assert _spend_row(first.message.id)["call_type"] == "anthropic_messages" + + +def test_r06_responses_non_stream_through_the_sync_openai_sdk(gateway: Gateway) -> None: + identity: Final = _identity("resp") + with gateway.scenario() as scenario, wire_server(_responses_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + with _openai(gateway) as client: + response: Final = client.responses.create(model=model, input=TEXT, extra_body=dict(NO_CACHE)) + assert identity in response_identities(response.id), response.id + assert response.output_text == TEXT, response + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/responses"),), served + upstream_body: Final = _body(served[0]) + assert upstream_body["model"] == UPSTREAM_MODEL, upstream_body + assert upstream_body["input"] == TEXT, upstream_body + assert _spend_row(response.id)["model"] == DEPLOYMENT_MODEL + + +def test_r07_responses_stream_through_the_async_openai_sdk(gateway: Gateway) -> None: + identity: Final = _identity("resp") + with gateway.scenario() as scenario, wire_server(_responses_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + + async def stream(client: openai.AsyncOpenAI) -> tuple[ResponseStreamEvent, ...]: + events: Final = await client.responses.create( + model=model, input=TEXT, stream=True, extra_body=dict(NO_CACHE) + ) + return tuple([event async for event in events]) + + received: Final = _async_openai(gateway, stream) + assert received[0].type == "response.created", received + last: Final = received[-1] + assert last.type == "response.completed", received + assert identity in response_identities(last.response.id), last + assert last.response.output_text == TEXT, last + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/responses"),), served + assert _body(served[0])["stream"] is True, served[0].body + assert _spend_row_keyed_by_any(response_identities(last.response.id))["model"] == DEPLOYMENT_MODEL + + +def test_r08_completions_non_stream_through_the_sync_openai_sdk(gateway: Gateway) -> None: + identity: Final = _identity("cmpl") + with gateway.scenario() as scenario, wire_server(_completion_peer(identity)) as wire: + model: Final = scenario.model(model=INSTRUCT_DEPLOYMENT_MODEL, api_base=f"{wire.url}/v1") + with _openai(gateway) as client: + completion: Final = client.completions.create(model=model, prompt=TEXT, extra_body=dict(NO_CACHE)) + assert completion.id == identity, completion + assert completion.choices[0].text == TEXT, completion + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/completions"),), served + upstream_body: Final = _body(served[0]) + assert upstream_body["model"] == INSTRUCT_MODEL, upstream_body + assert upstream_body["prompt"] == TEXT, upstream_body + assert _spend_row(identity)["model"] == INSTRUCT_DEPLOYMENT_MODEL + + +def test_r09_completions_stream_through_the_async_openai_sdk(gateway: Gateway) -> None: + identity: Final = _identity("cmpl") + with gateway.scenario() as scenario, wire_server(_completion_peer(identity)) as wire: + model: Final = scenario.model(model=INSTRUCT_DEPLOYMENT_MODEL, api_base=f"{wire.url}/v1") + + async def stream(client: openai.AsyncOpenAI) -> tuple[Completion, ...]: + chunks: Final = await client.completions.create( + model=model, prompt=TEXT, stream=True, extra_body=dict(NO_CACHE) + ) + return tuple([chunk async for chunk in chunks]) + + received: Final = _async_openai(gateway, stream) + assert {chunk.id for chunk in received} == {identity}, received + assert "".join(chunk.choices[0].text or "" for chunk in received if chunk.choices) == TEXT, received + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/completions"),), served + assert _body(served[0])["stream"] is True, served[0].body + assert _spend_row(identity)["model"] == INSTRUCT_DEPLOYMENT_MODEL + + +def test_r10_moderations_through_the_sync_openai_sdk(gateway: Gateway) -> None: + identity: Final = _identity("modr") + with gateway.scenario() as scenario, wire_server(_moderation_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + moderation: Final = _moderation_after_worker_sync(gateway, model) + assert isinstance(moderation, ModerationCreateResponse), moderation + assert moderation.id == identity, moderation + assert [result.flagged for result in moderation.results] == [False], moderation + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/moderations"),), served + assert served[0].headers["authorization"] == PROVIDER_AUTHORIZATION, served[0].headers + assert _body(served[0])["input"] == TEXT, served[0].body + assert _spend_row(moderation.id)["call_type"] == "amoderation" + + +def test_r11_files_create_retrieve_and_content_through_the_sync_openai_sdk(gateway: Gateway) -> None: + file_id: Final = _identity("file") + content: Final = _batch_line(file_id) + with gateway.scenario() as scenario, wire_server(_files_peer(file_id, content)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + with _openai(gateway) as client: + created: Final = _after_worker_sync( + lambda: client.files.create(file=("batch.jsonl", content), purpose="batch", extra_body={"model": model}) + ) + retrieved: Final = _after_worker_sync( + lambda: client.files.retrieve(created.id, extra_query={"model": model}) + ) + downloaded: Final = _after_worker_sync( + lambda: client.files.content(created.id, extra_query={"model": model}) + ).content + assert _upstream_id(created.id, "file-") == file_id, created + assert _upstream_id(retrieved.id, "file-") == file_id, retrieved + assert retrieved.bytes == len(content), retrieved + assert downloaded == content, downloaded + served: Final = _served(wire) + assert _routes(served) == ( + ("POST", "/v1/files"), + ("GET", f"/v1/files/{file_id}"), + ("GET", f"/v1/files/{file_id}/content"), + ), served + assert content in served[0].body, served[0].body + assert {request.headers["authorization"] for request in served} == {PROVIDER_AUTHORIZATION}, served + + +def test_r12_file_content_streams_through_raw_httpx(gateway: Gateway) -> None: + file_id: Final = _identity("file") + content: Final = STREAM_HEAD + STREAM_TAIL + with ( + gateway.scenario() as scenario, + wire_server(_files_peer(file_id, content, chunks=(STREAM_HEAD, STREAM_TAIL))) as wire, + ): + model: Final = scenario.model(api_base=f"{wire.url}/v1") + status, received = eventually( + lambda: _streamed_file_content(gateway, file_id, model), + lambda outcome: not (outcome[0] == 400 and b"Invalid model name" in b"".join(outcome[1])), + seconds=WORKER_SYNC_SECONDS + 10, + ) + assert status == 200, b"".join(received).decode() + assert b"".join(received) == content, (len(received), sum(len(chunk) for chunk in received)) + served: Final = _served(wire) + assert _routes(served) == (("GET", f"/v1/files/{file_id}/content"),), served + assert served[0].headers["authorization"] == PROVIDER_AUTHORIZATION, served[0].headers + + +def test_r13_batches_create_retrieve_list_and_cancel_through_the_async_openai_sdk(gateway: Gateway) -> None: + file_id: Final = _identity("file") + batch_id: Final = _identity("batch") + with gateway.scenario() as scenario, wire_server(_batches_peer(file_id, batch_id)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + uploaded: Final = _async_after_worker_sync( + gateway, + lambda client: client.files.create( + file=("batch.jsonl", _batch_line("r1")), purpose="batch", extra_body={"model": model} + ), + ) + created: Final = _async_after_worker_sync( + gateway, + lambda client: client.batches.create( + input_file_id=uploaded.id, + endpoint="/v1/chat/completions", + completion_window="24h", + extra_body={"model": model}, + ), + ) + retrieved: Final = _async_after_worker_sync(gateway, lambda client: client.batches.retrieve(created.id)) + listed: Final = _async_after_worker_sync( + gateway, lambda client: client.batches.list(extra_query={"model": model}) + ) + cancelled: Final = _async_after_worker_sync(gateway, lambda client: client.batches.cancel(created.id)) + listed_ids: Final = tuple(batch.id for batch in listed.data) + assert _upstream_id(uploaded.id, "file-") == file_id, uploaded + assert _upstream_id(created.id, "batch_") == batch_id, created + assert _upstream_id(retrieved.id, "batch_") == batch_id, retrieved + assert batch_id in [_upstream_id(listed_id, "batch_") for listed_id in listed_ids], listed_ids + assert _upstream_id(cancelled.id, "batch_") == batch_id, cancelled + assert cancelled.status == "cancelling", cancelled + served: Final = _served(wire) + assert _routes(served) == ( + ("POST", "/v1/files"), + ("POST", "/v1/batches"), + ("GET", f"/v1/batches/{batch_id}"), + ("POST", f"/v1/batches/{batch_id}/cancel"), + ), served + assert _body(served[1])["input_file_id"] == file_id, served[1].body + assert {request.headers["authorization"] for request in served} == {PROVIDER_AUTHORIZATION}, served + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_r14_fine_tuning_jobs_create_list_retrieve_and_cancel_through_the_sync_openai_sdk( + configured_proxy: ConfiguredProxy, +) -> None: + suffix: Final = uuid.uuid4().hex + job_id: Final = f"ftjob-{suffix}" + with _openai(configured_proxy.gateway) as client: + created: Final = client.fine_tuning.jobs.create( + model=UPSTREAM_MODEL, training_file=f"file-{suffix}", extra_body={"custom_llm_provider": "openai"} + ) + listed: Final = client.fine_tuning.jobs.list(extra_query={"custom_llm_provider": "openai"}) + retrieved: Final = client.fine_tuning.jobs.retrieve(created.id, extra_query={"custom_llm_provider": "openai"}) + cancelled: Final = client.fine_tuning.jobs.cancel(created.id, extra_body={"custom_llm_provider": "openai"}) + assert created.id == job_id, created + assert created.status == "queued", created + assert [job.id for job in listed.data] == [LISTED_JOB], listed + assert retrieved.id == job_id, retrieved + assert retrieved.training_file == f"file-{suffix}", retrieved + assert cancelled.id == job_id and cancelled.status == "cancelled", cancelled + served: Final = _served(configured_proxy.wire) + assert _routes(served) == ( + ("POST", "/v1/fine_tuning/jobs"), + ("GET", "/v1/fine_tuning/jobs"), + ("GET", f"/v1/fine_tuning/jobs/{job_id}"), + ("POST", f"/v1/fine_tuning/jobs/{job_id}/cancel"), + ), served + assert _body(served[0])["training_file"] == f"file-{suffix}", served[0].body + assert {request.headers["authorization"] for request in served} == {PROVIDER_AUTHORIZATION}, served + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_r15_assistants_list_through_raw_httpx(configured_proxy: ConfiguredProxy) -> None: + response: Final = configured_proxy.gateway.request("GET", "/v1/assistants") + assert response.status_code == 200, response.text + body: Final = JSON_OBJECT.validate_json(response.content) + assert body["object"] == "list", response.text + listed: Final = body["data"] + assert isinstance(listed, list), response.text + assert [string_value(object_value(entry)["id"]) for entry in listed] == [ASSISTANT], response.text + served: Final = _served(configured_proxy.wire) + assert _routes(served) == (("GET", "/v1/assistants"),), served + assert served[0].headers["authorization"] == PROVIDER_AUTHORIZATION, served[0].headers + + +def test_r17_chat_completions_on_an_azure_cloudflare_gateway_deployment_through_the_sync_openai_sdk( + gateway: Gateway, +) -> None: + identity: Final = _identity("chatcmpl") + deployment: Final = f"gpt-4o-mini-{uuid.uuid4().hex[:8]}" + with gateway.scenario() as scenario, wire_server(_chat_peer(identity)) as wire: + model: Final = scenario.model( + model=f"azure/{deployment}", api_base=f"{wire.url}{CLOUDFLARE_GATEWAY}", api_version=AZURE_API_VERSION + ) + with _openai(gateway) as client: + response: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": TEXT}], extra_body=dict(NO_CACHE) + ) + streamed: Final = tuple( + client.chat.completions.create( + model=model, messages=[{"role": "user", "content": TEXT}], stream=True, extra_body=dict(NO_CACHE) + ) + ) + assert response.id == identity, response + assert response.choices[0].message.content == TEXT, response + assert {chunk.id for chunk in streamed} == {identity}, streamed + assert "".join(_chunk_text(chunk) for chunk in streamed) == TEXT, streamed + served: Final = _served(wire) + assert _routes(served) == (("POST", f"{CLOUDFLARE_GATEWAY}/{deployment}/chat/completions"),) * 2, served + assert {request.target.split("?", 1)[1] for request in served} == {f"api-version={AZURE_API_VERSION}"}, served + assert {request.headers["api-key"] for request in served} == {PROVIDER_KEY}, served + assert [_streams(request) for request in served] == [False, True], served + assert _spend_row(identity)["model"] == f"azure/{deployment}" + + +def test_r18_audio_speech_answers_with_the_upstream_audio_type_through_the_sync_openai_sdk(gateway: Gateway) -> None: + with gateway.scenario() as scenario, wire_server(_speech_peer("audio/ogg")) as wire: + model: Final = scenario.model(model=f"openai/{SPEECH_UPSTREAM_MODEL}", api_base=f"{wire.url}/v1") + speech: Final = _speech_after_worker_sync(gateway, model) + assert not isinstance(speech, openai.APIStatusError), speech + assert speech.status_code == 200, speech.text + assert speech.headers["content-type"] == "audio/ogg", dict(speech.headers) + assert speech.content == SPEECH_AUDIO, speech.content[:16] + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/audio/speech"),), served + assert served[0].headers["authorization"] == PROVIDER_AUTHORIZATION, served[0].headers + upstream_body: Final = _body(served[0]) + assert (upstream_body["model"], upstream_body["input"], upstream_body["response_format"]) == ( + SPEECH_UPSTREAM_MODEL, + TEXT, + "mp3", + ), upstream_body + assert _spend_row(speech.headers["x-litellm-call-id"])["call_type"] == "aspeech" + + +def test_s01_chat_completions_upstream_400_reaches_the_sync_openai_sdk_as_a_bad_request(gateway: Gateway) -> None: + with gateway.scenario() as scenario, wire_server(_status_peer(openai_error(400))) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1", max_retries=0) + with _openai(gateway) as client, pytest.raises(openai.BadRequestError) as raised: + client.chat.completions.create(model=model, messages=[_user_message()], extra_body=dict(NO_CACHE)) + assert raised.value.status_code == 400, raised.value + assert "scripted 400" in raised.value.response.text, raised.value.response.text + assert _routes(_served(wire)) == (("POST", "/v1/chat/completions"),) + + +def test_s02_chat_completions_upstream_429_keeps_retry_after_through_raw_httpx(gateway: Gateway) -> None: + limited: Final = Reply(status=429, body=openai_error(429).body, headers={"retry-after": "7"}) + with gateway.scenario() as scenario, wire_server(_status_peer(limited)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1", max_retries=0) + response: Final = _chat(gateway, model) + assert response.status_code == 429, response.text + assert response.headers.get("llm_provider-retry-after") == "7", dict(response.headers) + assert "scripted 429" in _error_message(response), response.text + assert _routes(_served(wire)) == (("POST", "/v1/chat/completions"),) + + +def test_s03_chat_completions_upstream_500_reaches_the_sync_openai_sdk_as_a_server_error(gateway: Gateway) -> None: + with gateway.scenario() as scenario, wire_server(_status_peer(openai_error(500))) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1", max_retries=0) + with _openai(gateway) as client, pytest.raises(openai.InternalServerError) as raised: + client.chat.completions.create(model=model, messages=[_user_message()], extra_body=dict(NO_CACHE)) + assert raised.value.status_code == 500, raised.value + assert "scripted 500" in raised.value.response.text, raised.value.response.text + assert _routes(_served(wire)) == (("POST", "/v1/chat/completions"),) + + +def test_s04_chat_completions_body_timeout_shapes_answer_without_breaking_the_proxy(gateway: Gateway) -> None: + identity: Final = _identity("chatcmpl") + with gateway.scenario() as scenario, wire_server(_chat_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + for shape, timeout in REJECTED_CHAT_TIMEOUTS.items(): + _assert_rejected(_chat(gateway, model, {"timeout": timeout}), shape) + for shape, timeout in ACCEPTED_CHAT_TIMEOUTS.items(): + _assert_served(_chat(gateway, model, {"timeout": timeout}), identity, shape) + duplicated: Final = gateway.client.post( + "/v1/chat/completions", + content=_duplicated_timeout_body(model), + headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, + ) + _assert_served(duplicated, identity, "duplicated key") + _assert_served(_chat(gateway, model), identity, "follow-up") + assert len(_served(wire)) == len(ACCEPTED_CHAT_TIMEOUTS) + 2 + + +def test_s05_model_new_timeout_shapes_are_rejected_or_stored_without_breaking_the_proxy(gateway: Gateway) -> None: + identity: Final = _identity("chatcmpl") + with gateway.scenario() as scenario, wire_server(_chat_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + for shape, timeout in REJECTED_DEPLOYMENT_TIMEOUTS.items(): + _assert_rejected(_new_model(gateway, f"{wire.url}/v1", timeout), shape) + for shape, timeout in ACCEPTED_DEPLOYMENT_TIMEOUTS.items(): + _register_accepted_model(scenario, _new_model(gateway, f"{wire.url}/v1", timeout), shape) + _assert_served(_chat(gateway, model), identity, "follow-up") + assert _routes(_served(wire)) == (("POST", "/v1/chat/completions"),) + + +def test_s06_upstream_errors_on_completions_moderations_and_files_reach_the_sync_openai_sdk( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + with wire_server(_status_peer(openai_error(401))) as unauthorized: + model: Final = scenario.model( + model=INSTRUCT_DEPLOYMENT_MODEL, api_base=f"{unauthorized.url}/v1", max_retries=0 + ) + with _openai(gateway) as client, pytest.raises(openai.AuthenticationError) as completion_error: + client.completions.create(model=model, prompt=TEXT, extra_body=dict(NO_CACHE)) + assert completion_error.value.status_code == 401, completion_error.value + assert "scripted 401" in completion_error.value.response.text, completion_error.value.response.text + assert _routes(_served(unauthorized)) == (("POST", "/v1/completions"),) + with wire_server(_status_peer(openai_error(400))) as rejecting: + rejected_model: Final = scenario.model(api_base=f"{rejecting.url}/v1", max_retries=0) + moderation_error: Final = _moderation_after_worker_sync(gateway, rejected_model) + assert isinstance(moderation_error, openai.BadRequestError), moderation_error + assert "scripted 400" in moderation_error.response.text, moderation_error.response.text + with _openai(gateway) as client: + file_error: Final = _outcome_after_worker_sync( + lambda: client.files.create( + file=("batch.jsonl", _batch_line("rejected")), + purpose="batch", + extra_body={"model": rejected_model}, + ) + ) + assert isinstance(file_error, openai.BadRequestError), file_error + assert "scripted 400" in file_error.response.text, file_error.response.text + assert _routes(_served(rejecting)) == (("POST", "/v1/moderations"), ("POST", "/v1/files")) + + +@pytest.mark.parametrize("provider", tuple(PAYMENT_REQUIRED_DEPLOYMENTS)) +def test_s07_chat_completions_upstream_402_reaches_the_sync_openai_sdk_as_payment_required( + gateway: Gateway, provider: str +) -> None: + with gateway.scenario() as scenario, wire_server(_status_peer(openai_error(402))) as wire: + model: Final = scenario.model(model=PAYMENT_REQUIRED_DEPLOYMENTS[provider], api_base=wire.url, max_retries=0) + with _openai(gateway) as client, pytest.raises(openai.APIStatusError) as raised: + client.chat.completions.create(model=model, messages=[_user_message()], extra_body=dict(NO_CACHE)) + assert raised.value.status_code == 402, raised.value + assert "scripted 402" in raised.value.response.text, raised.value.response.text + assert len(_served(wire)) == 1 + + +def test_e01_deployment_timeout_answers_408_and_a_body_timeout_on_the_same_deployment_recovers( + gateway: Gateway, +) -> None: + identity: Final = _identity("chatcmpl") + gate: Final = threading.Event() + with gateway.scenario() as scenario, wire_server(_held_chat_peer(gate, identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1", timeout=0.5, max_retries=0) + timed_out: Final = _chat(gateway, model) + gate.set() + assert timed_out.status_code == 408, timed_out.text + assert "timeout" in _error_message(timed_out).lower(), timed_out.text + recovered: Final = _chat(gateway, model, {"timeout": 5}) + assert _response_id(recovered) == identity, recovered.text + assert _routes(_served(wire)) == (("POST", "/v1/chat/completions"),) * 2 + + +def test_e02_three_identical_no_cache_requests_through_raw_httpx_land_three_spend_rows(gateway: Gateway) -> None: + served_ids: Final[SimpleQueue[str]] = SimpleQueue() + with gateway.scenario() as scenario, wire_server(_fresh_chat_peer(served_ids)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + received: Final = ( + _response_id(_chat(gateway, model)), + _response_id(_chat(gateway, model)), + _response_id(_chat(gateway, model)), + ) + assert len(set(received)) == 3, received + assert _routes(_served(wire)) == (("POST", "/v1/chat/completions"),) * 3 + assert {served_ids.get_nowait() for _ in range(3)} == set(received), received + assert _spend_request_ids(received) == frozenset(received) + + +def test_e03_deployment_timeout_as_a_string_still_serves_through_raw_httpx(gateway: Gateway) -> None: + identity: Final = _identity("chatcmpl") + with gateway.scenario() as scenario, wire_server(_chat_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1", timeout="30") + assert _response_id(_chat(gateway, model)) == identity + assert _routes(_served(wire)) == (("POST", "/v1/chat/completions"),) + assert _spend_row(identity)["model"] == DEPLOYMENT_MODEL diff --git a/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py b/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py new file mode 100644 index 00000000000..0bcc9fc4340 --- /dev/null +++ b/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py @@ -0,0 +1,424 @@ +from __future__ import annotations + +import json +import signal +import subprocess +import uuid +from collections.abc import Mapping +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass +from functools import partial +from itertools import product +from pathlib import Path +from typing import Final, Literal, TypeAlias + +import httpx +import psutil +import pytest +from anthropic import Anthropic +from anthropic import APIError as AnthropicApiError +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy_process, owned_upstream +from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario +from openai import APIError as OpenAiApiError +from openai import OpenAI +from pydantic import JsonValue + +from tests.integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, SseResponse + +_Kind: TypeAlias = Literal["chat", "chat_stream", "messages", "responses", "completions", "moderations"] + +_ALL_KINDS: Final[tuple[_Kind, ...]] = ("chat", "chat_stream", "messages", "responses", "completions", "moderations") +_CHAT_KINDS: Final[tuple[_Kind, ...]] = ("chat", "chat_stream") +_NO_CACHE: Final = {"cache": {"no-cache": True}} +_SPEND_QUERY: Final = 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = %s' +_CHAT: Final = JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$UNIQUE_ID", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "scripted"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, +) +_MODERATION: Final = JsonResponse( + content_type="application/json", + body={ + "id": "modr-$UNIQUE_ID", + "model": "omni-moderation-latest", + "results": [{"flagged": False, "categories": {}, "category_scores": {}}], + }, +) +_RESPONSES: Final = JsonResponse( + content_type="application/json", + body={ + "id": "resp_$UNIQUE_ID", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_$UNIQUE_ID", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "scripted", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, +) +_PLAIN: Final = RoutedResponse( + content_type="application/x-routed", + routes={"POST /chat/completions": _CHAT, "POST /moderations": _MODERATION, "POST /responses": _RESPONSES}, +) +_STREAM: Final = SseResponse( + content_type="text/event-stream", + frames=( + 'data: {"id":"chatcmpl-$UNIQUE_ID","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",' + '"choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}', + 'data: {"id":"chatcmpl-$UNIQUE_ID","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",' + '"choices":[{"index":0,"delta":{"content":"streamed"},"finish_reason":null}]}', + 'data: {"id":"chatcmpl-$UNIQUE_ID","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",' + '"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ), + frame_delay_ms=150, +) + + +@dataclass(frozen=True, slots=True) +class _Call: + kind: _Kind + marker: str + + +@dataclass(frozen=True, slots=True) +class _Completed: + call: _Call + response_id: str + + +@dataclass(frozen=True, slots=True) +class _Failed: + call: _Call + error: str + + +@dataclass(frozen=True, slots=True) +class _Clients: + openai: OpenAI + anthropic: Anthropic + plain_model: str + stream_model: str + + +class _Observations: + def __init__(self, url: str) -> None: + self.url = url.rstrip("/") + self.items: tuple[Mapping[str, JsonValue], ...] = () + + def read(self) -> tuple[Mapping[str, JsonValue], ...]: + with httpx.Client(timeout=10, trust_env=False) as client: + payload: Final = JSON_OBJECT.validate_python(client.get(f"{self.url}/__observations").json()) + requests: Final = payload.get("requests") + assert isinstance(requests, list) + self.items = (*self.items, *(object_value(item) for item in requests if isinstance(item, dict))) + return self.items + + +def _register( + scenario: Scenario, upstream_url: str, label: str, response: RoutedResponse | SseResponse +) -> ScenarioHandle: + handle: Final = register_scenario(f"chaos-{label}-{uuid.uuid4().hex}", response, control_url=upstream_url) + scenario.cleanups.callback(delete_scenario, handle) + return handle + + +def _deployments(scenario: Scenario, upstream_url: str) -> tuple[str, str]: + plain_handle: Final = _register(scenario, upstream_url, "plain", _PLAIN) + stream_handle: Final = _register(scenario, upstream_url, "stream", _STREAM) + return ( + scenario.model(api_base=plain_handle.api_base(), api_key=plain_handle.scenario_id), + scenario.model(api_base=stream_handle.api_base(), api_key=stream_handle.scenario_id), + ) + + +def _clients(gateway: Gateway, deployments: tuple[str, str]) -> _Clients: + base: Final = str(gateway.client.base_url).rstrip("/") + clients: Final = _Clients( + openai=OpenAI(base_url=f"{base}/v1", api_key=gateway.key, max_retries=0), + anthropic=Anthropic(base_url=base, api_key=gateway.key, max_retries=0), + plain_model=deployments[0], + stream_model=deployments[1], + ) + _import_lazy_sdk_resources_on_this_thread(clients) + return clients + + +def _import_lazy_sdk_resources_on_this_thread(clients: _Clients) -> None: + resources: Final = ( + clients.openai.chat.completions, + clients.openai.completions, + clients.openai.responses, + clients.openai.moderations, + clients.anthropic.messages, + ) + assert all(resource is not None for resource in resources) + + +def _invoke(clients: _Clients, call: _Call) -> str: + match call.kind: + case "chat": + return clients.openai.chat.completions.create( + model=clients.plain_model, messages=[{"role": "user", "content": call.marker}], extra_body=_NO_CACHE + ).id + case "chat_stream": + chunks: Final = tuple( + clients.openai.chat.completions.create( + model=clients.stream_model, + messages=[{"role": "user", "content": call.marker}], + stream=True, + extra_body=_NO_CACHE, + ) + ) + assert chunks[-1].choices[0].finish_reason == "stop", chunks + assert len({chunk.id for chunk in chunks}) == 1, chunks + return chunks[0].id + case "messages": + return clients.anthropic.messages.create( + model=clients.plain_model, + max_tokens=64, + messages=[{"role": "user", "content": call.marker}], + extra_body=_NO_CACHE, + ).id + case "responses": + return clients.openai.responses.create( + model=clients.plain_model, input=call.marker, extra_body=_NO_CACHE + ).id + case "completions": + return clients.openai.completions.create( + model=clients.plain_model, prompt=call.marker, extra_body=_NO_CACHE + ).id + case "moderations": + return clients.openai.moderations.create( + model=clients.plain_model, input=call.marker, extra_body=_NO_CACHE + ).id + + +def _attempt(clients: _Clients, call: _Call) -> _Completed | _Failed: + try: + return _Completed(call, _invoke(clients, call)) + except (OpenAiApiError, AnthropicApiError, httpx.HTTPError) as error: + return _Failed(call, f"{type(error).__name__}: {str(error)[:200]}") + + +def _every_worker_serves_every_kind(gateway: Gateway, deployments: tuple[str, str]) -> bool: + calls: Final = _burst_calls(_ALL_KINDS, 2) + with ThreadPoolExecutor(max_workers=len(calls)) as pool: + futures: Final = _submit(pool, _clients(gateway, deployments), calls) + outcomes: Final = tuple(future.result(timeout=60) for future in futures) + return all(isinstance(outcome, _Completed) for outcome in outcomes) + + +def _burst_calls(kinds: tuple[_Kind, ...], per_kind: int) -> tuple[_Call, ...]: + return tuple(_Call(kind, f"burst-{kind}-{uuid.uuid4().hex}") for kind, _ in product(kinds, range(per_kind))) + + +def _submit( + pool: ThreadPoolExecutor, clients: _Clients, calls: tuple[_Call, ...] +) -> tuple[Future[_Completed | _Failed], ...]: + return tuple(pool.submit(_attempt, clients, call) for call in calls) + + +def _done_count(futures: tuple[Future[_Completed | _Failed], ...]) -> int: + return sum(future.done() for future in futures) + + +def _liveliness(gateway: Gateway) -> int: + return gateway.client.get("/health/liveliness").status_code + + +def _spend_rows(request_id: str) -> tuple[Mapping[str, JsonValue], ...]: + return tuple(read_rows(_SPEND_QUERY, (request_id,))) + + +def _marker_counts(items: tuple[Mapping[str, JsonValue], ...], markers: tuple[str, ...]) -> tuple[int, ...]: + bodies: Final = tuple(json.dumps(item.get("body")) for item in items) + return tuple(sum(marker in body for body in bodies) for marker in markers) + + +def _landed_once(response_ids: tuple[str, ...]) -> tuple[tuple[Mapping[str, JsonValue], ...], ...]: + rows: Final = tuple( + eventually(partial(_spend_rows, response_id), lambda values: len(values) == 1, seconds=90) + for response_id in response_ids + ) + assert all(found[0]["request_id"] == response_id for found, response_id in zip(rows, response_ids)), rows + return rows + + +def _is_worker(child: psutil.Process) -> bool: + try: + return "spawn_main" in " ".join(child.cmdline()) + except psutil.Error: + return False + + +def _worker_alive(worker: psutil.Process) -> bool: + try: + return worker.is_running() and worker.status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + +def _alive_workers(process: subprocess.Popen[bytes]) -> tuple[psutil.Process, ...]: + children: Final = tuple(child for child in psutil.Process(process.pid).children() if _is_worker(child)) + return tuple(child for child in children if _worker_alive(child)) + + +def _process_tree(process: subprocess.Popen[bytes]) -> str: + root: Final = psutil.Process(process.pid) + return "\n".join(_process_line(member) for member in (root, *root.children(recursive=True))) + + +def _process_line(member: psutil.Process) -> str: + try: + return f"{member.pid} {' '.join(member.cmdline())}" + except psutil.Error: + return f"{member.pid} " + + +def _kind_counts(outcomes: tuple[_Completed | _Failed, ...]) -> str: + kinds: Final = tuple(outcome.call.kind for outcome in outcomes) + return str({kind: kinds.count(kind) for kind in _ALL_KINDS if kind in kinds}) + + +def test_c01_upstream_pause_mid_burst_keeps_liveliness_and_lands_every_id_once( + gateway: Gateway, + tmp_path: Path, + record_property: pytest.RecordProperty, +) -> None: + with owned_upstream(tmp_path) as slot, gateway.scenario() as scenario: + upstream: Final = slot.process + assert upstream is not None + deployments: Final = _deployments(scenario, slot.url) + eventually(partial(_every_worker_serves_every_kind, gateway, deployments), lambda served: served, seconds=90) + clients: Final = _clients(gateway, deployments) + calls: Final = _burst_calls(_ALL_KINDS, 5) + observations: Final = _Observations(slot.url) + with ThreadPoolExecutor(max_workers=len(calls)) as pool: + futures: Final = _submit(pool, clients, calls) + eventually(partial(_done_count, futures), lambda done: done >= 1, seconds=60) + upstream.send_signal(signal.SIGSTOP) + try: + done_at_pause: Final = _done_count(futures) + paused_liveliness: Final = tuple(_liveliness(gateway) for _ in range(3)) + finally: + upstream.send_signal(signal.SIGCONT) + done_at_resume: Final = _done_count(futures) + outcomes: Final = tuple(future.result(timeout=180) for future in futures) + record_property("c01_burst_size", len(calls)) + record_property("c01_done_at_pause", done_at_pause) + record_property("c01_done_at_resume", done_at_resume) + record_property("c01_paused_liveliness_statuses", str(paused_liveliness)) + assert paused_liveliness == (200, 200, 200), paused_liveliness + assert done_at_resume < len(calls), done_at_resume + failures: Final = tuple(outcome for outcome in outcomes if isinstance(outcome, _Failed)) + assert not failures, failures + completed: Final = tuple(outcome for outcome in outcomes if isinstance(outcome, _Completed)) + record_property("c01_completed_by_kind", _kind_counts(completed)) + response_ids: Final = tuple(outcome.response_id for outcome in completed) + assert len(set(response_ids)) == len(calls), response_ids + spend_rows: Final = _landed_once(response_ids) + record_property("c01_spend_query", _SPEND_QUERY) + record_property("c01_spend_row_counts", str(tuple(len(rows) for rows in spend_rows))) + markers: Final = tuple(call.marker for call in calls) + eventually( + observations.read, + lambda items: _marker_counts(items, markers) == (1,) * len(markers), + seconds=60, + ) + record_property("c01_upstream_marker_counts", str(_marker_counts(observations.items, markers))) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_c02_worker_sigkill_mid_burst_keeps_serving_and_respawns( + gateway: Gateway, + tmp_path: Path, + record_property: pytest.RecordProperty, +) -> None: + with owned_upstream(tmp_path) as slot, gateway.scenario() as scenario: + deployments: Final = _deployments(scenario, slot.url) + with owned_proxy_process(gateway, tmp_path, {"INTEGRATION_UPSTREAM_URL": slot.url}, workers=2) as owned: + eventually( + partial(_every_worker_serves_every_kind, owned.gateway, deployments), lambda served: served, seconds=90 + ) + clients: Final = _clients(owned.gateway, deployments) + workers: Final = eventually( + partial(_alive_workers, owned.process), lambda found: len(found) == 2, seconds=60 + ) + record_property("c02_process_tree_before_kill", _process_tree(owned.process)) + victim: Final = workers[0] + calls: Final = _burst_calls(_CHAT_KINDS, 15) + observations: Final = _Observations(slot.url) + with ThreadPoolExecutor(max_workers=len(calls)) as pool: + futures: Final = _submit(pool, clients, calls) + eventually(partial(_done_count, futures), lambda done: done >= 1, seconds=60) + victim.kill() + eventually(partial(_worker_alive, victim), lambda alive: not alive, seconds=10) + done_at_kill: Final = _done_count(futures) + survivor: Final = eventually( + lambda: _attempt(clients, _Call("chat", f"after-kill-{uuid.uuid4().hex}")), + lambda outcome: isinstance(outcome, _Completed), + seconds=30, + ) + outcomes: Final = tuple(future.result(timeout=180) for future in futures) + record_property("c02_burst_size", len(calls)) + record_property("c02_killed_worker_pid", victim.pid) + record_property("c02_done_at_kill", done_at_kill) + assert isinstance(survivor, _Completed), survivor + completed: Final = tuple(outcome for outcome in outcomes if isinstance(outcome, _Completed)) + failed: Final = tuple(outcome for outcome in outcomes if isinstance(outcome, _Failed)) + record_property("c02_completed_by_kind", _kind_counts(completed)) + record_property("c02_inflight_failure_count", len(failed)) + record_property("c02_inflight_failure_errors", str(tuple(outcome.error for outcome in failed))) + assert completed, outcomes + respawned: Final = eventually( + partial(_alive_workers, owned.process), + lambda found: len(found) == 2 and victim.pid not in {worker.pid for worker in found}, + seconds=graceful_stop_seconds() + 60, + ) + record_property("c02_process_tree_after_respawn", _process_tree(owned.process)) + record_property("c02_respawned_worker_pids", str(tuple(worker.pid for worker in respawned))) + after_respawn: Final = eventually( + lambda: _attempt(clients, _Call("chat", f"after-respawn-{uuid.uuid4().hex}")), + lambda outcome: isinstance(outcome, _Completed), + seconds=60, + ) + assert isinstance(after_respawn, _Completed), after_respawn + _landed_once((survivor.response_id, after_respawn.response_id)) + response_ids: Final = tuple(outcome.response_id for outcome in completed) + assert len(set(response_ids)) == len(completed), response_ids + spend_row_counts: Final = tuple((outcome, len(_spend_rows(outcome.response_id))) for outcome in completed) + landed: Final = tuple(outcome for outcome, count in spend_row_counts if count == 1) + lost: Final = tuple(outcome for outcome, count in spend_row_counts if count == 0) + assert len(landed) + len(lost) == len(completed), spend_row_counts + record_property("c02_spend_query", _SPEND_QUERY) + record_property("c02_burst_ids_landed_once", len(landed)) + record_property("c02_burst_ids_lost", len(lost)) + record_property("c02_burst_ids_lost_by_kind", _kind_counts(lost)) + completed_markers: Final = tuple(outcome.call.marker for outcome in completed) + eventually( + observations.read, + lambda items: _marker_counts(items, completed_markers) == (1,) * len(completed_markers), + seconds=60, + ) + failed_markers: Final = tuple(outcome.call.marker for outcome in failed) + record_property( + "c02_failed_calls_seen_upstream", + sum(count == 1 for count in _marker_counts(observations.items, failed_markers)), + ) diff --git a/tests/integration/sdk/test_openai_sdk_client_sessions_wire.py b/tests/integration/sdk/test_openai_sdk_client_sessions_wire.py new file mode 100644 index 00000000000..e8ec1a37082 --- /dev/null +++ b/tests/integration/sdk/test_openai_sdk_client_sessions_wire.py @@ -0,0 +1,488 @@ +import asyncio +import base64 +import json +import struct +import threading +import zlib +from collections.abc import Awaitable, Callable, Iterator, Mapping +from queue import SimpleQueue +from typing import Final, TypeVar +from urllib.parse import urlsplit + +import httpx +import openai +import pytest +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +import litellm +from litellm import Router +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.utils import ImageResponse, ModelResponse, TextCompletionResponse + +_R: Final = TypeVar("_R") + +_PROVIDER_KEY: Final = "sk-scripted-provider" +_API_VERSION: Final = "2024-10-21" +_GATEWAY_PATH: Final = "/gateway.ai.cloudflare.com/v1/scripted-account/scripted-gateway/azure-openai/scripted-resource" +_DEPLOYMENT: Final = "gpt-4o-mini-gateway" +_ANSWER: Final = "wire answer" +_IMAGE_URL: Final = "https://images.example.invalid/variation.png" + + +def _png_chunk(kind: bytes, data: bytes) -> bytes: + return struct.pack(">I", len(data)) + kind + data + struct.pack(">I", zlib.crc32(kind + data)) + + +_PNG: Final = ( + b"\x89PNG\r\n\x1a\n" + + _png_chunk(b"IHDR", struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0)) + + _png_chunk(b"IDAT", zlib.compress(b"\x00\x00\x00\x00")) + + _png_chunk(b"IEND", b"") +) + + +def _json(body: Mapping[str, JsonValue]) -> Reply: + return Reply(body=json.dumps(body).encode()) + + +def _chat_completion() -> Mapping[str, JsonValue]: + return { + "id": "chatcmpl-wire", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": _ANSWER}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + } + + +def _text_completion() -> Mapping[str, JsonValue]: + return { + "id": "cmpl-wire", + "object": "text_completion", + "created": 1, + "model": "gpt-3.5-turbo-instruct", + "choices": [{"index": 0, "text": _ANSWER, "finish_reason": "stop", "logprobs": None}], + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + } + + +def _peer(expected_path: str, body: Mapping[str, JsonValue]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert urlsplit(request.target).path == expected_path, request.target + return _json(body) + + return respond + + +def _held_peer(gate: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert gate.wait(timeout=10), "the timeout cell never released its peer" + return _json(_chat_completion()) + + return respond + + +def _only_request(wire: Wire) -> Request: + received: Final = wire.drain() + assert len(received) == 1, received + return received[0] + + +def _drain(seen: SimpleQueue[str]) -> tuple[str, ...]: + return tuple(seen.get_nowait() for _ in range(seen.qsize())) + + +def _content(response: ModelResponse) -> str | None: + choice: Final = response.choices[0] + return choice.message.content if isinstance(choice, litellm.Choices) else None + + +@pytest.fixture +def sync_session(monkeypatch: pytest.MonkeyPatch) -> Iterator[tuple[httpx.Client, SimpleQueue[str]]]: + seen: Final[SimpleQueue[str]] = SimpleQueue() + + def record(request: httpx.Request) -> None: + seen.put(str(request.url)) + + with httpx.Client(event_hooks={"request": [record]}) as client: + monkeypatch.setattr(litellm, "client_session", client) + monkeypatch.setattr(litellm, "aclient_session", None) + yield client, seen + + +def _run_with_async_session( + monkeypatch: pytest.MonkeyPatch, call: Callable[[], Awaitable[_R]] +) -> tuple[_R, tuple[str, ...]]: + seen: Final[SimpleQueue[str]] = SimpleQueue() + + async def record(request: httpx.Request) -> None: + seen.put(str(request.url)) + + async def run() -> _R: + async with httpx.AsyncClient(event_hooks={"request": [record]}) as client: + monkeypatch.setattr(litellm, "client_session", None) + monkeypatch.setattr(litellm, "aclient_session", client) + return await call() + + result: Final = asyncio.run(run()) + return result, _drain(seen) + + +def _deployment(timeout: httpx.Timeout | openai.Timeout, api_base: str, deployment_id: str) -> Router: + return Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": api_base, + "api_key": _PROVIDER_KEY, + "timeout": timeout, + "max_retries": 0, + }, + "model_info": {"id": deployment_id}, + } + ], + num_retries=0, + ) + + +def test_f01_async_image_variation_with_only_a_sync_session_set_builds_its_own_async_client( + sync_session: tuple[httpx.Client, SimpleQueue[str]], +) -> None: + with wire_server(_peer("/images/variations", {"created": 1, "data": [{"url": _IMAGE_URL}]})) as wire: + response: Final = asyncio.run( + litellm.aimage_variation( + image=("probe.png", _PNG, "image/png"), + model="dall-e-2", + custom_llm_provider="openai", + api_base=wire.url, + api_key=_PROVIDER_KEY, + num_retries=0, + ) + ) + assert isinstance(response, ImageResponse), type(response) + assert response.data is not None and response.data[0].url == _IMAGE_URL, response + request: Final = _only_request(wire) + assert request.method == "POST", request.method + assert _PNG in request.body + assert request.headers.get("authorization") == f"Bearer {_PROVIDER_KEY}", request.headers + assert _drain(sync_session[1]) == () + + +def test_f02_async_azure_cloudflare_gateway_call_with_only_a_sync_session_set_builds_its_own_async_client( + sync_session: tuple[httpx.Client, SimpleQueue[str]], +) -> None: + with wire_server(_peer(f"{_GATEWAY_PATH}/{_DEPLOYMENT}/chat/completions", _chat_completion())) as wire: + response: Final = asyncio.run( + litellm.acompletion( + model=f"azure/{_DEPLOYMENT}", + messages=[{"role": "user", "content": "via the gateway"}], + api_base=f"{wire.url}{_GATEWAY_PATH}", + api_key=_PROVIDER_KEY, + api_version=_API_VERSION, + num_retries=0, + ) + ) + assert isinstance(response, ModelResponse), type(response) + assert _content(response) == _ANSWER, response + request: Final = _only_request(wire) + assert urlsplit(request.target).query == f"api-version={_API_VERSION}", request.target + assert request.headers.get("api-key") == _PROVIDER_KEY, request.headers + assert _drain(sync_session[1]) == () + + +def test_f03_async_text_completion_uses_the_callers_async_session(monkeypatch: pytest.MonkeyPatch) -> None: + with wire_server(_peer("/completions", _text_completion())) as wire: + response, seen = _run_with_async_session( + monkeypatch, + lambda: litellm.atext_completion( + model="gpt-3.5-turbo-instruct", + prompt="complete this", + custom_llm_provider="openai", + api_base=wire.url, + api_key=_PROVIDER_KEY, + num_retries=0, + ), + ) + assert isinstance(response, TextCompletionResponse), type(response) + assert response.choices[0].text == _ANSWER, response + assert _only_request(wire).method == "POST" + assert seen == (f"{wire.url}/completions",) + + +def test_f04_sync_text_completion_uses_the_callers_sync_session( + sync_session: tuple[httpx.Client, SimpleQueue[str]], +) -> None: + with wire_server(_peer("/completions", _text_completion())) as wire: + response: Final = litellm.text_completion( + model="gpt-3.5-turbo-instruct", + prompt="complete this", + custom_llm_provider="openai", + api_base=wire.url, + api_key=_PROVIDER_KEY, + num_retries=0, + ) + assert isinstance(response, TextCompletionResponse), type(response) + assert response.choices[0].text == _ANSWER, response + assert _only_request(wire).method == "POST" + assert _drain(sync_session[1]) == (f"{wire.url}/completions",) + + +def test_f05_sync_moderation_reaches_the_peer_and_returns_its_verdict() -> None: + verdict: Final[Mapping[str, JsonValue]] = { + "id": "modr-wire", + "model": "omni-moderation-latest", + "results": [{"flagged": True, "categories": {"harassment": True}, "category_scores": {"harassment": 0.91}}], + } + with wire_server(_peer("/moderations", verdict)) as wire: + response: Final = litellm.moderation( + input="moderate this", + model="omni-moderation-latest", + api_key=_PROVIDER_KEY, + api_base=wire.url, + ) + assert response.results[0].flagged is True, response + request: Final = _only_request(wire) + assert json.loads(request.body) == {"input": "moderate this", "model": "omni-moderation-latest"}, request.body + + +@pytest.mark.parametrize("timeout_type", (httpx.Timeout, openai.Timeout), ids=("httpx", "openai")) +def test_f06_sync_completion_accepts_a_structured_timeout( + timeout_type: type[httpx.Timeout] | type[openai.Timeout], +) -> None: + with wire_server(_peer("/chat/completions", _chat_completion())) as wire: + response: Final = litellm.completion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "with a structured timeout"}], + custom_llm_provider="openai", + api_base=wire.url, + api_key=_PROVIDER_KEY, + timeout=timeout_type(5.0, connect=2.0), + num_retries=0, + ) + assert isinstance(response, ModelResponse), type(response) + assert _content(response) == _ANSWER, response + assert _only_request(wire).method == "POST" + + +def test_f07_async_azure_cloudflare_gateway_call_uses_the_callers_async_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with wire_server(_peer(f"{_GATEWAY_PATH}/{_DEPLOYMENT}/chat/completions", _chat_completion())) as wire: + response, seen = _run_with_async_session( + monkeypatch, + lambda: litellm.acompletion( + model=f"azure/{_DEPLOYMENT}", + messages=[{"role": "user", "content": "via the gateway"}], + api_base=f"{wire.url}{_GATEWAY_PATH}", + api_key=_PROVIDER_KEY, + api_version=_API_VERSION, + num_retries=0, + ), + ) + assert isinstance(response, ModelResponse), type(response) + assert _content(response) == _ANSWER, response + assert _only_request(wire).headers.get("api-key") == _PROVIDER_KEY + assert seen == (f"{wire.url}{_GATEWAY_PATH}/{_DEPLOYMENT}/chat/completions?api-version={_API_VERSION}",) + + +def test_f08_sync_image_variation_uses_the_callers_sync_session( + sync_session: tuple[httpx.Client, SimpleQueue[str]], +) -> None: + with wire_server(_peer("/images/variations", {"created": 1, "data": [{"url": _IMAGE_URL}]})) as wire: + response: Final = litellm.image_variation( + image=("probe.png", _PNG, "image/png"), + model="dall-e-2", + custom_llm_provider="openai", + api_base=wire.url, + api_key=_PROVIDER_KEY, + num_retries=0, + ) + assert isinstance(response, ImageResponse), type(response) + assert response.data is not None and response.data[0].url == _IMAGE_URL, response + assert _PNG in _only_request(wire).body + assert _drain(sync_session[1]) == (f"{wire.url}/images/variations",) + + +@pytest.mark.parametrize("timeout_type", (httpx.Timeout, openai.Timeout), ids=("httpx", "openai")) +def test_f09_router_deployment_keeps_a_structured_timeout_as_httpx_and_serves( + timeout_type: type[httpx.Timeout] | type[openai.Timeout], +) -> None: + with wire_server(_peer("/chat/completions", _chat_completion())) as wire: + router: Final = _deployment(timeout_type(5.0, connect=2.0), wire.url, "deployment-f09") + response: Final = router.completion( + model="gpt-4o-mini", messages=[{"role": "user", "content": "through the router"}] + ) + assert isinstance(response, ModelResponse), type(response) + assert _content(response) == _ANSWER, response + assert _only_request(wire).method == "POST" + deployment: Final = router.get_deployment("deployment-f09") + assert deployment is not None + stored: Final = deployment.litellm_params.timeout + assert isinstance(stored, httpx.Timeout), type(stored) + assert (stored.read, stored.connect) == (5.0, 2.0), stored + + +def test_f10_router_deployment_structured_read_timeout_fires_once_as_a_408() -> None: + gate: Final = threading.Event() + with wire_server(_held_peer(gate)) as wire: + router: Final = _deployment(httpx.Timeout(0.5, connect=2.0), wire.url, "deployment-f10") + with pytest.raises(litellm.Timeout) as raised: + router.completion(model="gpt-4o-mini", messages=[{"role": "user", "content": "held upstream"}]) + gate.set() + assert raised.value.status_code == 408, raised.value + assert len(wire.drain()) == 1 + + +_SPEECH_AUDIO: Final = b"OggS" + bytes(range(60)) +_SPEECH_AUDIO_REPLY: Final = Reply(body=_SPEECH_AUDIO, content_type="audio/ogg") +_ELEVENLABS_VOICE: Final = "21m00Tcm4TlvDq8ikWAM" +_VERTEX_PROJECT: Final = "scripted-project" +_SPEECH_ROUTES: Final[Mapping[str, tuple[Mapping[str, str], str, Reply]]] = { + "openai": ( + {"model": "openai/gpt-4o-mini-tts", "voice": "alloy", "api_key": _PROVIDER_KEY}, + "/audio/speech", + _SPEECH_AUDIO_REPLY, + ), + "azure": ( + {"model": "azure/tts-deployment", "voice": "alloy", "api_key": _PROVIDER_KEY, "api_version": _API_VERSION}, + "/openai/deployments/tts-deployment/audio/speech", + _SPEECH_AUDIO_REPLY, + ), + "azure_ava": ( + {"model": "azure/speech/azure-tts", "voice": "alloy", "api_key": _PROVIDER_KEY}, + "/cognitiveservices/v1", + _SPEECH_AUDIO_REPLY, + ), + "elevenlabs": ( + {"model": "elevenlabs/eleven_multilingual_v2", "voice": _ELEVENLABS_VOICE, "api_key": _PROVIDER_KEY}, + f"/v1/text-to-speech/{_ELEVENLABS_VOICE}", + _SPEECH_AUDIO_REPLY, + ), + "edenai": ( + {"model": "edenai/openai", "voice": "alloy", "api_key": _PROVIDER_KEY}, + "/audio/speech", + _SPEECH_AUDIO_REPLY, + ), + "minimax": ( + {"model": "minimax/speech-02-hd", "voice": "alloy", "api_key": _PROVIDER_KEY}, + "/v1/t2a_v2", + _json({"data": {"audio": _SPEECH_AUDIO.hex()}, "base_resp": {"status_code": 0, "status_msg": "success"}}), + ), + "mistral": ( + {"model": "mistral/voxtral-mini-tts-2603", "voice": "alloy", "api_key": _PROVIDER_KEY}, + "/v1/audio/speech", + _json({"audio_data": base64.b64encode(_SPEECH_AUDIO).decode()}), + ), + "aws_polly": ( + { + "model": "aws_polly/neural", + "voice": "Joanna", + "aws_access_key_id": "AKIASCRIPTED", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + }, + "/v1/speech", + _SPEECH_AUDIO_REPLY, + ), +} +_PAYMENT_REQUIRED_MODELS: Final[Mapping[str, type[litellm.BadRequestError]]] = { + "openai/gpt-4o-mini": litellm.BadRequestError, + "anthropic/claude-sonnet-4-5": litellm.PaymentRequiredError, +} + + +def _speech_peer(expected_path: str, reply: Reply) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert urlsplit(request.target).path == expected_path, request.target + return reply + + return respond + + +def _vertex_speech_peer(request: Request) -> Reply: + if urlsplit(request.target).path == "/_oauth/token": + return _json({"access_token": "scripted-vertex-token", "expires_in": 3600, "token_type": "Bearer"}) + return _json({"audioContent": base64.b64encode(_SPEECH_AUDIO).decode()}) + + +def _payment_required_peer(request: Request) -> Reply: + return Reply( + status=402, + body=json.dumps({"error": {"message": "scripted 402", "type": "invalid_request_error"}}).encode(), + ) + + +@pytest.mark.parametrize("provider", tuple(_SPEECH_ROUTES)) +def test_f11_sync_speech_accepts_an_sdk_timeout_on_every_provider_branch(provider: str) -> None: + route, expected_path, reply = _SPEECH_ROUTES[provider] + with wire_server(_speech_peer(expected_path, reply)) as wire: + response: Final = litellm.speech( + input="speak this", + api_base=wire.url, + timeout=openai.Timeout(5.0, connect=2.0), + **route, + ) + assert response.content == _SPEECH_AUDIO, response.content[:16] + assert _only_request(wire).method == "POST" + + +def test_f13_sync_vertex_speech_accepts_an_sdk_timeout() -> None: + with wire_server(_vertex_speech_peer) as wire: + response: Final = litellm.speech( + model="vertex_ai/chirp", + voice="alloy", + input="speak this", + api_base=wire.url, + vertex_credentials=service_account_json(_VERTEX_PROJECT, wire.url), + vertex_project=_VERTEX_PROJECT, + vertex_location="us-central1", + timeout=openai.Timeout(5.0, connect=2.0), + ) + assert response.content == _SPEECH_AUDIO, response.content[:16] + assert tuple((seen.method, urlsplit(seen.target).path) for seen in wire.drain()) == ( + ("POST", "/_oauth/token"), + ("POST", "/"), + ) + + +def _runwayml_rejecting_peer(request: Request) -> Reply: + assert urlsplit(request.target).path == "/v1/text_to_speech", request.target + return Reply(status=400, body=json.dumps({"error": "scripted 400"}).encode()) + + +def test_f14_sync_runwayml_speech_sends_its_task_with_an_sdk_timeout() -> None: + with wire_server(_runwayml_rejecting_peer) as wire: + with pytest.raises(BaseLLMException) as raised: + litellm.speech( + model="runwayml/eleven_multilingual_v2", + voice="Maya", + input="speak this", + api_base=wire.url, + api_key=_PROVIDER_KEY, + timeout=openai.Timeout(5.0, connect=2.0), + ) + assert raised.value.status_code == 400, raised.value + assert "scripted 400" in raised.value.message, raised.value.message + assert _only_request(wire).method == "POST" + + +@pytest.mark.parametrize("model", tuple(_PAYMENT_REQUIRED_MODELS)) +def test_f12_sync_completion_maps_an_upstream_402_with_its_message(model: str) -> None: + with wire_server(_payment_required_peer) as wire: + with pytest.raises(_PAYMENT_REQUIRED_MODELS[model]) as raised: + litellm.completion( + model=model, + messages=[{"role": "user", "content": "out of credit"}], + api_base=wire.url, + api_key=_PROVIDER_KEY, + num_retries=0, + max_retries=0, + ) + assert raised.value.status_code == 402, raised.value + assert "scripted 402" in raised.value.message, raised.value.message + assert len(wire.drain()) == 1 diff --git a/tests/typing/openai_sdk_compat.py b/tests/typing/openai_sdk_compat.py new file mode 100644 index 00000000000..03f44db6a9f --- /dev/null +++ b/tests/typing/openai_sdk_compat.py @@ -0,0 +1,66 @@ +from typing import Final + +import httpx +from openai import DefaultAsyncHttpxClient, DefaultHttpxClient +from openai import HttpxBinaryResponseContent as SDKBinaryResponse +from openai import Timeout as SDKTimeout +from openai._types import Response as SDKResponse +from typing_extensions import assert_type + +import litellm +from litellm.exceptions import ( + APIConnectionError, + APIResponseValidationError, + BadGatewayError, + InternalServerError, + InvalidRequestError, + RateLimitError, + Timeout, +) +from litellm.litellm_core_utils.completion_timeout import CompletionTimeout +from litellm.types.llms.openai import HttpxBinaryResponseContent +from litellm.types.router import GenericLiteLLMParams, LiteLLMParamsTypedDict + + +def consume_binary_response(response: HttpxBinaryResponseContent) -> httpx.Response: + assert_type(response.response, httpx.Response) + assert_type(response.response.request, httpx.Request) + return response.response + + +def binary_response_types(legacy_response: httpx.Response, sdk_response: SDKResponse) -> None: + legacy: Final = HttpxBinaryResponseContent(legacy_response) + native: Final = HttpxBinaryResponseContent(sdk_response) + + assert_type(legacy.response, httpx.Response) + assert_type(native.response, SDKResponse) + assert_type(legacy.read(), bytes) + assert_type(native.read(), bytes) + assert_type(consume_binary_response(legacy), httpx.Response) + + sdk_wrapper: Final[SDKBinaryResponse] = legacy + assert_type(sdk_wrapper.read(), bytes) + + +def synthesized_response_types( + error: RateLimitError | BadGatewayError | InternalServerError | APIResponseValidationError | InvalidRequestError, +) -> None: + assert_type(error.response, httpx.Response) + assert_type(error.request, httpx.Request) + + +def synthesized_request_types(error: APIConnectionError | Timeout) -> None: + assert_type(error.request, httpx.Request) + + +def configured_clients( + sync_client: httpx.Client | DefaultHttpxClient, + async_client: httpx.AsyncClient | DefaultAsyncHttpxClient, + timeout: httpx.Timeout | SDKTimeout, +) -> None: + litellm.client_session = sync_client # test-quality-ok: [TQ005] Type-check-only fixture; never executed + litellm.aclient_session = async_client # test-quality-ok: [TQ005] Type-check-only fixture; never executed + assert_type(CompletionTimeout.normalize(timeout), httpx.Timeout) + params: Final[LiteLLMParamsTypedDict] = {"timeout": 1.0} + params["timeout"] = timeout + GenericLiteLLMParams(timeout=timeout) diff --git a/tests/typing/pyrightconfig.json b/tests/typing/pyrightconfig.json new file mode 100644 index 00000000000..0910ca64f5b --- /dev/null +++ b/tests/typing/pyrightconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../pyrightconfig.json", + "include": ["openai_sdk_compat.py"], + "exclude": [], + "extraPaths": ["../.."], + "pythonVersion": "3.10", + "reportMissingImports": true +} diff --git a/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py index 57f4557c6f7..ea0416cecbd 100644 --- a/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py @@ -1204,8 +1204,10 @@ def test_speech_response_without_a_byte_count_produces_no_output() -> None: assert data.choices_out == () -def test_speech_binary_response_is_logged_as_its_summary_not_dropped() -> None: - import httpx +@pytest.mark.parametrize("http_module", ("httpx", "httpx2")) +def test_speech_binary_response_is_logged_as_its_summary_not_dropped(http_module: str) -> None: + httpx: Final = pytest.importorskip(http_module) + from openai import HttpxBinaryResponseContent as SDKBinaryResponse from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.litellm_logging import _extract_response_obj_and_hidden_params @@ -1216,13 +1218,17 @@ def test_speech_binary_response_is_logged_as_its_summary_not_dropped() -> None: set_provider_response_headers_in_hidden_params(speech, raw.headers) response_obj, hidden_params = _extract_response_obj_and_hidden_params(speech, None) + assert isinstance(speech, SDKBinaryResponse) + assert speech.response is raw + assert speech.read() == raw.content assert response_obj == {"object": "binary", "content_type": "audio/mpeg", "num_bytes": 1234} assert hidden_params is not None assert hidden_params["headers"]["content-type"] == "audio/mpeg" -def test_speech_binary_response_still_streaming_reports_the_bytes_downloaded_so_far() -> None: - import httpx +@pytest.mark.parametrize("http_module", ("httpx", "httpx2")) +def test_speech_binary_response_still_streaming_reports_the_bytes_downloaded_so_far(http_module: str) -> None: + httpx: Final = pytest.importorskip(http_module) from litellm.types.llms.openai import HttpxBinaryResponseContent diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index 3774545488d..053ec7e141e 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -1,11 +1,16 @@ -import asyncio, importlib, os +import asyncio +import importlib +import os import json -from typing import Final -from unittest.mock import AsyncMock, Mock, patch +from collections.abc import Mapping +from types import ModuleType +from typing import Final, Protocol +from unittest.mock import Mock, patch import httpx import pytest from openai import AsyncOpenAI, OpenAI +from openai._types import Response as SDKResponse import litellm from litellm.llms.openai.openai import ( @@ -77,19 +82,48 @@ def test_get_stream_options_passes_caller_stream_options_through_on_any_host(api } +@pytest.fixture(params=("httpx", "sdk")) +def http_backend(request: pytest.FixtureRequest) -> ModuleType: + return importlib.import_module("httpx" if request.param == "httpx" else SDKResponse.__module__) + + +@pytest.fixture(params=("openai", "perplexity", "cerebras", "nvidia_nim")) +def sdk_provider(request: pytest.FixtureRequest) -> str: + return str(request.param) + + +class _Request(Protocol): + @property + def content(self) -> bytes: ... + + @property + def extensions(self) -> Mapping[str, object]: ... + + @property + def url(self) -> object: ... + + @property + def headers(self) -> Mapping[str, str]: ... + + @pytest.mark.asyncio -async def test_acompletion_returns_json_reply_over_injected_transport(): +async def test_acompletion_returns_json_reply_over_injected_transport( + http_backend: ModuleType, sdk_provider: str, monkeypatch: pytest.MonkeyPatch +): outbound: Final = asyncio.Queue() - def respond(request: httpx.Request) -> httpx.Response: + def respond(request: _Request) -> httpx.Response | SDKResponse: + assert str(request.url) == f"https://{sdk_provider}.example/v1/chat/completions" + assert request.headers["authorization"] == f"Bearer {sdk_provider}-key" + assert request.extensions["timeout"] == {"connect": 7, "read": 7, "write": 7, "pool": 7} outbound.put_nowait(json.loads(request.content)) - return httpx.Response( + return http_backend.Response( 200, json={ "id": "chatcmpl-smoke", "object": "chat.completion", "created": 0, - "model": "gpt-5.6", + "model": "sdk-compat", "choices": [ { "index": 0, @@ -101,31 +135,71 @@ async def test_acompletion_returns_json_reply_over_injected_transport(): }, ) - async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: - client: Final = AsyncOpenAI(api_key="transport-only", http_client=http_client) + async with http_backend.AsyncClient(transport=http_backend.MockTransport(respond)) as http_client: + monkeypatch.setattr(litellm, "aclient_session", http_client) response: Final = await asyncio.wait_for( litellm.acompletion( - model="openai/gpt-5.6", - api_key="transport-only", - client=client, + model=f"{sdk_provider}/sdk-compat", + api_key=f"{sdk_provider}-key", + api_base=f"https://{sdk_provider}.example/v1", messages=[{"role": "user", "content": "smoke-json-request"}], + timeout=7, num_retries=0, max_retries=0, ), timeout=10, ) request: Final = await asyncio.wait_for(outbound.get(), timeout=10) - assert request["model"] == "gpt-5.6" + assert request["model"] == "sdk-compat" assert request["messages"] == [{"role": "user", "content": "smoke-json-request"}] assert not request.get("stream") assert outbound.empty() assert response.choices[0].message.content == "smoke-json-reply" assert response.choices[0].finish_reason == "stop" assert response.usage.total_tokens == 15 + assert not http_client.is_closed @pytest.mark.asyncio -async def test_acompletion_streams_text_deltas_over_injected_transport(): +@pytest.mark.parametrize( + ("status", "expected_error"), + ( + (400, litellm.BadRequestError), + (429, litellm.RateLimitError), + (503, litellm.ServiceUnavailableError), + (None, litellm.Timeout), + ), +) +async def test_acompletion_preserves_public_errors_for_both_http_clients( + http_backend: ModuleType, + sdk_provider: str, + status: int | None, + expected_error: type[Exception], + monkeypatch: pytest.MonkeyPatch, +) -> None: + def respond(request: _Request) -> httpx.Response | SDKResponse: + if status is None: + raise http_backend.ReadTimeout("upstream timeout", request=request) + return http_backend.Response(status, json={"error": {"message": "upstream failure", "type": "api_error"}}) + + async with http_backend.AsyncClient(transport=http_backend.MockTransport(respond)) as http_client: + monkeypatch.setattr(litellm, "aclient_session", http_client) + with pytest.raises(expected_error, match="timed out" if status is None else "upstream failure"): + await litellm.acompletion( + model=f"{sdk_provider}/sdk-compat", + api_key=f"{sdk_provider}-key", + api_base=f"https://{sdk_provider}.example/v1", + messages=[{"role": "user", "content": "request"}], + num_retries=0, + max_retries=0, + ) + assert not http_client.is_closed + + +@pytest.mark.asyncio +async def test_acompletion_streams_text_deltas_over_injected_transport( + http_backend: ModuleType, sdk_provider: str, monkeypatch: pytest.MonkeyPatch +): outbound: Final = asyncio.Queue() def chunk(delta: dict, finish: str | None) -> bytes: @@ -133,18 +207,18 @@ async def test_acompletion_streams_text_deltas_over_injected_transport(): "id": "chatcmpl-smoke", "object": "chat.completion.chunk", "created": 0, - "model": "gpt-5.6", + "model": "sdk-compat", "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], } return f"data: {json.dumps(body)}\n\n".encode() - def respond(request: httpx.Request) -> httpx.Response: + def respond(request: _Request) -> httpx.Response | SDKResponse: outbound.put_nowait(json.loads(request.content)) usage: Final = { "id": "chatcmpl-smoke", "object": "chat.completion.chunk", "created": 0, - "model": "gpt-5.6", + "model": "sdk-compat", "choices": [], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, } @@ -156,14 +230,14 @@ async def test_acompletion_streams_text_deltas_over_injected_transport(): b"data: [DONE]\n\n", ) ) - return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=content) + return http_backend.Response(200, headers={"content-type": "text/event-stream"}, content=content) - async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: - client: Final = AsyncOpenAI(api_key="transport-only", http_client=http_client) + async with http_backend.AsyncClient(transport=http_backend.MockTransport(respond)) as http_client: + monkeypatch.setattr(litellm, "aclient_session", http_client) stream: Final = await litellm.acompletion( - model="openai/gpt-5.6", - api_key="transport-only", - client=client, + model=f"{sdk_provider}/sdk-compat", + api_key=f"{sdk_provider}-key", + api_base=f"https://{sdk_provider}.example/v1", messages=[{"role": "user", "content": "smoke-stream-request"}], stream=True, num_retries=0, @@ -190,7 +264,7 @@ async def test_acompletion_streams_text_deltas_over_injected_transport(): @pytest.mark.asyncio -async def test_acompletion_streams_tool_call_arguments_over_injected_transport(): +async def test_acompletion_streams_tool_call_arguments_over_injected_transport(http_backend: ModuleType): outbound: Final = asyncio.Queue() tools: Final = [ { @@ -217,7 +291,8 @@ async def test_acompletion_streams_tool_call_arguments_over_injected_transport() } return f"data: {json.dumps(body)}\n\n".encode() - def respond(request: httpx.Request) -> httpx.Response: + def respond(request: _Request) -> httpx.Response | SDKResponse: + assert request.extensions["timeout"] == {"connect": 1, "read": None, "write": 2, "pool": 3} outbound.put_nowait(json.loads(request.content)) content: Final = b"".join( ( @@ -239,9 +314,9 @@ async def test_acompletion_streams_tool_call_arguments_over_injected_transport() b"data: [DONE]\n\n", ) ) - return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=content) + return http_backend.Response(200, headers={"content-type": "text/event-stream"}, content=content) - async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + async with http_backend.AsyncClient(transport=http_backend.MockTransport(respond)) as http_client: client: Final = AsyncOpenAI(api_key="transport-only", http_client=http_client) messages: Final = [{"role": "user", "content": "weather in Paris"}] stream: Final = await litellm.acompletion( @@ -251,6 +326,7 @@ async def test_acompletion_streams_tool_call_arguments_over_injected_transport() messages=messages, tools=tools, stream=True, + timeout=http_backend.Timeout(connect=1, read=None, write=2, pool=3), num_retries=0, max_retries=0, ) diff --git a/tests/unit/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py index fb605c6da7a..3c7dff34296 100644 --- a/tests/unit/llms/openai/test_openai_common_utils.py +++ b/tests/unit/llms/openai/test_openai_common_utils.py @@ -1,9 +1,10 @@ -from unittest.mock import MagicMock, call, patch +from typing import Final +from unittest.mock import MagicMock, patch import httpx import openai import pytest - +import respx import litellm from litellm.litellm_core_utils.token_counter import token_counter @@ -413,3 +414,124 @@ def test_is_openai_backed_api_base_decides_by_hostname_only(api_base, expected): assert is_openai_backed_api_base(api_base) is expected +def _sdk_api_client( + api: str, + is_async: bool, + timeout: float | httpx.Timeout | openai.Timeout | None, + client: openai.OpenAI | openai.AsyncOpenAI | None = None, +) -> openai.OpenAI | openai.AsyncOpenAI | None: + from litellm.llms.azure.common_utils import BaseAzureLLM + from litellm.llms.openai.fine_tuning.handler import OpenAIFineTuningAPI + from litellm.llms.openai.image_variations.handler import OpenAIImageVariationsHandler + from litellm.llms.openai.openai import OpenAIAssistantsAPI, OpenAIBatchesAPI, OpenAIFilesAPI + + kwargs: Final = { + "api_key": "transport-only", + "api_base": "https://sdk-default.example/v1", + "timeout": timeout, + "max_retries": 0, + "organization": None, + "client": client, + } + if api == "assistants": + assistant_factory: Final = OpenAIAssistantsAPI() + return ( + assistant_factory.async_get_openai_client(**kwargs) + if is_async + else assistant_factory.get_openai_client(**kwargs) + ) + if api == "image_variations": + variation_factory: Final = OpenAIImageVariationsHandler() + params: Final = { + "api_key": "transport-only", + "base_url": kwargs["api_base"], + "timeout": timeout, + } + return ( + variation_factory.get_async_client(client=client, init_client_params=params) + if is_async + else variation_factory.get_sync_client(client=client, init_client_params=params) + ) + if api == "azure_gateway": + return BaseAzureLLM()._init_azure_client_for_cloudflare_ai_gateway( + api_base="https://sdk-default.example", + model="deployment", + api_version="2024-02-01", + max_retries=0, + timeout=timeout, + litellm_params={}, + api_key="transport-only", + azure_ad_token=None, + azure_ad_token_provider=None, + acompletion=is_async, + client=client, + ) + factory: Final = {"files": OpenAIFilesAPI, "batches": OpenAIBatchesAPI, "fine_tuning": OpenAIFineTuningAPI}[api]() + return factory.get_openai_client(**kwargs, _is_async=is_async) + + +async def _list_model_ids(sdk_client: openai.OpenAI | openai.AsyncOpenAI) -> list[str]: + page: Final = ( + await sdk_client.models.list() if isinstance(sdk_client, openai.AsyncOpenAI) else sdk_client.models.list() + ) + return [model.id for model in page.data] + + +@pytest.mark.parametrize("api", ["files", "batches", "assistants", "fine_tuning", "image_variations", "azure_gateway"]) +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_sdk_api_factories_keep_httpx_transport_and_request_timeouts(api: str, is_async: bool) -> None: + timeout: Final = openai.Timeout(connect=1, read=7, write=2, pool=3) + with respx.mock(assert_all_called=True) as mock: + mock.get(host="sdk-default.example", path__regex=r".*/models$").mock( + return_value=httpx.Response(307, headers={"location": "https://sdk-default.example/v1/moved-models"}) + ) + moved: Final = mock.get(host="sdk-default.example", path="/v1/moved-models").mock( + return_value=httpx.Response( + 200, + json={"object": "list", "data": [{"id": "sdk-compat", "object": "model", "created": 0, "owned_by": "t"}]}, + ) + ) + sdk_client: Final = _sdk_api_client(api, is_async, timeout) + assert sdk_client is not None + try: + assert _sdk_api_client(api, is_async, 19, client=sdk_client) is sdk_client + model_ids: Final = await _list_model_ids(sdk_client) + finally: + if is_async: + await sdk_client.close() + else: + sdk_client.close() + assert model_ids == ["sdk-compat"] + assert moved.calls.last.request.extensions["timeout"] == {"connect": 1.0, "read": 7.0, "write": 2.0, "pool": 3.0} + + +@pytest.mark.parametrize("backend", ["httpx", "sdk_default"]) +@pytest.mark.asyncio +async def test_azure_gateway_and_image_variations_use_the_callers_async_session( + backend: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from importlib import import_module + from io import BytesIO + + from openai._types import Response as SDKResponse + + from litellm.images.main import aimage_variation + + http_module: Final = httpx if backend == "httpx" else import_module(SDKResponse.__module__) + transport: Final = http_module.MockTransport( + lambda request: http_module.Response( + 200, json={"created": 0, "data": [{"url": "https://example.com/image.png"}]} + ) + ) + async with http_module.AsyncClient(transport=transport) as session: + monkeypatch.setattr(litellm, "aclient_session", session) + monkeypatch.setattr(litellm, "client_session", None) + azure_client: Final = _sdk_api_client("azure_gateway", True, 7) + assert azure_client is not None + assert azure_client._client is session + response: Final = await aimage_variation( + image=BytesIO(b"image-bytes"), api_key="transport-only", api_base="https://sdk-default.example/v1" + ) + assert response.data[0].url == "https://example.com/image.png" + assert not session.is_closed diff --git a/tests/unit/llms/openai/test_openai_workload_identity.py b/tests/unit/llms/openai/test_openai_workload_identity.py index 7d78a2c5238..d0cf52d0e28 100644 --- a/tests/unit/llms/openai/test_openai_workload_identity.py +++ b/tests/unit/llms/openai/test_openai_workload_identity.py @@ -1,5 +1,6 @@ import json import sys +from functools import partial from pathlib import Path from typing import Final @@ -36,6 +37,24 @@ CHAT_COMPLETION_BODY: Final = { } +@pytest.fixture(autouse=True) +def mock_sdk_token_exchange_transport(monkeypatch: pytest.MonkeyPatch) -> None: + native_client_factory: Final = getattr(sys.modules.get("openai.auth._workload"), "DefaultHttpx2Client", None) + if native_client_factory is not None: + import httpx2 + + def handle_exchange(request: httpx2.Request) -> httpx2.Response: + response: Final = respx.mock.handler( + httpx.Request(request.method, str(request.url), headers=dict(request.headers), content=request.content) + ) + return httpx2.Response(response.status_code, headers=dict(response.headers), content=response.content) + + monkeypatch.setattr( + "openai.auth._workload.DefaultHttpx2Client", + partial(native_client_factory, transport=httpx2.MockTransport(handle_exchange), trust_env=False), + ) + + @pytest.fixture def wif_env(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> OpenAIWorkloadIdentityConfig: token_file: Final = tmp_path / "subject_token.jwt" @@ -602,7 +621,6 @@ class TestDiscoverModels: assert models_route.calls.last.request.headers["Authorization"] == "Bearer None" assert not exchange_route.called - @respx.mock def test_empty_static_key_never_borrows_the_env_key(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("OPENAI_API_KEY", "sk-env-key-that-must-stay-home") diff --git a/tests/unit/test_completion_timeout_resolution.py b/tests/unit/test_completion_timeout_resolution.py index 7eb79e90e60..77166d6e947 100644 --- a/tests/unit/test_completion_timeout_resolution.py +++ b/tests/unit/test_completion_timeout_resolution.py @@ -2,8 +2,11 @@ import os import sys +from typing import Final, Literal import httpx +import pytest +from openai import Timeout as SDKTimeout sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))) @@ -142,3 +145,35 @@ def test_httpx_timeout_preserved_for_openai(): ) assert out is t assert isinstance(out, httpx.Timeout) + + +@pytest.mark.parametrize("source", ("model", "timeout", "request_timeout")) +def test_sdk_timeout_preserves_components_for_httpx_providers( + source: Literal["model", "timeout", "request_timeout"], +) -> None: + timeout: Final = SDKTimeout(connect=2.0, read=None, write=5.0, pool=7.0) + resolved: Final = CompletionTimeout.resolve( + timeout if source == "model" else None, + {} if source == "model" else {source: timeout}, + "openai", + global_timeout=None, + supports_httpx_timeout=supports_httpx_timeout, + ) + assert isinstance(resolved, httpx.Timeout) + assert resolved.as_dict() == timeout.as_dict() + + +@pytest.mark.parametrize("read_timeout, expected", ((23.0, 23.0), (None, 600.0))) +def test_sdk_timeout_coerces_read_timeout_for_providers_without_httpx_support( + read_timeout: float | None, expected: float +) -> None: + assert ( + CompletionTimeout.resolve( + SDKTimeout(connect=2.0, read=read_timeout, write=5.0, pool=7.0), + {}, + "azure_ai", + global_timeout=None, + supports_httpx_timeout=supports_httpx_timeout, + ) + == expected + ) diff --git a/tests/unit/test_exception_header_preservation.py b/tests/unit/test_exception_header_preservation.py index dd142d9d40b..7acacf05226 100644 --- a/tests/unit/test_exception_header_preservation.py +++ b/tests/unit/test_exception_header_preservation.py @@ -8,6 +8,8 @@ This is important for debugging and observability - headers like x-request-id, x-ms-region, rate limit headers, etc. should be available even when errors occur. """ +from typing import Final + import httpx import pytest @@ -17,6 +19,7 @@ from litellm.exceptions import ( ContextWindowExceededError, ImageFetchError, MidStreamFallbackError, + PaymentRequiredError, RateLimitError, ServiceUnavailableError, ) @@ -25,10 +28,11 @@ from litellm.exceptions import ( class TestExceptionHeaderPreservation: """Test that exception classes preserve headers from provider responses.""" - @pytest.fixture - def mock_response_with_headers(self) -> httpx.Response: + @pytest.fixture(params=("httpx", "httpx2")) + def mock_response_with_headers(self, request: pytest.FixtureRequest) -> httpx.Response: """Create a mock response with typical provider headers.""" - return httpx.Response( + transport: Final = pytest.importorskip(request.param) + return transport.Response( status_code=400, headers={ "x-request-id": "req-abc123", @@ -36,7 +40,7 @@ class TestExceptionHeaderPreservation: "x-ratelimit-remaining-requests": "99", "x-ratelimit-remaining-tokens": "9999", }, - request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions"), + request=transport.Request("POST", "https://api.openai.com/v1/chat/completions"), ) def test_bad_request_error_preserves_headers( @@ -50,11 +54,30 @@ class TestExceptionHeaderPreservation: response=mock_response_with_headers, ) - assert error.response is not None + assert error.response is mock_response_with_headers + assert error.request is mock_response_with_headers.request + assert error.request_id == mock_response_with_headers.headers["x-request-id"] assert error.response.headers.get("x-request-id") == "req-abc123" assert error.response.headers.get("x-ms-region") == "eastus" assert error.response.headers.get("x-ratelimit-remaining-requests") == "99" + def test_payment_required_error_preserves_headers( + self, mock_response_with_headers: httpx.Response + ): + """PaymentRequiredError should keep the provider response like its BadRequestError parent.""" + error = PaymentRequiredError( + message="Insufficient credit", + model="gpt-4", + llm_provider="openrouter", + response=mock_response_with_headers, + ) + + assert error.status_code == 402 + assert error.response is mock_response_with_headers + assert error.request is mock_response_with_headers.request + assert error.request_id == mock_response_with_headers.headers["x-request-id"] + assert error.response.headers.get("x-ms-region") == "eastus" + def test_content_policy_violation_error_preserves_headers( self, mock_response_with_headers: httpx.Response ): diff --git a/tests/unit/test_openai_embedding_encoding_format_default.py b/tests/unit/test_openai_embedding_encoding_format_default.py index 7a42eaf0f0a..9788eb30b1b 100644 --- a/tests/unit/test_openai_embedding_encoding_format_default.py +++ b/tests/unit/test_openai_embedding_encoding_format_default.py @@ -7,9 +7,19 @@ import respx import litellm +_SDK_EMBEDDING_MODELS: Final = ( + "openai/sdk-compat", + "mistral/sdk-compat", + "fireworks_ai/accounts/fireworks/models/sdk-compat", + "together_ai/sdk-compat", + "nvidia_nim/sdk-compat", +) -def _mock_openai_embedding_route(respx_mock: respx.MockRouter) -> respx.Route: - return respx_mock.post("https://api.openai.com/v1/embeddings").mock( + +def _mock_openai_embedding_route( + respx_mock: respx.MockRouter, api_base: str = "https://api.openai.com/v1" +) -> respx.Route: + return respx_mock.post(f"{api_base}/embeddings").mock( return_value=httpx.Response( 200, json={ @@ -27,13 +37,21 @@ def clear_default_encoding_format_env(monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", raising=False) -def test_embedding_openai_omits_encoding_format_when_client_omits_it(respx_mock: respx.MockRouter) -> None: - mock_route: Final = _mock_openai_embedding_route(respx_mock) +@pytest.mark.parametrize("model", _SDK_EMBEDDING_MODELS) +def test_embedding_sdk_providers_omit_encoding_format_when_client_omits_it( + respx_mock: respx.MockRouter, model: str +) -> None: + provider, upstream_model = model.split("/", 1) + api_base: Final = f"https://{provider.replace('_', '-')}.example/v1" + mock_route: Final = _mock_openai_embedding_route(respx_mock, api_base) - response: Final = litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + response: Final = litellm.embedding(model=model, input=["hello"], api_key="sk-test", api_base=api_base) request_body: Final = json.loads(mock_route.calls.last.request.read()) assert "encoding_format" not in request_body + assert request_body["model"] == upstream_model + assert request_body["input"] == ["hello"] + assert mock_route.calls.last.request.headers["authorization"] == "Bearer sk-test" assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] @@ -89,16 +107,22 @@ def test_embedding_openai_env_none_omits_encoding_format( @pytest.mark.asyncio -async def test_aembedding_openai_omits_encoding_format_when_client_omits_it( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +@pytest.mark.parametrize("model", _SDK_EMBEDDING_MODELS) +async def test_aembedding_sdk_providers_omit_encoding_format_when_client_omits_it( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, model: str ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - mock_route: Final = _mock_openai_embedding_route(respx_mock) + provider, upstream_model = model.split("/", 1) + api_base: Final = f"https://{provider.replace('_', '-')}.example/v1" + mock_route: Final = _mock_openai_embedding_route(respx_mock, api_base) - response: Final = await litellm.aembedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + response: Final = await litellm.aembedding(model=model, input=["hello"], api_key="sk-test", api_base=api_base) request_body: Final = json.loads(mock_route.calls.last.request.read()) assert "encoding_format" not in request_body + assert request_body["model"] == upstream_model + assert request_body["input"] == ["hello"] + assert mock_route.calls.last.request.headers["authorization"] == "Bearer sk-test" assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] diff --git a/tests/unit/types/test_router.py b/tests/unit/types/test_router.py index 89238e04a5a..6188bb966ee 100644 --- a/tests/unit/types/test_router.py +++ b/tests/unit/types/test_router.py @@ -1,6 +1,9 @@ import logging +from typing import Final +import httpx import pytest +from openai import Timeout as SDKTimeout from pydantic import ValidationError from litellm.types.router import ( @@ -26,6 +29,13 @@ from litellm.types.utils import ( ) +def test_sdk_timeout_is_normalized_for_provider_clients() -> None: + timeout: Final = SDKTimeout(connect=2.0, read=None, write=5.0, pool=7.0) + params: Final = GenericLiteLLMParams(timeout=timeout) + assert isinstance(params.timeout, httpx.Timeout) + assert params.timeout.as_dict() == timeout.as_dict() + + def test_model_info_declares_mirrored_pricing_fields(): """The pricing keys Deployment mirrors onto model_info must be declared fields, not extras that only survive because ModelInfo sets extra="allow".""" diff --git a/uv.lock b/uv.lock index 96fa27289fa..471460a32b5 100644 --- a/uv.lock +++ b/uv.lock @@ -4814,7 +4814,7 @@ requires-dist = [ { name = "numpy", marker = "extra == 'stt-nvidia-riva'", specifier = ">=1.26.0" }, { name = "numpydoc", marker = "extra == 'utils'", specifier = ">=1.8.0,<2.0" }, { name = "nvidia-riva-client", marker = "extra == 'stt-nvidia-riva'", specifier = ">=2.15.0" }, - { name = "openai", specifier = ">=2.20.0,<3.0.0" }, + { name = "openai", specifier = ">=2.20.0,<4.0.0" }, { name = "opentelemetry-api", marker = "extra == 'proxy-runtime'", specifier = "==1.33.1" }, { name = "opentelemetry-exporter-otlp", marker = "extra == 'proxy-runtime'", specifier = "==1.33.1" }, { name = "opentelemetry-instrumentation-fastapi", marker = "extra == 'proxy-runtime'", specifier = "==0.54b1" }, From 62aee8bee95a8841a2eb50320b8f369252cf20c7 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 20:03:01 -0700 Subject: [PATCH 022/104] feat(ui): build System One playground requests with a form (#45721) * feat(ui): build System One playground requests with a form * test(ui): fill System One form fields with single change events --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../components/systemOneUI/JsonEditor.tsx | 4 +- .../components/systemOneUI/SystemOneForm.tsx | 337 ++++++++++++++++++ .../SystemOneUI.integration.test.tsx | 178 ++++++++- .../components/systemOneUI/SystemOneUI.tsx | 37 +- .../systemOneUI/lib/formPayload.test.ts | 124 +++++++ .../components/systemOneUI/lib/formPayload.ts | 89 +++++ 6 files changed, 754 insertions(+), 15 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneForm.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/formPayload.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/formPayload.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/JsonEditor.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/JsonEditor.tsx index 1effe55ef41..fce617b4c47 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/JsonEditor.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/JsonEditor.tsx @@ -31,7 +31,7 @@ interface JsonEditorProps { validation: SystemOnePayloadValidation; } -function ValidationStatus({ validation }: { validation: SystemOnePayloadValidation }) { +export function ValidationStatus({ validation }: { validation: SystemOnePayloadValidation }) { const errorCount = validation.issues.filter((issue) => issue.severity === "error").length; if (errorCount > 0) { return ( @@ -43,7 +43,7 @@ function ValidationStatus({ validation }: { validation: SystemOnePayloadValidati return Valid payload; } -function IssueList({ id, validation }: { id: string; validation: SystemOnePayloadValidation }) { +export function IssueList({ id, validation }: { id: string; validation: SystemOnePayloadValidation }) { if (validation.issues.length === 0) { return (

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneForm.tsx new file mode 100644 index 00000000000..d799afd13ea --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneForm.tsx @@ -0,0 +1,337 @@ +"use client"; + +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Textarea } from "@/components/ui/textarea"; +import { cn } from "@/lib/cva.config"; +import { Plus, Trash2 } from "lucide-react"; +import { useId, useMemo, useState, type ReactNode } from "react"; +import { IssueList, ValidationStatus } from "./JsonEditor"; +import { + blankQuestion, + freeName, + readForm, + removeAt, + renameKey, + replaceAt, + retype, + withNoulCriterion, + withoutKey, + type ChoiceQuestion, + type FormPayload, + type FormQuestion, + type NoulQuestion, + type QuestionType, + type ScoreQuestion, +} from "./lib/formPayload"; +import type { SystemOnePayloadValidation } from "./lib/validatePayload"; + +const QUESTION_TYPES: readonly QuestionType[] = ["choice", "noul", "score"]; +const TYPE_LABELS: Record = { choice: "Choice", noul: "Yes / no", score: "Score" }; + +const isQuestionType = (value: unknown): value is QuestionType => QUESTION_TYPES.some((type) => type === value); + +interface SystemOneFormProps { + value: string; + onChange: (value: string) => void; + validation: SystemOnePayloadValidation; + onOpenJson: () => void; +} + +interface QuestionEditorProps { + question: Q; + onChange: (question: FormQuestion) => void; +} + +interface KeyInputProps { + value: string; + taken: readonly string[]; + onCommit: (value: string) => void; + label: string; + placeholder?: string; +} + +function keyError(next: string, current: string, taken: readonly string[]): string | undefined { + if (next === "") { + return "Enter a name"; + } + if (next !== current && taken.includes(next)) { + return `"${next}" is already used`; + } + return undefined; +} + +function KeyInput({ value, taken, onCommit, label, placeholder }: KeyInputProps) { + const errorId = useId(); + const [draft, setDraft] = useState(value); + const next = draft.trim(); + const error = keyError(next, value, taken); + + function commit() { + if (error === undefined && next !== value) { + onCommit(next); + return; + } + setDraft(value); + } + + return ( +

+ setDraft(event.target.value)} + onBlur={commit} + /> + {error !== undefined && ( +

+ {error} +

+ )} +
+ ); +} + +function RemoveButton({ label, onClick }: { label: string; onClick: () => void }) { + return ( + + ); +} + +function AddButton({ children, onClick }: { children: ReactNode; onClick: () => void }) { + return ( + + ); +} + +function ChoiceOptions({ question, onChange }: QuestionEditorProps) { + const labels = Object.keys(question.criteria); + const writeCriteria = (criteria: Record) => onChange({ ...question, criteria }); + + return ( +
+ Options + {Object.entries(question.criteria).map(([label, description], index) => ( +
+ writeCriteria(renameKey(question.criteria, label, next))} + /> + writeCriteria({ ...question.criteria, [label]: event.target.value })} + /> + writeCriteria(withoutKey(question.criteria, label))} + /> +
+ ))} + writeCriteria({ ...question.criteria, [freeName(labels, "option")]: "" })}> + Add option + +
+ ); +} + +function NoulCriteria({ question, onChange }: QuestionEditorProps) { + const yesId = useId(); + const noId = useId(); + + return ( +
+
+ + onChange(withNoulCriterion(question, "true", event.target.value))} + /> +
+
+ + onChange(withNoulCriterion(question, "false", event.target.value))} + /> +
+
+ ); +} + +function ScoreLevels({ question, onChange }: QuestionEditorProps) { + const writeCriteria = (criteria: string[]) => onChange({ ...question, criteria }); + + return ( +
+ Levels, lowest first + {question.criteria.map((level, index) => ( +
+ {index} + writeCriteria(replaceAt(question.criteria, index, event.target.value))} + /> + writeCriteria(removeAt(question.criteria, index))} + /> +
+ ))} + writeCriteria([...question.criteria, ""])}>Add level +
+ ); +} + +function CriteriaEditor({ question, onChange }: QuestionEditorProps) { + switch (question.type) { + case "choice": + return ; + case "noul": + return ; + case "score": + return ; + } +} + +interface QuestionCardProps extends QuestionEditorProps { + name: string; + taken: readonly string[]; + onRename: (name: string) => void; + onRemove: () => void; +} + +function QuestionCard({ name, question, taken, onChange, onRename, onRemove }: QuestionCardProps) { + const instructionsId = useId(); + + return ( +
+
+ + + +
+
+ +