diff --git a/litellm-rust/.agents/skills/rust-string-enums/SKILL.md b/litellm-rust/.agents/skills/rust-string-enums/SKILL.md index fc6988639d0..5fc067790a0 100644 --- a/litellm-rust/.agents/skills/rust-string-enums/SKILL.md +++ b/litellm-rust/.agents/skills/rust-string-enums/SKILL.md @@ -35,7 +35,7 @@ pub enum EventType { } ``` -Use `#[strum(serialize_all = "snake_case")]` or another supported case style when it exactly matches the contract. Use explicit variant spellings otherwise. With multiple accepted aliases, set `to_string` to the existing canonical output: Strum Display otherwise selects the longest `serialize` spelling +Spell each variant explicitly with `#[strum(serialize = "...")]`; do not use `serialize_all`. With multiple accepted aliases, set `to_string` to the existing canonical output: Strum Display otherwise selects the longest `serialize` spelling Use only the derives the contract needs. A deserialize-only type should remain deserialize-only. Do not add an unknown variant to a closed enum, or derive Serde for a type that currently has no serialization contract @@ -43,7 +43,7 @@ Plain Serde derives with `rename` or `rename_all` remain appropriate for closed ## Preserve behavior during migration -Read the type, its callers, and existing tests before changing it. Preserve canonical output, accepted aliases, case sensitivity, whitespace handling, unknown values, malformed-input rejection, and existing public conversion APIs. Keep conversions needed by callers or compatibility even when Serde no longer uses them +Read the type, its callers, and existing tests before changing it. Preserve canonical output, accepted aliases, case sensitivity, whitespace handling, unknown values, malformed-input rejection, and existing public conversion contracts. Wrappers that only delegate to a Strum-derived trait are removed rather than kept: callers use `FromStr` and `From for &'static str` directly. Keep conversions that add behavior needed by callers or compatibility even when Serde no longer uses them Use the workspace dependencies and enable `serde_with.workspace = true` in a crate only when needed. Check the versions and enabled features in `Cargo.toml` and `Cargo.lock` rather than upgrading dependencies for this refactor diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index 442dfd8e957..c22f9575f74 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -4,6 +4,10 @@ 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 + +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` + ## Test placement - Never create a `tests.rs` (or `test.rs`) file under `src/`, and never `#[path = "tests.rs"] mod tests;` @@ -17,6 +21,9 @@ For string-valued enums and their Serde conversions, follow [.agents/skills/rust 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 +- 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` diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 057faa279e3..9d511774b04 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3485,6 +3485,7 @@ dependencies = [ "rstest", "serde", "serde_json", + "strum", "thiserror 2.0.19", "tokio", ] @@ -3692,7 +3693,9 @@ dependencies = [ "litellm-auth-types", "rstest", "serde", + "serde_json", "serde_yaml_ng", + "strum", "tempfile", "thiserror 2.0.19", ] @@ -3990,6 +3993,7 @@ dependencies = [ "rustls-native-certs", "serde", "serde_json", + "strum", "tempfile", "thiserror 2.0.19", "tokio", @@ -4539,6 +4543,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.10.9", + "strum", "tar", "target-lexicon", "tempfile", diff --git a/litellm-rust/crates/cache/Cargo.toml b/litellm-rust/crates/cache/Cargo.toml index f18dbd9cb26..51bbe12a85e 100644 --- a/litellm-rust/crates/cache/Cargo.toml +++ b/litellm-rust/crates/cache/Cargo.toml @@ -8,6 +8,7 @@ repository.workspace = true [dependencies] serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } +strum.workspace = true thiserror.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/cache/src/cache_type.rs b/litellm-rust/crates/cache/src/cache_type.rs index 22d8c8c7cb5..80d040db487 100644 --- a/litellm-rust/crates/cache/src/cache_type.rs +++ b/litellm-rust/crates/cache/src/cache_type.rs @@ -1,57 +1,44 @@ use serde::{Deserialize, Serialize}; -#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, Hash)] +#[derive( + Clone, + Copy, + Debug, + Deserialize, + Serialize, + PartialEq, + Eq, + Hash, + strum::EnumString, + strum::IntoStaticStr, + strum::VariantArray, +)] pub enum CacheType { #[serde(rename = "local")] + #[strum(serialize = "local")] Local, #[serde(rename = "redis")] + #[strum(serialize = "redis")] Redis, #[serde(rename = "redis-semantic")] + #[strum(serialize = "redis-semantic")] RedisSemantic, #[serde(rename = "valkey-semantic")] + #[strum(serialize = "valkey-semantic")] ValkeySemantic, #[serde(rename = "s3")] + #[strum(serialize = "s3")] S3, #[serde(rename = "disk")] + #[strum(serialize = "disk")] Disk, #[serde(rename = "qdrant-semantic")] + #[strum(serialize = "qdrant-semantic")] QdrantSemantic, #[serde(rename = "azure-blob")] + #[strum(serialize = "azure-blob")] AzureBlob, #[serde(rename = "gcs")] + #[strum(serialize = "gcs")] Gcs, } - -impl CacheType { - pub const ALL: [Self; 9] = [ - Self::Local, - Self::Redis, - Self::RedisSemantic, - Self::ValkeySemantic, - Self::S3, - Self::Disk, - Self::QdrantSemantic, - Self::AzureBlob, - Self::Gcs, - ]; - - pub const fn as_python_name(self) -> &'static str { - match self { - Self::Local => "local", - Self::Redis => "redis", - Self::RedisSemantic => "redis-semantic", - Self::ValkeySemantic => "valkey-semantic", - Self::S3 => "s3", - Self::Disk => "disk", - Self::QdrantSemantic => "qdrant-semantic", - Self::AzureBlob => "azure-blob", - Self::Gcs => "gcs", - } - } - - pub fn from_python_name(value: &str) -> Option { - Self::ALL - .into_iter() - .find(|cache_type| cache_type.as_python_name() == value) - } -} diff --git a/litellm-rust/crates/cache/tests/cache_type.rs b/litellm-rust/crates/cache/tests/cache_type.rs index 24aaba8e5fd..49b3e8b27ad 100644 --- a/litellm-rust/crates/cache/tests/cache_type.rs +++ b/litellm-rust/crates/cache/tests/cache_type.rs @@ -1,5 +1,6 @@ use litellm_cache::CacheType; use rstest::rstest; +use strum::VariantArray; #[rstest] #[case(CacheType::Local, "local")] @@ -15,16 +16,16 @@ fn every_python_cache_type_has_one_round_trip_identity( #[case] cache_type: CacheType, #[case] name: &str, ) { - assert_eq!(cache_type.as_python_name(), name); - assert_eq!(CacheType::from_python_name(name), Some(cache_type)); + assert_eq!(<&'static str>::from(cache_type), name); + assert_eq!(name.parse::().ok(), Some(cache_type)); assert_eq!( serde_json::to_value(cache_type).unwrap(), serde_json::Value::from(name) ); assert_eq!( - CacheType::ALL + CacheType::VARIANTS .iter() - .filter(|candidate| candidate.as_python_name() == name) + .filter(|candidate| <&'static str>::from(**candidate) == name) .count(), 1 ); @@ -33,7 +34,10 @@ fn every_python_cache_type_has_one_round_trip_identity( #[rstest] fn python_cache_types_are_listed_in_python_order() { assert_eq!( - CacheType::ALL.map(CacheType::as_python_name), + CacheType::VARIANTS + .iter() + .map(|cache_type| <&'static str>::from(*cache_type)) + .collect::>(), [ "local", "redis", @@ -52,5 +56,5 @@ fn python_cache_types_are_listed_in_python_order() { #[case::unknown("memcached")] #[case::case_sensitive("Redis")] fn unknown_python_names_have_no_cache_type(#[case] name: &str) { - assert_eq!(CacheType::from_python_name(name), None); + assert!(name.parse::().is_err()); } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 70b7ecb45fd..dd3da8991cc 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -936,9 +936,6 @@ mod payload_tests { /// The payload phases of `Logging` on top of `StubLogger`, with `pre_call` handing the /// payload to the case's `on_pre_call`. const PAYLOAD_LOGGER: &CStr = c" -class Request: - pass - class PayloadLogger(StubLogger): def update_from_kwargs(self, **update): self.update = update @@ -954,7 +951,7 @@ class PayloadLogger(StubLogger): self.record('post_call', None) self.post = (original_response, api_key, additional_args) -request = Request() +bound = {} kwargs = {} logger = PayloadLogger() on_pre_call = lambda additional_args: None @@ -1225,10 +1222,10 @@ on_pre_call = lambda args: observed.append( def check(): assert observed == [(True, True)], observed ")] - #[case::request_attribute_behind_an_omitted_keyword(c" + #[case::bound_value_behind_an_omitted_keyword(c" document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} pages = [0] -request.document = document +bound['document'] = document kwargs = {'pages': pages} observed = [] on_pre_call = lambda args: observed.append( diff --git a/litellm-rust/crates/callbacks-legacy-python/src/call.rs b/litellm-rust/crates/callbacks-legacy-python/src/call.rs index 8e7645cae7b..8ea67292ae7 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/call.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/call.rs @@ -13,21 +13,21 @@ use pyo3::{ pub struct PublicCall { args: Py, kwargs: Py, - request: Py, + bound: Py, } impl PublicCall { /// Copies the keyword arguments once, so the legacy path's rewrites never reach the /// caller's own dict while every value keeps its identity. pub fn capture( - request: &Bound<'_, PyAny>, + bound: &Bound<'_, PyDict>, args: &Bound<'_, PyTuple>, kwargs: &Bound<'_, PyDict>, ) -> PyResult { Ok(Self { args: args.clone().unbind(), kwargs: kwargs.copy()?.unbind(), - request: request.clone().unbind(), + bound: bound.clone().unbind(), }) } @@ -55,31 +55,30 @@ impl PublicCall { py: Python<'py>, name: &str, ) -> PyResult>> { - lookup(self.kwargs.bind(py), self.request.bind(py), name) + lookup(self.kwargs.bind(py), self.bound.bind(py), name) } pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { visit.call(&self.args)?; visit.call(&self.kwargs)?; - visit.call(&self.request) + visit.call(&self.bound) } } #[cfg(test)] mod tests { use super::*; + use crate::test_support::{local, local_dict}; fn capture<'py>(py: Python<'py>, source: &std::ffi::CStr) -> (PublicCall, Bound<'py, PyDict>) { let locals = PyDict::new(py); py.run(source, Some(&locals), Some(&locals)).unwrap(); - let request = locals.get_item("request").unwrap().unwrap(); - let kwargs = locals - .get_item("kwargs") - .unwrap() - .unwrap() - .cast_into::() - .unwrap(); - let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap(); + let call = PublicCall::capture( + &local_dict(&locals, "bound"), + &PyTuple::empty(py), + &local_dict(&locals, "kwargs"), + ) + .unwrap(); (call, locals) } @@ -91,25 +90,22 @@ mod tests { py, c" pages = [0] -class Request: - pass -request = Request() +bound = {'pages': [1]} kwargs = {'pages': pages} ", ); - let caller = locals - .get_item("kwargs") - .unwrap() - .unwrap() - .cast_into::() - .unwrap(); + let caller = local_dict(&locals, "kwargs"); call.kwargs() .bind(py) .set_item("litellm_call_id", "call") .unwrap(); assert!(!caller.contains("litellm_call_id").unwrap()); - let pages = locals.get_item("pages").unwrap().unwrap(); - assert!(call.lookup(py, "pages").unwrap().unwrap().is(&pages)); + assert!( + call.lookup(py, "pages") + .unwrap() + .unwrap() + .is(local(&locals, "pages")) + ); }); } } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index f252a45b562..a7b28f1b0ca 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -173,21 +173,23 @@ pub(crate) fn local<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, locals.get_item(name).unwrap().unwrap() } -/// A legacy call over the namespace's `kwargs` (or none) and `request` (or `None`). +pub(crate) fn local_dict<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyDict> { + local(locals, name).cast_into().unwrap() +} + +/// A legacy call over the namespace's `kwargs` and `bound` dicts, each empty when absent. pub(crate) fn legacy_call( py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool, ) -> LegacyLogging { - let request = locals - .get_item("request") - .unwrap() - .unwrap_or_else(|| py.None().into_bound(py)); - let kwargs = locals - .get_item("kwargs") - .unwrap() - .map(|kwargs| kwargs.cast_into::().unwrap()) - .unwrap_or_else(|| PyDict::new(py)); - let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap(); + let dict = |name: &str| { + locals + .get_item(name) + .unwrap() + .map(|value| value.cast_into::().unwrap()) + .unwrap_or_else(|| PyDict::new(py)) + }; + let call = PublicCall::capture(&dict("bound"), &PyTuple::empty(py), &dict("kwargs")).unwrap(); LegacyLogging::new(py, crate::LoggingOperation::Ocr, call, asynchronous) } diff --git a/litellm-rust/crates/config/Cargo.toml b/litellm-rust/crates/config/Cargo.toml index 36bd68fe2a0..e23739997e2 100644 --- a/litellm-rust/crates/config/Cargo.toml +++ b/litellm-rust/crates/config/Cargo.toml @@ -9,8 +9,10 @@ repository.workspace = true litellm-auth-types.workspace = true serde.workspace = true serde_yaml_ng = "0.10.0" +strum.workspace = true thiserror.workspace = true [dev-dependencies] rstest.workspace = true +serde_json.workspace = true tempfile.workspace = true diff --git a/litellm-rust/crates/config/src/mcp.rs b/litellm-rust/crates/config/src/mcp.rs index 522744e4fe4..cbe992f9956 100644 --- a/litellm-rust/crates/config/src/mcp.rs +++ b/litellm-rust/crates/config/src/mcp.rs @@ -41,57 +41,43 @@ impl fmt::Debug for McpServer { } } -#[derive(Clone, Copy, Debug, Default, Deserialize, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, Default, Deserialize, PartialEq, Eq, strum::IntoStaticStr)] #[serde(rename_all = "snake_case")] pub enum McpTransport { #[default] + #[strum(serialize = "http")] Http, + #[strum(serialize = "sse")] Sse, + #[strum(serialize = "stdio")] Stdio, } -impl McpTransport { - pub fn as_str(self) -> &'static str { - match self { - Self::Http => "http", - Self::Sse => "sse", - Self::Stdio => "stdio", - } - } -} - -#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, strum::IntoStaticStr)] #[serde(rename_all = "snake_case")] pub enum McpAuth { + #[strum(serialize = "none")] None, + #[strum(serialize = "api_key")] ApiKey, + #[strum(serialize = "bearer_token")] BearerToken, + #[strum(serialize = "basic")] Basic, + #[strum(serialize = "authorization")] Authorization, + #[strum(serialize = "token")] Token, + #[strum(serialize = "oauth2")] Oauth2, + #[strum(serialize = "aws_sigv4")] AwsSigv4, + #[strum(serialize = "oauth2_token_exchange")] Oauth2TokenExchange, + #[strum(serialize = "oauth2_id_jag")] Oauth2IdJag, + #[strum(serialize = "true_passthrough")] TruePassthrough, + #[strum(serialize = "oauth_delegate")] OauthDelegate, } - -impl McpAuth { - pub fn as_str(self) -> &'static str { - match self { - Self::None => "none", - Self::ApiKey => "api_key", - Self::BearerToken => "bearer_token", - Self::Basic => "basic", - Self::Authorization => "authorization", - Self::Token => "token", - Self::Oauth2 => "oauth2", - Self::AwsSigv4 => "aws_sigv4", - Self::Oauth2TokenExchange => "oauth2_token_exchange", - Self::Oauth2IdJag => "oauth2_id_jag", - Self::TruePassthrough => "true_passthrough", - Self::OauthDelegate => "oauth_delegate", - } - } -} diff --git a/litellm-rust/crates/config/tests/mcp.rs b/litellm-rust/crates/config/tests/mcp.rs new file mode 100644 index 00000000000..9ac6ad02bbb --- /dev/null +++ b/litellm-rust/crates/config/tests/mcp.rs @@ -0,0 +1,39 @@ +use litellm_config::{McpAuth, McpTransport}; +use rstest::rstest; +use serde_json::json; + +#[rstest] +#[case::http(McpTransport::Http, "http")] +#[case::sse(McpTransport::Sse, "sse")] +#[case::stdio(McpTransport::Stdio, "stdio")] +fn mcp_transport_as_str_matches_the_serde_spelling( + #[case] transport: McpTransport, + #[case] name: &str, +) { + assert_eq!(<&'static str>::from(transport), name); + assert_eq!( + serde_json::from_value::(json!(name)).unwrap(), + transport + ); +} + +#[rstest] +#[case::none(McpAuth::None, "none")] +#[case::api_key(McpAuth::ApiKey, "api_key")] +#[case::bearer_token(McpAuth::BearerToken, "bearer_token")] +#[case::basic(McpAuth::Basic, "basic")] +#[case::authorization(McpAuth::Authorization, "authorization")] +#[case::token(McpAuth::Token, "token")] +#[case::oauth2(McpAuth::Oauth2, "oauth2")] +#[case::aws_sigv4(McpAuth::AwsSigv4, "aws_sigv4")] +#[case::oauth2_token_exchange(McpAuth::Oauth2TokenExchange, "oauth2_token_exchange")] +#[case::oauth2_id_jag(McpAuth::Oauth2IdJag, "oauth2_id_jag")] +#[case::true_passthrough(McpAuth::TruePassthrough, "true_passthrough")] +#[case::oauth_delegate(McpAuth::OauthDelegate, "oauth_delegate")] +fn mcp_auth_as_str_matches_the_serde_spelling(#[case] auth: McpAuth, #[case] name: &str) { + assert_eq!(<&'static str>::from(auth), name); + assert_eq!( + serde_json::from_value::(json!(name)).unwrap(), + auth + ); +} 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 e71ecdffa44..a2bb462c2a9 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 @@ -7,17 +7,26 @@ pub struct CustomLlmProvider<'a> { } #[derive(Debug, Clone, Copy, PartialEq, Eq, EnumString, IntoStaticStr)] -#[strum(serialize_all = "snake_case")] pub enum LlmProviders { + #[strum(serialize = "anthropic")] Anthropic, + #[strum(serialize = "aws_textract")] AwsTextract, + #[strum(serialize = "azure_ai")] AzureAi, + #[strum(serialize = "bedrock")] Bedrock, + #[strum(serialize = "cohere")] Cohere, + #[strum(serialize = "mistral")] Mistral, + #[strum(serialize = "openai")] Openai, + #[strum(serialize = "openai_like")] OpenaiLike, + #[strum(serialize = "reducto")] Reducto, + #[strum(serialize = "vertex_ai")] VertexAi, } diff --git a/litellm-rust/crates/core-utils/src/prompt_templates/factory.rs b/litellm-rust/crates/core-utils/src/prompt_templates/factory.rs index 10a0d719e9d..dcb98ed17f9 100644 --- a/litellm-rust/crates/core-utils/src/prompt_templates/factory.rs +++ b/litellm-rust/crates/core-utils/src/prompt_templates/factory.rs @@ -17,18 +17,13 @@ pub const EMPTY_TEXT_PLACEHOLDER: &str = "[System: Empty message content sanitised to satisfy protocol]"; #[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq)] -#[strum(serialize_all = "snake_case")] pub enum TurnRole { + #[strum(serialize = "user")] User, + #[strum(serialize = "assistant")] Assistant, } -impl TurnRole { - pub fn as_str(self) -> &'static str { - self.into() - } -} - #[derive(Clone, Debug, PartialEq, Eq)] pub struct Turn { pub role: TurnRole, diff --git a/litellm-rust/crates/gateway-mcp/src/configured.rs b/litellm-rust/crates/gateway-mcp/src/configured.rs index 06a1779a274..bb4f5594f1f 100644 --- a/litellm-rust/crates/gateway-mcp/src/configured.rs +++ b/litellm-rust/crates/gateway-mcp/src/configured.rs @@ -49,8 +49,8 @@ fn info(name: &str, config: &McpServer) -> ServerInfo { let identity = format!( "{name}|{}|{}|{}|{}", config.url.as_ref().map_or("", SecretValue::expose), - config.transport.as_str(), - config.auth_type.map_or("", McpAuth::as_str), + <&'static str>::from(config.transport), + config.auth_type.map_or("", <&'static str>::from), config.alias.as_deref().unwrap_or("") ); ServerInfo { diff --git a/litellm-rust/crates/host-python/src/argument.rs b/litellm-rust/crates/host-python/src/argument.rs index 13214e0e9ed..0b89f94cd74 100644 --- a/litellm-rust/crates/host-python/src/argument.rs +++ b/litellm-rust/crates/host-python/src/argument.rs @@ -2,78 +2,97 @@ use pyo3::{prelude::*, types::PyDict}; pub fn lookup<'py>( kwargs: &Bound<'py, PyDict>, - request: &Bound<'py, PyAny>, + bound: &Bound<'py, PyDict>, name: &str, ) -> PyResult>> { - if let Some(value) = kwargs.get_item(name)? { - return Ok(Some(value)); + match kwargs.get_item(name)? { + Some(value) => Ok(Some(value)), + None => bound.get_item(name), } - if let Ok(bound) = request.cast::() { - return bound.get_item(name); - } - request.getattr_opt(name) +} + +pub fn present<'py>( + kwargs: &Bound<'py, PyDict>, + bound: &Bound<'py, PyDict>, + name: &str, +) -> PyResult>> { + Ok(lookup(kwargs, bound, name)?.filter(|value| !value.is_none())) } #[cfg(test)] mod tests { + use rstest::rstest; + use super::*; - #[test] - fn lookup_prefers_the_keyword_even_when_none_and_falls_back_to_the_request() { + fn dicts<'py>( + py: Python<'py>, + kwargs: &str, + bound: &str, + ) -> (Bound<'py, PyDict>, Bound<'py, PyDict>) { + let eval = |source: &str| { + py.eval(&std::ffi::CString::new(source).unwrap(), None, None) + .unwrap() + .cast_into::() + .unwrap() + }; + (eval(kwargs), eval(bound)) + } + + #[rstest] + #[case::keyword_wins( + "{'api_key': 'keyword'}", + "{'api_key': 'bound'}", + Some(Some("keyword")) + )] + #[case::explicit_none_wins("{'api_key': None}", "{'api_key': 'bound'}", Some(None))] + #[case::bound_fallback("{}", "{'api_key': 'bound'}", Some(Some("bound")))] + #[case::missing("{}", "{}", None)] + fn lookup_prefers_the_keyword_and_falls_back_to_bound( + #[case] kwargs: &str, + #[case] bound: &str, + #[case] expected: Option>, + ) { crate::initialize_python(); Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - c" -key = object() -document = {'type': 'document_url'} -class Request: - api_key = 'from-request' - api_base = 'from-request' - document = document -request = Request() -kwargs = {'api_key': key, 'api_base': None} -", - Some(&locals), - Some(&locals), - ) - .unwrap(); - let item = |name: &str| locals.get_item(name).unwrap().unwrap(); - let kwargs = item("kwargs").cast_into::().unwrap(); - let request = item("request"); - let find = |name: &str| lookup(&kwargs, &request, name).unwrap(); - assert!(find("api_key").unwrap().is(item("key"))); - assert!(find("api_base").unwrap().is_none()); - assert!(find("document").unwrap().is(item("document"))); - assert!(find("model").is_none()); + let (kwargs, bound) = dicts(py, kwargs, bound); + let value = lookup(&kwargs, &bound, "api_key") + .unwrap() + .map(|value| value.extract::>().unwrap()); + assert_eq!(value, expected.map(|value| value.map(str::to_owned))); }); } - #[rstest::rstest] - #[case::prepared_value("{'api_key': 'replacement'}", Some("replacement"))] - #[case::explicit_none("{'api_key': None}", None)] - #[case::bound_fallback("{}", Some("original"))] - fn prepared_mapping_overrides_bound_values( - #[case] source: &str, + #[rstest] + #[case::explicit_none_hides_bound("{'api_key': None}", "{'api_key': 'bound'}", None)] + #[case::bound_none("{}", "{'api_key': None}", None)] + #[case::bound_value("{}", "{'api_key': 'bound'}", Some("bound"))] + fn present_treats_none_as_unset( + #[case] kwargs: &str, + #[case] bound: &str, #[case] expected: Option<&str>, ) { crate::initialize_python(); Python::attach(|py| { + let (kwargs, bound) = dicts(py, kwargs, bound); + let value = present(&kwargs, &bound, "api_key") + .unwrap() + .map(|value| value.extract::().unwrap()); + assert_eq!(value.as_deref(), expected); + }); + } + + #[rstest] + fn lookup_returns_the_callers_object() { + crate::initialize_python(); + Python::attach(|py| { + let document = PyDict::new(py); let bound = PyDict::new(py); - bound.set_item("api_key", "original").unwrap(); - let source = std::ffi::CString::new(source).unwrap(); - let prepared = py - .eval(&source, None, None) - .unwrap() - .cast_into::() - .unwrap(); - let value = lookup(&prepared, bound.as_any(), "api_key") + bound.set_item("document", &document).unwrap(); + let found = lookup(&PyDict::new(py), &bound, "document") .unwrap() .unwrap(); - assert_eq!( - value.extract::>().unwrap().as_deref(), - expected - ); + assert!(found.is(&document)); }); } } diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 00543f64085..323f5b7a79b 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -19,7 +19,7 @@ mod owned; mod runtime; mod services; -pub use argument::lookup; +pub use argument::{lookup, present}; pub use binding::PythonBinding; pub use conversion_cache::{FromPythonCache, ToPythonCache}; pub use driver::{CallOptions, run_call}; diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml index 3d880be49f9..7d64467f2f7 100644 --- a/litellm-rust/crates/http/Cargo.toml +++ b/litellm-rust/crates/http/Cargo.toml @@ -21,6 +21,7 @@ rustls-native-certs.workspace = true tokio-tungstenite.workspace = true serde.workspace = true serde_json.workspace = true +strum.workspace = true thiserror.workspace = true tokio.workspace = true tracing.workspace = true diff --git a/litellm-rust/crates/http/src/tls.rs b/litellm-rust/crates/http/src/tls.rs index c58076607e4..0b1c44bf6fe 100644 --- a/litellm-rust/crates/http/src/tls.rs +++ b/litellm-rust/crates/http/src/tls.rs @@ -42,30 +42,28 @@ impl KeyExchangeGroup { } } -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, strum::EnumString)] +#[strum( + parse_err_ty = Unsupported, + parse_err_fn = unsupported_cipher_token +)] pub enum Tls12CipherSuite { + #[strum(serialize = "ECDHE-ECDSA-AES128-GCM-SHA256")] EcdheEcdsaAes128Gcm, + #[strum(serialize = "ECDHE-ECDSA-AES256-GCM-SHA384")] EcdheEcdsaAes256Gcm, + #[strum(serialize = "ECDHE-ECDSA-CHACHA20-POLY1305")] EcdheEcdsaChacha20, + #[strum(serialize = "ECDHE-RSA-AES128-GCM-SHA256")] EcdheRsaAes128Gcm, + #[strum(serialize = "ECDHE-RSA-AES256-GCM-SHA384")] EcdheRsaAes256Gcm, + #[strum(serialize = "ECDHE-RSA-CHACHA20-POLY1305")] EcdheRsaChacha20, } -impl FromStr for Tls12CipherSuite { - type Err = Unsupported; - - fn from_str(name: &str) -> Result { - match name { - "ECDHE-ECDSA-AES128-GCM-SHA256" => Ok(Self::EcdheEcdsaAes128Gcm), - "ECDHE-ECDSA-AES256-GCM-SHA384" => Ok(Self::EcdheEcdsaAes256Gcm), - "ECDHE-ECDSA-CHACHA20-POLY1305" => Ok(Self::EcdheEcdsaChacha20), - "ECDHE-RSA-AES128-GCM-SHA256" => Ok(Self::EcdheRsaAes128Gcm), - "ECDHE-RSA-AES256-GCM-SHA384" => Ok(Self::EcdheRsaAes256Gcm), - "ECDHE-RSA-CHACHA20-POLY1305" => Ok(Self::EcdheRsaChacha20), - _ => Err(Unsupported::CipherToken(name.to_owned())), - } - } +fn unsupported_cipher_token(name: &str) -> Unsupported { + Unsupported::CipherToken(name.to_owned()) } impl Tls12CipherSuite { @@ -367,6 +365,39 @@ mod tests { ); } + #[rstest] + #[case::ecdhe_ecdsa_aes128( + "ECDHE-ECDSA-AES128-GCM-SHA256", + Tls12CipherSuite::EcdheEcdsaAes128Gcm + )] + #[case::ecdhe_ecdsa_aes256( + "ECDHE-ECDSA-AES256-GCM-SHA384", + Tls12CipherSuite::EcdheEcdsaAes256Gcm + )] + #[case::ecdhe_ecdsa_chacha20( + "ECDHE-ECDSA-CHACHA20-POLY1305", + Tls12CipherSuite::EcdheEcdsaChacha20 + )] + #[case::ecdhe_rsa_aes128("ECDHE-RSA-AES128-GCM-SHA256", Tls12CipherSuite::EcdheRsaAes128Gcm)] + #[case::ecdhe_rsa_aes256("ECDHE-RSA-AES256-GCM-SHA384", Tls12CipherSuite::EcdheRsaAes256Gcm)] + #[case::ecdhe_rsa_chacha20("ECDHE-RSA-CHACHA20-POLY1305", Tls12CipherSuite::EcdheRsaChacha20)] + fn tls12_cipher_suite_parses_each_openssl_name( + #[case] name: &str, + #[case] expected: Tls12CipherSuite, + ) { + assert_eq!(name.parse::().unwrap(), expected); + } + + #[rstest] + #[case::unknown("AES128-SHA")] + #[case::case_sensitive("ecdhe-rsa-aes256-gcm-sha384")] + fn unsupported_cipher_token_keeps_the_verbatim_name(#[case] name: &str) { + assert_eq!( + name.parse::().unwrap_err(), + Unsupported::CipherToken(name.to_owned()) + ); + } + #[test] fn named_suites_are_the_only_tls12_suites_offered_and_tls13_stays() { let tls = ClientConfig::try_from(&config(HttpSettings { diff --git a/litellm-rust/crates/llms-types/src/formats/chat_completions.rs b/litellm-rust/crates/llms-types/src/formats/chat_completions.rs index ca4ffa3c718..8f154dd7fb5 100644 --- a/litellm-rust/crates/llms-types/src/formats/chat_completions.rs +++ b/litellm-rust/crates/llms-types/src/formats/chat_completions.rs @@ -9,41 +9,25 @@ pub use content::{ /// Reasoning effort level accepted or applied by the model. #[macro_rules_attribute::apply(crate::wire_type)] -#[derive(Copy, Eq, IntoStaticStr)] +#[derive(Copy, Eq, IntoStaticStr, strum::EnumString, strum::VariantArray)] #[serde(rename_all = "snake_case")] -#[strum(serialize_all = "snake_case")] pub enum ReasoningEffort { + #[strum(serialize = "none")] None, + #[strum(serialize = "minimal")] Minimal, + #[strum(serialize = "low")] Low, + #[strum(serialize = "medium")] Medium, + #[strum(serialize = "high")] High, + #[strum(serialize = "xhigh")] Xhigh, + #[strum(serialize = "max")] Max, } -impl ReasoningEffort { - pub const ALL: [Self; 7] = [ - Self::None, - Self::Minimal, - Self::Low, - Self::Medium, - Self::High, - Self::Xhigh, - Self::Max, - ]; - - pub fn as_str(self) -> &'static str { - self.into() - } - - pub fn parse(value: &str) -> Option { - Self::ALL - .into_iter() - .find(|effort| effort.as_str() == value) - } -} - #[macro_rules_attribute::apply(crate::wire_type)] #[serde(untagged)] pub enum ChatMessageContent { @@ -195,6 +179,7 @@ pub struct ChatCompletionChunk { #[cfg(test)] mod tests { use rstest::rstest; + use strum::VariantArray; use super::*; @@ -213,10 +198,10 @@ mod tests { ) { assert_eq!( serde_json::to_value(effort).unwrap(), - Value::String(effort.as_str().to_string()) + Value::String(<&'static str>::from(effort).to_string()) ); - assert_eq!(ReasoningEffort::parse(effort.as_str()), Some(effort)); - assert!(ReasoningEffort::ALL.contains(&effort)); + assert_eq!(<&'static str>::from(effort).parse(), Ok(effort)); + assert!(ReasoningEffort::VARIANTS.contains(&effort)); } #[rstest] @@ -224,6 +209,6 @@ mod tests { #[case::uppercase("HIGH")] #[case::empty("")] fn reasoning_effort_parse_rejects(#[case] value: &str) { - assert_eq!(ReasoningEffort::parse(value), None); + assert!(value.parse::().is_err()); } } diff --git a/litellm-rust/crates/llms-types/src/formats/messages/metadata.rs b/litellm-rust/crates/llms-types/src/formats/messages/metadata.rs index 9db13cab498..6e5ebb2b98d 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/metadata.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/metadata.rs @@ -165,10 +165,12 @@ pub struct Safeguard { )] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] #[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))] -#[strum(serialize_all = "snake_case")] pub enum MessageRole { + #[strum(serialize = "user")] User, + #[strum(serialize = "assistant")] Assistant, + #[strum(serialize = "system")] System, #[strum(default, transparent)] Other(String), @@ -205,8 +207,8 @@ impl From for String { )] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] #[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))] -#[strum(serialize_all = "snake_case")] pub enum MessageType { + #[strum(serialize = "message")] Message, #[strum(default, transparent)] Other(String), @@ -243,15 +245,22 @@ impl From for String { )] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] #[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))] -#[strum(serialize_all = "snake_case")] pub enum StopReason { + #[strum(serialize = "end_turn")] EndTurn, + #[strum(serialize = "max_tokens")] MaxTokens, + #[strum(serialize = "stop_sequence")] StopSequence, + #[strum(serialize = "tool_use")] ToolUse, + #[strum(serialize = "refusal")] Refusal, + #[strum(serialize = "compaction")] Compaction, + #[strum(serialize = "pause_turn")] PauseTurn, + #[strum(serialize = "model_context_window_exceeded")] ModelContextWindowExceeded, #[strum(default, transparent)] Other(String), diff --git a/litellm-rust/crates/llms-types/src/formats/messages/request.rs b/litellm-rust/crates/llms-types/src/formats/messages/request.rs index c4a31142fb6..d065f45d41c 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/request.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/request.rs @@ -30,16 +30,24 @@ pub enum MessageContent { )] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] #[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))] -#[strum(serialize_all = "snake_case")] pub enum ContentBlockType { + #[strum(serialize = "text")] Text, + #[strum(serialize = "thinking")] Thinking, + #[strum(serialize = "redacted_thinking")] RedactedThinking, + #[strum(serialize = "tool_use")] ToolUse, + #[strum(serialize = "server_tool_use")] ServerToolUse, + #[strum(serialize = "tool_result")] ToolResult, + #[strum(serialize = "compaction")] Compaction, + #[strum(serialize = "advisor_tool_result")] AdvisorToolResult, + #[strum(serialize = "web_search_tool_result")] WebSearchToolResult, #[strum(default, transparent)] Other(String), @@ -124,23 +132,21 @@ pub struct Message { } #[macro_rules_attribute::apply(crate::wire_type)] -#[derive(Copy, Hash, IntoStaticStr, Eq)] +#[derive(Copy, Hash, IntoStaticStr, Eq, strum::VariantArray)] #[serde(rename_all = "lowercase")] -#[strum(serialize_all = "lowercase")] pub enum EffortLevel { + #[strum(serialize = "low")] Low, + #[strum(serialize = "medium")] Medium, + #[strum(serialize = "high")] High, + #[strum(serialize = "xhigh")] Xhigh, + #[strum(serialize = "max")] Max, } -impl EffortLevel { - pub fn as_str(self) -> &'static str { - self.into() - } -} - impl From for ReasoningEffort { fn from(level: EffortLevel) -> Self { match level { @@ -156,18 +162,13 @@ impl From for ReasoningEffort { #[macro_rules_attribute::apply(crate::wire_type)] #[derive(Copy, IntoStaticStr, Eq)] #[serde(rename_all = "lowercase")] -#[strum(serialize_all = "lowercase")] pub enum Speed { + #[strum(serialize = "fast")] Fast, + #[strum(serialize = "standard")] Standard, } -impl Speed { - pub fn as_str(self) -> &'static str { - self.into() - } -} - /// The tools whose presence changes how the request is sent. Every other tool, custom or /// server, deserializes as `Recognized::Unrecognized` and passes through verbatim. #[macro_rules_attribute::apply(crate::wire_type)] @@ -691,7 +692,10 @@ mod tests { #[rstest] fn speed_names_match_the_wire(#[values(Speed::Fast, Speed::Standard)] speed: Speed) { - assert_eq!(serde_json::to_value(speed).unwrap(), json!(speed.as_str())); + assert_eq!( + serde_json::to_value(speed).unwrap(), + json!(<&'static str>::from(speed)) + ); } #[rstest] @@ -705,10 +709,13 @@ mod tests { )] level: EffortLevel, ) { - assert_eq!(serde_json::to_value(level).unwrap(), json!(level.as_str())); + assert_eq!( + serde_json::to_value(level).unwrap(), + json!(<&'static str>::from(level)) + ); assert_eq!( serde_json::to_value(ReasoningEffort::from(level)).unwrap(), - json!(level.as_str()) + json!(<&'static str>::from(level)) ); } } diff --git a/litellm-rust/crates/llms-types/src/providers/AGENTS.md b/litellm-rust/crates/llms-types/src/providers/AGENTS.md index 69712edbcbb..128d29652d1 100644 --- a/litellm-rust/crates/llms-types/src/providers/AGENTS.md +++ b/litellm-rust/crates/llms-types/src/providers/AGENTS.md @@ -4,7 +4,8 @@ These links describe the provider-specific wire types in this directory. Shared ## anthropic.rs -`AnthropicBeta` and `BetaSet` represent the `anthropic-beta` header values and their wire spelling +`AnthropicBeta`, `BetaSet` and `BetaProvider` represent the `anthropic-beta` header values, their +wire spelling and per-host support - https://platform.claude.com/docs/en/api/beta-headers.md diff --git a/litellm-rust/crates/llms-types/src/providers/anthropic.rs b/litellm-rust/crates/llms-types/src/providers/anthropic.rs index 3e0c4369b11..97bb26a4088 100644 --- a/litellm-rust/crates/llms-types/src/providers/anthropic.rs +++ b/litellm-rust/crates/llms-types/src/providers/anthropic.rs @@ -7,6 +7,39 @@ use std::{ str::FromStr, }; +use strum::VariantArray; + +/// A provider column of `litellm/anthropic_beta_headers_config.json`: which betas a host accepts +/// and under which name. +#[derive( + Clone, + Copy, + Debug, + PartialEq, + Eq, + Hash, + strum::AsRefStr, + strum::Display, + strum::EnumString, + strum::VariantArray, +)] +pub enum BetaProvider { + #[strum(serialize = "anthropic")] + Anthropic, + #[strum(serialize = "azure_ai")] + AzureAi, + #[strum(serialize = "bedrock_converse")] + BedrockConverse, + #[strum(serialize = "bedrock")] + Bedrock, + #[strum(serialize = "bedrock_mantle")] + BedrockMantle, + #[strum(serialize = "vertex_ai")] + VertexAi, + #[strum(serialize = "databricks")] + Databricks, +} + /// One value of the `anthropic-beta` header. Equality, ordering and hashing follow the wire /// string, so a value parsed from a caller's header never disagrees with the matching variant. #[derive(Clone, Debug, strum::AsRefStr, strum::Display, strum::EnumString)] @@ -35,12 +68,68 @@ pub enum AnthropicBeta { PerTurnControl20260701, #[strum(serialize = "dangerous-tool-use-2026-09-03")] DangerousToolUse20260903, + #[strum(serialize = "bash_20241022")] + Bash20241022, + #[strum(serialize = "bash_20250124")] + Bash20250124, + #[strum(serialize = "claude-code-20250219")] + ClaudeCode20250219, + #[strum(serialize = "code-execution-2025-08-25")] + CodeExecution20250825, + #[strum(serialize = "computer-use-2025-01-24")] + ComputerUse20250124, + #[strum(serialize = "computer-use-2025-11-24")] + ComputerUse20251124, + #[strum(serialize = "context-1m-2025-08-07")] + Context1m20250807, + #[strum(serialize = "effort-2025-11-24")] + Effort20251124, + #[strum(serialize = "files-api-2025-04-14")] + FilesApi20250414, + #[strum(serialize = "fine-grained-tool-streaming-2025-05-14")] + FineGrainedToolStreaming20250514, + #[strum(serialize = "interleaved-thinking-2025-05-14")] + InterleavedThinking20250514, + #[strum(serialize = "mcp-client-2025-04-04")] + McpClient20250404, + #[strum(serialize = "mcp-client-2025-11-20")] + McpClient20251120, + #[strum(serialize = "mcp-servers-2025-12-04")] + McpServers20251204, + #[strum(serialize = "mid-conversation-output-config-2026-07-01")] + MidConversationOutputConfig20260701, + #[strum(serialize = "mid-conversation-tool-changes-2026-07-01")] + MidConversationToolChanges20260701, + #[strum(serialize = "output-128k-2025-02-19")] + Output128k20250219, + #[strum(serialize = "prompt-caching-scope-2026-01-05")] + PromptCachingScope20260105, + #[strum(serialize = "skills-2025-10-02")] + Skills20251002, + #[strum(serialize = "structured-output-2024-03-01")] + StructuredOutput20240301, + #[strum(serialize = "text_editor_20241022")] + TextEditor20241022, + #[strum(serialize = "text_editor_20250124")] + TextEditor20250124, + #[strum(serialize = "thinking-binding-controls-2026-08-01")] + ThinkingBindingControls20260801, + #[strum(serialize = "thinking-display-updates-2026-08-18")] + ThinkingDisplayUpdates20260818, + #[strum(serialize = "token-efficient-tools-2025-02-19")] + TokenEfficientTools20250219, + #[strum(serialize = "tool-examples-2025-10-29")] + ToolExamples20251029, + #[strum(serialize = "tool-search-tool-2025-10-19")] + ToolSearchTool20251019, + #[strum(serialize = "inline-tools-2026-09-15")] + InlineTools20260915, #[strum(default, transparent)] Other(String), } impl AnthropicBeta { - pub const KNOWN: [Self; 12] = [ + pub const KNOWN: [Self; 40] = [ Self::Oauth20250420, Self::WebFetch20250910, Self::WebSearch20250305, @@ -53,11 +142,153 @@ impl AnthropicBeta { Self::AdvisorTool20260301, Self::PerTurnControl20260701, Self::DangerousToolUse20260903, + Self::Bash20241022, + Self::Bash20250124, + Self::ClaudeCode20250219, + Self::CodeExecution20250825, + Self::ComputerUse20250124, + Self::ComputerUse20251124, + Self::Context1m20250807, + Self::Effort20251124, + Self::FilesApi20250414, + Self::FineGrainedToolStreaming20250514, + Self::InterleavedThinking20250514, + Self::McpClient20250404, + Self::McpClient20251120, + Self::McpServers20251204, + Self::MidConversationOutputConfig20260701, + Self::MidConversationToolChanges20260701, + Self::Output128k20250219, + Self::PromptCachingScope20260105, + Self::Skills20251002, + Self::StructuredOutput20240301, + Self::TextEditor20241022, + Self::TextEditor20250124, + Self::ThinkingBindingControls20260801, + Self::ThinkingDisplayUpdates20260818, + Self::TokenEfficientTools20250219, + Self::ToolExamples20251029, + Self::ToolSearchTool20251019, + Self::InlineTools20260915, ]; pub fn as_str(&self) -> &str { self.as_ref() } + + pub fn on(&self, provider: BetaProvider) -> Option { + let resolved = match self { + Self::Other(raw) => raw.parse().unwrap(), + _ => self.clone(), + }; + if resolved.rejected_by().contains(&provider) { + return None; + } + Some(match (resolved, provider) { + ( + Self::AdvancedToolUse20251120, + BetaProvider::Bedrock | BetaProvider::BedrockMantle | BetaProvider::VertexAi, + ) => Self::ToolSearchTool20251019, + (beta, _) => beta, + }) + } + + fn rejected_by(&self) -> &'static [BetaProvider] { + match self { + Self::AdvancedToolUse20251120 | Self::ContextManagement20250627 => { + &[BetaProvider::BedrockConverse] + } + Self::AdvisorTool20260301 | Self::Compact20260904 => &[ + BetaProvider::AzureAi, + BetaProvider::BedrockConverse, + BetaProvider::Bedrock, + BetaProvider::BedrockMantle, + BetaProvider::VertexAi, + BetaProvider::Databricks, + ], + Self::Bash20241022 + | Self::Bash20250124 + | Self::McpServers20251204 + | Self::StructuredOutput20240301 + | Self::TextEditor20241022 + | Self::TextEditor20250124 => BetaProvider::VARIANTS, + Self::ClaudeCode20250219 | Self::ToolExamples20251029 => &[ + BetaProvider::Anthropic, + BetaProvider::AzureAi, + BetaProvider::BedrockConverse, + BetaProvider::Bedrock, + BetaProvider::VertexAi, + BetaProvider::Databricks, + ], + Self::CodeExecution20250825 + | Self::FilesApi20250414 + | Self::McpClient20250404 + | Self::McpClient20251120 + | Self::PromptCachingScope20260105 + | Self::Skills20251002 + | Self::WebFetch20250910 => &[ + BetaProvider::BedrockConverse, + BetaProvider::Bedrock, + BetaProvider::BedrockMantle, + BetaProvider::VertexAi, + ], + Self::Compact20260112 => &[BetaProvider::AzureAi, BetaProvider::BedrockConverse], + Self::ComputerUse20250124 | Self::ComputerUse20251124 | Self::Context1m20250807 => &[], + Self::DangerousToolUse20260903 => { + &[BetaProvider::BedrockConverse, BetaProvider::Databricks] + } + Self::Effort20251124 => &[BetaProvider::VertexAi], + Self::FastMode20260201 | Self::Oauth20250420 => &[ + BetaProvider::AzureAi, + BetaProvider::BedrockConverse, + BetaProvider::Bedrock, + BetaProvider::BedrockMantle, + BetaProvider::VertexAi, + ], + Self::FineGrainedToolStreaming20250514 => { + &[BetaProvider::AzureAi, BetaProvider::VertexAi] + } + Self::InterleavedThinking20250514 | Self::WebSearch20250305 => { + &[BetaProvider::BedrockConverse, BetaProvider::Bedrock] + } + Self::MidConversationOutputConfig20260701 + | Self::MidConversationToolChanges20260701 => &[ + BetaProvider::AzureAi, + BetaProvider::BedrockConverse, + BetaProvider::BedrockMantle, + BetaProvider::VertexAi, + BetaProvider::Databricks, + ], + Self::Output128k20250219 | Self::TokenEfficientTools20250219 => &[ + BetaProvider::AzureAi, + BetaProvider::BedrockConverse, + BetaProvider::Bedrock, + BetaProvider::VertexAi, + ], + Self::PerTurnControl20260701 => &[ + BetaProvider::BedrockConverse, + BetaProvider::Bedrock, + BetaProvider::Databricks, + ], + Self::StructuredOutputs20251113 => &[BetaProvider::Bedrock, BetaProvider::VertexAi], + Self::ThinkingBindingControls20260801 => &[BetaProvider::AzureAi], + Self::ThinkingDisplayUpdates20260818 => &[BetaProvider::Databricks], + Self::ToolSearchTool20251019 => &[ + BetaProvider::Anthropic, + BetaProvider::AzureAi, + BetaProvider::BedrockConverse, + BetaProvider::Databricks, + ], + Self::InlineTools20260915 => &[ + BetaProvider::AzureAi, + BetaProvider::BedrockConverse, + BetaProvider::Bedrock, + BetaProvider::BedrockMantle, + BetaProvider::Databricks, + ], + Self::Other(_) => BetaProvider::VARIANTS, + } + } } impl PartialEq for AnthropicBeta { @@ -149,10 +380,20 @@ impl fmt::Display for BetaSet { #[cfg(test)] mod tests { + use std::collections::BTreeSet; + + use indexmap::IndexMap; use rstest::rstest; use super::*; + fn beta_headers_config() -> IndexMap { + serde_json::from_str(include_str!( + "../../../../../litellm/anthropic_beta_headers_config.json" + )) + .unwrap() + } + fn set(header: &str) -> BetaSet { header.parse().unwrap_or_else(|never| match never {}) } @@ -171,7 +412,35 @@ mod tests { AnthropicBeta::FastMode20260201, AnthropicBeta::AdvisorTool20260301, AnthropicBeta::PerTurnControl20260701, - AnthropicBeta::DangerousToolUse20260903 + AnthropicBeta::DangerousToolUse20260903, + AnthropicBeta::Bash20241022, + AnthropicBeta::Bash20250124, + AnthropicBeta::ClaudeCode20250219, + AnthropicBeta::CodeExecution20250825, + AnthropicBeta::ComputerUse20250124, + AnthropicBeta::ComputerUse20251124, + AnthropicBeta::Context1m20250807, + AnthropicBeta::Effort20251124, + AnthropicBeta::FilesApi20250414, + AnthropicBeta::FineGrainedToolStreaming20250514, + AnthropicBeta::InterleavedThinking20250514, + AnthropicBeta::McpClient20250404, + AnthropicBeta::McpClient20251120, + AnthropicBeta::McpServers20251204, + AnthropicBeta::MidConversationOutputConfig20260701, + AnthropicBeta::MidConversationToolChanges20260701, + AnthropicBeta::Output128k20250219, + AnthropicBeta::PromptCachingScope20260105, + AnthropicBeta::Skills20251002, + AnthropicBeta::StructuredOutput20240301, + AnthropicBeta::TextEditor20241022, + AnthropicBeta::TextEditor20250124, + AnthropicBeta::ThinkingBindingControls20260801, + AnthropicBeta::ThinkingDisplayUpdates20260818, + AnthropicBeta::TokenEfficientTools20250219, + AnthropicBeta::ToolExamples20251029, + AnthropicBeta::ToolSearchTool20251019, + AnthropicBeta::InlineTools20260915 )] beta: AnthropicBeta, ) { @@ -181,6 +450,119 @@ mod tests { assert!(AnthropicBeta::KNOWN.contains(&beta)); } + #[rstest] + fn known_betas_are_exactly_the_config_keys() { + let config = beta_headers_config(); + let config_keys: BTreeSet = config + .iter() + .filter(|(provider, _)| provider.as_str() != "description") + .flat_map(|(_, column)| { + column + .as_object() + .into_iter() + .flat_map(|values| values.keys().cloned()) + }) + .collect(); + let known_keys: BTreeSet = AnthropicBeta::KNOWN + .iter() + .map(|beta| beta.as_str().to_string()) + .collect(); + assert_eq!(known_keys, config_keys); + } + + #[rstest] + #[case::anthropic(BetaProvider::Anthropic)] + #[case::azure_ai(BetaProvider::AzureAi)] + #[case::bedrock_converse(BetaProvider::BedrockConverse)] + #[case::bedrock(BetaProvider::Bedrock)] + #[case::bedrock_mantle(BetaProvider::BedrockMantle)] + #[case::vertex_ai(BetaProvider::VertexAi)] + #[case::databricks(BetaProvider::Databricks)] + fn provider_columns_match_beta_provider_variants(#[case] provider: BetaProvider) { + let config = beta_headers_config(); + let config_columns: Vec = config + .keys() + .filter(|column| column.as_str() != "description") + .cloned() + .collect(); + let beta_provider_columns: Vec = BetaProvider::VARIANTS + .iter() + .map(ToString::to_string) + .collect(); + assert_eq!(config_columns, beta_provider_columns); + assert_eq!( + provider.to_string().parse::().unwrap(), + provider + ); + } + + #[rstest] + fn on_matches_every_config_cell() { + let config = beta_headers_config(); + for provider in BetaProvider::VARIANTS { + for beta in &AnthropicBeta::KNOWN { + let expected = config + .get(provider.as_ref()) + .and_then(serde_json::Value::as_object) + .and_then(|column| column.get(beta.as_str())) + .and_then(serde_json::Value::as_str) + .map(str::to_string); + let actual = beta.on(*provider).map(|name| name.to_string()); + assert_eq!(actual, expected, "provider {provider}, beta {beta}"); + if matches!(beta, AnthropicBeta::AdvancedToolUse20251120) + && matches!( + provider, + BetaProvider::Bedrock + | BetaProvider::BedrockMantle + | BetaProvider::VertexAi + ) + { + let parsed: AnthropicBeta = actual.as_deref().unwrap().parse().unwrap(); + assert!(AnthropicBeta::KNOWN.contains(&parsed)); + assert!(!matches!(parsed, AnthropicBeta::Other(_))); + } + } + } + } + + #[rstest] + #[case::renamed( + "advanced-tool-use-2025-11-20", + BetaProvider::Bedrock, + Some(AnthropicBeta::ToolSearchTool20251019) + )] + #[case::kept( + "oauth-2025-04-20", + BetaProvider::Anthropic, + Some(AnthropicBeta::Oauth20250420) + )] + #[case::rejected("effort-2025-11-24", BetaProvider::VertexAi, None)] + #[case::unknown("example-beta-2099-01-01", BetaProvider::Anthropic, None)] + fn on_resolves_known_spellings_held_as_other( + #[case] raw: &str, + #[case] provider: BetaProvider, + #[case] expected: Option, + ) { + let actual = AnthropicBeta::Other(raw.to_string()).on(provider); + assert_eq!(actual, expected); + assert!(!matches!(actual, Some(AnthropicBeta::Other(_)))); + } + + #[rstest] + #[case::anthropic(BetaProvider::Anthropic)] + #[case::azure_ai(BetaProvider::AzureAi)] + #[case::bedrock_converse(BetaProvider::BedrockConverse)] + #[case::bedrock(BetaProvider::Bedrock)] + #[case::bedrock_mantle(BetaProvider::BedrockMantle)] + #[case::vertex_ai(BetaProvider::VertexAi)] + #[case::databricks(BetaProvider::Databricks)] + fn unknown_beta_is_rejected_by_every_provider(#[case] provider: BetaProvider) { + assert_eq!( + AnthropicBeta::Other("example-beta-2099-01-01".into()).on(provider), + None + ); + } + #[test] fn unknown_values_are_kept_verbatim() { let parsed: AnthropicBeta = "claude-code-20250219".parse().unwrap(); diff --git a/litellm-rust/crates/llms/src/anthropic/beta_headers.rs b/litellm-rust/crates/llms/src/anthropic/beta_headers.rs new file mode 100644 index 00000000000..3b83b12d727 --- /dev/null +++ b/litellm-rust/crates/llms/src/anthropic/beta_headers.rs @@ -0,0 +1,36 @@ +use litellm_http::request::{with_header, without_headers}; +use litellm_llms_types::providers::anthropic::{BetaProvider, BetaSet}; + +use crate::{anthropic::common_utils::existing_betas, base_llm::auth::Headers}; + +const BETA_HEADER: &str = "anthropic-beta"; + +/// How the `anthropic-beta` header is treated before a request leaves for a host. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum BetaPolicy { + /// Send the header untouched and let the host decide. + Forward, + /// Keep only the betas the provider accepts, under the names it expects. + Filter(BetaProvider), + /// Remove every `anthropic-beta` header. + Drop, +} + +impl BetaPolicy { + pub fn apply(self, headers: Headers) -> Headers { + let provider = match self { + Self::Forward => return headers, + Self::Drop => return without_headers(headers, &[BETA_HEADER]), + Self::Filter(provider) => provider, + }; + let accepted: BetaSet = existing_betas(&headers) + .iter() + .filter_map(|beta| beta.on(provider)) + .collect(); + let headers = without_headers(headers, &[BETA_HEADER]); + if accepted.is_empty() { + return headers; + } + with_header(headers, BETA_HEADER, accepted.to_string()) + } +} diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index ecc6cfaf83d..e1fd1956167 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -241,7 +241,7 @@ fn anthropic_body( .iter() .map(|turn| { json!({ - "role": turn.role.as_str(), + "role": <&'static str>::from(turn.role), "content": turn.texts.iter().map(|text| text_block(text)).collect::>(), }) }) diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 6e2fa3b8785..ccf3a6d15f2 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -596,17 +596,10 @@ mod tests { use crate::base_llm::messages::context::SupportedEffortTiers; use rstest::{fixture, rstest}; use serde_json::json; + use strum::VariantArray; use super::*; - const ALL_LEVELS: [EffortLevel; 5] = [ - EffortLevel::Low, - EffortLevel::Medium, - EffortLevel::High, - EffortLevel::Xhigh, - EffortLevel::Max, - ]; - fn apply(sanitizer: fn(Vec) -> Vec, messages: Value) -> Value { let parsed: Vec = serde_json::from_value(messages).unwrap(); serde_json::to_value(sanitizer(parsed)).unwrap() @@ -1712,7 +1705,10 @@ mod tests { ..unmapped }; assert_eq!( - ALL_LEVELS.map(|level| supports_effort_tier(&capabilities, level)), + EffortLevel::VARIANTS + .iter() + .map(|level| supports_effort_tier(&capabilities, *level)) + .collect::>(), expected ); } diff --git a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs index 2162c39c229..df1ae13fd28 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs @@ -10,6 +10,7 @@ use litellm_llms_types::{ }; use litellm_python_compat::{json::from_json, repr::repr, truthy::truthy}; use serde_json::Value; +use strum::VariantArray; use crate::base_llm::messages::context::{ MessagesModelCapabilities, ThinkingBudgets, ThinkingContext, @@ -54,14 +55,17 @@ fn unmapped_effort(effort: &Value) -> Error { Error::InvalidRequest(crate::ErrorDetail::InvalidChoice { field: "reasoning effort", actual: repr(&from_json(effort.clone())), - choices: ReasoningEffort::ALL.map(|effort| effort.as_str()).into(), + choices: ReasoningEffort::VARIANTS + .iter() + .map(|effort| <&'static str>::from(*effort)) + .collect(), }) } fn unsupported_effort(level: EffortLevel, model: &str) -> Error { Error::InvalidRequest(crate::ErrorDetail::UnsupportedValue { field: "effort", - value: level.as_str(), + value: <&'static str>::from(level), model: model.into(), }) } @@ -134,7 +138,7 @@ fn legacy_reasoning_effort( Some(Recognized::Known(level)) => Ok((*level).into()), Some(Recognized::Unrecognized(value)) if truthy(&from_json(value.clone())) => value .as_str() - .and_then(ReasoningEffort::parse) + .and_then(|text| text.parse().ok()) .ok_or_else(|| unmapped_effort(value)), None | Some(Recognized::Unrecognized(_)) => Ok(ReasoningEffort::Medium), } diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index b9cf6c37272..74ac35752ab 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -270,7 +270,7 @@ fn drop_unsupported_params( fn speed_text(speed: &Recognized) -> String { match speed { - Recognized::Known(speed) => speed.as_str().to_string(), + Recognized::Known(speed) => <&'static str>::from(*speed).to_string(), Recognized::Unrecognized(Value::String(text)) => text.clone(), Recognized::Unrecognized(other) => other.to_string(), } diff --git a/litellm-rust/crates/llms/src/anthropic/mod.rs b/litellm-rust/crates/llms/src/anthropic/mod.rs index a884c146dca..72c390cfe21 100644 --- a/litellm-rust/crates/llms/src/anthropic/mod.rs +++ b/litellm-rust/crates/llms/src/anthropic/mod.rs @@ -1,3 +1,4 @@ +pub mod beta_headers; pub mod common_utils; pub mod batches; 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 90f5b97322f..4139ee860cf 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 @@ -23,9 +23,11 @@ const HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAA /// Textract has operations rather than models; the model slot of /// `aws_textract/` names the one to call. #[derive(Clone, Copy, Debug, EnumString, IntoStaticStr, VariantNames, PartialEq, Eq)] -#[strum(serialize_all = "kebab-case", ascii_case_insensitive)] +#[strum(ascii_case_insensitive)] pub enum TextractOperation { + #[strum(serialize = "detect-document-text")] DetectDocumentText, + #[strum(serialize = "analyze-document")] AnalyzeDocument, } 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 cee906e77a4..c0b2c7aa3b7 100644 --- a/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs +++ b/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs @@ -30,20 +30,17 @@ pub struct BedrockAudioTranscriptionConfig; #[derive(Clone, Copy, Deserialize, IntoStaticStr)] #[serde(rename_all = "lowercase")] -#[strum(serialize_all = "lowercase")] enum AudioFormat { + #[strum(serialize = "wav")] Wav, + #[strum(serialize = "mp3")] Mp3, + #[strum(serialize = "flac")] Flac, + #[strum(serialize = "ogg")] Ogg, } -impl AudioFormat { - fn as_str(self) -> &'static str { - self.into() - } -} - struct AudioInput { data: String, format: AudioFormat, @@ -118,7 +115,7 @@ impl BaseAudioTranscriptionConfig for BedrockAudioTranscriptionConfig { "messages": [{ "role": "user", "content": [ - {"audio": {"format": audio.format.as_str(), "source": {"bytes": audio.data}}}, + {"audio": {"format": <&'static str>::from(audio.format), "source": {"bytes": audio.data}}}, {"text": instruction} ] }], 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 3a13e388a4b..bcd1077b830 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -385,7 +385,7 @@ fn converse_body(conversation: &Conversation, optional_params: &Map::from(turn.role), "content": turn.texts.iter().map(|text| json!({"text": text})).collect::>(), }) }) diff --git a/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs b/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs index ad7025ab6af..373ad4ca4f5 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs @@ -64,27 +64,23 @@ pub fn invoke_anthropic_event_stream(bytes: ByteStream) -> EventStream { Box::pin(invoke_chunk_stream(bytes).map(|chunk| decode_invoke_anthropic_chunk(chunk?))) } -#[derive(Clone, Copy)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString)] enum InvokeProvider { + #[strum(serialize = "anthropic")] Anthropic, + #[strum(serialize = "deepseek_r1")] DeepseekR1, + #[strum(serialize = "moonshot")] Moonshot, + #[strum(disabled)] Unsupported, } -impl From<&str> for InvokeProvider { - fn from(value: &str) -> Self { - match value { - "anthropic" => Self::Anthropic, - "deepseek_r1" => Self::DeepseekR1, - "moonshot" => Self::Moonshot, - _ => Self::Unsupported, - } - } -} - pub fn invoke_chat_stream(invoke_provider: &str, shape: StreamShape) -> Result { - match InvokeProvider::from(invoke_provider) { + match invoke_provider + .parse() + .unwrap_or(InvokeProvider::Unsupported) + { InvokeProvider::Anthropic => Ok(ChatStream::new( invoke_anthropic_event_stream, ModelResponseIterator::new(shape), @@ -116,6 +112,22 @@ mod tests { const TEXT_DELTA: &str = r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#; + #[rstest::rstest] + #[case::anthropic("anthropic", InvokeProvider::Anthropic)] + #[case::deepseek_r1("deepseek_r1", InvokeProvider::DeepseekR1)] + #[case::moonshot("moonshot", InvokeProvider::Moonshot)] + #[case::unknown("qwen", InvokeProvider::Unsupported)] + #[case::case_sensitive("Anthropic", InvokeProvider::Unsupported)] + fn invoke_provider_maps_each_model_family_name( + #[case] name: &str, + #[case] expected: InvokeProvider, + ) { + assert_eq!( + name.parse().unwrap_or(InvokeProvider::Unsupported), + expected + ); + } + fn aws_wire(chunk: &str) -> Vec { let payload = serde_json::json!({"bytes": STANDARD.encode(chunk)}); let message = Message::new(Bytes::from(serde_json::to_vec(&payload).unwrap())).add_header( diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index 67b81ec6feb..dadab413fe4 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -25,11 +25,13 @@ const COHERE_API_KEY_ENV: &str = "COHERE_API_KEY"; const COHERE_PARSE_HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC"; -#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)] +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, strum::IntoStaticStr)] #[serde(rename_all = "lowercase")] pub enum OutputFormat { #[default] + #[strum(serialize = "markdown")] Markdown, + #[strum(serialize = "blocks")] Blocks, } @@ -267,11 +269,7 @@ fn build_request(model: &str, image_url: String, params: &CohereOptions) -> Cohe CohereRequest { model: model.into(), document: CohereParseDocument::ImageUrl { image_url }, - output_format: match params.output_format.unwrap_or_default() { - OutputFormat::Markdown => "markdown", - OutputFormat::Blocks => "blocks", - } - .into(), + output_format: <&'static str>::from(params.output_format.unwrap_or_default()).into(), } } diff --git a/litellm-rust/crates/llms/tests/anthropic_beta_headers.rs b/litellm-rust/crates/llms/tests/anthropic_beta_headers.rs new file mode 100644 index 00000000000..422c3d4b5f0 --- /dev/null +++ b/litellm-rust/crates/llms/tests/anthropic_beta_headers.rs @@ -0,0 +1,111 @@ +use litellm_llms::anthropic::beta_headers::BetaPolicy; +use litellm_llms_types::providers::anthropic::BetaProvider; +use rstest::rstest; + +fn headers(pairs: &[(&str, &str)]) -> Vec<(String, String)> { + pairs + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() +} + +#[rstest] +fn forward_keeps_headers_as_is() { + let input = headers(&[ + ( + "Anthropic-Beta", + "example-beta-2099-01-01,web-search-2025-03-05", + ), + ("x-api-key", "k"), + ]); + assert_eq!(BetaPolicy::Forward.apply(input.clone()), input); +} + +#[rstest] +fn drop_removes_every_casing() { + assert_eq!( + BetaPolicy::Drop.apply(headers(&[ + ("anthropic-beta", "a"), + ("x-api-key", "k"), + ("ANTHROPIC-BETA", "b"), + ])), + headers(&[("x-api-key", "k")]) + ); +} + +#[rstest] +#[case::bedrock(BetaProvider::Bedrock, Some("tool-search-tool-2025-10-19"))] +#[case::bedrock_mantle(BetaProvider::BedrockMantle, Some("tool-search-tool-2025-10-19"))] +#[case::vertex_ai(BetaProvider::VertexAi, Some("tool-search-tool-2025-10-19"))] +#[case::anthropic(BetaProvider::Anthropic, Some("advanced-tool-use-2025-11-20"))] +#[case::azure_ai(BetaProvider::AzureAi, Some("advanced-tool-use-2025-11-20"))] +#[case::databricks(BetaProvider::Databricks, Some("advanced-tool-use-2025-11-20"))] +#[case::bedrock_converse(BetaProvider::BedrockConverse, None)] +fn filter_renames_per_host(#[case] provider: BetaProvider, #[case] expected_beta: Option<&str>) { + let expected = match expected_beta { + Some(beta) => headers(&[("x-api-key", "k"), ("anthropic-beta", beta)]), + None => headers(&[("x-api-key", "k")]), + }; + assert_eq!( + BetaPolicy::Filter(provider).apply(headers(&[ + ("Anthropic-Beta", "advanced-tool-use-2025-11-20"), + ("x-api-key", "k"), + ])), + expected + ); +} + +#[rstest] +#[case::azure_ai(BetaProvider::AzureAi, "fast-mode-2026-02-01,example-beta-2099-01-01")] +#[case::anthropic(BetaProvider::Anthropic, "bash_20241022")] +#[case::vertex_ai(BetaProvider::VertexAi, " , ")] +fn filter_removes_header_when_nothing_survives( + #[case] provider: BetaProvider, + #[case] beta_value: &str, +) { + assert_eq!( + BetaPolicy::Filter(provider).apply(headers(&[ + ("x-api-key", "k"), + ("anthropic-beta", beta_value) + ])), + headers(&[("x-api-key", "k")]) + ); +} + +#[rstest] +#[case::anthropic( + BetaProvider::Anthropic, + "web-search-2025-03-05", + "oauth-2025-04-20,web-search-2025-03-05,example-beta-2099-01-01", + "oauth-2025-04-20,web-search-2025-03-05" +)] +#[case::vertex_ai_renames_and_deduplicates( + BetaProvider::VertexAi, + "advanced-tool-use-2025-11-20", + "tool-search-tool-2025-10-19", + "tool-search-tool-2025-10-19" +)] +fn filter_deduplicates_and_sorts( + #[case] provider: BetaProvider, + #[case] first_beta: &str, + #[case] second_beta: &str, + #[case] expected_beta: &str, +) { + assert_eq!( + BetaPolicy::Filter(provider).apply(headers(&[ + ("ANTHROPIC-BETA", first_beta), + ("x-api-key", "k"), + ("anthropic-beta", second_beta), + ])), + headers(&[("x-api-key", "k"), ("anthropic-beta", expected_beta)]) + ); +} + +#[rstest] +#[case::forward(BetaPolicy::Forward)] +#[case::drop(BetaPolicy::Drop)] +#[case::filter(BetaPolicy::Filter(BetaProvider::Bedrock))] +fn headers_without_a_beta_are_untouched(#[case] policy: BetaPolicy) { + let input = headers(&[("x-api-key", "k"), ("anthropic-version", "2023")]); + assert_eq!(policy.apply(input.clone()), input); +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index cf1080a3881..92acf796e6b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -17,6 +17,10 @@ mod tokenizer; #[pymodule(gil_used = true)] mod _native { + #[pymodule_export] + use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking}; + use pyo3::{prelude::*, types::PyModule}; + #[cfg(feature = "panic-test")] #[pymodule_export] use crate::diagnostics::_panic_for_test; @@ -29,9 +33,7 @@ mod _native { #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; #[pymodule_export] - use crate::routes::chat_completions::{ - achat_completions, acompletion, chat_completions, completion, - }; + use crate::routes::chat_completions::{acompletion, completion}; #[pymodule_export] use crate::routes::embeddings::{aembedding, embedding}; #[pymodule_export] @@ -51,9 +53,6 @@ mod _native { use crate::tokenizer::HuggingFaceEncoding; #[pymodule_export] use crate::tokenizer::Tokenizer; - #[pymodule_export] - use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking}; - use pyo3::{prelude::*, types::PyModule}; #[pymodule_init] fn init(module: &Bound<'_, PyModule>) -> PyResult<()> { @@ -101,8 +100,6 @@ mod tests { "atranscription", "messages", "amessages", - "chat_completions", - "achat_completions", "completion", "acompletion", "responses", diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index df15a9591a1..4a7efeffdd2 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -19,13 +19,6 @@ pub(crate) struct RouteOptions { pub(crate) timeout: Option, } -pub(crate) fn messages_argument(value: &Bound<'_, PyAny>) -> PyResult> { - match from_py_argument(value)? { - Value::Array(values) => Ok(values), - _ => Err(PyValueError::new_err("messages must be a list")), - } -} - fn required_object(name: &'static str, value: Value) -> PyResult> { match value { Value::Object(values) => Ok(values), @@ -449,18 +442,6 @@ module.__getattr__ = fail fn argument_converters_keep_nested_values_and_accept_explicit_none() { Python::initialize(); Python::attach(|py| { - let messages = py - .eval( - c"[{'role': 'user', 'content': [{'type': 'text', 'text': 'hi'}]}]", - None, - None, - ) - .unwrap(); - assert_eq!( - Value::Array(messages_argument(&messages).unwrap()), - json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) - ); - let params = py.eval(c"{'temperature': 0.2}", None, None).unwrap(); assert_eq!( optional_object("optional_params", ¶ms).unwrap(), diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 8f49444f730..99afc772b99 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -6,13 +6,16 @@ use crate::coercion::{FieldSpec, ProjectionError}; const MODULE: &str = "litellm.rust_bridge.settings"; #[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq)] -#[strum(serialize_all = "snake_case")] pub(crate) enum PythonSettings { #[strum(serialize = "http_settings")] Http, + #[strum(serialize = "url_policy")] UrlPolicy, + #[strum(serialize = "provider_defaults")] ProviderDefaults, + #[strum(serialize = "secret_manager")] SecretManager, + #[strum(serialize = "secret_manager_binding")] SecretManagerBinding, } @@ -23,17 +26,16 @@ pub(crate) struct Snapshot<'py> { impl Snapshot<'_> { pub(crate) fn read(&self, spec: &FieldSpec) -> Result { - spec.read(&self.value, self.group.name()) + spec.read(&self.value, self.group.into()) } } impl PythonSettings { - pub(crate) fn name(self) -> &'static str { - self.into() - } - pub(crate) fn read(self, py: Python<'_>) -> PyResult> { - let value = py.import(MODULE)?.getattr(self.name())?.call0()?; + let value = py + .import(MODULE)? + .getattr(<&'static str>::from(self))? + .call0()?; Ok(Snapshot { group: self, value }) } diff --git a/litellm-rust/crates/python-bridge/src/routes/AGENTS.md b/litellm-rust/crates/python-bridge/src/routes/AGENTS.md index c2afed49b45..948c67eb647 100644 --- a/litellm-rust/crates/python-bridge/src/routes/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/routes/AGENTS.md @@ -11,3 +11,7 @@ The host driver owns sequencing and terminal events; the bridge supplies fallibl 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 Regression tests must observe that an unstarted call does no setup, hook and preflight rewrites affect resource configuration, setup failures reach the selected failure handler once, and provider work is not replayed. Retain existing read-point and object-identity guarantees while changing setup timing + +## Layout + +`messages/` is the reference shape. Each lifecycle route is a folder: `mod.rs` holds `run_` and the thin sync and async `#[pyfunction]` entrypoints that take `NativeCall` and call it, `host.rs` holds the `PythonBinding` host, and anything else route-specific (projection, error mapping, extra pyclasses) gets its own sibling file. Scaffolded routes without a lifecycle (embeddings, audio transcription) stay a single `.rs` diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index fc7bfb08b75..4c0c2760875 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -5,6 +5,8 @@ use litellm_inference_transcription::{ use pyo3::prelude::*; use serde_json::{Map, Value}; +use super::NativeCall; + use crate::{ errors::route_error_to_pyerr, marshal::{RouteOptions, optional_object_field, required_field, value_route_options}, @@ -40,8 +42,7 @@ async fn execute( } #[pyfunction] -pub(crate) fn transcription(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult> { - let call = super::NativeCall::extract(&call)?; +pub(crate) fn transcription(py: Python<'_>, call: NativeCall<'_>) -> PyResult> { let audio: Value = litellm_host_python::from_py_argument(&required_field(&call.bound, "audio")?)?; let options = value_route_options(&call.bound)?; @@ -59,9 +60,8 @@ pub(crate) fn transcription(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult< #[pyfunction] pub(crate) fn atranscription<'py>( py: Python<'py>, - call: Bound<'py, PyAny>, + call: NativeCall<'py>, ) -> PyResult> { - let call = super::NativeCall::extract(&call)?; let audio: Value = litellm_host_python::from_py_argument(&required_field(&call.bound, "audio")?)?; let options = value_route_options(&call.bound)?; diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs deleted file mode 100644 index 27f076b7cce..00000000000 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ /dev/null @@ -1,150 +0,0 @@ -mod host; - -use pyo3::types::{PyDict, PyTuple}; - -use crate::execution::{run_async, run_sync}; -use litellm_inference_chat::{ChatCompletionsRoute, Error, types::ChatCompletionsRequest}; -use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse; -use pyo3::prelude::*; -use serde_json::{Map, Value}; - -use crate::{ - errors::route_error_to_pyerr, - marshal::{ - RouteOptions, messages_argument, optional_object_field, required_field, value_route_options, - }, -}; - -async fn execute( - http: Result, - secrets: std::sync::Arc, - messages: Vec, - optional_params: Map, - options: RouteOptions, -) -> Result { - let RouteOptions { - model, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout, - } = options; - ChatCompletionsRoute::new(http?, crate::http::resources().auth.clone(), secrets) - .execute( - ChatCompletionsRequest { - model: &model, - messages: Value::Array(messages), - optional_params, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - }, - &(), - None, - ) - .await -} - -#[pyfunction] -pub(crate) fn chat_completions(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult> { - let call = super::NativeCall::extract(&call)?; - let messages: Vec = messages_argument(&required_field(&call.bound, "messages")?)?; - let optional_params = - optional_object_field(&call.bound, "optional_params")?.unwrap_or_default(); - let options = value_route_options(&call.bound)?; - let http = crate::http::provider_client(py, &call.kwargs, false)?; - let secrets = crate::secrets::source(py)?; - run_sync( - py, - execute(http, secrets, messages, optional_params, options), - route_error_to_pyerr, - ) -} - -#[pyfunction] -pub(crate) fn achat_completions<'py>( - py: Python<'py>, - call: Bound<'py, PyAny>, -) -> PyResult> { - let call = super::NativeCall::extract(&call)?; - let messages: Vec = messages_argument(&required_field(&call.bound, "messages")?)?; - let optional_params = - optional_object_field(&call.bound, "optional_params")?.unwrap_or_default(); - let options = value_route_options(&call.bound)?; - let http = crate::http::provider_client(py, &call.kwargs, true)?; - let secrets = crate::secrets::source(py)?; - run_async( - py, - execute(http, secrets, messages, optional_params, options), - route_error_to_pyerr, - ) -} - -fn run_public( - py: Python<'_>, - request: Bound<'_, PyAny>, - args: Bound<'_, PyTuple>, - kwargs: Bound<'_, PyDict>, - asynchronous: bool, -) -> PyResult> { - use super::inference::InferenceHost; - use litellm_callbacks_legacy_python::LoggingOperation; - let host = InferenceHost::new( - request.clone().unbind(), - "litellm.rust_bridge.chat_completions.route_host", - ); - let cache_call_type = if asynchronous { - "acompletion" - } else { - "completion" - }; - crate::cache::admit_native(py, &kwargs, cache_call_type)?; - let (arguments, hooks) = crate::routes::call_hooks( - py, - LoggingOperation::Completion, - &request, - &args, - &kwargs, - asynchronous, - )?; - crate::routes::run_public_call( - py, - arguments, - move |py, arguments, request| { - let route = ChatCompletionsRoute::new( - crate::http::provider_client(py, arguments, asynchronous)? - .map_err(crate::http::client_error)?, - crate::http::resources().auth.clone(), - crate::secrets::source(py)?, - ); - let (cache, cache_options) = - crate::cache::configured_native(py, arguments, cache_call_type)?; - let route = match cache { - Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new( - cache, - litellm_cache_response::CacheScope::Shared, - )), - None => route, - }; - Ok(route.machine(request, cache_options.policy)) - }, - host::ChatCompletionsPythonHost(host), - hooks, - asynchronous, - ) -} - -#[pyfunction] -pub(crate) fn completion(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult> { - let call = super::NativeCall::extract(&call)?; - run_public(py, call.bound.into_any(), call.args, call.kwargs, false) -} - -#[pyfunction] -pub(crate) fn acompletion(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult> { - let call = super::NativeCall::extract(&call)?; - run_public(py, call.bound.into_any(), call.args, call.kwargs, true) -} 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 e7525cd3a91..1547a4daa93 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 @@ -1,6 +1,6 @@ use std::convert::Infallible; -use super::super::inference::InferenceHost; +use crate::routes::inference::InferenceHost; use litellm_host_python::{InvokeError, PythonBinding, PythonHostCalls, PythonOwned}; use litellm_inference_chat::{Error, route::ChatCompletions, types::ChatCompletionsCall}; use pyo3::{ 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 new file mode 100644 index 00000000000..1146436bfe6 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs @@ -0,0 +1,71 @@ +mod host; + +use std::sync::Arc; + +use host::ChatCompletionsPythonHost; +use litellm_auth::AuthServices; +use litellm_cache_response::{CachePolicy, ScopedCache}; +use litellm_callbacks_legacy_python::LoggingOperation; +use litellm_host::{call::HostedMachine, protocol::Protocol}; +use litellm_inference_chat::{ChatCompletionsRoute, route::ChatCompletions}; +use litellm_secrets::source::SecretSource; +use pyo3::prelude::*; + +use super::{ + NativeCall, + inference::{InferenceHost, InferenceRoute, run_inference}, +}; + +fn run_chat_completions( + py: Python<'_>, + call: NativeCall<'_>, + asynchronous: bool, +) -> PyResult> { + let host = InferenceHost::new( + call.bound.clone().unbind(), + "litellm.rust_bridge.chat_completions.route_host", + ); + run_inference::( + py, + call, + asynchronous, + ChatCompletionsPythonHost(host), + ) +} + +#[pyfunction] +pub(crate) fn completion(py: Python<'_>, call: NativeCall<'_>) -> PyResult> { + run_chat_completions(py, call, false) +} + +#[pyfunction] +pub(crate) fn acompletion(py: Python<'_>, call: NativeCall<'_>) -> PyResult> { + run_chat_completions(py, call, true) +} + +impl InferenceRoute for ChatCompletionsRoute { + type Protocol = ChatCompletions; + const OPERATION: LoggingOperation = LoggingOperation::Completion; + const SYNC_CALL_TYPE: &'static str = "completion"; + const ASYNC_CALL_TYPE: &'static str = "acompletion"; + + fn new( + http: litellm_http::Client, + auth: Arc, + secrets: Arc, + ) -> Self { + Self::new(http, auth, secrets) + } + + fn with_cache(self, cache: ScopedCache) -> Self { + self.with_cache(cache) + } + + fn machine( + self, + call: ::Request, + policy: CachePolicy, + ) -> HostedMachine { + self.machine(call, policy) + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/embeddings.rs b/litellm-rust/crates/python-bridge/src/routes/embeddings.rs index 1d34e3c21ff..2b184274041 100644 --- a/litellm-rust/crates/python-bridge/src/routes/embeddings.rs +++ b/litellm-rust/crates/python-bridge/src/routes/embeddings.rs @@ -1,17 +1,18 @@ use pyo3::prelude::*; +use super::NativeCall; use crate::errors::RustBridgeDeclined; #[pyfunction] -pub(crate) fn embedding(call: Bound<'_, PyAny>) -> PyResult> { - drop(super::NativeCall::extract(&call)?); +pub(crate) fn embedding(call: NativeCall<'_>) -> PyResult> { + drop(call); Err(RustBridgeDeclined::new_err( "native embeddings route is not implemented", )) } #[pyfunction] -pub(crate) fn aembedding(call: Bound<'_, PyAny>) -> PyResult> { +pub(crate) fn aembedding(call: NativeCall<'_>) -> PyResult> { embedding(call) } @@ -36,7 +37,8 @@ call = SimpleNamespace(args=(), kwargs={}, bound={'model':'test-model','input':' Some(&locals), ) .unwrap(); - let call = locals.get_item("call").unwrap().unwrap(); + let call: super::NativeCall<'_> = + locals.get_item("call").unwrap().unwrap().extract().unwrap(); let error = if asynchronous { super::aembedding(call) } else { diff --git a/litellm-rust/crates/python-bridge/src/routes/inference.rs b/litellm-rust/crates/python-bridge/src/routes/inference.rs index 28025b80b0f..0333aaeb04a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/inference.rs +++ b/litellm-rust/crates/python-bridge/src/routes/inference.rs @@ -1,11 +1,19 @@ +use std::sync::Arc; + +use litellm_auth::AuthServices; +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_python::{from_py, lookup}; +use litellm_host::{call::HostedMachine, protocol::Protocol}; +use litellm_host_python::{PythonBinding, PythonHostCalls, from_py, present}; use litellm_http::transport::Error as TransportError; use litellm_inference::RouteError; +use litellm_secrets::source::SecretSource; use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; use serde::Serialize; use serde_json::{Map, Value}; +use super::NativeCall; use crate::{ errors::{RustUpstreamError, route_error_to_pyerr}, marshal::{ @@ -15,7 +23,7 @@ use crate::{ }; pub(super) struct InferenceHost { - pub request: Py, + pub request: Py, module: &'static str, } @@ -26,7 +34,7 @@ pub(super) struct ProjectedCall { } impl InferenceHost { - pub fn new(request: Py, module: &'static str) -> Self { + pub fn new(request: Py, module: &'static str) -> Self { Self { request, module } } @@ -85,21 +93,7 @@ impl InferenceHost { arguments: &Bound<'py, PyDict>, name: &str, ) -> PyResult>> { - let request = self.request.bind(py); - if let Some(value) = lookup(arguments, request, name)? { - return Ok((!value.is_none()).then_some(value)); - } - if request.is_instance_of::() { - return Ok(None); - } - let parameter = request - .getattr("parameters")? - .call_method1("get", (name,))?; - if !parameter.is_none() { - return Ok(Some(parameter)); - } - let extra = request.getattr("kwargs")?.call_method1("get", (name,))?; - Ok((!extra.is_none()).then_some(extra)) + present(arguments, self.request.bind(py), name) } pub fn parameters( @@ -140,3 +134,62 @@ impl InferenceHost { Ok(PyErr::from_value(mapped)) } } + +pub(super) trait InferenceRoute: Sized + 'static { + type Protocol: Protocol; + const OPERATION: LoggingOperation; + const SYNC_CALL_TYPE: &'static str; + const ASYNC_CALL_TYPE: &'static str; + + fn new( + http: litellm_http::Client, + auth: Arc, + secrets: Arc, + ) -> Self; + fn with_cache(self, cache: ScopedCache) -> Self; + fn machine( + self, + call: ::Request, + policy: CachePolicy, + ) -> HostedMachine; +} + +pub(super) fn run_inference( + py: Python<'_>, + call: NativeCall<'_>, + asynchronous: bool, + host: H, +) -> PyResult> +where + R: InferenceRoute, + H: PythonBinding + PythonHostCalls + 'static, +{ + let call_type = if asynchronous { + R::ASYNC_CALL_TYPE + } else { + R::SYNC_CALL_TYPE + }; + crate::cache::admit_native(py, &call.kwargs, call_type)?; + let (arguments, hooks) = super::call_hooks(py, R::OPERATION, &call, asynchronous)?; + super::run_public_call( + py, + arguments, + move |py, arguments, request| { + let route = R::new( + crate::http::provider_client(py, arguments, asynchronous)? + .map_err(crate::http::client_error)?, + crate::http::resources().auth.clone(), + crate::secrets::source(py)?, + ); + let (cache, cache_options) = crate::cache::configured_native(py, arguments, call_type)?; + let route = match cache { + Some(cache) => route.with_cache(ScopedCache::new(cache, CacheScope::Shared)), + None => route, + }; + Ok(route.machine(request, cache_options.policy)) + }, + host, + hooks, + asynchronous, + ) +} 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 bde1cc7269c..841ef3081e3 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, lookup, to_py}; +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, @@ -89,12 +89,12 @@ 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, + request: Py, cache: PythonCache, } impl MessagesPythonHost { - pub(super) fn new(request: Py, asynchronous: bool) -> Self { + pub(super) fn new(request: Py, asynchronous: bool) -> Self { Self { request, cache: PythonCache::new(asynchronous), @@ -107,9 +107,7 @@ impl MessagesPythonHost { arguments: &Bound<'_, PyDict>, ) -> PyResult> { let request = self.request.bind(py); - let argument = |name: &str| -> PyResult>> { - Ok(lookup(arguments, request, name)?.filter(|value| !value.is_none())) - }; + let argument = |name: &str| present(arguments, request, name); let string = |name: &str| -> PyResult> { argument(name)?.map(|value| value.extract()).transpose() }; @@ -153,8 +151,7 @@ impl MessagesPythonHost { ) -> PyResult>> { let request = self.request.bind(py); let mapping = |name: &str| -> PyResult>> { - lookup(arguments, request, name)? - .filter(|value| !value.is_none()) + present(arguments, request, name)? .map(|value| from_py(&value)) .transpose() }; @@ -169,8 +166,7 @@ impl MessagesPythonHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> PyResult> { - lookup(arguments, self.request.bind(py), "provider_specific_header")? - .filter(|value| !value.is_none()) + present(arguments, self.request.bind(py), "provider_specific_header")? .map(|value| from_py(&value)) .transpose() } @@ -200,9 +196,9 @@ impl MessagesPythonHost { self.request .bind(py) .get_item("custom_llm_provider") - .and_then(|value| value.extract::>()) .ok() .flatten() + .and_then(|value| value.extract::>().ok().flatten()) .unwrap_or_else(|| "anthropic".into()) } 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 318aae27121..6ed6656cd21 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -2,26 +2,13 @@ mod host; use host::MessagesPythonHost; use litellm_callbacks_legacy_python::LoggingOperation; -use pyo3::{ - prelude::*, - types::{PyDict, PyTuple}, -}; +use pyo3::prelude::*; -fn run_messages( - py: Python<'_>, - request: Bound<'_, PyAny>, - args: Bound<'_, PyTuple>, - kwargs: Bound<'_, PyDict>, - asynchronous: bool, -) -> PyResult> { - let (arguments, hooks) = crate::routes::call_hooks( - py, - LoggingOperation::Messages, - &request, - &args, - &kwargs, - asynchronous, - )?; +use super::NativeCall; + +fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult> { + let (arguments, hooks) = + crate::routes::call_hooks(py, LoggingOperation::Messages, &call, asynchronous)?; crate::routes::run_public_call( py, arguments, @@ -57,20 +44,18 @@ fn run_messages( }, )) }, - MessagesPythonHost::new(request.unbind(), asynchronous), + MessagesPythonHost::new(call.bound.unbind(), asynchronous), hooks, asynchronous, ) } #[pyfunction] -pub(crate) fn messages(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult> { - let call = super::NativeCall::extract(&call)?; - run_messages(py, call.bound.into_any(), call.args, call.kwargs, false) +pub(crate) fn messages(py: Python<'_>, call: NativeCall<'_>) -> PyResult> { + run_messages(py, call, false) } #[pyfunction] -pub(crate) fn amessages(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult> { - let call = super::NativeCall::extract(&call)?; - run_messages(py, call.bound.into_any(), call.args, call.kwargs, true) +pub(crate) fn amessages(py: Python<'_>, call: NativeCall<'_>) -> PyResult> { + run_messages(py, call, true) } diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index bf52fe3bfe8..646b7799730 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -8,8 +8,7 @@ pub(crate) mod responses; pub(crate) mod token_counter; pub(crate) mod traces; -use litellm_callbacks_legacy_python::LoggingOperation; -use litellm_callbacks_legacy_python::{LegacyLogging, PublicCall}; +use litellm_callbacks_legacy_python::{LegacyLogging, LoggingOperation, PublicCall}; use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol}; use litellm_host_python::{HookChain, PythonBinding, PythonCallHooks, PythonHostCalls}; use pyo3::{ @@ -17,14 +16,16 @@ use pyo3::{ types::{PyDict, PyMapping, PyTuple}, }; -struct NativeCall<'py> { +pub(crate) struct NativeCall<'py> { args: Bound<'py, PyTuple>, kwargs: Bound<'py, PyDict>, bound: Bound<'py, PyDict>, } -impl<'py> NativeCall<'py> { - fn extract(call: &Bound<'py, PyAny>) -> PyResult { +impl<'py> FromPyObject<'_, 'py> for NativeCall<'py> { + type Error = PyErr; + + fn extract(call: Borrowed<'_, 'py, PyAny>) -> PyResult { Ok(Self { args: call.getattr("args")?.cast_into()?, kwargs: mapping_dict(&call.getattr("kwargs")?)?, @@ -46,12 +47,10 @@ fn mapping_dict<'py>(value: &Bound<'py, PyAny>) -> PyResult> fn call_hooks( py: Python<'_>, operation: LoggingOperation, - request: &Bound<'_, PyAny>, - args: &Bound<'_, PyTuple>, - kwargs: &Bound<'_, PyDict>, + call: &NativeCall<'_>, asynchronous: bool, ) -> PyResult<(Py, impl PythonCallHooks + use<>)> { - let call = PublicCall::capture(request, args, kwargs)?; + let call = PublicCall::capture(&call.bound, &call.args, &call.kwargs)?; let arguments = call.arguments(py); Ok(( arguments, @@ -99,6 +98,7 @@ mod tests { prelude::*, types::{PyDict, PyList}, }; + use rstest::rstest; fn value_call<'py>( py: Python<'py>, @@ -126,8 +126,10 @@ mod tests { .unwrap() } - #[test] - fn route_arguments_that_fail_to_convert_raise_value_error() { + #[rstest] + #[case::sync("transcription")] + #[case::asynchronous("atranscription")] + fn route_arguments_that_fail_to_convert_raise_value_error(#[case] name: &str) { Python::initialize(); Python::attach(|py| { let module = crate::native_module(py); @@ -151,48 +153,24 @@ value = Broken() .expect("locals should be readable") .expect("helper value should exist"); - for name in ["chat_completions", "achat_completions"] { - let error = module - .getattr(name) - .and_then(|function| { - function.call1((value_call(py, "messages", &broken, None),)) - }) - .expect_err("route should reject a value it cannot convert"); + let error = module + .getattr(name) + .and_then(|function| function.call1((value_call(py, "audio", &broken, None),))) + .expect_err("route should reject a value it cannot convert"); - assert!( - error.is_instance_of::(py), - "{name} surfaced {error} instead of ValueError" - ); - } + assert!( + error.is_instance_of::(py), + "{name} surfaced {error} instead of ValueError" + ); }); } - #[test] + #[rstest] fn sync_and_async_routes_apply_the_same_input_validation() { Python::initialize(); Python::attach(|py| { let module = crate::native_module(py); - let invalid_messages = PyDict::new(py); - let sync_chat_error = module - .getattr("chat_completions") - .and_then(|function| { - function.call1((value_call(py, "messages", &invalid_messages, None),)) - }) - .expect_err("sync chat should reject a non-list messages value"); - let async_chat_error = module - .getattr("achat_completions") - .and_then(|function| { - function.call1((value_call(py, "messages", &invalid_messages, None),)) - }) - .expect_err("async chat should reject a non-list messages value"); - - assert_eq!( - sync_chat_error.to_string(), - "ValueError: messages must be a list" - ); - assert_eq!(async_chat_error.to_string(), sync_chat_error.to_string()); - let invalid_headers = PyList::empty(py); let kwargs = PyDict::new(py); kwargs @@ -221,51 +199,43 @@ value = Broken() }); } - #[test] + #[rstest] + fn missing_and_explicit_none_optional_params_share_the_next_error() { + Python::initialize(); + Python::attach(|py| { + let module = crate::native_module(py); + let audio = PyDict::new(py); + let transcribe = |optional_params: Option>| { + let kwargs = PyDict::new(py); + if let Some(optional_params) = optional_params { + kwargs.set_item("optional_params", optional_params).unwrap(); + } + module + .getattr("transcription") + .and_then(|function| { + function.call1((value_call(py, "audio", &audio, Some(&kwargs)),)) + }) + .expect_err("model 'model' has no provider") + .to_string() + }; + + let omitted = transcribe(None); + assert_eq!(transcribe(Some(py.None().into_bound(py))), omitted); + assert_ne!(omitted, "ValueError: optional_params must be a dict"); + assert_eq!( + transcribe(Some(PyList::empty(py).into_any())), + "ValueError: optional_params must be a dict" + ); + }); + } + + #[rstest] fn route_input_validation_preserves_left_to_right_order() { Python::initialize(); Python::attach(|py| { let module = crate::native_module(py); let invalid = PyList::empty(py); - let chat_kwargs = PyDict::new(py); - chat_kwargs - .set_item("optional_params", &invalid) - .expect("kwargs should accept optional_params"); - chat_kwargs - .set_item("extra_headers", &invalid) - .expect("kwargs should accept extra_headers"); - let invalid_messages = PyDict::new(py); - let error = module - .getattr("chat_completions") - .and_then(|function| { - function.call1((value_call( - py, - "messages", - &invalid_messages, - Some(&chat_kwargs), - ),)) - }) - .expect_err("messages should be validated first"); - assert_eq!(error.to_string(), "ValueError: messages must be a list"); - - let valid_messages = PyList::empty(py); - let error = module - .getattr("chat_completions") - .and_then(|function| { - function.call1((value_call( - py, - "messages", - &valid_messages, - Some(&chat_kwargs), - ),)) - }) - .expect_err("optional_params should be validated before headers"); - assert_eq!( - error.to_string(), - "ValueError: optional_params must be a dict" - ); - let headers_kwargs = PyDict::new(py); headers_kwargs .set_item("extra_headers", &invalid) @@ -286,43 +256,4 @@ value = Broken() assert!(!error.to_string().contains("extra_headers")); }); } - - #[test] - fn missing_and_explicit_none_optional_params_share_the_next_error() { - Python::initialize(); - Python::attach(|py| { - let module = crate::native_module(py); - let messages = PyList::empty(py); - let headers = PyList::empty(py); - let omitted = PyDict::new(py); - omitted - .set_item("extra_headers", &headers) - .expect("kwargs should accept extra_headers"); - let explicit = PyDict::new(py); - explicit - .set_item("optional_params", py.None()) - .expect("kwargs should accept optional_params"); - explicit - .set_item("extra_headers", &headers) - .expect("kwargs should accept extra_headers"); - - let omitted_error = module - .getattr("chat_completions") - .and_then(|function| { - function.call1((value_call(py, "messages", &messages, Some(&omitted)),)) - }) - .expect_err("omitted optional_params should reach header validation"); - let explicit_error = module - .getattr("chat_completions") - .and_then(|function| { - function.call1((value_call(py, "messages", &messages, Some(&explicit)),)) - }) - .expect_err("None optional_params should reach header validation"); - assert_eq!( - omitted_error.to_string(), - "ValueError: extra_headers must be a dict" - ); - assert_eq!(explicit_error.to_string(), omitted_error.to_string()); - }); - } } 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 7ddaba70937..ca683600d71 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -27,12 +27,12 @@ enum OcrHostData { /// document as it goes), acquires Azure AD tokens, and builds the public response and /// exception. pub(super) struct OcrPythonHost { - request: Py, + request: Py, data: OcrHostData, } impl OcrPythonHost { - pub(super) fn new(request: Py) -> Self { + pub(super) fn new(request: Py) -> Self { Self { request, data: OcrHostData::Unprojected, @@ -209,7 +209,7 @@ del provider .unwrap() .cast_into::() .unwrap(); - let mut host = OcrPythonHost::new(py.None()); + let mut host = OcrPythonHost::new(PyDict::new(py).unbind()); 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 71b1054fc3c..4dcb2f4c53e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -9,10 +9,9 @@ use litellm_core_utils::settings::ProcessEnvironment; use litellm_host_python::to_py; use litellm_inference_ocr::provider_config; use litellm_llms::base_llm::ocr::settings::OcrSettings; -use pyo3::{ - prelude::*, - types::{PyDict, PyTuple}, -}; +use pyo3::prelude::*; + +use super::NativeCall; use crate::{ coercion::FieldSpec, @@ -29,21 +28,9 @@ const ENABLE_AZURE_AD_TOKEN_REFRESH: FieldSpec = Ok(field.exact_true()) }); -fn run_ocr( - py: Python<'_>, - request: Bound<'_, PyAny>, - args: Bound<'_, PyTuple>, - kwargs: Bound<'_, PyDict>, - asynchronous: bool, -) -> PyResult> { - let (arguments, hooks) = crate::routes::call_hooks( - py, - LoggingOperation::Ocr, - &request, - &args, - &kwargs, - asynchronous, - )?; +fn run_ocr(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult> { + let (arguments, hooks) = + crate::routes::call_hooks(py, LoggingOperation::Ocr, &call, asynchronous)?; crate::routes::run_public_call( py, arguments, @@ -61,7 +48,7 @@ fn run_ocr( let route = litellm_inference_ocr::OcrRoute::new(client); Ok(route.machine(request, None)) }, - OcrPythonHost::new(request.unbind()), + OcrPythonHost::new(call.bound.unbind()), hooks, asynchronous, ) @@ -81,15 +68,13 @@ fn project_provider_defaults(snapshot: &Snapshot<'_>) -> PyResult { } #[pyfunction] -pub(crate) fn ocr(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult> { - let call = super::NativeCall::extract(&call)?; - run_ocr(py, call.bound.into_any(), call.args, call.kwargs, false) +pub(crate) fn ocr(py: Python<'_>, call: NativeCall<'_>) -> PyResult> { + run_ocr(py, call, false) } #[pyfunction] -pub(crate) fn aocr(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult> { - let call = super::NativeCall::extract(&call)?; - run_ocr(py, call.bound.into_any(), call.args, call.kwargs, true) +pub(crate) fn aocr(py: Python<'_>, call: NativeCall<'_>) -> PyResult> { + run_ocr(py, call, true) } #[pyfunction] 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 ee7f6238510..7ae5a1f3ce2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -1,5 +1,5 @@ use litellm_auth::SecretValue; -use litellm_host_python::from_py; +use litellm_host_python::{from_py, present}; use litellm_inference_ocr::{ types::{LiteLLMOcrRequest, OcrDocumentInput}, wire::{OcrWireRequest, consumed_optional_params, decode_document, decode_request_input}, @@ -22,13 +22,13 @@ pub(super) struct OcrHostHandles { } struct OcrArguments<'a, 'py> { - request: &'a Bound<'py, PyAny>, + bound: &'a Bound<'py, PyDict>, kwargs: &'a Bound<'py, PyDict>, } impl<'py> OcrArguments<'_, 'py> { fn lookup(&self, name: &str) -> PyResult> { - litellm_host_python::lookup(self.kwargs, self.request, name)? + litellm_host_python::lookup(self.kwargs, self.bound, name)? .ok_or_else(|| PyValueError::new_err(format!("missing argument: {name}"))) } @@ -58,7 +58,7 @@ impl<'py> OcrArguments<'_, 'py> { fn extra_headers(&self) -> PyResult>> { self.lookup("extra_headers")? .extract::>>()? - .map(|value| from_py(value.bind(self.request.py()))) + .map(|value| from_py(value.bind(self.bound.py()))) .transpose() } @@ -66,7 +66,7 @@ impl<'py> OcrArguments<'_, 'py> { Ok(self .lookup("timeout")? .extract::>>()? - .map(|value| python_timeout_seconds(self.request.py(), value)) + .map(|value| python_timeout_seconds(self.bound.py(), value)) .transpose()? .flatten()) } @@ -110,10 +110,10 @@ impl ProjectedDocument { } pub(super) fn project_request( - request: &Bound<'_, PyAny>, + bound: &Bound<'_, PyDict>, kwargs: &Bound<'_, PyDict>, ) -> PyResult<(LiteLLMOcrRequest, OcrHostHandles)> { - let arguments = OcrArguments { request, kwargs }; + let arguments = OcrArguments { bound, kwargs }; let model = arguments.model()?; let custom_llm_provider = arguments.custom_llm_provider()?; let document = ProjectedDocument::project(&arguments.document()?)?; @@ -122,7 +122,7 @@ pub(super) fn project_request( .map_err(ocr_error_to_pyerr)?; let names = specs.iter().map(|spec| spec.name).collect::>(); let optional_params = - project_optional_fields(names.iter().copied(), |name| kwargs.get_item(name))?; + project_optional_fields(names.iter().copied(), |name| present(kwargs, bound, name))?; let input_sources = request_input_sources( kwargs, names @@ -136,7 +136,7 @@ pub(super) fn project_request( let timeout_seconds = arguments.timeout_seconds()?; let wire = OcrWireRequest { model, - document: document.resolve(request.py())?, + document: document.resolve(bound.py())?, api_key, api_base, custom_llm_provider, @@ -169,11 +169,20 @@ mod tests { locals } + fn dict<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyDict> { + locals + .get_item(name) + .unwrap() + .unwrap() + .cast_into::() + .unwrap() + } + fn arguments<'a, 'py>( - request: &'a Bound<'py, PyAny>, + bound: &'a Bound<'py, PyDict>, kwargs: &'a Bound<'py, PyDict>, ) -> OcrArguments<'a, 'py> { - OcrArguments { request, kwargs } + OcrArguments { bound, kwargs } } fn project_document(document: &Bound<'_, PyAny>) -> PyResult { @@ -204,141 +213,26 @@ sys.modules['litellm.rust_bridge.timeouts'] = timeouts } #[test] - fn kwargs_override_request_attributes_including_explicit_none() { + fn kwargs_override_bound_values_including_explicit_none() { Python::initialize(); Python::attach(|py| { let locals = eval( py, c" -class Request: - def __init__(self): - self.accesses = [] - def __getattribute__(self, name): - if name != 'accesses': - object.__getattribute__(self, 'accesses').append(name) - return object.__getattribute__(self, name) -request = Request() -request.model = 'from-request' -request.custom_llm_provider = 'mistral' +bound = {'model': 'from-bound', 'custom_llm_provider': 'mistral'} kwargs = {'model': 'from-kwargs', 'custom_llm_provider': None} ", ); - let request = locals.get_item("request").unwrap().unwrap(); - let kwargs = locals - .get_item("kwargs") - .unwrap() - .unwrap() - .cast_into::() - .unwrap(); - let arguments = arguments(&request, &kwargs); + let (bound, kwargs) = (dict(&locals, "bound"), dict(&locals, "kwargs")); + let arguments = arguments(&bound, &kwargs); assert_eq!(arguments.model().unwrap(), "from-kwargs"); assert_eq!(arguments.custom_llm_provider().unwrap(), None); - let accesses: Vec = request.getattr("accesses").unwrap().extract().unwrap(); - assert_eq!(accesses, Vec::::new()); }); } - #[test] - fn missing_kwargs_read_the_request_property_once() { - Python::initialize(); - Python::attach(|py| { - let locals = eval( - py, - c" -class Request: - def __init__(self): - self.reads = 0 - @property - def model(self): - self.reads += 1 - return 'mistral-ocr-latest' -request = Request() -kwargs = {} -", - ); - let request = locals.get_item("request").unwrap().unwrap(); - let kwargs = locals - .get_item("kwargs") - .unwrap() - .unwrap() - .cast_into::() - .unwrap(); - assert_eq!( - arguments(&request, &kwargs).model().unwrap(), - "mistral-ocr-latest" - ); - assert_eq!( - request.getattr("reads").unwrap().extract::().unwrap(), - 1 - ); - }); - } - - #[test] - fn request_property_exceptions_keep_their_identity() { - Python::initialize(); - Python::attach(|py| { - let locals = eval( - py, - c" -failure = LookupError('model failed') -class Request: - @property - def model(self): - raise failure -request = Request() -kwargs = {} -", - ); - let request = locals.get_item("request").unwrap().unwrap(); - let kwargs = locals - .get_item("kwargs") - .unwrap() - .unwrap() - .cast_into::() - .unwrap(); - let error = arguments(&request, &kwargs).model().unwrap_err(); - assert!( - error - .value(py) - .is(locals.get_item("failure").unwrap().unwrap()) - ); - }); - } - - #[test] - fn unused_raising_property_is_never_inspected() { - Python::initialize(); - Python::attach(|py| { - let locals = eval( - py, - c" -class Request: - @property - def unused(self): - raise RuntimeError('unused') - model = 'mistral-ocr-latest' - custom_llm_provider = None -request = Request() -kwargs = {} -", - ); - let request = locals.get_item("request").unwrap().unwrap(); - let kwargs = locals - .get_item("kwargs") - .unwrap() - .unwrap() - .cast_into::() - .unwrap(); - let arguments = arguments(&request, &kwargs); - assert_eq!(arguments.model().unwrap(), "mistral-ocr-latest"); - assert_eq!(arguments.custom_llm_provider().unwrap(), None); - }); - } - - /// A reader that rewrites the request while it runs shows which arguments projection - /// read before it and which after: every other argument is read first, and the read - /// happens exactly once. + /// A reader that rewrites the bound arguments while it runs shows which arguments + /// projection read before it and which after: every other argument is read first, and + /// the read happens exactly once. #[test] fn document_readers_are_read_once_after_every_other_argument() { Python::initialize(); @@ -347,37 +241,28 @@ kwargs = {} let locals = eval( py, c" -class Request: - model = 'mistral/mistral-ocr-latest' - custom_llm_provider = None - api_key = None - api_base = 'https://original.example.com' - extra_headers = {'x-source': 'original'} - timeout = 1 - @property - def document(self): - return document class Reader: reads = 0 def read(self): Reader.reads += 1 - Request.api_base = 'https://mutated.example.com' - Request.extra_headers = {'x-source': 'mutated'} - Request.timeout = 9 + bound['api_base'] = 'https://mutated.example.com' + bound['extra_headers'] = {'x-source': 'mutated'} + bound['timeout'] = 9 return b'abc' -document = {'type': 'file', 'file': Reader(), 'mime_type': 'application/pdf'} -request = Request() +bound = { + 'model': 'mistral/mistral-ocr-latest', + 'custom_llm_provider': None, + 'api_key': None, + 'api_base': 'https://original.example.com', + 'extra_headers': {'x-source': 'original'}, + 'timeout': 1, + 'document': {'type': 'file', 'file': Reader(), 'mime_type': 'application/pdf'}, +} kwargs = {} ", ); - let request = locals.get_item("request").unwrap().unwrap(); - let kwargs = locals - .get_item("kwargs") - .unwrap() - .unwrap() - .cast_into::() - .unwrap(); - let (projected, _) = project_request(&request, &kwargs).unwrap(); + let (projected, _) = + project_request(&dict(&locals, "bound"), &dict(&locals, "kwargs")).unwrap(); assert_eq!( py.eval(c"Reader.reads", Some(&locals), Some(&locals)) .unwrap() @@ -524,34 +409,26 @@ document = Document() }); } - fn request_and_kwargs<'py>( + fn bound_and_kwargs<'py>( py: Python<'py>, kwargs: &std::ffi::CStr, - ) -> (Bound<'py, PyAny>, Bound<'py, PyDict>) { + ) -> (Bound<'py, PyDict>, Bound<'py, PyDict>) { let locals = eval( py, c" -class Request: - model = 'mistral/mistral-ocr-latest' - custom_llm_provider = 'mistral' - document = {'type': 'document_url', 'document_url': 'https://example.com/request.pdf'} - api_key = None - api_base = 'https://request.example.com' - extra_headers = {'x-source': 'request'} - timeout = 1 -request = Request() +bound = { + 'model': 'mistral/mistral-ocr-latest', + 'custom_llm_provider': 'mistral', + 'document': {'type': 'document_url', 'document_url': 'https://example.com/bound.pdf'}, + 'api_key': None, + 'api_base': 'https://bound.example.com', + 'extra_headers': {'x-source': 'bound'}, + 'timeout': 1, +} ", ); py.run(kwargs, Some(&locals), Some(&locals)).unwrap(); - ( - locals.get_item("request").unwrap().unwrap(), - locals - .get_item("kwargs") - .unwrap() - .unwrap() - .cast_into::() - .unwrap(), - ) + (dict(&locals, "bound"), dict(&locals, "kwargs")) } #[test] @@ -559,7 +436,7 @@ request = Request() Python::initialize(); Python::attach(|py| { stub_timeout_conversion(py); - let (request, kwargs) = request_and_kwargs( + let (bound, kwargs) = bound_and_kwargs( py, c" kwargs = { @@ -575,7 +452,7 @@ kwargs = { } ", ); - let (projected, _) = project_request(&request, &kwargs).unwrap(); + let (projected, _) = project_request(&bound, &kwargs).unwrap(); assert_eq!( projected.optional_params.keys().collect::>(), ["pages"] @@ -584,12 +461,29 @@ kwargs = { }); } + #[rstest::rstest] + #[case::explicit_none_is_unset(c"bound['pages'] = [1]\nkwargs = {'pages': None}", None)] + #[case::bound_fallback(c"bound['pages'] = [1]\nkwargs = {}", Some(serde_json::json!([1])))] + #[case::keyword_wins(c"bound['pages'] = [1]\nkwargs = {'pages': [0]}", Some(serde_json::json!([0])))] + fn optional_params_read_through_bound_and_drop_none( + #[case] script: &std::ffi::CStr, + #[case] expected: Option, + ) { + Python::initialize(); + Python::attach(|py| { + stub_timeout_conversion(py); + let (bound, kwargs) = bound_and_kwargs(py, script); + let (projected, _) = project_request(&bound, &kwargs).unwrap(); + assert_eq!(projected.optional_params.get("pages").cloned(), expected); + }); + } + #[test] fn replacement_kwargs_project_provider_connection_and_timeout() { Python::initialize(); Python::attach(|py| { stub_timeout_conversion(py); - let (request, kwargs) = request_and_kwargs( + let (bound, kwargs) = bound_and_kwargs( py, c" kwargs = { @@ -602,7 +496,7 @@ kwargs = { } ", ); - let (projected, handles) = project_request(&request, &kwargs).unwrap(); + let (projected, handles) = project_request(&bound, &kwargs).unwrap(); assert_eq!(handles.provider, "azure_ai"); assert_eq!(projected.model, "mistral-ocr-latest"); assert_eq!( 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 25d3fc07332..b22bbc6192b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs @@ -1,6 +1,6 @@ use std::convert::Infallible; -use super::super::inference::InferenceHost; +use crate::routes::inference::InferenceHost; use litellm_host_python::{InvokeError, PythonBinding, PythonHostCalls, PythonOwned}; use litellm_inference_responses::{Error, route::Responses, types::ResponsesCall}; use pyo3::{ diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs new file mode 100644 index 00000000000..a5b81b91752 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs @@ -0,0 +1,102 @@ +mod host; +mod websocket; + +use std::sync::Arc; + +use host::ResponsesPythonHost; +use litellm_auth::AuthServices; +use litellm_cache_response::{CachePolicy, ScopedCache}; +use litellm_callbacks_legacy_python::LoggingOperation; +use litellm_host::{call::HostedMachine, protocol::Protocol}; +use litellm_host_python::present; +use litellm_inference_responses::{ResponsesRoute, route::Responses}; +use litellm_secrets::source::SecretSource; +use pyo3::prelude::*; +use serde_json::Value; +pub(crate) use websocket::ResponsesWebSocketConnection; + +use super::{ + NativeCall, + inference::{InferenceHost, InferenceRoute, run_inference}, +}; +use crate::errors::RustBridgeDeclined; + +const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.responses.route_host"; + +fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult> { + if let Some(reason) = py + .import(ROUTE_HOST_MODULE)? + .getattr("decline_reason")? + .call1((&call.bound,))? + .extract::>()? + { + return Err(RustBridgeDeclined::new_err(reason)); + } + let argument = |name: &str| present(&call.kwargs, &call.bound, name); + let model = argument("model")? + .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("model is required"))? + .extract::()?; + let provider = argument("custom_llm_provider")? + .map(|value| value.extract::()) + .transpose()?; + if provider + .as_deref() + .is_some_and(|provider| provider != "openai") + || model + .strip_prefix("openai/") + .unwrap_or(&model) + .contains('/') + { + return Err(RustBridgeDeclined::new_err( + "native HTTP responses provider", + )); + } + if argument("stream")? + .map(|value| litellm_host_python::from_py::(&value)) + .transpose()? + .is_some_and(|value| value == Value::Bool(true)) + { + return Err(RustBridgeDeclined::new_err( + "native Python responses streaming", + )); + } + let host = InferenceHost::new(call.bound.clone().unbind(), ROUTE_HOST_MODULE); + run_inference::(py, call, asynchronous, ResponsesPythonHost(host)) +} + +#[pyfunction] +pub(crate) fn responses(py: Python<'_>, call: NativeCall<'_>) -> PyResult> { + run_responses(py, call, false) +} + +#[pyfunction] +pub(crate) fn aresponses(py: Python<'_>, call: NativeCall<'_>) -> PyResult> { + run_responses(py, call, true) +} + +impl InferenceRoute for ResponsesRoute { + type Protocol = Responses; + const OPERATION: LoggingOperation = LoggingOperation::Responses; + const SYNC_CALL_TYPE: &'static str = "responses"; + const ASYNC_CALL_TYPE: &'static str = "aresponses"; + + fn new( + http: litellm_http::Client, + auth: Arc, + secrets: Arc, + ) -> Self { + Self::new(http, auth, secrets) + } + + fn with_cache(self, cache: ScopedCache) -> Self { + self.with_cache(cache) + } + + fn machine( + self, + call: ::Request, + policy: CachePolicy, + ) -> HostedMachine { + self.machine(call, policy) + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/responses.rs b/litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs similarity index 57% rename from litellm-rust/crates/python-bridge/src/routes/responses.rs rename to litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs index 780f2e6929b..615a838a597 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs @@ -1,121 +1,12 @@ -mod host; - use litellm_inference_responses::websocket::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; -use pyo3::{ - prelude::*, - types::{PyDict, PyTuple}, -}; +use pyo3::prelude::*; use serde_json::Value; use crate::{ - errors::{RustBridgeDeclined, route_error_to_pyerr}, + errors::route_error_to_pyerr, marshal::{marshal_headers, optional_timeout}, }; -fn run_public( - py: Python<'_>, - request: Bound<'_, PyAny>, - args: Bound<'_, PyTuple>, - kwargs: Bound<'_, PyDict>, - asynchronous: bool, -) -> PyResult> { - use super::inference::InferenceHost; - use litellm_callbacks_legacy_python::LoggingOperation; - let host = InferenceHost::new( - request.clone().unbind(), - "litellm.rust_bridge.responses.route_host", - ); - if let Some(reason) = py - .import("litellm.rust_bridge.responses.route_host")? - .getattr("decline_reason")? - .call1((&request,))? - .extract::>()? - { - return Err(RustBridgeDeclined::new_err(reason)); - } - let model = host - .argument(py, &kwargs, "model")? - .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("model is required"))? - .extract::()?; - let provider = host - .argument(py, &kwargs, "custom_llm_provider")? - .map(|value| value.extract::()) - .transpose()?; - if provider - .as_deref() - .is_some_and(|provider| provider != "openai") - || model - .strip_prefix("openai/") - .unwrap_or(&model) - .contains('/') - { - return Err(RustBridgeDeclined::new_err( - "native HTTP responses provider", - )); - } - if host - .argument(py, &kwargs, "stream")? - .map(|value| litellm_host_python::from_py::(&value)) - .transpose()? - .is_some_and(|value| value == Value::Bool(true)) - { - return Err(RustBridgeDeclined::new_err( - "native Python responses streaming", - )); - } - let cache_call_type = if asynchronous { - "aresponses" - } else { - "responses" - }; - crate::cache::admit_native(py, &kwargs, cache_call_type)?; - let (arguments, hooks) = crate::routes::call_hooks( - py, - LoggingOperation::Responses, - &request, - &args, - &kwargs, - asynchronous, - )?; - crate::routes::run_public_call( - py, - arguments, - move |py, arguments, request| { - let route = litellm_inference_responses::ResponsesRoute::new( - crate::http::provider_client(py, arguments, asynchronous)? - .map_err(crate::http::client_error)?, - crate::http::resources().auth.clone(), - crate::secrets::source(py)?, - ); - let (cache, cache_options) = - crate::cache::configured_native(py, arguments, cache_call_type)?; - let route = match cache { - Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new( - cache, - litellm_cache_response::CacheScope::Shared, - )), - None => route, - }; - Ok(route.machine(request, cache_options.policy)) - }, - host::ResponsesPythonHost(host), - hooks, - asynchronous, - ) -} - -#[pyfunction] -pub(crate) fn responses(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult> { - let call = super::NativeCall::extract(&call)?; - run_public(py, call.bound.into_any(), call.args, call.kwargs, false) -} - -#[pyfunction] -pub(crate) fn aresponses(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult> { - let call = super::NativeCall::extract(&call)?; - run_public(py, call.bound.into_any(), call.args, call.kwargs, true) -} - #[pyclass] pub(crate) struct ResponsesWebSocketConnection { inner: RustResponsesWebSocketConnection, diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs index 61066e322a7..413245c1ec2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/traces.rs +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -172,7 +172,7 @@ impl NativeTraceStorage { BTreeMap, >, ) -> PyResult> { - let table = InsertTable::parse(table).map_err(map_error)?; + let table = table.parse::().map_err(map_error)?; let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; let connection = self.config.storage().writer().clone(); let database = self.config.storage().database().to_owned(); diff --git a/litellm-rust/crates/secrets-types/src/config.rs b/litellm-rust/crates/secrets-types/src/config.rs index 82d48f7b2e1..9b6e82c3a6b 100644 --- a/litellm-rust/crates/secrets-types/src/config.rs +++ b/litellm-rust/crates/secrets-types/src/config.rs @@ -5,18 +5,37 @@ use strum::IntoStaticStr; use crate::SecretValue; -#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, IntoStaticStr, PartialEq, Serialize)] +#[derive( + Clone, + Copy, + Debug, + Deserialize, + Eq, + Hash, + IntoStaticStr, + PartialEq, + Serialize, + strum::VariantArray, +)] #[serde(rename_all = "snake_case")] -#[strum(serialize_all = "snake_case")] pub enum KeyManagementSystem { + #[strum(serialize = "google_kms")] GoogleKms, + #[strum(serialize = "azure_key_vault")] AzureKeyVault, + #[strum(serialize = "aws_secret_manager")] AwsSecretManager, + #[strum(serialize = "google_secret_manager")] GoogleSecretManager, + #[strum(serialize = "hashicorp_vault")] HashicorpVault, + #[strum(serialize = "cyberark")] Cyberark, + #[strum(serialize = "local")] Local, + #[strum(serialize = "aws_kms")] AwsKms, + #[strum(serialize = "custom")] Custom, } diff --git a/litellm-rust/crates/secrets-types/tests/context.rs b/litellm-rust/crates/secrets-types/tests/context.rs index a167e87954f..8f8eeabdb62 100644 --- a/litellm-rust/crates/secrets-types/tests/context.rs +++ b/litellm-rust/crates/secrets-types/tests/context.rs @@ -128,17 +128,21 @@ fn rotation_write_context_preserves_the_operation_context(aws_context: SecretOpe fn provider_context_accepts_only_its_owner( #[case] owner: KeyManagementSystem, #[case] context: SecretOperationContext, -) { - for system in [ - KeyManagementSystem::AwsSecretManager, + #[values( + KeyManagementSystem::GoogleKms, KeyManagementSystem::AzureKeyVault, + KeyManagementSystem::AwsSecretManager, KeyManagementSystem::GoogleSecretManager, KeyManagementSystem::HashicorpVault, KeyManagementSystem::Cyberark, - ] { - assert_eq!(context.validate_for(system).is_ok(), system == owner); - assert!(SecretOperationContext::Default.validate_for(system).is_ok()); - } + KeyManagementSystem::Local, + KeyManagementSystem::AwsKms, + KeyManagementSystem::Custom + )] + system: KeyManagementSystem, +) { + assert_eq!(context.validate_for(system).is_ok(), system == owner); + assert!(SecretOperationContext::Default.validate_for(system).is_ok()); if matches!( owner, KeyManagementSystem::AzureKeyVault | KeyManagementSystem::GoogleSecretManager diff --git a/litellm-rust/crates/secrets/src/oidc.rs b/litellm-rust/crates/secrets/src/oidc.rs index b6e8dbc123b..adeb6c3db00 100644 --- a/litellm-rust/crates/secrets/src/oidc.rs +++ b/litellm-rust/crates/secrets/src/oidc.rs @@ -23,17 +23,22 @@ const OIDC_ALLOWED_CREDENTIAL_DIRS: &str = "LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS const DEFAULT_CREDENTIAL_DIRS: &str = "/var/run/secrets,/run/secrets"; #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, strum::EnumString, strum::AsRefStr)] -#[strum(serialize_all = "snake_case")] pub enum OidcProvider { + #[strum(serialize = "google")] Google, #[strum(serialize = "circleci")] CircleCi, #[strum(serialize = "circleci_v2")] CircleCiV2, + #[strum(serialize = "github")] Github, + #[strum(serialize = "azure")] Azure, + #[strum(serialize = "file")] File, + #[strum(serialize = "env")] Env, + #[strum(serialize = "env_path")] EnvPath, } diff --git a/litellm-rust/crates/testkit/Cargo.toml b/litellm-rust/crates/testkit/Cargo.toml index 98a36a1e87f..97efb753384 100644 --- a/litellm-rust/crates/testkit/Cargo.toml +++ b/litellm-rust/crates/testkit/Cargo.toml @@ -13,6 +13,7 @@ serde.workspace = true semver.workspace = true serde_json.workspace = true sha2.workspace = true +strum.workspace = true tar.workspace = true target-lexicon.workspace = true thiserror.workspace = true diff --git a/litellm-rust/crates/testkit/src/agent/configure.rs b/litellm-rust/crates/testkit/src/agent/configure.rs index a095aacc92f..a41bec9379a 100644 --- a/litellm-rust/crates/testkit/src/agent/configure.rs +++ b/litellm-rust/crates/testkit/src/agent/configure.rs @@ -5,7 +5,7 @@ use semver::Version; use crate::Error; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] pub enum Wire { ChatCompletions, Messages, diff --git a/litellm-rust/crates/testkit/tests/configure.rs b/litellm-rust/crates/testkit/tests/configure.rs index ca3587c3474..87f53d2f9c0 100644 --- a/litellm-rust/crates/testkit/tests/configure.rs +++ b/litellm-rust/crates/testkit/tests/configure.rs @@ -2,6 +2,7 @@ use std::path::Path; use litellm_testkit::{ClaudeCode, Codex, Configure, Error, Opencode, Settings, Version, Wire}; use rstest::rstest; +use strum::VariantArray; fn settings(wire: Wire) -> Settings { Settings { @@ -121,7 +122,10 @@ fn opencode_uses_a_different_provider_package_for_every_wire() { .unwrap() .to_owned() }; - let packages = [Wire::ChatCompletions, Wire::Responses, Wire::Messages].map(package); + let packages = Wire::VARIANTS + .iter() + .map(|wire| package(*wire)) + .collect::>(); assert_eq!( packages diff --git a/litellm-rust/crates/traces-clickhouse/src/insert.rs b/litellm-rust/crates/traces-clickhouse/src/insert.rs index ed8db188c14..18c077977cc 100644 --- a/litellm-rust/crates/traces-clickhouse/src/insert.rs +++ b/litellm-rust/crates/traces-clickhouse/src/insert.rs @@ -30,29 +30,19 @@ fn max_insert_bytes() -> Result { pub type InsertRow = BTreeMap>; +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString, strum::IntoStaticStr)] +#[strum(parse_err_ty = Error, parse_err_fn = invalid_table)] pub enum InsertTable { + #[strum(serialize = "otel_traces")] OtelTraces, + #[strum(serialize = "spend_logs")] SpendLogs, + #[strum(serialize = "lens_feedback")] LensFeedback, } -impl InsertTable { - pub fn parse(value: &str) -> Result { - match value { - "otel_traces" => Ok(Self::OtelTraces), - "spend_logs" => Ok(Self::SpendLogs), - "lens_feedback" => Ok(Self::LensFeedback), - _ => Err(Error::InvalidTable), - } - } - - fn name(&self) -> &'static str { - match self { - Self::OtelTraces => "otel_traces", - Self::SpendLogs => "spend_logs", - Self::LensFeedback => "lens_feedback", - } - } +fn invalid_table(_name: &str) -> Error { + Error::InvalidTable } pub async fn insert_rows( @@ -81,7 +71,7 @@ pub async fn insert_shared_rows( client, connection, database, - table.name(), + <&'static str>::from(table), &token, body, ) @@ -242,7 +232,25 @@ mod tests { use rstest::rstest; use serde_json::json; - use super::{Error, shared_rows, write_rows}; + use super::{Error, InsertTable, shared_rows, write_rows}; + + #[rstest] + #[case::otel_traces("otel_traces", InsertTable::OtelTraces)] + #[case::spend_logs("spend_logs", InsertTable::SpendLogs)] + #[case::lens_feedback("lens_feedback", InsertTable::LensFeedback)] + fn insert_table_parses_each_table_name(#[case] name: &str, #[case] expected: InsertTable) { + assert_eq!(name.parse::().unwrap(), expected); + } + + #[rstest] + #[case::unknown("events")] + #[case::case_sensitive("OTEL_TRACES")] + fn insert_table_rejects_unknown_names(#[case] name: &str) { + assert!(matches!( + name.parse::(), + Err(Error::InvalidTable) + )); + } #[rstest] fn encoded_limit_counts_utf8_bytes_across_rows() { diff --git a/litellm-rust/crates/traces-clickhouse/src/query.rs b/litellm-rust/crates/traces-clickhouse/src/query.rs index a08affb0cad..cd1a96eef31 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query.rs @@ -56,15 +56,21 @@ enum PathPart { #[macro_rules_attribute::apply(crate::response_type)] #[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, strum::Display)] #[serde(rename_all = "lowercase")] -#[strum(serialize_all = "lowercase")] #[cfg_attr(feature = "schema", schemars(rename = "MetadataValueType"))] enum JsonKind { + #[strum(serialize = "array")] Array, + #[strum(serialize = "boolean")] Boolean, + #[strum(serialize = "integer")] Integer, + #[strum(serialize = "null")] Null, + #[strum(serialize = "number")] Number, + #[strum(serialize = "object")] Object, + #[strum(serialize = "string")] String, } diff --git a/litellm-rust/crates/traces-clickhouse/src/table.rs b/litellm-rust/crates/traces-clickhouse/src/table.rs index 60ab4943f5e..c9039ad8bed 100644 --- a/litellm-rust/crates/traces-clickhouse/src/table.rs +++ b/litellm-rust/crates/traces-clickhouse/src/table.rs @@ -4,9 +4,11 @@ Clone, Copy, Debug, strum::Display, strum::AsRefStr, strum::EnumIter, strum::IntoStaticStr, )] #[serde(rename_all = "snake_case")] -#[strum(serialize_all = "snake_case")] pub enum TraceTable { + #[strum(serialize = "otel_traces")] OtelTraces, + #[strum(serialize = "agent_traces_by_key")] AgentTracesByKey, + #[strum(serialize = "spend_logs")] SpendLogs, } diff --git a/litellm-rust/crates/traces-clickhouse/tests/queries.rs b/litellm-rust/crates/traces-clickhouse/tests/queries.rs index 6e63adfc347..5cb24dc5548 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/queries.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/queries.rs @@ -126,10 +126,12 @@ async fn lens_sample_keeps_spans_before_window_start_and_excludes_old_only_trace } #[derive(Clone, Copy, strum::AsRefStr)] -#[strum(serialize_all = "snake_case")] enum ScopeCase { + #[strum(serialize = "admin")] Admin, + #[strum(serialize = "team")] Team, + #[strum(serialize = "other_team")] OtherTeam, } diff --git a/litellm-rust/crates/traces/src/normalize/format/claude_code.rs b/litellm-rust/crates/traces/src/normalize/format/claude_code.rs index 37fb2d01330..d849d0fea87 100644 --- a/litellm-rust/crates/traces/src/normalize/format/claude_code.rs +++ b/litellm-rust/crates/traces/src/normalize/format/claude_code.rs @@ -15,14 +15,23 @@ use crate::{ /// Claude Code's built-in tracing, identified by its instrumentation scope. pub(crate) struct ClaudeCode; +#[derive(Debug, PartialEq, Eq, strum::EnumString)] enum SpanType { + #[strum(serialize = "assistant_response")] AssistantResponse, + #[strum(serialize = "tool_result")] ToolResult, + #[strum(serialize = "api_request_body")] ApiRequestBody, + #[strum(serialize = "compaction")] Compaction, + #[strum(serialize = "interaction")] Interaction, + #[strum(serialize = "llm_request")] LlmRequest, + #[strum(serialize = "tool")] Tool, + #[strum(disabled)] Other, } @@ -33,16 +42,7 @@ fn span_type(name: &str, attributes: &BTreeMap) -> SpanType { } else { kind }; - match kind { - "assistant_response" => SpanType::AssistantResponse, - "tool_result" => SpanType::ToolResult, - "api_request_body" => SpanType::ApiRequestBody, - "compaction" => SpanType::Compaction, - "interaction" => SpanType::Interaction, - "llm_request" => SpanType::LlmRequest, - "tool" => SpanType::Tool, - _ => SpanType::Other, - } + kind.parse().unwrap_or(SpanType::Other) } /// `agent:custom:search_agent` -> `search_agent`: the subagent a request ran for. @@ -318,7 +318,7 @@ mod tests { use rstest::rstest; use serde_json::Value; - use super::CLAUDE_CODE_SCOPE; + use super::{CLAUDE_CODE_SCOPE, SpanType, span_type}; use crate::{ Error, normalize::{Normalization, NormalizedSpan, ObservationType}, @@ -355,6 +355,22 @@ mod tests { .collect() } + #[rstest] + #[case::assistant_response("assistant_response", SpanType::AssistantResponse)] + #[case::tool_result("tool_result", SpanType::ToolResult)] + #[case::api_request_body("api_request_body", SpanType::ApiRequestBody)] + #[case::compaction("compaction", SpanType::Compaction)] + #[case::interaction("interaction", SpanType::Interaction)] + #[case::llm_request("llm_request", SpanType::LlmRequest)] + #[case::tool("tool", SpanType::Tool)] + #[case::unknown("surprise", SpanType::Other)] + fn span_type_maps_each_recorded_kind(#[case] kind: &str, #[case] expected: SpanType) { + assert_eq!( + span_type("anything", &attributes(&[("span.type", kind)])), + expected + ); + } + #[rstest] fn notification_prompts_keep_user_provenance_and_compaction_is_system() { let prompt_text = diff --git a/litellm-rust/crates/traces/src/normalize/format/genai.rs b/litellm-rust/crates/traces/src/normalize/format/genai.rs index 5ba9b449737..2bfd6ce5370 100644 --- a/litellm-rust/crates/traces/src/normalize/format/genai.rs +++ b/litellm-rust/crates/traces/src/normalize/format/genai.rs @@ -11,18 +11,24 @@ use crate::{ pub(crate) struct GenAi; #[derive(strum::EnumString)] -#[strum(serialize_all = "snake_case")] pub(crate) enum Operation { + #[strum(serialize = "create_agent")] CreateAgent, + #[strum(serialize = "invoke_agent")] InvokeAgent, + #[strum(serialize = "invoke_workflow")] InvokeWorkflow, + #[strum(serialize = "chat")] Chat, #[strum(serialize = "text_completion", serialize = "completion")] TextCompletion, + #[strum(serialize = "generate_content")] GenerateContent, + #[strum(serialize = "execute_tool")] ExecuteTool, #[strum(serialize = "embeddings", serialize = "embedding")] Embeddings, + #[strum(serialize = "retrieval")] Retrieval, } diff --git a/litellm-rust/crates/traces/src/normalize/metadata.rs b/litellm-rust/crates/traces/src/normalize/metadata.rs index 100008714c5..600bfce4287 100644 --- a/litellm-rust/crates/traces/src/normalize/metadata.rs +++ b/litellm-rust/crates/traces/src/normalize/metadata.rs @@ -24,32 +24,56 @@ pub enum AgentType { serde_with::DeserializeFromStr, serde_with::SerializeDisplay, )] -#[strum(serialize_all = "kebab-case")] pub enum Integration { + #[strum(serialize = "claude-code")] ClaudeCode, + #[strum(serialize = "claude-agent-sdk")] ClaudeAgentSdk, + #[strum(serialize = "openai-codex")] OpenaiCodex, + #[strum(serialize = "deepagents-code")] DeepagentsCode, + #[strum(serialize = "cursor")] Cursor, + #[strum(serialize = "pi")] Pi, + #[strum(serialize = "opencode")] Opencode, + #[strum(serialize = "copilot")] Copilot, + #[strum(serialize = "langchain")] Langchain, + #[strum(serialize = "langgraph")] Langgraph, + #[strum(serialize = "deepagents")] Deepagents, + #[strum(serialize = "autogen")] Autogen, + #[strum(serialize = "crewai")] Crewai, + #[strum(serialize = "google-adk")] GoogleAdk, + #[strum(serialize = "llama-index")] LlamaIndex, + #[strum(serialize = "mastra")] Mastra, + #[strum(serialize = "microsoft-agent-framework")] MicrosoftAgentFramework, + #[strum(serialize = "openai-agents")] OpenaiAgents, + #[strum(serialize = "pydantic-ai")] PydanticAi, + #[strum(serialize = "semantic-kernel")] SemanticKernel, + #[strum(serialize = "strands")] Strands, + #[strum(serialize = "vercel-ai-sdk")] VercelAiSdk, + #[strum(serialize = "instructor")] Instructor, + #[strum(serialize = "n8n")] N8n, + #[strum(serialize = "temporal")] Temporal, #[strum(default)] Other(String), @@ -138,23 +162,36 @@ impl AgentMetadata { } #[derive(strum::EnumString, strum::IntoStaticStr)] -#[strum(serialize_all = "snake_case")] enum MetadataField { + #[strum(serialize = "lc_agent_name")] LcAgentName, + #[strum(serialize = "ls_integration")] LsIntegration, + #[strum(serialize = "ls_agent_type")] LsAgentType, + #[strum(serialize = "ls_agent_purpose")] LsAgentPurpose, + #[strum(serialize = "ls_agent_runtime")] LsAgentRuntime, #[strum(serialize = "ls_agent_runtime_version", to_string = "ls_agent_version")] LsAgentVersion, + #[strum(serialize = "ls_trace_schema_version")] LsTraceSchemaVersion, + #[strum(serialize = "thread_id")] ThreadId, + #[strum(serialize = "ls_subagent_id")] LsSubagentId, + #[strum(serialize = "ls_subagent_type")] LsSubagentType, + #[strum(serialize = "ls_tool_name")] LsToolName, + #[strum(serialize = "ls_model_name")] LsModelName, + #[strum(serialize = "ls_provider")] LsProvider, + #[strum(serialize = "git_branch")] GitBranch, + #[strum(serialize = "git_commit_sha")] GitCommitSha, #[strum(serialize = "repository_url", to_string = "git_repo_url")] GitRepoUrl, diff --git a/litellm-rust/crates/traces/src/normalize/mod.rs b/litellm-rust/crates/traces/src/normalize/mod.rs index 4d11e76e19b..894fcd1bb03 100644 --- a/litellm-rust/crates/traces/src/normalize/mod.rs +++ b/litellm-rust/crates/traces/src/normalize/mod.rs @@ -32,20 +32,32 @@ pub use metadata::{AgentMetadata, AgentType, Integration}; #[macro_rules_attribute::apply(crate::wire_type)] #[derive(Clone, Copy, Debug, Eq, PartialEq, strum::EnumString)] #[serde(rename_all = "lowercase")] -#[strum(serialize_all = "lowercase", ascii_case_insensitive)] +#[strum(ascii_case_insensitive)] #[cfg_attr(feature = "schema", schemars(rename = "SpanType"))] pub enum ObservationType { + #[strum(serialize = "agent")] Agent, + #[strum(serialize = "llm")] Llm, + #[strum(serialize = "tool")] Tool, + #[strum(serialize = "chain")] Chain, + #[strum(serialize = "framework")] Framework, + #[strum(serialize = "retriever")] Retriever, + #[strum(serialize = "embedding")] Embedding, + #[strum(serialize = "reranker")] Reranker, + #[strum(serialize = "guardrail")] Guardrail, + #[strum(serialize = "evaluator")] Evaluator, + #[strum(serialize = "prompt")] Prompt, + #[strum(serialize = "decision")] Decision, } diff --git a/litellm-rust/crates/traces/src/query.rs b/litellm-rust/crates/traces/src/query.rs index 26eaa33d5ac..33dcbb31abf 100644 --- a/litellm-rust/crates/traces/src/query.rs +++ b/litellm-rust/crates/traces/src/query.rs @@ -2,23 +2,38 @@ pub mod guide; pub mod named; #[derive(Clone, Copy, Debug, Eq, PartialEq, strum::EnumString, strum::Display, strum::AsRefStr)] -#[strum(serialize_all = "snake_case")] pub enum ReadQuery { + #[strum(serialize = "list_traces")] ListTraces, + #[strum(serialize = "trace_agents")] TraceAgents, + #[strum(serialize = "trace_spans")] TraceSpans, + #[strum(serialize = "trace_page_spans")] TracePageSpans, + #[strum(serialize = "trace_identity")] TraceIdentity, + #[strum(serialize = "span_detail")] SpanDetail, + #[strum(serialize = "span_error")] SpanError, + #[strum(serialize = "spend_by_response_ids")] SpendByResponseIds, + #[strum(serialize = "availability")] Availability, + #[strum(serialize = "agents")] Agents, + #[strum(serialize = "sample")] Sample, + #[strum(serialize = "content")] Content, + #[strum(serialize = "evidence")] Evidence, + #[strum(serialize = "feedback_target")] FeedbackTarget, + #[strum(serialize = "feedback")] Feedback, + #[strum(serialize = "feedback_summary")] FeedbackSummary, } diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 85ef5a93937..b928f33e961 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -71,15 +71,25 @@ class CacheMode(str, Enum): #### LiteLLM.Completion / Embedding Cache #### +def _is_conversation_item(item: object) -> bool: + if isinstance(item, BaseModel): + return True + if not isinstance(item, Mapping): + return False + block: Final = cast(Mapping[str, object], item) # cast-ok: isinstance leaves the key and value types unknown + return block.get("type") != "file" + + def _request_message_count(kwargs: Mapping[str, object]) -> int: - """Chat and Messages API `messages`, else Responses API `input` items; embedding `input` strings count as none""" + """Chat and Messages API `messages`, else Responses API `input` items; embedding strings and file blocks count as none""" messages: Final = kwargs.get("messages") if isinstance(messages, list): return len(messages) input_items: Final = kwargs.get("input") if not isinstance(input_items, list): return 0 - return sum(1 for item in input_items if isinstance(item, (Mapping, BaseModel))) + items: Final = cast(list[object], input_items) # cast-ok: isinstance leaves the element type unknown + return sum(1 for item in items if _is_conversation_item(item)) class Cache: diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index f9215864d1e..9ffc95c6260 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -21,7 +21,7 @@ import time from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Mapping from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar -from pydantic import ConfigDict, ValidationError +from pydantic import ConfigDict, SkipValidation, ValidationError import litellm from litellm._internal_context import post_response_phase @@ -39,7 +39,7 @@ from litellm.litellm_core_utils.logging_utils import ( from litellm.types.caching import CACHED_STREAM_EVENTS_KEY, EMBEDDING_CACHE_FORMAT_VERSION, CachedEmbedding from litellm.types.integrations.custom_logger import converted_stream_requested from litellm.types.llms.base import LiteLLMBaseModel -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import ChatCompletionFileObject, ResponsesAPIResponse from litellm.types.rerank import RerankResponse from litellm.types.utils import ( CachingDetails, @@ -71,6 +71,8 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +EmbeddingCacheInputElement = str | list[int] | ChatCompletionFileObject + class CachingHandlerResponse(LiteLLMBaseModel): """ @@ -82,7 +84,7 @@ class CachingHandlerResponse(LiteLLMBaseModel): cached_result: object | None = None final_embedding_cached_response: EmbeddingResponse | None = None embedding_all_elements_cache_hit: bool = False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call - embedding_uncached_input: list[str | list[int]] | None = None + embedding_uncached_input: SkipValidation[list[EmbeddingCacheInputElement]] | None = None in_memory_cache_obj: Final = InMemoryCache() @@ -508,7 +510,7 @@ class LLMCachingHandler: _sync_get_cache = sync_get_cache - def handle_kwargs_input_list_or_str(self, kwargs: dict[str, object]) -> list[str]: + def handle_kwargs_input_list_or_str(self, kwargs: dict[str, object]) -> list[EmbeddingCacheInputElement]: """ Handles the input of kwargs['input'] being a list or a string """ @@ -517,7 +519,11 @@ class LLMCachingHandler: elif isinstance(kwargs["input"], list): return kwargs["input"] else: - raise ValueError("input must be a string or a list") + raise litellm.BadRequestError( + message="input must be a string or a list of strings and content blocks", + model=str(kwargs.get("model")), + llm_provider=str(kwargs.get("custom_llm_provider")), + ) def _extract_model_from_cached_results(self, non_null_list: list[tuple[int, CachedEmbedding]]) -> str | None: """ @@ -851,10 +857,7 @@ class LLMCachingHandler: self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs) cached_result: object | None = None if call_type == CallTypes.aembedding.value: - if isinstance(new_kwargs["input"], str): - new_kwargs["input"] = [new_kwargs["input"]] - elif not isinstance(new_kwargs["input"], list): - raise ValueError("input must be a string or a list") + new_kwargs["input"] = self.handle_kwargs_input_list_or_str(new_kwargs) tasks: Final[list[Awaitable[object]]] = [] for idx, i in enumerate(new_kwargs["input"]): preset_cache_key = litellm.cache.get_cache_key(**{**new_kwargs, "input": i}) diff --git a/litellm/constants.py b/litellm/constants.py index c4adc0de22f..b8048ad30b4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1,10 +1,20 @@ import os import sys +from enum import Enum from types import MappingProxyType from typing import Final, Literal from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_in_range, get_env_int_or_none +SERVER_STREAMING_CLASSIFICATION_KEY: Final = "litellm_server_streaming_classification" + + +class ServerStreamingClassification(str, Enum): + MARKER = "litellm-server-streaming" + + +SERVER_STREAMING_CLASSIFICATION_MARKER: Final = ServerStreamingClassification.MARKER + DEFER_PYDANTIC_BUILD: Final = os.getenv("DEFER_PYDANTIC_BUILD", "true") in ("true", "1", "on") DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm")) AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview")) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 94f1207b9ae..383a09f0931 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -6,7 +6,7 @@ import secrets from collections.abc import Mapping, Sequence from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args import httpx @@ -21,11 +21,14 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.secret_managers.main import str_to_bool from litellm.types.guardrails import ( + DEFAULT_GUARDRAIL_STREAM_SCOPE, DynamicGuardrailParams, GuardrailEventHooks, + GuardrailStreamScope, LitellmParams, LoggingOnlyScope, Mode, + runtime_stream_scope, ) from litellm.types.llms.openai import AllMessageValues from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -49,6 +52,8 @@ from litellm.constants import ( GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, LOGS_GUARDRAIL_INFORMATION_MARKER, PRE_CALL_EXECUTED_GUARDRAILS_KEY, + SERVER_STREAMING_CLASSIFICATION_KEY, + SERVER_STREAMING_CLASSIFICATION_MARKER, ) from litellm.exceptions import ( BlockedPiiEntityError, @@ -173,6 +178,42 @@ def get_session_id_from_request_data(request_data: dict[str, Any]) -> str | None return None +_REALTIME_STREAMING_HOOKS: Final = frozenset({GuardrailEventHooks.realtime_input_transcription}) + + +def without_server_streaming_classification(data: Mapping[str, object]) -> dict[str, object]: + return { + key: value + for key, value in data.items() + if key != SERVER_STREAMING_CLASSIFICATION_KEY or value != SERVER_STREAMING_CLASSIFICATION_MARKER + } + + +def guardrail_request_data_with_streaming( + data: Mapping[str, object], + *, + is_streaming: bool, +) -> dict[str, object]: + data_without_server_classification: Final = without_server_streaming_classification(data) + if not is_streaming: + return data_without_server_classification + return { + **data_without_server_classification, + SERVER_STREAMING_CLASSIFICATION_KEY: SERVER_STREAMING_CLASSIFICATION_MARKER, + } + + +def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = None) -> bool: + if event_type in _REALTIME_STREAMING_HOOKS: + return True + if not isinstance(data, Mapping): + return False + return ( + data.get("stream") is True + or data.get(SERVER_STREAMING_CLASSIFICATION_KEY) is SERVER_STREAMING_CLASSIFICATION_MARKER + ) + + class CustomGuardrail(CustomLogger): # If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path. use_native_during_call_hook: ClassVar[bool] = False @@ -183,6 +224,9 @@ class CustomGuardrail(CustomLogger): records_own_guardrail_information: ClassVar[bool] = False logging_only_scope: LoggingOnlyScope | None + stream_scope_default: GuardrailStreamScope = DEFAULT_GUARDRAIL_STREAM_SCOPE + stream_scope_by_hook: tuple[tuple[str, GuardrailStreamScope], ...] = () + timeout: float | httpx.Timeout | None = None def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: forwarded to cooperative __init_subclass__ hooks @@ -258,6 +302,8 @@ class CustomGuardrail(CustomLogger): self.run_in_parallel: bool = run_in_parallel self.scan_raw_request: bool = scan_raw_request self.only_scan_new_messages: bool = only_scan_new_messages + stream_scope_arg: Final[object] = cast(object, kwargs.pop("stream_scope", None)) # cast-ok: config + self.apply_stream_scope(stream_scope_arg) self.logging_only_scope = None if timeout is not None: self.timeout = timeout @@ -1099,6 +1145,23 @@ class CustomGuardrail(CustomLogger): return name in suppressed_compression_guardrails() + def apply_stream_scope(self, stream_scope: object) -> None: + default, by_hook = runtime_stream_scope(stream_scope) + self.stream_scope_default = default + self.stream_scope_by_hook = tuple(by_hook.items()) + + def stream_scope_allows(self, data: object, event_type: GuardrailEventHooks) -> bool: + scope: Final = next( + (scope for hook, scope in self.stream_scope_by_hook if hook == event_type.value), + self.stream_scope_default, + ) + if scope == "both": + return True + is_streaming: Final = _request_is_streaming(data, event_type) + if scope == "streaming": + return is_streaming + return not is_streaming + def should_run_guardrail( self, data, @@ -1142,8 +1205,10 @@ class CustomGuardrail(CustomLogger): data, self.event_hook, event_type ) if result is not None: - return result - return True + tagged_result: Final[bool] = bool(cast(object, result)) # cast-ok: helper return + data_obj: Final[object] = cast(object, data) # cast-ok: data param + return tagged_result and self.stream_scope_allows(data_obj, event_type) + return self.stream_scope_allows(cast(object, data), event_type) # cast-ok: data param return False if ( @@ -1167,8 +1232,9 @@ class CustomGuardrail(CustomLogger): ) result = EnterpriseCustomGuardrailHelper._should_run_if_mode_by_tag(data, self.event_hook, event_type) if result is not None: - return result - return True + mode_tag_result: Final[bool] = bool(cast(object, result)) # cast-ok: helper return + return mode_tag_result and self.stream_scope_allows(cast(object, data), event_type) # cast-ok: data + return self.stream_scope_allows(cast(object, data), event_type) # cast-ok: data param def _event_hook_is_event_type(self, event_type: GuardrailEventHooks) -> bool: """ diff --git a/litellm/integrations/otel/emitter.py b/litellm/integrations/otel/emitter.py index e9441ee2a9a..a4cb0a84541 100644 --- a/litellm/integrations/otel/emitter.py +++ b/litellm/integrations/otel/emitter.py @@ -91,13 +91,32 @@ def span_attribute_limit(span: Span) -> int | None: return span._limits.max_span_attributes # pyright: ignore[reportPrivateUsage] # SDK has no public getter -def attribute_budget(span: Span, reserved: int) -> int | None: - """How many mapped attributes fit on ``span`` next to what it already carries and ``reserved`` more.""" +def _carried_keys( + attributes: Mapping[str, AttrValue] | tuple[tuple[str, AttrValue], ...], +) -> frozenset[str]: + """The keys a span already carries: a live span exposes a ``Mapping``, an ended one a tuple of pairs.""" + if isinstance(attributes, Mapping): + return frozenset(attributes) + return frozenset(key for key, _value in attributes) + + +def attribute_budget(span: Span, reserved: int, overwrites: frozenset[str] = frozenset()) -> int | None: + """How many mapped attributes fit on ``span`` next to what it already carries and ``reserved`` more. + + ``overwrites`` are the mapped keys already present on ``span``: setting one + replaces the value in place and consumes no slot against the limit, so only + the genuinely new pre-existing keys reduce the budget. Counting the + overwritten ones too reserves slots the fit can never spend and sheds + indexed message attributes for nothing. + """ limit: Final = span_attribute_limit(span) if limit is None: return None - on_span: Final = len(span.attributes or ()) if isinstance(span, ReadableSpan) else 0 - return limit - on_span - reserved + if not isinstance(span, ReadableSpan): + return limit - reserved + existing_keys: Final = _carried_keys(span.attributes or ()) + fresh: Final = len(existing_keys - overwrites) if overwrites else len(existing_keys) + return limit - fresh - reserved def stamp_error( @@ -284,7 +303,7 @@ class SpanEmitter: ) stamped_later: Final = error_attributes(error) if error else _NO_ATTRIBUTES reserved: Final = len(stamped_later.keys() - mapped.keys()) - for key, value in fit_indexed_messages(mapped, attribute_budget(span, reserved)).items(): + for key, value in fit_indexed_messages(mapped, attribute_budget(span, reserved, frozenset(mapped))).items(): span.set_attribute(key, value) if error: stamped: Final = stamp_error(span, error) diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index 888c90b7661..498ebe655da 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -575,6 +575,7 @@ class OpenTelemetryV2(CustomLogger): request_purpose=call.purpose, trace=call.trace, session_id=call.session_id, + metadata_keys=tuple(self.config.baggage_metadata_keys), ) end_time_ns: Final = to_ns(end_time) if carrier is not None and carrier.span is not None: diff --git a/litellm/integrations/otel/mappers/openinference.py b/litellm/integrations/otel/mappers/openinference.py index a064c2c7e61..8d7b4d7cf4a 100644 --- a/litellm/integrations/otel/mappers/openinference.py +++ b/litellm/integrations/otel/mappers/openinference.py @@ -7,7 +7,7 @@ Phoenix + any other OpenInference-aware backend simultaneously. """ import json -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from itertools import accumulate, chain, groupby from types import MappingProxyType from typing import Final @@ -15,10 +15,12 @@ from typing import Final from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData from litellm.integrations.otel.mappers.utils import ( MAX_TOOL_DEFINITION_ATTRS_PER_SPAN, + MessageToolCall, collect, - drop_none, + drop_none_pairs, json_if, message_content, + message_tool_calls, output_messages, tool_definition_attrs, ) @@ -31,49 +33,146 @@ from litellm.integrations.otel.model.payloads import ( _INPUT_MESSAGES: Final = "llm.input_messages" _OUTPUT_MESSAGES: Final = "llm.output_messages" _MESSAGE_FAMILIES: Final = (_INPUT_MESSAGES, _OUTPUT_MESSAGES) +_MESSAGE_BASE: Final = -1 + +_ParsedMessage = tuple[object, str | None, tuple[MessageToolCall, ...]] -def _message_key_groups(attrs: Mapping[str, AttrValue]) -> Mapping[tuple[str, int], tuple[str, ...]]: - """Per-index message keys in ``attrs`` grouped by ``(family, index)``.""" - tagged: Final = sorted( - (family, int(key.split(".")[2]), key) - for key in attrs - for family in _MESSAGE_FAMILIES - if key.startswith(f"{family}.") +def _parse_message(message: object) -> _ParsedMessage: + role: Final = message.get("role") if isinstance(message, dict) else None + return role, message_content(message), message_tool_calls(message) + + +def _tool_call_attribute_pairs( + prefix: str, idx: int, tool_calls: tuple[MessageToolCall, ...] +) -> Iterator[tuple[str, str | None]]: + for tool_idx, tool_call in enumerate(tool_calls): + yield f"{prefix}.{idx}.message.tool_calls.{tool_idx}.tool_call.id", tool_call.id + yield f"{prefix}.{idx}.message.tool_calls.{tool_idx}.tool_call.function.name", tool_call.name + yield f"{prefix}.{idx}.message.tool_calls.{tool_idx}.tool_call.function.arguments", tool_call.arguments + + +def _message_attribute_pairs( + prefix: str, + messages: Sequence[_ParsedMessage], + *, + with_tool_call_attrs: bool, +) -> Iterator[tuple[str, str | None]]: + for idx, (role, content, tool_calls) in enumerate(messages): + yield f"{prefix}.{idx}.message.role", role if isinstance(role, str) else None + yield f"{prefix}.{idx}.message.content", content + if with_tool_call_attrs: + yield from _tool_call_attribute_pairs(prefix, idx, tool_calls) + + +def _message_value(messages: Sequence[_ParsedMessage]) -> str: + return json.dumps( + [ + { + "role": role, + "content": content, + **({"tool_calls": [tool_call.to_openai_dict() for tool_call in tool_calls]} if tool_calls else {}), + } + for role, content, tool_calls in messages + ] + ) + + +def _message_key_group(key: str) -> tuple[str, int, int, str] | None: + family: Final = next( + (family for family in _MESSAGE_FAMILIES if key.startswith(f"{family}.")), + None, + ) + if family is None: + return None + parts: Final = key.split(".") + message_idx: Final = int(parts[2]) + tool_idx: Final = int(parts[5]) if parts[4] == "tool_calls" else _MESSAGE_BASE + return family, message_idx, tool_idx, key + + +def _message_key_groups(attrs: Mapping[str, AttrValue]) -> Mapping[tuple[str, int, int], tuple[str, ...]]: + """Message and tool-call keys in ``attrs`` grouped by family, message index, and tool index.""" + tagged: Final = tuple(tag for key in attrs if (tag := _message_key_group(key)) is not None) + return MappingProxyType( + {group: tuple(key for _, _, _, key in keys) for group, keys in groupby(sorted(tagged), key=lambda tag: tag[:3])} + ) + + +def _tool_call_groups_by_message( + groups: Mapping[tuple[str, int, int], tuple[str, ...]], +) -> Mapping[tuple[str, int], tuple[tuple[str, int, int], ...]]: + """Tool-call groups indexed by ``(family, message index)``, each tuple highest tool index first. + + Indexing once keeps the shed order linear in the group count: rescanning the + full group map per message made attribute fitting quadratic on long prompts. + """ + ordered: Final = sorted( + (group for group in groups if group[2] != _MESSAGE_BASE), + key=lambda group: (group[0], group[1], group[2]), ) return MappingProxyType( - {group: tuple(key for _, _, key in keys) for group, keys in groupby(tagged, key=lambda tag: tag[:2])} + { + message: tuple(reversed(tuple(message_tool_groups))) + for message, message_tool_groups in groupby(ordered, key=lambda group: (group[0], group[1])) + } ) -def _shed_order(groups: Mapping[tuple[str, int], tuple[str, ...]]) -> tuple[tuple[str, int], ...]: - """Message groups least valuable first: middle prompt turns, extra choices, then the opener, the newest turn - and the first choice.""" - inputs: Final = sorted(idx for family, idx in groups if family == _INPUT_MESSAGES) - outputs: Final = sorted(idx for family, idx in groups if family == _OUTPUT_MESSAGES) +def _message_shed_groups( + groups: Mapping[tuple[str, int, int], tuple[str, ...]], + tool_call_groups: Mapping[tuple[str, int], tuple[tuple[str, int, int], ...]], + family: str, + message_idx: int, +) -> Iterator[tuple[str, int, int]]: + yield from tool_call_groups.get((family, message_idx), ()) + base_group: Final = (family, message_idx, _MESSAGE_BASE) + if base_group in groups: + yield base_group + + +def _shed_order(groups: Mapping[tuple[str, int, int], tuple[str, ...]]) -> tuple[tuple[str, int, int], ...]: + """Middle inputs, extra choices, pinned inputs, then the first choice, with tool calls before message keys.""" + tool_call_groups: Final = _tool_call_groups_by_message(groups) + inputs: Final = sorted(frozenset(idx for family, idx, _ in groups if family == _INPUT_MESSAGES)) + outputs: Final = sorted(frozenset(idx for family, idx, _ in groups if family == _OUTPUT_MESSAGES)) pinned_inputs: Final = tuple(dict.fromkeys((*inputs[:1], *inputs[-1:]))) - return ( + message_order: Final = ( *((_INPUT_MESSAGES, idx) for idx in inputs[1:-1]), *((_OUTPUT_MESSAGES, idx) for idx in reversed(outputs[1:])), *((_INPUT_MESSAGES, idx) for idx in pinned_inputs), *((_OUTPUT_MESSAGES, idx) for idx in outputs[:1]), ) + return tuple( + chain.from_iterable( + _message_shed_groups(groups, tool_call_groups, family, message_idx) for family, message_idx in message_order + ) + ) + + +_METADATA_KEY: Final = "metadata" def fit_indexed_messages(attrs: Mapping[str, AttrValue], budget: int | None) -> Mapping[str, AttrValue]: - """``attrs`` with whole per-index messages shed, least valuable first, until at most ``budget`` keys remain. + """``attrs`` with indexed message attributes shed, least valuable first, until at most ``budget`` keys remain. ``None`` means the span has no attribute count limit. Every message still rides the ``input.value`` and - ``output.value`` blobs, so shedding a per-index pair loses no content. + ``output.value`` blobs, so shedding a per-index pair loses no content. The ``metadata`` blob sheds only + after every indexed message attribute: message attributes are the indexed, queryable view (the blobs + carry no per-index keys), so a count-squeezed span keeps them and the single metadata key absorbs only + the residual shortfall. Shedding happens here, before ``span.set_attribute``, so the fit is exact and + the SDK's dropped-attributes counter never silently masks the choice. """ if budget is None or len(attrs) <= budget: return attrs groups: Final = _message_key_groups(attrs) - order: Final = _shed_order(groups) - running: Final = tuple(accumulate(len(groups[group]) for group in order)) + sheddable: Final = (*(groups[group] for group in _shed_order(groups)),) + flex: Final = (_METADATA_KEY,) if _METADATA_KEY in attrs else () + candidates: Final = (*sheddable, flex) if flex else sheddable + running: Final = tuple(accumulate(len(candidate) for candidate in candidates)) excess: Final = len(attrs) - budget - shed_count: Final = next((n + 1 for n, total in enumerate(running) if total >= excess), len(order)) - shed: Final = frozenset(chain.from_iterable(groups[group] for group in order[:shed_count])) + shed_count: Final = next((n + 1 for n, total in enumerate(running) if total >= excess), len(candidates)) + shed: Final = frozenset(chain.from_iterable(candidates[:shed_count])) return MappingProxyType({key: value for key, value in attrs.items() if key not in shed}) @@ -85,7 +184,9 @@ class OpenInferenceMapper: - ``llm.model_name`` / ``llm.provider`` / ``llm.invocation_parameters`` - ``llm.input_messages.{i}.message.role`` / ``...content`` - ``llm.output_messages.{i}.message.role`` / ``...content`` + - ``llm.output_messages.{i}.message.tool_calls.{j}.tool_call.*`` - ``llm.token_count.prompt`` / ``...completion`` / ``...total`` + - ``metadata`` — JSON object of allowlisted promoted request metadata - ``input.value`` / ``output.value`` — JSON-serialized request / response """ @@ -121,6 +222,7 @@ class OpenInferenceMapper: "llm.invocation_parameters": lambda d: json_if( collect(OpenInferenceMapper._INVOCATION_PARAMS, d.request_params) ), + "metadata": lambda d: json_if(dict(sorted(d.promoted_metadata.items()))), } def __init__(self, tool_attr_budget: int = MAX_TOOL_DEFINITION_ATTRS_PER_SPAN) -> None: @@ -137,27 +239,36 @@ class OpenInferenceMapper: return { **collect(self._LLM_CALL_ATTRS, data), **collect(self._BLOB_ATTRS, data), - **self._messages(_INPUT_MESSAGES, "input.value", data.messages_in), - **self._messages(_OUTPUT_MESSAGES, "output.value", output_messages(data)), + **self._messages( + _INPUT_MESSAGES, + "input.value", + data.messages_in, + with_tool_call_attrs=False, + ), + **self._messages( + _OUTPUT_MESSAGES, + "output.value", + output_messages(data), + with_tool_call_attrs=True, + ), **self._tools(data), } @staticmethod - def _messages(prefix: str, value_key: str, messages: Sequence[object]) -> AttributeMap: + def _messages( + prefix: str, + value_key: str, + messages: Sequence[object], + *, + with_tool_call_attrs: bool, + ) -> AttributeMap: """``{prefix}.{idx}.message.*`` keys for every message + the ``value_key`` blob of all of them.""" - parsed: Final = [(m.get("role") if isinstance(m, dict) else None, message_content(m)) for m in messages] - attrs: Final = drop_none( - { - key: value - for idx, (role, content) in enumerate(parsed) - for key, value in ( - (f"{prefix}.{idx}.message.role", role if isinstance(role, str) else None), - (f"{prefix}.{idx}.message.content", content), - ) - } + parsed: Final = tuple(_parse_message(message) for message in messages) + attrs: Final = drop_none_pairs( + _message_attribute_pairs(prefix, parsed, with_tool_call_attrs=with_tool_call_attrs) ) if parsed: - attrs[value_key] = json.dumps([{"role": role, "content": content} for role, content in parsed]) + attrs[value_key] = _message_value(parsed) return attrs def _tools(self, data: LLMCallSpanData) -> AttributeMap: diff --git a/litellm/integrations/otel/mappers/utils.py b/litellm/integrations/otel/mappers/utils.py index 5582734585f..80ef2dd553a 100644 --- a/litellm/integrations/otel/mappers/utils.py +++ b/litellm/integrations/otel/mappers/utils.py @@ -7,11 +7,28 @@ they live in one place. import json from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass from typing import Final from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue from litellm.integrations.otel.model.payloads import LLMCallSpanData, ToolDefinition + +@dataclass(frozen=True, slots=True) +class MessageToolCall: + id: str | None + type: str + name: str | None + arguments: str | None + + def to_openai_dict(self) -> Mapping[str, object]: + return { + "id": self.id, + "type": self.type, + "function": {"name": self.name, "arguments": self.arguments}, + } + + DEFAULT_SPAN_ATTRIBUTE_LIMIT: Final = 128 """The OTel SDK's default per-span attribute count limit.""" @@ -121,3 +138,50 @@ def message_content(message: object) -> str | None: def output_messages(data: LLMCallSpanData) -> list: """The ``message`` payload of each response choice.""" return [c.get("message") for c in data.choices_out if isinstance(c, dict)] + + +def message_tool_calls(message: object) -> tuple[MessageToolCall, ...]: + if not isinstance(message, dict): + return () + tool_calls: Final = message.get("tool_calls") + if not isinstance(tool_calls, (list, tuple)) or not tool_calls: + return () + return tuple(tool_call for value in tool_calls if (tool_call := _message_tool_call(value)) is not None) + + +def _message_tool_call(value: object) -> MessageToolCall | None: + if not isinstance(value, dict): + return None + function: Final = value.get("function") + function_data: Final = function if isinstance(function, dict) else {} + raw_type: Final = value.get("type") + raw_arguments: Final = function_data.get("arguments") + arguments: Final = ( + raw_arguments + if isinstance(raw_arguments, str) + else _stringify_tool_arguments(raw_arguments) + if raw_arguments is not None + else None + ) + name: Final = function_data.get("name") + identifier: Final = value.get("id") + return MessageToolCall( + id=identifier if isinstance(identifier, str) else None, + type=raw_type if isinstance(raw_type, str) else "function", + name=name if isinstance(name, str) else None, + arguments=arguments, + ) + + +def _stringify_tool_arguments(value: object) -> str: + """Serialize non-string tool-call arguments, falling back to ``repr``. + + Arguments normally arrive as already-JSON strings, but provider adapters and + ``model_construct`` responses hand over raw Python objects. ``json.dumps`` + raises on those (tuple-keyed dicts, cycles) and the escaping exception would + lose the whole span, so keep a readable ``repr`` instead. + """ + try: + return json.dumps(value, default=str) + except (TypeError, ValueError): + return repr(value) diff --git a/litellm/integrations/otel/model/baggage.py b/litellm/integrations/otel/model/baggage.py index 131848e1380..606fd79f218 100644 --- a/litellm/integrations/otel/model/baggage.py +++ b/litellm/integrations/otel/model/baggage.py @@ -18,7 +18,7 @@ from collections.abc import Callable, Mapping from types import MappingProxyType from typing import Final -from litellm.integrations.otel.model.metadata import REQUESTER_METADATA_PATH, RequestIdentity +from litellm.integrations.otel.model.metadata import RequestIdentity, allowlisted_metadata from litellm.integrations.otel.model.semconv import GenAI, LiteLLM # Attribute key -> value extractor over (identity, request_model, @@ -92,9 +92,8 @@ def promoted_metadata(metadata: Mapping[str, str], metadata_keys: tuple[str, ... """Allowlisted entries of a flattened metadata mapping under ``litellm.metadata.*``.""" return MappingProxyType( { - f"{LiteLLM.METADATA_PREFIX}{meta_key.removeprefix(REQUESTER_METADATA_PATH)}": value - for meta_key in metadata_keys - if (value := metadata.get(meta_key)) + f"{LiteLLM.METADATA_PREFIX}{meta_key}": value + for meta_key, value in allowlisted_metadata(metadata, metadata_keys).items() } ) diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index 809fc794461..8524df71b44 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -61,6 +61,16 @@ REQUESTER_METADATA_KEY: Final = "requester_metadata" REQUESTER_METADATA_PATH: Final = f"{REQUESTER_METADATA_KEY}." +def allowlisted_metadata(metadata: Mapping[str, str], metadata_keys: tuple[str, ...]) -> Mapping[str, str]: + return MappingProxyType( + { + meta_key.removeprefix(REQUESTER_METADATA_PATH): value + for meta_key in metadata_keys + if (value := metadata.get(meta_key)) + } + ) + + @dataclass(frozen=True) class RequestIdentity: call_id: str | None = None diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index ccec93920c9..101e6285946 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -12,7 +12,7 @@ from urllib.parse import urlsplit from typing_extensions import ReadOnly, TypedDict -from litellm.integrations.otel.model.metadata import RequestContext, RequestIdentity +from litellm.integrations.otel.model.metadata import RequestContext, RequestIdentity, allowlisted_metadata from litellm.integrations.otel.model.semconv import ( GenAIOperation, GenAIOutputType, @@ -61,6 +61,8 @@ if TYPE_CHECKING: StandardLoggingPayload, ) +_EMPTY_METADATA: Final[Mapping[str, str]] = MappingProxyType({}) + # --- typed sub-structures ---------------------------------------------------- # @@ -437,6 +439,7 @@ class LLMCallSpanData: trace: TraceControls = field(default_factory=TraceControls) session_id: str | None = None embedding_output: EmbeddingOutput | None = None + promoted_metadata: Mapping[str, str] = field(default_factory=lambda: _EMPTY_METADATA) routing_attributes: Mapping[str, RoutingAttributeValue] = field(default_factory=lambda: MappingProxyType({})) @classmethod @@ -449,6 +452,8 @@ class LLMCallSpanData: request_purpose: str | None = None, trace: TraceControls | None = None, session_id: str | None = None, + *, + metadata_keys: tuple[str, ...] = (), ) -> LLMCallSpanData: params: Final = cast(Mapping[str, object], payload.get("model_parameters") or {}) # The single parse of the request's metadata — the request-vs-provider @@ -487,6 +492,7 @@ class LLMCallSpanData: cost=LLMCost.from_breakdown(cast("Mapping[str, object] | None", payload.get("cost_breakdown"))), server=ServerInfo.from_api_base(context.api_base), identity=context.identity, + promoted_metadata=allowlisted_metadata(context.identity.metadata, metadata_keys), is_streaming=as_bool(payload.get("stream")), tools=_extract_tools(params), messages_in=_dicts(payload.get("messages")) if capture_content else (), diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index a781be610a6..473be107cbb 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -782,7 +782,7 @@ class RealTimeStreaming: isinstance(cb, CustomGuardrail) and any( cb.should_run_guardrail( - data=self.request_data, + data={**self.request_data, "stream": True}, event_type=et, ) for et in event_hooks @@ -847,7 +847,7 @@ class RealTimeStreaming: if event_hooks is None: event_hooks = [GuardrailEventHooks.realtime_input_transcription] _realtime_event_types: Final = event_hooks - _check_data: Final = {**self.request_data, "transcript": transcript} + _check_data: Final = {**self.request_data, "transcript": transcript, "stream": True} _already_run: Final[set] = set() for callback in litellm.callbacks: diff --git a/litellm/llms/base_llm/decisions/systemone.py b/litellm/llms/base_llm/decisions/systemone.py deleted file mode 100644 index 43be31b0c30..00000000000 --- a/litellm/llms/base_llm/decisions/systemone.py +++ /dev/null @@ -1,247 +0,0 @@ -"""The Jev / System One wire shape and its translation to and from the OpenAI Decisions shape. - -System One (TypeSafe, Perplexity, OpenRouter, Cloudflare Clef, Strands Decider) takes -{"model", "state", "questions": {name: question}} and answers with {"model", "answers": {name: answer}, "usage"}. -Predicates are `noul` questions, choice options are a `criteria` map, score levels are a `criteria` list. -""" - -import itertools -from collections.abc import Mapping, Sequence -from typing import Final, Literal, TypeAlias - -from pydantic import ConfigDict, TypeAdapter -from typing_extensions import assert_never - -from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.types.llms.base import LiteLLMPydanticObjectBase -from litellm.types.openai_decisions import ( - ChoiceAnswer, - ChoiceProbability, - ChoiceQuestion, - DecisionAnswer, - DecisionChoice, - DecisionInput, - DecisionInputMessage, - DecisionInputPart, - DecisionInputTokensDetails, - DecisionOutputTokensDetails, - DecisionQuestion, - DecisionsRequest, - DecisionsRequestBody, - DecisionsResponse, - DecisionUsage, - PredicateAnswer, - PredicateQuestion, - ScoreAnswer, - ScoreProbability, - ScoreQuestion, -) - - -class SystemOneObjectBase(LiteLLMPydanticObjectBase): - model_config = ConfigDict(extra="allow", frozen=True) - - -class SystemOneNoulAnswer(SystemOneObjectBase): - type: Literal["noul"] - noul: float - - -class SystemOneChoiceAnswer(SystemOneObjectBase): - type: Literal["choice"] - choice: str - confidence: float - probabilities: Mapping[str, float] - - -class SystemOneScoreAnswer(SystemOneObjectBase): - type: Literal["score"] - score: float - confidence: float - probabilities: Mapping[str, float] - - -SystemOneAnswer: TypeAlias = SystemOneNoulAnswer | SystemOneChoiceAnswer | SystemOneScoreAnswer - - -class SystemOneUsage(SystemOneObjectBase): - input_tokens: int = 0 - output_tokens: int = 0 - - -class SystemOneResponse(SystemOneObjectBase): - model: str | None = None - answers: Mapping[str, SystemOneAnswer] - usage: SystemOneUsage | None = None - - -SYSTEM_ONE_RESPONSE_ADAPTER: Final[TypeAdapter[SystemOneResponse]] = TypeAdapter(SystemOneResponse) - - -def _unsupported(what: str, custom_llm_provider: str) -> BaseLLMException: - return BaseLLMException( - status_code=400, - message=f"Decisions provider '{custom_llm_provider}' does not support {what}", - ) - - -def to_system_one_request(model: str, body: DecisionsRequestBody, custom_llm_provider: str) -> dict[str, object]: - keys: Final = question_keys(body.questions, custom_llm_provider) - return { - "model": model, - "state": _state(body.input, custom_llm_provider), - "questions": { - key: _question(question, custom_llm_provider) for key, question in zip(keys, body.questions, strict=True) - }, - } - - -def question_keys(questions: Sequence[DecisionQuestion], custom_llm_provider: str) -> tuple[str, ...]: - """System One keys questions and answers by name, so unnamed questions get a positional key.""" - names: Final = tuple(question.name for question in questions if question.name is not None) - if len(set(names)) != len(names): - raise BaseLLMException( - status_code=400, - message=f"Decisions provider '{custom_llm_provider}' requires a unique name per question", - ) - taken: Final = frozenset(names) - return tuple( - question.name if question.name is not None else _positional_key(index, taken) - for index, question in enumerate(questions) - ) - - -def _positional_key(index: int, taken: frozenset[str]) -> str: - candidates: Final = (f"{'_' * depth}q{index}" for depth in itertools.count()) - return next(key for key in candidates if key not in taken) - - -def _state(input_value: DecisionInput, custom_llm_provider: str) -> str: - if isinstance(input_value, str): - return input_value - return "\n".join(_message_text(message, custom_llm_provider) for message in input_value) - - -def _message_text(message: DecisionInputMessage, custom_llm_provider: str) -> str: - if isinstance(message.content, str): - return message.content - return "\n".join(_part_text(part, custom_llm_provider) for part in message.content) - - -def _part_text(part: DecisionInputPart, custom_llm_provider: str) -> str: - if part.type != "input_text": - raise _unsupported("input_image parts", custom_llm_provider) - return part.text - - -def _question(question: DecisionQuestion, custom_llm_provider: str) -> dict[str, object]: - match question: - case PredicateQuestion(): - return {"type": "noul", "instructions": question.instructions} - case ChoiceQuestion(): - return { - "type": "choice", - "instructions": question.instructions, - "criteria": _choice_criteria(question, custom_llm_provider), - } - case ScoreQuestion(): - return { - "type": "score", - "instructions": question.instructions, - "criteria": [_level_text(level.label, level.description) for level in question.levels], - } - case _: - assert_never(question) - - -def _choice_criteria(question: ChoiceQuestion, custom_llm_provider: str) -> dict[str, str | None]: - values: Final = tuple(_choice_key(choice, custom_llm_provider) for choice in question.choices) - if len(set(values)) != len(values): - raise _unsupported("repeated choice values", custom_llm_provider) - return {value: choice.description for value, choice in zip(values, question.choices, strict=True)} - - -def _choice_key(choice: DecisionChoice, custom_llm_provider: str) -> str: - if not isinstance(choice.value, str): - raise _unsupported("boolean choice values", custom_llm_provider) - return choice.value - - -def _level_text(label: str, description: str | None) -> str: - return description if description is not None else label - - -def to_decisions_response( - system_one: SystemOneResponse, - request: DecisionsRequest, - custom_llm_provider: str, -) -> DecisionsResponse: - keys: Final = question_keys(request.body.questions, custom_llm_provider) - return DecisionsResponse( - model=system_one.model if system_one.model is not None else request.model, - answers=[ - _answer(key, question, system_one.answers, custom_llm_provider) - for key, question in zip(keys, request.body.questions, strict=True) - ], - usage=_usage(system_one.usage), - ) - - -def _answer( - key: str, - question: DecisionQuestion, - answers: Mapping[str, SystemOneAnswer], - custom_llm_provider: str, -) -> DecisionAnswer: - match question, answers.get(key): - case PredicateQuestion(), SystemOneNoulAnswer() as answer: - return PredicateAnswer(type="predicate", name=question.name, probability=answer.noul) - case ChoiceQuestion(), SystemOneChoiceAnswer() as answer: - return ChoiceAnswer( - type="choice", - name=question.name, - choice=answer.choice, - probabilities=_choice_probabilities(question, answer), - confidence=answer.confidence, - ) - case ScoreQuestion(), SystemOneScoreAnswer() as answer: - return ScoreAnswer( - type="score", - name=question.name, - score=answer.score, - probabilities=_score_probabilities(question, answer), - confidence=answer.confidence, - ) - case _: - raise BaseLLMException( - status_code=500, - message=( - f"Decisions provider '{custom_llm_provider}' returned no {question.type} answer " - f"for question '{key}'" - ), - ) - - -def _choice_probabilities(question: ChoiceQuestion, answer: SystemOneChoiceAnswer) -> list[ChoiceProbability]: - return [ - ChoiceProbability(value=choice.value, probability=answer.probabilities.get(str(choice.value), 0.0)) - for choice in question.choices - ] - - -def _score_probabilities(question: ScoreQuestion, answer: SystemOneScoreAnswer) -> list[ScoreProbability]: - return [ - ScoreProbability(value=index, label=level.label, probability=answer.probabilities.get(str(index), 0.0)) - for index, level in enumerate(question.levels) - ] - - -def _usage(usage: SystemOneUsage | None) -> DecisionUsage: - counted: Final = usage if usage is not None else SystemOneUsage() - return DecisionUsage( - input_tokens=counted.input_tokens, - input_tokens_details=DecisionInputTokensDetails(cached_tokens=0, cache_write_tokens=0), - output_tokens=counted.output_tokens, - output_tokens_details=DecisionOutputTokensDetails(reasoning_tokens=0), - total_tokens=counted.input_tokens + counted.output_tokens, - ) diff --git a/litellm/llms/bedrock/chat/chat_completions/transformation.py b/litellm/llms/bedrock/chat/chat_completions/transformation.py index ad5f0da8542..f7a81082f80 100644 --- a/litellm/llms/bedrock/chat/chat_completions/transformation.py +++ b/litellm/llms/bedrock/chat/chat_completions/transformation.py @@ -36,6 +36,7 @@ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token from litellm.llms.bedrock.common_utils import ( BedrockError, bedrock_model_is_openai_gpt, + bedrock_runtime_chat_completions_serves_reasoning_inline, split_bedrock_region_path, ) from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler @@ -213,7 +214,12 @@ def split_reasoning_tag(content: str) -> tuple[str | None, str]: class BedrockRuntimeChatCompletionsStreamingHandler(OpenAIChatCompletionStreamingHandler): - """OpenAI chunk parsing plus the ```` split, tracked per choice index.""" + """OpenAI chunk parsing plus the inline ```` split, tracked per choice index. + + Every chunk echoes the model id litellm sent, so the split engages only when that id's price-map + row carries ``supports_bedrock_runtime_chat_completions_inline_reasoning`` (gpt-oss); a GPT 5.6 or Grok + answer that starts with a literal ```` tag streams as content. + """ def __init__( self, @@ -226,6 +232,8 @@ class BedrockRuntimeChatCompletionsStreamingHandler(OpenAIChatCompletionStreamin def chunk_parser(self, chunk: dict) -> ModelResponseStream: # mutable-ok: BaseModelResponseIterator signature parsed: Final = super().chunk_parser(chunk) + if not bedrock_runtime_chat_completions_serves_reasoning_inline(parsed.model or ""): + return parsed for choice in parsed.choices: next_state, reasoning, content = _split_streamed_content( self._splitters.get(choice.index, ReasoningTagSplitter()), @@ -475,6 +483,8 @@ class AmazonBedrockRuntimeChatCompletionsConfig(OpenAILikeChatConfig): json_mode=json_mode, ) set_provider_response_headers_in_hidden_params(response, raw_response.headers) + if not bedrock_runtime_chat_completions_serves_reasoning_inline(model): + return response for choice in response.choices: if not isinstance(choice, Choices) or not isinstance(choice.message.content, str): continue diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 7c637bf6724..c21cb61ce60 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -909,6 +909,17 @@ def bedrock_model_is_openai_gpt(model: str) -> bool: return _openai_gpt_version(model) is not None +def bedrock_runtime_chat_completions_serves_reasoning_inline(model: str) -> bool: + """Whether AWS's native Chat Completions writes this model's reasoning inline in the answer text. + + Data-driven from the price-map ``supports_bedrock_runtime_chat_completions_inline_reasoning`` flag (gpt-oss). + A flagged model opens its answer with a ``...`` block instead of a + ``reasoning_content`` field, so litellm splits that block out for it and keeps every other model's + text as sent. + """ + return _bedrock_price_map_flag(model, "supports_bedrock_runtime_chat_completions_inline_reasoning") + + BEDROCK_CONVERSE_ONLY_REQUEST_KEYS: Final = frozenset( ( "guardrailConfig", diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 2d02b152c61..ed61707f99e 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -3,27 +3,100 @@ Bedrock Token Counter implementation using the CountTokens API. """ from collections.abc import Mapping, Sequence +from dataclasses import dataclass from typing import Any, Final +from pydantic import JsonValue +from typing_extensions import assert_never + from litellm._logging import verbose_logger from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler +from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountResponse + +RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api" +MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api" + + +class _CountTokensReply(LiteLLMPydanticObjectBase): + input_tokens: int + + +@dataclass(frozen=True, slots=True) +class CountedTokens: + input_tokens: int + original_response: Mapping[str, JsonValue] + tokenizer_type: str + + +@dataclass(frozen=True, slots=True) +class CountTokensFailure: + status_code: int + message: str + tokenizer_type: str + + +CountTokensOutcome = CountedTokens | CountTokensFailure + + +def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool: + return ( + isinstance(outcome, CountTokensFailure) + and outcome.status_code == 400 + and is_mantle_claude_model(resolved_model) + ) class BedrockTokenCounter(BaseTokenCounter): - """Token counter implementation for AWS Bedrock provider using the CountTokens API.""" + """Token counter for AWS Bedrock: bedrock-runtime CountTokens first, and Anthropic's count_tokens + on bedrock-mantle for the Claude models bedrock-runtime answers 400 for""" + + def __init__( + self, + runtime_handler: BedrockCountTokensHandler | None = None, + mantle_handler: BedrockMantleCountTokensHandler | None = None, + client: AsyncHTTPHandler | None = None, + ) -> None: + self._runtime_handler: Final = runtime_handler or BedrockCountTokensHandler() + self._mantle_handler: Final = mantle_handler or BedrockMantleCountTokensHandler() + self._client: Final = client def should_use_token_counting_api( self, custom_llm_provider: str | None = None, ) -> bool: - """ - Returns True if we should use the Bedrock CountTokens API for token counting. - """ return custom_llm_provider == LlmProviders.BEDROCK.value + async def _count_with( + self, + handler: BedrockCountTokensHandler | BedrockMantleCountTokensHandler, + tokenizer_type: str, + request_data: dict[str, object], + litellm_params: dict[str, object], + resolved_model: str, + ) -> CountTokensOutcome: + try: + result: Final = await handler.handle_count_tokens_request( + request_data=request_data, + litellm_params=litellm_params, + resolved_model=resolved_model, + client=self._client, + ) + reply: Final = _CountTokensReply.model_validate(result) + except BedrockError as e: + verbose_logger.debug( + "%s CountTokens API error: status=%s, message=%s", tokenizer_type, e.status_code, e.message + ) + return CountTokensFailure(status_code=e.status_code, message=e.message, tokenizer_type=tokenizer_type) + except Exception as e: + verbose_logger.debug("Error calling %s CountTokens API: %s", tokenizer_type, e) + return CountTokensFailure(status_code=500, message=str(e), tokenizer_type=tokenizer_type) + return CountedTokens(input_tokens=reply.input_tokens, original_response=result, tokenizer_type=tokenizer_type) + async def count_tokens( self, model_to_use: str, @@ -34,81 +107,52 @@ class BedrockTokenCounter(BaseTokenCounter): tools: Sequence[Mapping[str, object]] | None = None, system: object | None = None, ) -> TokenCountResponse | None: - """ - Count tokens using AWS Bedrock's CountTokens API. - - This method calls the existing BedrockCountTokensHandler to make an API call - to Bedrock's token counting endpoint, bypassing the local tiktoken-based counting. - - Args: - model_to_use: The model identifier - messages: The messages to count tokens for - contents: Alternative content format (not used for Bedrock) - deployment: Deployment configuration containing litellm_params - request_model: The original request model name - - Returns: - TokenCountResponse with token count, or None if counting fails - """ if not messages: return None - deployment = deployment or {} - litellm_params: Final = deployment.get("litellm_params", {}) - - # Build request data in the format expected by BedrockCountTokensHandler + litellm_params: Final = (deployment or {}).get("litellm_params", {}) request_data: Final[dict[str, object]] = { "model": model_to_use, "messages": messages, + **({"tools": tools} if tools else {}), + **({"system": system} if system else {}), } - - if tools: - request_data["tools"] = tools - - if system: - request_data["system"] = system - - # Get the resolved model (strip prefixes like bedrock/, converse/, etc.) resolved_model: Final = get_bedrock_base_model(model_to_use) - try: - handler: Final = BedrockCountTokensHandler() - result: Final = await handler.handle_count_tokens_request( - request_data=request_data, - litellm_params=litellm_params, - resolved_model=resolved_model, + runtime: Final = await self._count_with( + self._runtime_handler, RUNTIME_TOKENIZER_TYPE, request_data, litellm_params, resolved_model + ) + outcome: Final = ( + await self._count_with( + self._mantle_handler, MANTLE_TOKENIZER_TYPE, request_data, litellm_params, resolved_model ) - - # Transform response to TokenCountResponse - if result is not None: + if _runtime_rejected_claude_model(runtime, resolved_model) + else runtime + ) + match outcome: + case CountedTokens(): return TokenCountResponse( - total_tokens=result.get("input_tokens", 0), + total_tokens=outcome.input_tokens, request_model=request_model, model_used=model_to_use, - tokenizer_type="bedrock_api", - original_response=result, + tokenizer_type=outcome.tokenizer_type, + original_response=dict(outcome.original_response), ) - except BedrockError as e: - verbose_logger.warning("Bedrock CountTokens API error: status=%s, message=%s", e.status_code, e.message) - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type="bedrock_api", - error=True, - error_message=e.message, - status_code=e.status_code, - ) - except Exception as e: - verbose_logger.warning("Error calling Bedrock CountTokens API: %s", e) - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type="bedrock_api", - error=True, - error_message=str(e), - status_code=500, - ) - - return None + case CountTokensFailure(): + verbose_logger.warning( + "%s CountTokens API error: status=%s, message=%s", + outcome.tokenizer_type, + outcome.status_code, + outcome.message, + ) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type=outcome.tokenizer_type, + error=True, + error_message=outcome.message, + status_code=outcome.status_code, + ) + case _: + assert_never(outcome) diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index d7fb510f057..be8243e48ef 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -125,7 +125,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): raise except httpx.HTTPStatusError as e: # HTTP errors - preserve the actual status code - verbose_logger.error("HTTP error in CountTokens handler: %s", e) + verbose_logger.debug("HTTP error in CountTokens handler: %s", e) raise BedrockError( status_code=e.response.status_code, message=e.response.text, diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py new file mode 100644 index 00000000000..23c3dc54c95 --- /dev/null +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -0,0 +1,99 @@ +from typing import Final + +import httpx +from pydantic import JsonValue, TypeAdapter +from typing_extensions import NotRequired, ReadOnly, TypedDict + +import litellm +from litellm._logging import verbose_logger +from litellm.llms.anthropic.count_tokens.transformation import AnthropicCountTokensConfig +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing +from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client + +MANTLE_COUNT_TOKENS_SUFFIX: Final = "/count_tokens" +MANTLE_ANTHROPIC_VERSION: Final = "2023-06-01" + + +class MantleCountTokensRequest(TypedDict): + messages: ReadOnly[list[dict[str, JsonValue]]] + system: ReadOnly[NotRequired[JsonValue]] + tools: ReadOnly[NotRequired[list[dict[str, JsonValue]]]] + + +_COUNT_REQUEST: Final = TypeAdapter(MantleCountTokensRequest) +_COUNT_RESPONSE: Final = TypeAdapter(dict[str, JsonValue]) + + +class BedrockMantleCountTokensHandler(BaseAWSLLM): + """Counts tokens through Anthropic's count_tokens on the bedrock-mantle endpoint. + + Claude models that Bedrock offers only through cross-region inference answer 400 on + bedrock-runtime's CountTokens; AWS documents Mantle's /anthropic/v1/messages/count_tokens + as the way to count them, with the base model id and the deployment's AWS credentials + """ + + def __init__(self, anthropic_config: AnthropicCountTokensConfig | None = None) -> None: + super().__init__() + self._anthropic_config: Final = anthropic_config or AnthropicCountTokensConfig() + + def get_mantle_count_tokens_endpoint(self, aws_region_name: str) -> str: + messages_url: Final = build_mantle_messages_url( + api_base=None, aws_bedrock_runtime_endpoint=None, region=aws_region_name + ) + return f"{messages_url}{MANTLE_COUNT_TOKENS_SUFFIX}" + + async def handle_count_tokens_request( + self, + request_data: dict[str, object], + litellm_params: dict[str, object], + resolved_model: str, + client: AsyncHTTPHandler | None = None, + ) -> dict[str, JsonValue]: + try: + request: Final = _COUNT_REQUEST.validate_python(request_data) + aws_region_name: Final = self._get_aws_region_name( + optional_params=litellm_params, model=resolved_model, model_id=None + ) + endpoint_url: Final = self.get_mantle_count_tokens_endpoint(aws_region_name) + body: Final = self._anthropic_config.transform_request_to_count_tokens( + model=resolved_model, + messages=request["messages"], + tools=request.get("tools"), + system=request.get("system"), + ) + verbose_logger.debug("Making bedrock-mantle count_tokens request to: %s", endpoint_url) + api_key: Final = litellm_params.get("api_key") + signed_headers, signed_body = await run_aws_signing( + self._sign_request, + service_name="bedrock", + headers={"Content-Type": "application/json", "anthropic-version": MANTLE_ANTHROPIC_VERSION}, + optional_params=litellm_params, + request_data=body, + api_base=endpoint_url, + model=resolved_model, + api_key=api_key if isinstance(api_key, str) else None, + ) + async_client: Final = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) + response: Final = await async_client.post( + endpoint_url, headers=signed_headers, data=signed_body, timeout=30.0 + ) + if response.status_code != 200: + raise BedrockError( + status_code=response.status_code, + message=response.text, + headers=response.headers, + response=response, + ) + return _COUNT_RESPONSE.validate_json(response.content) + except BedrockError: + raise + except httpx.HTTPStatusError as e: + raise BedrockError( + status_code=e.response.status_code, + message=e.response.text, + headers=e.response.headers, + response=e.response, + ) + except Exception as e: + raise BedrockError(status_code=500, message=f"bedrock-mantle count_tokens processing error: {e}") diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index c61576fa7b2..0d7f7cd3562 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -22,6 +22,13 @@ if TYPE_CHECKING: from litellm.types.utils import CostResponseTypes +BEDROCK_STREAMING_ACTIONS: Final = frozenset({"invoke-with-response-stream", "converse-stream"}) + + +def is_bedrock_streaming_endpoint(endpoint: str) -> bool: + return endpoint.partition("?")[0].rstrip("/").rsplit("/", 1)[-1] in BEDROCK_STREAMING_ACTIONS + + _TEXT_ONLY_DELTA_FIELDS: Final = frozenset({"content", "role"}) diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 1f295a6e656..fd58ca2ddbb 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -10,6 +10,7 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Optional from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import without_server_streaming_classification from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.proxy._types import PassThroughGuardrailSettings from litellm.types.utils import GenericGuardrailAPIInputs @@ -80,7 +81,9 @@ class PassThroughEndpointHandler(BaseTranslation): from litellm.litellm_core_utils.safe_json_dumps import safe_dumps payload_to_check: Final = { - k: v for k, v in data.items() if not k.startswith("_") and k not in ("metadata", "litellm_logging_obj") + k: v + for k, v in without_server_streaming_classification(data).items() + if not k.startswith("_") and k not in ("metadata", "litellm_logging_obj") } verbose_proxy_logger.debug("PassThroughEndpointHandler: Using full payload for guardrail") return safe_dumps(payload_to_check) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 8c77413bace..6b759065706 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -25,6 +25,22 @@ from litellm.types.llms.vertex_ai import ( from litellm.types.utils import TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages +GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType( + {"fps": "fps", "start_offset": "startOffset", "end_offset": "endOffset"} +) + + +GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/" + + +def gemini_video_metadata_from_openai(video_metadata: Mapping[str, object]) -> dict[str, object]: + return { + gemini_key: video_metadata[openai_key] + for openai_key, gemini_key in GEMINI_VIDEO_METADATA_KEYS.items() + if openai_key in video_metadata + } + + VERTEX_SELF_DEPLOYED_ENDPOINT_UNSUPPORTED_PARAMS: Final = frozenset( { "audio", diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index da5b974f186..09ec2ba6ee4 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -67,7 +67,7 @@ from litellm.types.llms.openai import ( OpenAIFilesPurpose, PathLike, ) -from litellm.types.llms.vertex_ai import GcsBucketResponse, GeminiEmbeddingInput +from litellm.types.llms.vertex_ai import GcsBucketResponse, GeminiEmbeddingElement, GeminiEmbeddingInput from litellm.types.utils import ( Embedding, EmbeddingResponse, @@ -555,18 +555,27 @@ def _is_responses_batch_entry(openai_entry: Mapping[str, object]) -> bool: return path == "responses" or path.endswith("/responses") +def _own_embedding_input( + element: GeminiEmbeddingElement | list[str] | list[GeminiEmbeddingElement], +) -> GeminiEmbeddingInput: + if isinstance(element, (str, list)): + return element + file_block_alone: Final[list[GeminiEmbeddingElement]] = [element] + return file_block_alone + + def _openai_embedding_input_elements( embedding_input: GeminiEmbeddingInput, -) -> tuple[str | list[str], ...]: +) -> tuple[GeminiEmbeddingInput, ...]: """ Split an OpenAI `input` into the elements that each get their own embedding. - A string is one embedding, a flat array is one embedding per element, and a nested - array is one combined embedding per inner array, matching the online - `batchEmbedContents` path. + A string or a file content block is one embedding, a flat array is one embedding + per element, and a nested array is one combined embedding per inner array, + matching the online `batchEmbedContents` path. """ if isinstance(embedding_input, list): - return tuple(embedding_input) + return tuple(_own_embedding_input(element) for element in embedding_input) return (embedding_input,) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 8749b81d7f1..15e5e440b91 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -57,7 +57,9 @@ from litellm.types.llms.vertex_ai import ( from litellm.types.utils import GenericImageParsingChunk, LlmProviders from ..common_utils import ( + GEMINI_FILES_API_URI_PREFIX, check_text_in_content, + gemini_video_metadata_from_openai, get_supports_response_schema, get_supports_system_message, ) @@ -69,7 +71,6 @@ _GCS_METADATA_VERTEX_BASE: object | None = None # Shared sync client for GCS JSON API metadata reads so proxy/SSL settings # from litellm's HTTP stack apply (see Greptile review on PR #27278). _GCS_METADATA_HTTP_HANDLER: HTTPHandler | None = None -GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/" _GEMINI_MIME_TYPE_ALIASES: Final[dict[str, str]] = { "image/jpg": "image/jpeg", } @@ -197,13 +198,7 @@ def _apply_gemini_metadata( part_dict["media_resolution"] = media_resolution_enum if video_metadata is not None: - gemini_video_metadata: Final = {} - if "fps" in video_metadata: - gemini_video_metadata["fps"] = video_metadata["fps"] - if "start_offset" in video_metadata: - gemini_video_metadata["startOffset"] = video_metadata["start_offset"] - if "end_offset" in video_metadata: - gemini_video_metadata["endOffset"] = video_metadata["end_offset"] + gemini_video_metadata: Final = gemini_video_metadata_from_openai(video_metadata) if gemini_video_metadata: part_dict["video_metadata"] = gemini_video_metadata diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index d15cf89cf47..0a800497de1 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -22,6 +22,8 @@ from litellm.types.utils import EmbeddingResponse from ..gemini.vertex_and_google_ai_studio_gemini import VertexLLM from .batch_embed_content_transformation import ( + file_reference_name, + flatten_media_sources, is_file_reference, process_embed_content_response, process_response, @@ -38,11 +40,8 @@ class GoogleBatchEmbeddings(VertexLLM): def _flatten_and_detect_file_refs( input: GeminiEmbeddingInput, ) -> tuple[list[str], bool]: - """Flatten nested input lists and detect file references.""" - input_list: Final = [input] if isinstance(input, str) else input - flat_elements: Final = [ - e for item in input_list for e in (item if isinstance(item, list) else [item]) if isinstance(e, str) - ] + """Flatten nested input lists and file content blocks into their sources and detect file references.""" + flat_elements: Final = list(flatten_media_sources(input)) has_file_refs: Final = any(is_file_reference(e) for e in flat_elements) return flat_elements, has_file_refs @@ -68,7 +67,7 @@ class GoogleBatchEmbeddings(VertexLLM): for element in input_list: if isinstance(element, str) and is_file_reference(element): - url = f"https://generativelanguage.googleapis.com/v1beta/{element}" + url = f"https://generativelanguage.googleapis.com/v1beta/{file_reference_name(element)}" headers = {"x-goog-api-key": api_key} response = sync_handler.get(url=url, headers=headers) @@ -105,7 +104,7 @@ class GoogleBatchEmbeddings(VertexLLM): for element in input_list: if isinstance(element, str) and is_file_reference(element): - url = f"https://generativelanguage.googleapis.com/v1beta/{element}" + url = f"https://generativelanguage.googleapis.com/v1beta/{file_reference_name(element)}" headers = {"x-goog-api-key": api_key} response = await async_handler.get(url=url, headers=headers) @@ -139,6 +138,7 @@ class GoogleBatchEmbeddings(VertexLLM): timeout=300, client=None, extra_headers: dict | None = None, + drop_params: bool = False, ) -> EmbeddingResponse: _auth_header, vertex_project = self._ensure_access_token( credentials=vertex_credentials, @@ -207,22 +207,26 @@ class GoogleBatchEmbeddings(VertexLLM): api_key=api_key, optional_params=optional_params, logging_obj=logging_obj, + drop_params=drop_params, ) ### TRANSFORMATION (sync path) ### request_data: VertexAIBatchEmbeddingsRequestBody | dict[str, object] + flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input) if use_embed_content: resolved_files = {} if api_key: - resolved_files = self._resolve_file_references(input=input, api_key=api_key, sync_handler=sync_handler) + resolved_files = self._resolve_file_references( + input=flat_elements, api_key=api_key, sync_handler=sync_handler + ) request_data = transform_openai_input_gemini_embed_content( input=input, model=model, optional_params=optional_params, resolved_files=resolved_files, + drop_params=drop_params, ) else: - flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input) if has_file_refs and not api_key: raise ValueError( "An API key is required to resolve Gemini file references (files/...). " @@ -238,6 +242,7 @@ class GoogleBatchEmbeddings(VertexLLM): model=model, optional_params=optional_params, resolved_files=resolved_files, + drop_params=drop_params, ) ## LOGGING @@ -294,6 +299,7 @@ class GoogleBatchEmbeddings(VertexLLM): api_key: str | None = None, optional_params: dict | None = None, logging_obj: "LiteLLMLoggingObj | None" = None, + drop_params: bool = False, ) -> EmbeddingResponse: if client is None: _params: Final = {} @@ -312,20 +318,21 @@ class GoogleBatchEmbeddings(VertexLLM): async_handler = client ### TRANSFORMATION (async path) ### + flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input) if use_embed_content: resolved_files = {} if api_key: resolved_files = await self._async_resolve_file_references( - input=input, api_key=api_key, async_handler=async_handler + input=flat_elements, api_key=api_key, async_handler=async_handler ) data = transform_openai_input_gemini_embed_content( input=input, model=model, optional_params=optional_params or {}, resolved_files=resolved_files, + drop_params=drop_params, ) else: - flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input) if has_file_refs and not api_key: raise ValueError( "An API key is required to resolve Gemini file references (files/...). " @@ -341,6 +348,7 @@ class GoogleBatchEmbeddings(VertexLLM): model=model, optional_params=optional_params or {}, resolved_files=resolved_files, + drop_params=drop_params, ) ## LOGGING diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index 9ee2b71d30c..96a785c06c9 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -4,22 +4,27 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batc Why separate file? Make it easy to see how transformation works """ -from collections.abc import Mapping, Sequence -from typing import Final +from collections.abc import Iterator, Mapping, Sequence +from typing import Annotated, Final, Literal, cast -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, Field, TypeAdapter, ValidationError +from litellm.exceptions import BadRequestError +from litellm.llms.vertex_ai.common_utils import GEMINI_FILES_API_URI_PREFIX, gemini_video_metadata_from_openai +from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.vertex_ai import ( BlobType, ContentType, EmbedContentRequest, FileDataType, + GeminiEmbeddingElement, GeminiEmbeddingInput, PartType, PromptTokensDetails, UsageMetadata, VertexAIBatchEmbeddingsRequestBody, VertexAIBatchEmbeddingsResponseObject, + VideoMetadataType, ) from litellm.types.utils import ( Embedding, @@ -40,9 +45,16 @@ SUPPORTED_EMBEDDING_MIME_TYPES: Final = { } +_GEMINI_API_V1BETA: Final = GEMINI_FILES_API_URI_PREFIX.removesuffix("files/") + + def is_file_reference(s: str) -> bool: - """Check if string is a Gemini file reference (files/...).""" - return isinstance(s, str) and s.startswith("files/") + """A `files/...` name or the Files API URI that `/v1/files` returns as the file id.""" + return isinstance(s, str) and (s.startswith("files/") or s.startswith(GEMINI_FILES_API_URI_PREFIX)) + + +def file_reference_name(reference: str) -> str: + return reference.removeprefix(_GEMINI_API_V1BETA) _is_file_reference = is_file_reference @@ -88,12 +100,13 @@ def _infer_mime_type_from_gcs_url(gcs_url: str) -> str: ) -def _parse_data_url(data_url: str) -> tuple[str, str]: +def _parse_data_url(data_url: str, mime_type_override: str | None = None) -> tuple[str, str]: """ Parse a data URL to extract the media type and base64 data. Args: data_url: Data URL in format: data:image/jpeg;base64,/9j/4AAQ... + mime_type_override: An explicit media type that replaces the declared one, skipping the allowlist Returns: tuple: (media_type, base64_data) @@ -110,50 +123,225 @@ def _parse_data_url(data_url: str) -> tuple[str, str]: raise ValueError(f"Invalid data URL format (missing comma): {data_url[:50]}...") metadata, base64_data = data_url.split(",", 1) + declared_media_type: Final = metadata[5:].split(";")[0] - metadata = metadata[5:] + if mime_type_override is not None: + return mime_type_override, base64_data - if ";" in metadata: - media_type = metadata.split(";")[0] - else: - media_type = metadata - - if media_type not in SUPPORTED_EMBEDDING_MIME_TYPES: + if declared_media_type not in SUPPORTED_EMBEDDING_MIME_TYPES: raise ValueError( - f"Unsupported MIME type for embedding: {media_type}. " + f"Unsupported MIME type for embedding: {declared_media_type}. " f"Supported types: {', '.join(sorted(SUPPORTED_EMBEDDING_MIME_TYPES))}" ) - return media_type, base64_data + return declared_media_type, base64_data + + +def _is_data_url(s: str) -> bool: + return s.startswith("data:") and ";base64," in s + + +class _EmbeddingVideoMetadata(LiteLLMBaseModel): + model_config = ConfigDict(extra="forbid", strict=True) + + fps: float | None = None + start_offset: str | None = None + end_offset: str | None = None + + +class _EmbeddingFile(LiteLLMBaseModel): + model_config = ConfigDict(extra="forbid") + + file_id: str | None = None + file_data: str | None = None + filename: str | None = None + format: Annotated[str, Field(min_length=1)] | None = None + video_metadata: _EmbeddingVideoMetadata | None = None + + +class _EmbeddingFileBlock(LiteLLMBaseModel): + model_config = ConfigDict(extra="forbid") + + type: Literal["file"] + file: _EmbeddingFile + + +_file_block_adapter: Final = TypeAdapter(_EmbeddingFileBlock) +_video_metadata_adapter: Final = TypeAdapter(VideoMetadataType) +_input_shape_adapter: Final[TypeAdapter[str | list[object]]] = TypeAdapter(str | list[object]) +_mapping_adapter: Final = TypeAdapter(dict[str, object]) +_BLOCK_FIELDS: Final = frozenset(_EmbeddingFileBlock.model_fields) +_FILE_FIELDS: Final = frozenset(_EmbeddingFile.model_fields) +_VIDEO_METADATA_FIELDS: Final = frozenset(_EmbeddingVideoMetadata.model_fields) +_FILE_SOURCE_FORMS: Final = "a data: URI, a gs:// URL, a files/ reference, or a Gemini Files API URI" + + +def _invalid_input(message: str) -> BadRequestError: + return BadRequestError(message=message, model=None, llm_provider="gemini") + + +def _validation_error_summary(error: ValidationError) -> str: + return "; ".join(f"{'.'.join(str(loc) for loc in detail['loc'])}: {detail['msg']}" for detail in error.errors()) + + +def _as_mapping(value: object) -> dict[str, object] | None: + try: + return _mapping_adapter.validate_python(value) + except ValidationError: + return None + + +def _only_fields(mapping: Mapping[str, object], fields: frozenset[str]) -> dict[str, object]: + return {key: value for key, value in mapping.items() if key in fields} + + +def _dropping_unsupported_keys(block: Mapping[str, object]) -> dict[str, object]: + """What `drop_params` keeps of a file content block: the keys this surface understands, at every level.""" + kept_block: Final = _only_fields(block, _BLOCK_FIELDS) + file: Final = _as_mapping(block.get("file")) + if file is None: + return kept_block + kept_file: Final = _only_fields(file, _FILE_FIELDS) + video_metadata: Final = _as_mapping(file.get("video_metadata")) + if video_metadata is None: + return {**kept_block, "file": kept_file} + return { + **kept_block, + "file": {**kept_file, "video_metadata": _only_fields(video_metadata, _VIDEO_METADATA_FIELDS)}, + } + + +def _parse_file_block(element: object, drop_params: bool) -> _EmbeddingFileBlock: + if not isinstance(element, Mapping): + raise _invalid_input( + f"Embedding input elements must be strings or file content blocks, got {type(element).__name__}" + ) + block: Final = cast(Mapping[str, object], element) # cast-ok: isinstance leaves the key and value types unknown + try: + return _file_block_adapter.validate_python(_dropping_unsupported_keys(block) if drop_params else block) + except ValidationError as error: + raise _invalid_input( + f"Invalid file content block in embedding input: {_validation_error_summary(error)}" + ) from error + + +def _file_block_source(block: _EmbeddingFileBlock) -> str: + match (block.file.file_id, block.file.file_data): + case (str() as file_id, None): + return file_id + case (None, str() as file_data): + return file_data + case (None, None): + raise _invalid_input("A file content block in embedding input needs file.file_id or file.file_data") + case _: + raise _invalid_input( + "A file content block in embedding input takes file.file_id or file.file_data, not both" + ) + + +def _gemini_video_metadata(video_metadata: _EmbeddingVideoMetadata) -> VideoMetadataType: + return _video_metadata_adapter.validate_python( + gemini_video_metadata_from_openai(video_metadata.model_dump(exclude_none=True)) + ) + + +def _source_mime_type( + source: str, + mime_type_override: str | None, + resolved_files: Mapping[str, Mapping[str, str]], +) -> str | None: + if mime_type_override is not None: + return mime_type_override + if _is_data_url(source): + try: + return _parse_data_url(source)[0] + except ValueError: + return None + if _is_gcs_url(source): + try: + return _infer_mime_type_from_gcs_url(source) + except ValueError: + return None + if is_file_reference(source): + file_info: Final = resolved_files.get(source) + return None if file_info is None else file_info.get("mime_type") + return None + + +def _media_part( + source: str, + mime_type_override: str | None, + resolved_files: Mapping[str, Mapping[str, str]], +) -> PartType: + if _is_data_url(source): + mime_type, base64_data = _parse_data_url(source, mime_type_override) + return PartType(inline_data=BlobType(mime_type=mime_type, data=base64_data)) + if _is_gcs_url(source): + gcs_mime_type: Final = mime_type_override or _infer_mime_type_from_gcs_url(source) + return PartType(file_data=FileDataType(mime_type=gcs_mime_type, file_uri=source)) + if is_file_reference(source): + file_info: Final = resolved_files.get(source) + if file_info is None: + raise _invalid_input( + f"File reference {source!r} could not be resolved: " + "Gemini Files API references are only supported through the gemini/ provider" + ) + return PartType( + file_data=FileDataType(mime_type=mime_type_override or file_info["mime_type"], file_uri=file_info["uri"]) + ) + raise _invalid_input(f"A file content block source must be {_FILE_SOURCE_FORMS}, got {source[:50]!r}") + + +def _top_level_elements( + input: GeminiEmbeddingInput, +) -> Sequence[GeminiEmbeddingElement | list[str] | list[GeminiEmbeddingElement]]: + try: + _input_shape_adapter.validate_python(input) + except ValidationError as error: + raise _invalid_input( + f"Embedding input must be a string or a list of strings and file content blocks, got {type(input).__name__}" + ) from error + return [input] if isinstance(input, str) else input + + +def _elements(input: GeminiEmbeddingInput) -> Iterator[GeminiEmbeddingElement]: + for element in _top_level_elements(input): + if isinstance(element, list): + yield from element + else: + yield element + + +def _element_source(element: GeminiEmbeddingElement) -> str: + if isinstance(element, str): + return element + return _file_block_source(_parse_file_block(element, drop_params=True)) + + +def flatten_media_sources(input: GeminiEmbeddingInput) -> tuple[str, ...]: + """Every string element plus every file content block's source, in input order.""" + return tuple(_element_source(element) for element in _elements(input)) def _is_multimodal_input(input: GeminiEmbeddingInput) -> bool: """ Check if the input contains multimodal data (data URIs, file references, - GCS URLs, or nested lists for combined embeddings). + GCS URLs, file content blocks, or nested lists for combined embeddings). Args: - input: GeminiEmbeddingInput — str, List[str], or List[List[str]] for combined embeddings + input: GeminiEmbeddingInput — str, List[element], or List[List[element]] for combined embeddings Returns: - bool: True if any element is multimodal or a nested list + bool: True if any element is multimodal """ - if isinstance(input, str): - return _is_multimodal_element(input) - - for element in input: - if isinstance(element, list): - if any(_is_multimodal_element(sub) for sub in element if isinstance(sub, str)): - return True - elif isinstance(element, str) and _is_multimodal_element(element): - return True - - return False + return any(_is_multimodal_element(element) for element in _elements(input)) -def _is_multimodal_element(element: str) -> bool: - """Check if a single string element is multimodal.""" - if element.startswith("data:") and ";base64," in element: +def _is_multimodal_element(element: GeminiEmbeddingElement) -> bool: + """Check if a single element is multimodal.""" + if not isinstance(element, str): + return True + if _is_data_url(element): return True if is_file_reference(element): return True @@ -163,37 +351,29 @@ def _is_multimodal_element(element: str) -> bool: def _build_part_for_input( - element: str, - resolved_files: dict[str, dict[str, str]] | None = None, + element: GeminiEmbeddingElement, + resolved_files: Mapping[str, Mapping[str, str]] | None = None, + drop_params: bool = False, ) -> PartType: """ Build a single PartType for an input element, handling text, data URIs, - file references, and GCS URLs. + file references, GCS URLs, and file content blocks carrying a mime type + and video_metadata. """ - resolved_files = resolved_files or {} + files: Final = resolved_files or {} - if element.startswith("data:") and ";base64," in element: - mime_type, base64_data = _parse_data_url(element) - blob: Final[BlobType] = {"mime_type": mime_type, "data": base64_data} - return PartType(inline_data=blob) - elif _is_gcs_url(element): - mime_type = _infer_mime_type_from_gcs_url(element) - file_data: Final[FileDataType] = { - "mime_type": mime_type, - "file_uri": element, - } - return PartType(file_data=file_data) - elif is_file_reference(element): - if element not in resolved_files: - raise ValueError(f"File reference {element} not resolved") - file_info: Final = resolved_files[element] - file_data_ref: Final[FileDataType] = { - "mime_type": file_info["mime_type"], - "file_uri": file_info["uri"], - } - return PartType(file_data=file_data_ref) - else: - return PartType(text=element) + if isinstance(element, str): + return _media_part(element, None, files) if _is_multimodal_element(element) else PartType(text=element) + + block: Final = _parse_file_block(element, drop_params) + part: Final = _media_part(_file_block_source(block), block.file.format, files) + if block.file.video_metadata is None: + return part + video_metadata: Final = _gemini_video_metadata(block.file.video_metadata) + if not video_metadata: + return part + part_with_metadata: Final[PartType] = {**part, "video_metadata": video_metadata} + return part_with_metadata _SUPPORTED_EMBED_PARAMS: Final = {"outputDimensionality", "taskType", "title"} @@ -214,6 +394,7 @@ def transform_openai_input_gemini_content( model: str, optional_params: dict, resolved_files: dict[str, dict[str, str]] | None = None, + drop_params: bool = False, ) -> VertexAIBatchEmbeddingsRequestBody: """ Transform OpenAI embedding input to Gemini batchEmbedContents format. @@ -234,19 +415,18 @@ def transform_openai_input_gemini_content( gemini_params: Final = _filter_embed_params(optional_params) - input_list: Final = [input] if isinstance(input, str) else input + input_list: Final = _top_level_elements(input) requests: Final[list[EmbedContentRequest]] = [] for element in input_list: if isinstance(element, list): if not element: raise ValueError("Nested input list must not be empty") - for sub in element: - if not isinstance(sub, str): - raise ValueError(f"Elements inside a nested input list must be strings, got {type(sub)}") - parts = [_build_part_for_input(sub, resolved_files=resolved_files) for sub in element] + parts = [ + _build_part_for_input(sub, resolved_files=resolved_files, drop_params=drop_params) for sub in element + ] else: - parts = [_build_part_for_input(element, resolved_files=resolved_files)] + parts = [_build_part_for_input(element, resolved_files=resolved_files, drop_params=drop_params)] request = EmbedContentRequest( model=gemini_model_name, content=ContentType(parts=parts), @@ -262,6 +442,7 @@ def transform_openai_input_gemini_embed_content( model: str, optional_params: dict, resolved_files: dict[str, dict[str, str]] | None = None, + drop_params: bool = False, ) -> dict: """ Transform OpenAI embedding input to Gemini embedContent format (multimodal). @@ -279,7 +460,7 @@ def transform_openai_input_gemini_embed_content( gemini_params: Final = _filter_embed_params(optional_params) - input_list: Final = [input] if isinstance(input, str) else input + input_list: Final = _top_level_elements(input) parts: Final[list[PartType]] = [] for element in input_list: @@ -288,9 +469,7 @@ def transform_openai_input_gemini_embed_content( "Nested (combined) embeddings are not supported on the embedContent path. " "Use the batchEmbedContents path or pass a flat list instead." ) - if not isinstance(element, str): - raise ValueError(f"Unsupported input type: {type(element)}") - parts.append(_build_part_for_input(element, resolved_files=resolved_files)) + parts.append(_build_part_for_input(element, resolved_files=resolved_files, drop_params=drop_params)) request_body: Final[dict] = { "content": ContentType(parts=parts), @@ -313,31 +492,18 @@ def _parse_usage_metadata(raw_usage_metadata: object) -> UsageMetadata | None: return None -def _flatten_input(input: GeminiEmbeddingInput) -> tuple[str, ...]: - if isinstance(input, str): - return (input,) - return tuple(sub for element in input for sub in (element if isinstance(element, list) else [element])) +def _flatten_input(input: GeminiEmbeddingInput) -> tuple[GeminiEmbeddingElement, ...]: + return tuple(_elements(input)) def _is_image_element( - element: str, + element: GeminiEmbeddingElement, resolved_files: Mapping[str, Mapping[str, str]], ) -> bool: - if element.startswith("data:") and ";base64," in element: - try: - mime_type, _ = _parse_data_url(element) - except ValueError: - return False - return mime_type in _IMAGE_MIME_TYPES - if _is_gcs_url(element): - try: - return _infer_mime_type_from_gcs_url(element) in _IMAGE_MIME_TYPES - except ValueError: - return False - if is_file_reference(element): - file_info: Final = resolved_files.get(element) - return file_info is not None and file_info.get("mime_type") in _IMAGE_MIME_TYPES - return False + if isinstance(element, str): + return _source_mime_type(element, None, resolved_files) in _IMAGE_MIME_TYPES + block: Final = _parse_file_block(element, drop_params=True) + return _source_mime_type(_file_block_source(block), block.file.format, resolved_files) in _IMAGE_MIME_TYPES def _is_image_only_input( diff --git a/litellm/main.py b/litellm/main.py index 25c73ea281e..b02ccca9ecb 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6379,6 +6379,10 @@ def embedding( """ azure: Final = kwargs.get("azure", None) client: Final = kwargs.pop("client", None) + drop_params_kwarg: Final = ( + cast(object, kwargs["drop_params"]) if "drop_params" in kwargs else None # cast-ok: untyped request kwargs + ) + drop_unsupported_params: Final = litellm.drop_params is True or normalize_drop_params(drop_params_kwarg) is True shared_session: Final = kwargs.get("shared_session", None) max_retries: Final = kwargs.get("max_retries", None) litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") @@ -6859,6 +6863,7 @@ def embedding( api_base=api_base, client=client, extra_headers=headers, + drop_params=drop_unsupported_params, ) elif custom_llm_provider == "vertex_ai": @@ -6911,6 +6916,7 @@ def embedding( api_base=api_base, client=client, extra_headers=headers, + drop_params=drop_unsupported_params, ) elif ( "image" in optional_params diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 663c6d1a2a7..bbd0d9ab989 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -42320,6 +42320,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "input_cost_per_token": 1.5e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, @@ -42338,6 +42339,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "input_cost_per_token": 7e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, @@ -42363,6 +42365,7 @@ "output_cost_per_token_flex": 3e-07, "source": "https://aws.amazon.com/bedrock/pricing/", "supports_audio_input": false, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -42380,6 +42383,7 @@ "output_cost_per_token_flex": 1e-07, "source": "https://aws.amazon.com/bedrock/pricing/", "supports_audio_input": false, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -48105,6 +48109,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -48123,6 +48128,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -65951,6 +65957,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -65969,6 +65976,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -66219,6 +66227,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -66237,6 +66246,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -80458,29 +80468,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "xai.grok-4.7": { - "supports_regex_lookaround": false, - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_flex": 2.5e-07, - "cache_read_input_token_cost_priority": 8.75e-07, - "input_cost_per_token": 2e-06, - "input_cost_per_token_flex": 1e-06, - "input_cost_per_token_priority": 3.5e-06, - "litellm_provider": "bedrock_converse", - "max_input_tokens": 500000, - "max_output_tokens": 500000, - "max_tokens": 500000, - "mode": "chat", - "output_cost_per_token": 6e-06, - "output_cost_per_token_flex": 3e-06, - "output_cost_per_token_priority": 1.05e-05, - "source": "https://b0.p.awsstatic.com/pricing/2.0/meteredUnitMaps/bedrock/USD/current/bedrock.json", - "supports_function_calling": true, - "supports_prompt_caching": false, - "supports_reasoning": true, - "supports_tool_choice": true, - "supports_vision": true - }, "openrouter/anthropic/claude-sonnet-5.5:batch": { "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 9322fb77615..7dcca8f9ec0 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -11906,6 +11906,34 @@ "description": "When True (default), after sensitive data is detected and routed, all subsequent requests in the same session will continue routing to the same model.", "title": "Sticky Session Routing" }, + "stream_scope": { + "anyOf": [ + { + "enum": [ + "streaming", + "non_streaming", + "both" + ], + "type": "string" + }, + { + "additionalProperties": { + "enum": [ + "streaming", + "non_streaming", + "both" + ], + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "description": "Whether this guardrail runs on streaming requests, non-streaming requests, or both. A string applies to every configured mode. A map overrides named modes (pre_call, during_call, post_call, ...); omitted keys default to both. Unset means both, matching historical behavior.", + "title": "Stream Scope" + }, "template_id": { "anyOf": [ { @@ -14831,6 +14859,34 @@ "description": "When True (default), after sensitive data is detected and routed, all subsequent requests in the same session will continue routing to the same model.", "title": "Sticky Session Routing" }, + "stream_scope": { + "anyOf": [ + { + "enum": [ + "streaming", + "non_streaming", + "both" + ], + "type": "string" + }, + { + "additionalProperties": { + "enum": [ + "streaming", + "non_streaming", + "both" + ], + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "description": "Whether this guardrail runs on streaming requests, non-streaming requests, or both. A string applies to every configured mode. A map overrides named modes (pre_call, during_call, post_call, ...); omitted keys default to both. Unset means both, matching historical behavior.", + "title": "Stream Scope" + }, "template_id": { "anyOf": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4b486b7d0d3..8753aae72d5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple, Ty import httpx from pydantic import ( + AfterValidator, BaseModel, BeforeValidator, ConfigDict, @@ -2743,6 +2744,41 @@ class ScheduledJobStaggerSettings(LiteLLMPydanticObjectBase): ) +SPEND_LOGS_METADATA_ALWAYS_KEPT_FIELDS: Final = frozenset({"status", "cold_storage_object_key"}) + + +def _known_spend_logs_metadata_field(name: str) -> str: + if name not in SpendLogsMetadata.__annotations__: + raise ValueError(f"{name!r} is not a LiteLLM_SpendLogs.metadata field") + return name + + +SpendLogsMetadataFieldName: TypeAlias = Annotated[str, AfterValidator(_known_spend_logs_metadata_field)] + + +class SpendLogsMetadataFields(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + include: tuple[SpendLogsMetadataFieldName, ...] | None = None + exclude: tuple[SpendLogsMetadataFieldName, ...] | None = None + + @model_validator(mode="after") + def _exactly_one_list(self) -> "SpendLogsMetadataFields": + if (self.include is None) == (self.exclude is None): + raise ValueError("set exactly one of 'include' or 'exclude'") + always_kept_excluded: Final = sorted(SPEND_LOGS_METADATA_ALWAYS_KEPT_FIELDS.intersection(self.exclude or ())) + if always_kept_excluded: + raise ValueError(f"{always_kept_excluded} are always kept and cannot be excluded") + return self + + def keeps(self, name: str) -> bool: + if name in SPEND_LOGS_METADATA_ALWAYS_KEPT_FIELDS: + return True + if self.include is not None: + return name in self.include + return name not in (self.exclude or ()) + + DEFAULT_RESPONSES_WEBSOCKET_SESSION_LIMIT_SECONDS: Final[float] = 3600.0 @@ -3095,6 +3131,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="If True, a completed batch logged via aretrieve_batch also emits one callback event per JSONL line item (request paired with its response or error). The aggregate batch callback is unchanged. Default is False.", ) + spend_logs_metadata_fields: SpendLogsMetadataFields | None = Field( + None, + description="Which keys of LiteLLM_SpendLogs.metadata are written to the database. Set exactly one of 'include' (write only these keys) or 'exclude' (drop these keys). 'status' and 'cold_storage_object_key' are always written. Daily spend tables, budgets and logging callbacks still see every key. Unset writes every key", + ) disable_auto_add_proxy_admin_to_teams: bool | None = Field( None, description="By default, the user calling /team/new is automatically added to the new team as a team admin. If True, proxy admins are no longer auto-added; members explicitly listed in members_with_roles are unaffected. Default is False.", diff --git a/litellm/proxy/db/baseline_accounting.py b/litellm/proxy/db/baseline_accounting.py index 486717f9b93..f7a2be654c8 100644 --- a/litellm/proxy/db/baseline_accounting.py +++ b/litellm/proxy/db/baseline_accounting.py @@ -243,18 +243,19 @@ WHERE scope = $1 AND revision > $2::bigint AND ($5::float8 IS NULL OR publication::jsonb->>'status' = 'estimated') ORDER BY started_at, request_id """ +_PUBLISHED_LOG_FIELDS: Final = ("autorouter_savings_estimate", "autorouter_savings") _UPDATE_LOGS: Final = """ WITH changes AS ( SELECT request_id, publication::jsonb AS publication FROM jsonb_to_recordset($1::jsonb) AS x(request_id text, publication jsonb) ) UPDATE "LiteLLM_SpendLogs" AS logs -SET metadata = (COALESCE(logs.metadata::jsonb, '{}'::jsonb) - 'autorouter_baseline_observation') || jsonb_build_object( +SET metadata = (COALESCE(logs.metadata::jsonb, '{}'::jsonb) - 'autorouter_baseline_observation') || (jsonb_build_object( 'autorouter_savings_estimate', changes.publication, 'autorouter_savings', CASE WHEN changes.publication->>'status' = 'estimated' THEN (changes.publication->>'baseline_spend')::float8 - (changes.publication->>'actual_spend')::float8 ELSE NULL END -) +) - ARRAY(SELECT jsonb_array_elements_text($2::jsonb))) FROM changes WHERE logs.request_id = changes.request_id """ _UPDATE_PUBLICATIONS: Final = """ @@ -397,7 +398,11 @@ async def _publish(db: SupportsRawQueries, changes: Sequence[_Change]) -> None: if not changes: return serialized: Final = json.dumps(tuple(change.model_dump(mode="json") for change in changes), separators=(",", ":")) - await db.execute_raw(_UPDATE_LOGS, serialized) + from litellm.proxy.spend_tracking.spend_tracking_utils import configured_spend_logs_metadata_fields + + fields: Final = configured_spend_logs_metadata_fields() + unstored: Final = tuple(name for name in _PUBLISHED_LOG_FIELDS if fields is not None and not fields.keeps(name)) + await db.execute_raw(_UPDATE_LOGS, serialized, json.dumps(unstored)) await db.execute_raw(_UPDATE_SESSIONS, serialized) if any(change.user_id for change in changes): await db.execute_raw(_UPDATE_USER_SESSIONS, serialized) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 9fb8cc20981..5bafe19cd0b 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -654,7 +654,14 @@ class DBSpendUpdateWriter: from litellm.repositories.table_repositories import SpendLogsRepository request_id: Final = payload["request_id"] - row: Final = _batch_cost_row_to_write(payload, disable_spend_logs) + from litellm.proxy.spend_tracking.spend_tracking_utils import ( + configured_spend_logs_metadata_fields, + spend_log_row_with_retained_metadata, + ) + + row: Final = spend_log_row_with_retained_metadata( + _batch_cost_row_to_write(payload, disable_spend_logs), configured_spend_logs_metadata_fields() + ) spend_logs: Final = SpendLogsRepository(prisma_client).table try: claimed: Final = await spend_logs.create_many( diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 87d2ee07818..8638bc57615 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -191,26 +191,12 @@ async def google_count_tokens(request: Request, model_name: str): request=token_request, call_endpoint=True, ) - if token_response is not None: - # cast the response to the well known format - original_response: Final[dict] = token_response.original_response or {} - if original_response: - return TokenCountDetailsResponse( - totalTokens=original_response.get("totalTokens", 0), - promptTokensDetails=original_response.get("promptTokensDetails", []), - ) - else: - return TokenCountDetailsResponse( - totalTokens=token_response.total_tokens or 0, - promptTokensDetails=[], - ) - - ######################################################### - # Return the response in the well known format - ######################################################### + if token_response is None: + return TokenCountDetailsResponse(totalTokens=0, promptTokensDetails=[]) + original_response: Final[dict] = token_response.original_response or {} return TokenCountDetailsResponse( - totalTokens=0, - promptTokensDetails=[], + totalTokens=original_response.get("totalTokens") or token_response.total_tokens or 0, + promptTokensDetails=original_response.get("promptTokensDetails", []), ) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 8e1209401ab..6072d8bf22a 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -69,6 +69,7 @@ from litellm.types.guardrails import ( PresidioPresidioConfigModelUserInterface, SupportedGuardrailIntegrations, ToolPermissionGuardrailConfigModel, + with_tolerated_stream_scope, ) from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.proxy.guardrails.guardrail_hooks.hide_secrets import ( @@ -146,7 +147,7 @@ def _get_guardrails_list_response( GuardrailInfoResponse( guardrail_id=guardrail.get("guardrail_id"), guardrail_name=guardrail.get("guardrail_name"), - litellm_params=masked_params, + litellm_params=with_tolerated_stream_scope(masked_params), guardrail_info=guardrail.get("guardrail_info"), ) ) @@ -289,7 +290,7 @@ async def list_guardrails_v2( ) masked_litellm_params = ( parse_tolerant_litellm_params( - masked_litellm_params_dict, + with_tolerated_stream_scope(masked_litellm_params_dict), guardrail.get("guardrail_name") or "Unknown", params_model=BaseLitellmParams, ) @@ -336,7 +337,7 @@ async def list_guardrails_v2( ) masked_in_memory_litellm_params_typed = ( parse_tolerant_litellm_params( - masked_in_memory_litellm_params, + with_tolerated_stream_scope(masked_in_memory_litellm_params), guardrail.get("guardrail_name") or "Unknown", params_model=BaseLitellmParams, ) @@ -1263,7 +1264,7 @@ async def patch_guardrail( # Update litellm_params if default_on is provided or pii_entities_config is provided existing_litellm_params: Final = _as_str_object_mapping(dict(existing_guardrail.get("litellm_params", {}))) current_litellm_params: Final = parse_tolerant_litellm_params( - existing_litellm_params, + with_tolerated_stream_scope(existing_litellm_params), existing_guardrail.get("guardrail_name") or "Unknown", ) requested_litellm_params: Final[Mapping[str, object]] = ( @@ -1275,7 +1276,7 @@ async def patch_guardrail( MappingProxyType({**current_litellm_params.model_dump(exclude_unset=True), **requested_litellm_params}) ) try: - parsed_litellm_params: Final = LitellmParams(**merged_litellm_params) + parsed_litellm_params: Final = LitellmParams(**with_tolerated_stream_scope(merged_litellm_params)) except ValidationError as validation_error: raise HTTPException( status_code=422, @@ -1436,7 +1437,7 @@ async def get_guardrail_info(guardrail_id: str): ) masked_litellm_params = ( parse_tolerant_litellm_params( - masked_litellm_params_dict, + with_tolerated_stream_scope(masked_litellm_params_dict), result.get("guardrail_name") or "Unknown", params_model=BaseLitellmParams, ) diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index d14cc940885..439d17b8644 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -54,6 +54,7 @@ from litellm.types.guardrails import ( LakeraCategoryThresholds, LitellmParams, SupportedGuardrailIntegrations, + with_tolerated_stream_scope, ) from .guardrail_hooks.llm_as_a_judge import ( @@ -634,6 +635,8 @@ def _configure_callback_scoping( "skip_tool_message_in_guardrail are enabled together, which excludes every message from " "scanning, so no request content would ever be scanned. Remove one of the two." ) + if isinstance(custom_guardrail_callback, CustomGuardrail): # pyright: ignore[reportUnnecessaryIsInstance] # module-path classes may only subclass CustomLogger + custom_guardrail_callback.apply_stream_scope(litellm_params.stream_scope) _apply_configured_bool_overrides(custom_guardrail_callback, litellm_params) @@ -718,9 +721,11 @@ class InMemoryGuardrailHandler: if isinstance(litellm_params_data, dict): if reject_invalid_logging_only_scope: - litellm_params = LitellmParams(**litellm_params_data) + litellm_params = LitellmParams(**with_tolerated_stream_scope(litellm_params_data)) else: - litellm_params = parse_tolerant_litellm_params(litellm_params_data, guardrail["guardrail_name"]) + litellm_params = parse_tolerant_litellm_params( + with_tolerated_stream_scope(litellm_params_data), guardrail["guardrail_name"] + ) else: litellm_params = litellm_params_data @@ -863,14 +868,17 @@ class InMemoryGuardrailHandler: # Extract additional params from litellm_params to pass to custom guardrail # This matches the behavior of other guardrail initializers (e.g., initialize_lakera) # and aligns with the documented behavior for custom guardrails - if hasattr(litellm_params, "model_dump"): - extra_params = litellm_params.model_dump(exclude_none=True) - else: - extra_params = dict(litellm_params) if litellm_params else {} - - # Remove params that are handled explicitly or are internal - for key in ["guardrail", "mode", "default_on"]: - extra_params.pop(key, None) + excluded_extra_param_keys: Final = frozenset(("guardrail", "mode", "default_on", "stream_scope")) + extra_params_items: Final = ( + litellm_params.model_dump(exclude_none=True).items() + if hasattr(litellm_params, "model_dump") + else iter(litellm_params) + if litellm_params + else () + ) + extra_params: Final = MappingProxyType( + {key: value for key, value in extra_params_items if key not in excluded_extra_param_keys} + ) _guardrail_callback: Final = _guardrail_class( guardrail_name=guardrail["guardrail_name"], @@ -1007,7 +1015,7 @@ class InMemoryGuardrailHandler: return params.model_dump() if isinstance(params, dict): try: - return parse_tolerant_litellm_params(params, guardrail_name).model_dump() + return parse_tolerant_litellm_params(with_tolerated_stream_scope(params), guardrail_name).model_dump() except ValidationError as e: verbose_proxy_logger.warning( "Could not normalize guardrail litellm_params for comparison; treating the guardrail as changed. Error: %s", diff --git a/litellm/proxy/lens/inference.py b/litellm/proxy/lens/inference.py index 599c88cf4d2..fbb661c027c 100644 --- a/litellm/proxy/lens/inference.py +++ b/litellm/proxy/lens/inference.py @@ -325,7 +325,7 @@ def renew_reservation(lens: Lens, reservation_id: str, now: datetime) -> Lens: async def wait_for_reservation( repo: LensRepository, lens_id: str, reservation_id: str, reserve: Callable[[Lens], Lens] ) -> None: - while (reserved := await repo.update(lens_id, reserve)) is not None: + while (reserved := await repo.update_locked(lens_id, reserve)) is not None: if any(held.id == reservation_id for held in reserved.reservations): return await asyncio.sleep(0.25) @@ -341,7 +341,7 @@ async def renew_budget_reservation( try: async with timeout(BUDGET_RENEW_INTERVAL): if ( - await repo.update( + await repo.update_locked( lens_id, lambda e: renew_reservation(e, reservation_id, datetime.now(timezone.utc)) ) is None diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 46cdf328af6..66dca57e782 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -32,6 +32,7 @@ from litellm.constants import ( SESSION_ID_OMITTED_METADATA_KEY, X_LITELLM_DISABLE_CALLBACKS, ) +from litellm.integrations.custom_guardrail import without_server_streaming_classification from litellm.litellm_core_utils.core_helpers import is_codex_user_agent from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -2094,7 +2095,9 @@ def refresh_proxy_server_request_body_snapshot( | _TRANSPORT_ONLY_CREDENTIAL_KEYS | _CALLBACK_CREDENTIAL_KEYS ) - body: Final = {k: v for k, v in data.items() if k not in _body_snapshot_exclude} + body: Final = { + k: v for k, v in without_server_streaming_classification(data).items() if k not in _body_snapshot_exclude + } proxy_server_request["body"] = body if guardrails_applied and isinstance(logging_obj, Logging): metadata: Final = data.get(get_metadata_variable_name_from_kwargs(data)) diff --git a/litellm/proxy/management_endpoints/fallback_management_endpoints.py b/litellm/proxy/management_endpoints/fallback_management_endpoints.py index 543538f5edd..460b282a9f6 100644 --- a/litellm/proxy/management_endpoints/fallback_management_endpoints.py +++ b/litellm/proxy/management_endpoints/fallback_management_endpoints.py @@ -11,15 +11,21 @@ DELETE /fallback/{model} - Delete fallbacks for a specific model # pyright: reportMissingImports=false import json -from typing import TYPE_CHECKING, Final, Literal +from collections.abc import Mapping +from typing import TYPE_CHECKING, Annotated, Final, Literal + +from pydantic import Field, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.model_checks import get_all_fallbacks from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.utils import PrismaClient, evict_config_param, invalidate_config_param if TYPE_CHECKING: from fastapi import APIRouter, Depends, HTTPException, status + + from litellm.proxy.proxy_server import ProxyConfig else: try: from fastapi import APIRouter, Depends, HTTPException, status @@ -37,6 +43,33 @@ from litellm.types.management_endpoints.router_settings_endpoints import ( router: Final = APIRouter() +ROUTER_SETTINGS_PARAM: Final = "router_settings" +FallbackRule = dict[str, list[str]] +StoredFallback = Annotated[FallbackRule | dict[str, object] | str, Field(union_mode="left_to_right")] +_STORED_FALLBACKS: Final = TypeAdapter(list[StoredFallback]) + + +def _rule_covers(entry: StoredFallback, model: str) -> bool: + return isinstance(entry, dict) and model in entry + + +async def _router_settings_fresh_from_db(proxy_config: "ProxyConfig") -> dict[str, object]: + await evict_config_param(ROUTER_SETTINGS_PARAM) + config: Final = await proxy_config.get_config() + return config.get(ROUTER_SETTINGS_PARAM, {}) + + +async def _persist_router_settings(prisma_client: PrismaClient, router_settings: Mapping[str, object]) -> None: + router_settings_json: Final = json.dumps(router_settings) + await ConfigRepository(prisma_client).table.upsert( + where={"param_name": ROUTER_SETTINGS_PARAM}, + data={ + "create": {"param_name": ROUTER_SETTINGS_PARAM, "param_value": router_settings_json}, + "update": {"param_value": router_settings_json}, + }, + ) + await invalidate_config_param(ROUTER_SETTINGS_PARAM) + @router.post( "/fallback", @@ -85,24 +118,24 @@ async def create_fallback( ) # Validate that the model exists in the router - model_names: Final = llm_router.model_names - if data.model not in model_names: + known_model_names: Final = frozenset(llm_router.model_names) | llm_router.team_public_model_names + if data.model not in known_model_names: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail={ "error": f"Model '{data.model}' not found in router", - "available_models": list(model_names), + "available_models": sorted(known_model_names), }, ) # Validate that all fallback models exist in the router - invalid_fallback_models: Final = [m for m in data.fallback_models if m not in model_names] + invalid_fallback_models: Final = [m for m in data.fallback_models if m not in known_model_names] if invalid_fallback_models: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail={ "error": f"Invalid fallback models: {invalid_fallback_models}", - "available_models": list(model_names), + "available_models": sorted(known_model_names), }, ) @@ -122,9 +155,7 @@ async def create_fallback( }, ) - # Load existing config - config: Final = await proxy_config.get_config() - router_settings: Final = config.get("router_settings", {}) + router_settings: Final = await _router_settings_fresh_from_db(proxy_config) # Get the appropriate fallback list based on type fallback_key = "fallbacks" @@ -134,12 +165,12 @@ async def create_fallback( fallback_key = "content_policy_fallbacks" # Get existing fallbacks - existing_fallbacks: Final[list[dict[str, list[str]]]] = router_settings.get(fallback_key, []) + existing_fallbacks: Final = _STORED_FALLBACKS.validate_python(router_settings.get(fallback_key) or []) # Update or add the fallback configuration fallback_updated = False - for i, fallback_dict in enumerate(existing_fallbacks): - if data.model in fallback_dict: + for i, rule in enumerate(existing_fallbacks): + if _rule_covers(rule, data.model): # Update existing fallback existing_fallbacks[i] = {data.model: data.fallback_models} fallback_updated = True @@ -152,18 +183,7 @@ async def create_fallback( # Update router settings router_settings[fallback_key] = existing_fallbacks - # Save to database - convert router_settings to JSON string - router_settings_json: Final = json.dumps(router_settings) - await ConfigRepository(prisma_client).table.upsert( - where={"param_name": "router_settings"}, - data={ - "create": { - "param_name": "router_settings", - "param_value": router_settings_json, - }, - "update": {"param_value": router_settings_json}, - }, - ) + await _persist_router_settings(prisma_client, router_settings) # Update the in-memory router configuration setattr(llm_router, fallback_key, existing_fallbacks) @@ -291,9 +311,7 @@ async def delete_fallback( }, ) - # Load existing config - config: Final = await proxy_config.get_config() - router_settings: Final = config.get("router_settings", {}) + router_settings: Final = await _router_settings_fresh_from_db(proxy_config) # Get the appropriate fallback list based on type fallback_key = "fallbacks" @@ -303,14 +321,14 @@ async def delete_fallback( fallback_key = "content_policy_fallbacks" # Get existing fallbacks - existing_fallbacks: Final[list[dict[str, list[str]]]] = router_settings.get(fallback_key, []) + existing_fallbacks: Final = _STORED_FALLBACKS.validate_python(router_settings.get(fallback_key) or []) # Find and remove the fallback configuration fallback_found = False updated_fallbacks: Final = [] - for fallback_dict in existing_fallbacks: - if model not in fallback_dict: - updated_fallbacks.append(fallback_dict) + for rule in existing_fallbacks: + if not _rule_covers(rule, model): + updated_fallbacks.append(rule) else: fallback_found = True @@ -323,18 +341,7 @@ async def delete_fallback( # Update router settings router_settings[fallback_key] = updated_fallbacks - # Save to database - convert router_settings to JSON string - router_settings_json: Final = json.dumps(router_settings) - await ConfigRepository(prisma_client).table.upsert( - where={"param_name": "router_settings"}, - data={ - "create": { - "param_name": "router_settings", - "param_value": router_settings_json, - }, - "update": {"param_value": router_settings_json}, - }, - ) + await _persist_router_settings(prisma_client, router_settings) # Update the in-memory router configuration setattr(llm_router, fallback_key, updated_fallbacks) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index ca142088383..1eefb58d7fa 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -44,12 +44,14 @@ from litellm.constants import ( AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER, BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES, ) +from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers from litellm.llms.azure.passthrough.transformation import ( foreign_azure_deployment, is_azure_body_model_inference_endpoint, ) +from litellm.llms.bedrock.passthrough.transformation import is_bedrock_streaming_endpoint from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.deepgram.common_utils import ( deepgram_listen_callback_params, @@ -946,8 +948,6 @@ BEDROCK_ENDPOINT_ACTIONS: Final = { "count-tokens", } -BEDROCK_STREAMING_ACTIONS: Final = {"invoke-with-response-stream", "converse-stream"} - def is_bedrock_count_tokens_endpoint(endpoint: str) -> bool: return "count_tokens" in endpoint or "count-tokens" in endpoint @@ -1077,7 +1077,7 @@ async def handle_bedrock_passthrough_router_model( from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing # Detect streaming based on endpoint - is_streaming: Final = any(action in endpoint for action in BEDROCK_STREAMING_ACTIONS) + is_streaming: Final = is_bedrock_streaming_endpoint(endpoint) verbose_proxy_logger.debug( "Bedrock router passthrough: model='%s', endpoint='%s', streaming=%s", model, endpoint, is_streaming @@ -1085,15 +1085,19 @@ async def handle_bedrock_passthrough_router_model( # Use the common processing path (same as non-router models) # This ensures all metadata, hooks, and logging are properly initialized - data: Final[dict[str, object]] = {} + bedrock_payload: Final[dict[str, object]] = { + "model": model, + "method": request.method, + "endpoint": endpoint, + "data": request_body, + "custom_llm_provider": "bedrock", + } + data: Final[dict[str, object]] = guardrail_request_data_with_streaming( + MappingProxyType(bedrock_payload), + is_streaming=is_streaming, + ) base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) - data["model"] = model - data["method"] = request.method - data["endpoint"] = endpoint - data["data"] = request_body - data["custom_llm_provider"] = "bedrock" - # Use the common passthrough processing to handle metadata and hooks # This also handles all response formatting (streaming/non-streaming) and exceptions try: @@ -1285,14 +1289,19 @@ async def bedrock_llm_proxy_route( "Bedrock passthrough: Using direct Bedrock model '%s' for endpoint '%s'", model, endpoint ) - data: Final[dict[str, object]] = {} + is_streaming: Final = is_bedrock_streaming_endpoint(endpoint) + passthrough_payload: Final[dict[str, object]] = { + "method": request.method, + "endpoint": endpoint, + "data": request_body, + "custom_llm_provider": "bedrock", + } + data: Final[dict[str, object]] = guardrail_request_data_with_streaming( + MappingProxyType(passthrough_payload), + is_streaming=is_streaming, + ) base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) - data["method"] = request.method - data["endpoint"] = endpoint - data["data"] = request_body - data["custom_llm_provider"] = "bedrock" - try: result: Final = await base_llm_response_processor.base_passthrough_process_llm_request( request=request, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 07053ad0f32..ee958679c1f 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -48,7 +48,11 @@ from litellm.constants import ( SESSION_ID_OMITTED_METADATA_KEY, WEBSOCKET_CLOSE_REASON_MAX_BYTES, ) -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + guardrail_request_data_with_streaming, + without_server_streaming_classification, +) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( bind_budget_reservation_to_callbacks, @@ -603,6 +607,12 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): from litellm.proxy.proxy_server import llm_router _parsed_body = _parsed_body or {} + parsed_body_typed: Final[Mapping[str, object]] = cast(Mapping[str, object], _parsed_body) # cast-ok: json + server_marker_free_body: Final = without_server_streaming_classification(parsed_body_typed) + # The marker-free body must propagate through the caller's request dict, so + # downstream guardrail scans and snapshots never observe the server streaming marker. + _parsed_body.clear() + _parsed_body.update(server_marker_free_body) managed_model: Final = get_model_from_request( request_data=_parsed_body, route=get_request_route(request), @@ -828,7 +838,7 @@ def _build_passthrough_failure_request_payload( error response. Spend tracking only attributes a recovered cost when it comes paired with a usage object, so both keys are written together. """ - request_payload: Final[dict] = dict(parsed_body or {}) + request_payload: Final[dict] = dict(cast(Mapping[str, object], parsed_body or {})) # cast-ok: json body if kwargs: request_payload.update(kwargs) if logging_obj is not None: @@ -1163,7 +1173,7 @@ async def pass_through_request( is_multipart: Final = HttpPassThroughEndpointHelpers.is_multipart(request) and not custom_body if custom_body: - _parsed_body = custom_body + _parsed_body = dict(custom_body) elif is_multipart: # Don't parse multipart body here - it will be handled by make_multipart_http_request _parsed_body = {} @@ -1232,6 +1242,14 @@ async def pass_through_request( if _parsed_body is None: _parsed_body = {} _parsed_body["litellm_logging_obj"] = logging_obj + is_streaming_pass_through: Final = bool( + HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body( + parsed_body=_parsed_body, + stream=stream, + ) + ) + typed_body: Final[Mapping[str, object]] = cast(Mapping[str, object], _parsed_body) # cast-ok: json + _parsed_body = guardrail_request_data_with_streaming(typed_body, is_streaming=is_streaming_pass_through) ### CALL HOOKS ### - modify incoming data / reject request before calling the model _parsed_body = await proxy_logging_obj.pre_call_hook( @@ -2518,10 +2536,10 @@ async def websocket_passthrough_request( ) ### CALL HOOKS ### - modify incoming data / reject request before calling the model - websocket_data: dict[str, object] = {} - websocket_data = await proxy_logging_obj.pre_call_hook( + websocket_hook_data: Final = guardrail_request_data_with_streaming(MappingProxyType({}), is_streaming=True) + await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, - data=websocket_data, + data=websocket_hook_data, call_type="pass_through_endpoint", ) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index d80eb56d8e3..5ccf2920030 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -8,7 +8,8 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. import copy import time from collections.abc import Callable, Mapping, Sequence -from typing import TYPE_CHECKING, Final, Literal, TypeVar +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeVar, cast from pydantic import BaseModel @@ -29,6 +30,7 @@ from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_g from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.policy_engine.pipeline_types import ( PipelineExecutionResult, PipelineStep, @@ -265,6 +267,30 @@ class _LegacyHookStreamAdapter(CustomGuardrail): return recorder.inputs +_PIPELINE_EVENT_HOOKS: Final = MappingProxyType( + { + "pre_call": GuardrailEventHooks.pre_call, + "post_call": GuardrailEventHooks.post_call, + "during_call": GuardrailEventHooks.during_call, + } +) + + +def _pipeline_stream_scope_allows( + callback: CustomGuardrail, + hook_input: Mapping[str, object], + mode: str, + streaming_chunks: list[object] | None, +) -> bool: + event_type: Final = _PIPELINE_EVENT_HOOKS.get(mode) + if event_type is None: + return True + return callback.stream_scope_allows( + hook_input if streaming_chunks is None else {**hook_input, "stream": True}, + event_type, + ) + + def _prepare_hook_input( step: PipelineStep, callback: CustomGuardrail, @@ -494,7 +520,7 @@ class PipelineExecutor: streaming_chunks: list[object] | None = None, # mutable-ok: shared buffered-stream chunks, read per step endpoint_translation: "BaseTranslation | None" = None, ) -> tuple[ - Literal["pass", "fail", "error"], + Literal["pass", "fail", "error", "skip"], dict | None, str | None, Exception | None, @@ -504,7 +530,7 @@ class PipelineExecutor: Returns: Tuple of (outcome, modified_data, error_detail, original_exception): - - outcome: "pass", "fail", or "error" + - outcome: "pass", "fail", "error", or "skip" - modified_data: dict if guardrail returned modified data, else None - error_detail: error message string if fail/error, else None - original_exception: the exception the guardrail raised, so the @@ -516,6 +542,10 @@ class PipelineExecutor: verbose_proxy_logger.warning("Pipeline: guardrail '%s' not found in callbacks", step.guardrail) return ("error", None, f"Guardrail '{step.guardrail}' not found", None) + hook_data: Final[Mapping[str, object]] = cast(Mapping[str, object], data) # cast-ok: payload + if not _pipeline_stream_scope_allows(callback, hook_data, mode, streaming_chunks): + return ("skip", None, None, None) + hook_input, scans_raw_request = _prepare_hook_input(step, callback, data, raw_request_snapshot) snapshot_entries_before: Final = len(_recorded_guardrail_information(hook_input)) @@ -702,10 +732,13 @@ def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str: """ Map pipeline step outcome to the configured action. + - skip -> next (stream_scope mismatch; do not apply on_pass/on_fail) - pass -> on_pass - fail -> on_fail (content/policy intervention) - error -> on_error if set, else on_fail (backward compatible) """ + if outcome == "skip": + return "next" if outcome == "pass": return step.on_pass if outcome == "fail": diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f07656bc51d..887aa4a0ae9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6958,6 +6958,13 @@ class ProxyConfig: ).user_api_key_cache_max_size ) + if "spend_logs_metadata_fields" in general_settings: + _ = ConfigGeneralSettings.model_validate( + MappingProxyType( + {"spend_logs_metadata_fields": typed_general_settings["spend_logs_metadata_fields"]} + ) + ) + ### PKCE MULTI-INSTANCE PREREQUISITE CHECK ### # PKCE verifiers are stored in redis_usage_cache when available so they can # be read back by any instance (not just the one that started the auth flow). diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 9d52f7d71bb..99bd2fbc480 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -9,7 +9,7 @@ from functools import reduce from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Protocol, cast, runtime_checkable -from pydantic import BaseModel, JsonValue +from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -47,7 +47,12 @@ from litellm.litellm_core_utils.ptu_pricing import azure_spillover from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.llms.anthropic.common_utils import resolve_used_client_oauth_token -from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload, SpendLogsRouterMetadata +from litellm.proxy._types import ( + SpendLogsMetadata, + SpendLogsMetadataFields, + SpendLogsPayload, + SpendLogsRouterMetadata, +) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.proxy.utils import PrismaClient, hash_token @@ -1726,6 +1731,35 @@ def should_store_prompts_and_responses_in_spend_logs() -> bool: return get_secret_bool("STORE_PROMPTS_IN_SPEND_LOGS") is True +_SPEND_LOGS_METADATA_FIELDS_ADAPTER: Final[TypeAdapter[SpendLogsMetadataFields | None]] = TypeAdapter( + SpendLogsMetadataFields | None +) +_SPEND_LOGS_METADATA_ADAPTER: Final = TypeAdapter(dict[str, JsonValue]) + + +def configured_spend_logs_metadata_fields() -> SpendLogsMetadataFields | None: + from litellm.proxy.proxy_server import general_settings_view + + try: + return _SPEND_LOGS_METADATA_FIELDS_ADAPTER.validate_python( + general_settings_view().get("spend_logs_metadata_fields") + ) + except ValidationError as e: + verbose_proxy_logger.error("Ignoring invalid general_settings.spend_logs_metadata_fields: %s", e) + return None + + +def spend_log_row_with_retained_metadata( + row: Mapping[str, object], fields: SpendLogsMetadataFields | None +) -> Mapping[str, object]: + metadata_json: Final = row.get("metadata") + if fields is None or not isinstance(metadata_json, str): + return row + metadata: Final = _SPEND_LOGS_METADATA_ADAPTER.validate_json(metadata_json) + retained: Final = {name: value for name, value in metadata.items() if fields.keeps(name)} + return {**row, "metadata": safe_dumps(retained)} + + def _get_status_for_spend_log( metadata: dict, ) -> Literal["success", "failure"]: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8f796f065a6..4743970a526 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4396,9 +4396,13 @@ async def get_config_param(prisma_client: "PrismaClient", param_name: str) -> An return row -async def evict_config_param(param_name: str) -> None: - with service_target(CONFIG_PARAMS_TARGET): - await litellm_config_cache.async_delete_cache(_config_cache_key(param_name)) +async def evict_config_param(param_name: str, cache: DualCache | None = None) -> None: + target: Final = cache if cache is not None else litellm_config_cache + try: + with service_target(CONFIG_PARAMS_TARGET): + await target.async_delete_cache(_config_cache_key(param_name)) + except Exception as e: # noqa: BLE001 # best-effort eviction; config writes must never fail on redis errors + verbose_proxy_logger.warning("config cache eviction of %s failed: %s", param_name, e) async def invalidate_config_param(param_name: str) -> None: @@ -7446,6 +7450,13 @@ class ProxyUpdateSpend: "Spend tracking - processing %d spend logs for DB write", len(logs_to_process), ) + from litellm.proxy.spend_tracking.spend_tracking_utils import ( + configured_spend_logs_metadata_fields, + spend_log_row_with_retained_metadata, + ) + + retention: Final = configured_spend_logs_metadata_fields() + rows_to_write: Final = [spend_log_row_with_retained_metadata(row, retention) for row in logs_to_process] start_time: Final = time.time() try: for i in range(n_retry_times + 1): @@ -7455,7 +7466,7 @@ class ProxyUpdateSpend: if not base_url.endswith("/"): base_url += "/" verbose_proxy_logger.debug("base_url: %s", base_url) - json_data = json.dumps(logs_to_process) + json_data = json.dumps(rows_to_write) response = await db_writer_client.post( url=base_url + "spend/update", data=json_data, @@ -7466,8 +7477,8 @@ class ProxyUpdateSpend: # Items already removed from queue at start of function pass else: - for j in range(0, len(logs_to_process), BATCH_SIZE): - batch = logs_to_process[j : j + BATCH_SIZE] + for j in range(0, len(rows_to_write), BATCH_SIZE): + batch = rows_to_write[j : j + BATCH_SIZE] batch_with_dates = [prisma_client.jsonify_object({**entry}) for entry in batch] isolation_budget = MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH for statement_rows in spend_log_write_batches( diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 5ca7b79d127..c83e4e5318c 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -110,12 +110,6 @@ def messages( def amessages( call: NativeCall, ) -> Coroutine[object, object, AnthropicMessagesResponse | AsyncIterator[bytes]]: ... -def chat_completions( - call: NativeCall, -) -> dict[str, object]: ... -def achat_completions( - call: NativeCall, -) -> Future[dict[str, object]]: ... @final class ResponsesWebSocketConnection: @@ -251,14 +245,12 @@ __all__ = [ "RustUpstreamError", "TokenCounter", "Tokenizer", - "achat_completions", "acompletion", "aembedding", "amessages", "aocr", "aresponses", "atranscription", - "chat_completions", "completion", "embedding", "gil_stats", diff --git a/litellm/rust_bridge/chat_completions/route_host.py b/litellm/rust_bridge/chat_completions/route_host.py index cb06ea8e213..d600495c919 100644 --- a/litellm/rust_bridge/chat_completions/route_host.py +++ b/litellm/rust_bridge/chat_completions/route_host.py @@ -37,16 +37,12 @@ def response(value: Mapping[str, object]) -> ModelResponse: return ModelResponse(**value) -def arguments(request: Mapping[str, object]) -> Mapping[str, object]: - return request - - def map_failure(error: Exception, request: Mapping[str, object]) -> Exception: provider: Final = optional_str(request.get("custom_llm_provider")) or str(request["model"]).partition("/")[0] return failures.map_native_failure( error, str(request["model"]), provider, - arguments(request), + request, optional_str(request.get("api_base")) or optional_str(request.get("base_url")), ) diff --git a/litellm/rust_bridge/messages/route_host.py b/litellm/rust_bridge/messages/route_host.py index 88a81fdae4a..69e27e78223 100644 --- a/litellm/rust_bridge/messages/route_host.py +++ b/litellm/rust_bridge/messages/route_host.py @@ -39,10 +39,6 @@ def stream_hidden_params(headers: Sequence[tuple[str, str]]) -> Mapping[str, obj return anthropic_messages_stream_hidden_params(httpx.Headers(list(headers))) -def arguments(request: Mapping[str, object]) -> Mapping[str, object]: - return request - - def map_failure(error: Exception, request: Mapping[str, object], request_provider: str) -> Exception: if getattr(error, "messages_request_error", False): return litellm.BadRequestError( @@ -51,7 +47,7 @@ def map_failure(error: Exception, request: Mapping[str, object], request_provide llm_provider=request_provider, ) return failures.map_native_failure( - error, str(request["model"]), request_provider, arguments(request), optional_str(request.get("api_base")) + error, str(request["model"]), request_provider, request, optional_str(request.get("api_base")) ) diff --git a/litellm/rust_bridge/ocr/route_host.py b/litellm/rust_bridge/ocr/route_host.py index a7eb0829f5c..8f563a381ef 100644 --- a/litellm/rust_bridge/ocr/route_host.py +++ b/litellm/rust_bridge/ocr/route_host.py @@ -12,7 +12,7 @@ from litellm.rust_bridge import failures from litellm.rust_bridge.failures import UpstreamFailure from litellm.rust_bridge.public_call import optional_str -__all__ = ("UpstreamFailure", "arguments", "map_failure", "response") +__all__ = ("UpstreamFailure", "map_failure", "response") _RESPONSE_ADAPTER: Final = TypeAdapter(dict[str, object]) @@ -27,10 +27,6 @@ def response(value: Mapping[str, object]) -> OCRResponse: return normalized -def arguments(request: Mapping[str, object]) -> Mapping[str, object]: - return request - - def map_failure(error: Exception, request: Mapping[str, object], request_provider: str) -> Exception: if getattr(error, "ocr_request_format_error", False): return litellm.UnsupportedParamsError( @@ -39,5 +35,5 @@ def map_failure(error: Exception, request: Mapping[str, object], request_provide llm_provider=request_provider, ) return failures.map_native_failure( - error, str(request["model"]), request_provider, arguments(request), optional_str(request.get("api_base")) + error, str(request["model"]), request_provider, request, optional_str(request.get("api_base")) ) diff --git a/litellm/rust_bridge/responses/route_host.py b/litellm/rust_bridge/responses/route_host.py index 00f24244f49..d8ad2a94b75 100644 --- a/litellm/rust_bridge/responses/route_host.py +++ b/litellm/rust_bridge/responses/route_host.py @@ -20,17 +20,13 @@ def response(value: Mapping[str, object]) -> ResponsesAPIResponse: return ResponsesAPIResponse.model_validate(value) -def arguments(request: Mapping[str, object]) -> Mapping[str, object]: - return request - - def map_failure(error: Exception, request: Mapping[str, object]) -> Exception: provider: Final = optional_str(request.get("custom_llm_provider")) or "openai" return failures.map_native_failure( error, str(request["model"]), provider, - arguments(request), + request, optional_str(request.get("api_base")) or optional_str(request.get("base_url")), ) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 201b92f8107..ecc3d0d747c 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -2,11 +2,12 @@ from collections.abc import Mapping from datetime import datetime from enum import Enum from types import MappingProxyType -from typing import Final, Literal +from typing import Final, Literal, cast from pydantic import ConfigDict, Field, field_validator, model_validator from typing_extensions import ReadOnly, Required, TypedDict +from litellm._logging import verbose_logger from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.proxy.guardrails.guardrail_hooks.agent_365 import ( @@ -901,6 +902,93 @@ class ContentFilterConfigModel(LiteLLMBaseModel): MCP_SECURITY_ON_VIOLATION: Final = frozenset({"block", "alert"}) +GuardrailStreamScope = Literal["streaming", "non_streaming", "both"] +DEFAULT_GUARDRAIL_STREAM_SCOPE: Final[GuardrailStreamScope] = "both" + + +class GuardrailEventHooks(str, Enum): + pre_call = "pre_call" + post_call = "post_call" + during_call = "during_call" + logging_only = "logging_only" + pre_mcp_call = "pre_mcp_call" + during_mcp_call = "during_mcp_call" + post_mcp_call = "post_mcp_call" + realtime_input_transcription = "realtime_input_transcription" + + +GUARDRAIL_EVENT_HOOK_VALUES: Final = frozenset(member.value for member in GuardrailEventHooks) + +_GUARDRAIL_STREAM_SCOPES: Final[Mapping[str, GuardrailStreamScope]] = MappingProxyType( + { + "streaming": "streaming", + "non_streaming": "non_streaming", + "both": "both", + } +) + + +def _as_guardrail_stream_scope(value: object) -> GuardrailStreamScope: + if not isinstance(value, str): + raise ValueError(f"stream_scope values must be strings, got {type(value).__name__}") + scope: Final = _GUARDRAIL_STREAM_SCOPES.get(value.lower()) + if scope is None: + raise ValueError(f"stream_scope must be one of both, streaming, non_streaming, got {value!r}") + return scope + + +def _validated_stream_scope_hook(key: object) -> str: + if not isinstance(key, str): + raise ValueError(f"stream_scope keys must be strings, got {type(key).__name__}") + hook: Final = key.lower() + if hook not in GUARDRAIL_EVENT_HOOK_VALUES: + raise ValueError( + f"stream_scope keys must be guardrail modes ({sorted(GUARDRAIL_EVENT_HOOK_VALUES)}), got {key!r}" + ) + return hook + + +def coerce_stream_scope(value: object) -> GuardrailStreamScope | dict[str, GuardrailStreamScope] | None: + if value is None: + return None + if isinstance(value, str): + return _as_guardrail_stream_scope(value) + if isinstance(value, Mapping): + scope_map: Final[Mapping[str, object]] = cast(Mapping[str, object], value) # cast-ok: keys validated below + return { + _validated_stream_scope_hook(key): _as_guardrail_stream_scope(scope) for key, scope in scope_map.items() + } + raise ValueError(f"stream_scope must be a string or mapping, got {type(value).__name__}") + + +def stored_stream_scope(value: object) -> GuardrailStreamScope | dict[str, GuardrailStreamScope] | None: + try: + return coerce_stream_scope(value) + except ValueError: + verbose_logger.warning("Ignoring invalid stored stream_scope value of type %s", type(value).__name__) + return None + + +def with_tolerated_stream_scope(params: Mapping[str, object]) -> dict[str, object]: + if "stream_scope" not in params: + return dict(params) + return { + **params, + "stream_scope": stored_stream_scope(params["stream_scope"]), + } + + +def runtime_stream_scope( + stream_scope: object, +) -> tuple[GuardrailStreamScope, MappingProxyType[str, GuardrailStreamScope]]: + coerced: Final = coerce_stream_scope(stream_scope) + if coerced is None: + return DEFAULT_GUARDRAIL_STREAM_SCOPE, MappingProxyType({}) + if isinstance(coerced, str): + return coerced, MappingProxyType({}) + return DEFAULT_GUARDRAIL_STREAM_SCOPE, MappingProxyType(coerced) + + LoggingOnlyScope = Literal["input", "output", "both"] @@ -1144,6 +1232,21 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up ), ) + stream_scope: GuardrailStreamScope | dict[str, GuardrailStreamScope] | None = Field( + default=None, + description=( + "Whether this guardrail runs on streaming requests, non-streaming requests, or both. " + "A string applies to every configured mode. A map overrides named modes " + "(pre_call, during_call, post_call, ...); omitted keys default to both. " + "Unset means both, matching historical behavior." + ), + ) + + @field_validator("stream_scope", mode="before") + @classmethod + def normalize_stream_scope(cls, v: object) -> GuardrailStreamScope | dict[str, GuardrailStreamScope] | None: + return coerce_stream_scope(v) + logging_only_scope: LoggingOnlyScope | None = Field( default=None, description=( @@ -1278,17 +1381,6 @@ class guardrailConfig(TypedDict): guardrails: list[Guardrail] -class GuardrailEventHooks(str, Enum): - pre_call = "pre_call" - post_call = "post_call" - during_call = "during_call" - logging_only = "logging_only" - pre_mcp_call = "pre_mcp_call" - during_mcp_call = "during_mcp_call" - post_mcp_call = "post_mcp_call" - realtime_input_transcription = "realtime_input_transcription" - - class DynamicGuardrailParams(TypedDict): extra_body: ReadOnly[dict[str, object]] diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index ce51e46ef15..116fdda0b5d 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -2,16 +2,20 @@ from enum import Enum from typing import Any, Final, Literal, Protocol from typing_extensions import ( + ReadOnly, Required, TypedDict, ) -from litellm.types.llms.openai import EmbeddingInput +from litellm.types.llms.openai import ChatCompletionFileObject, EmbeddingInput # Gemini supports nested-list inputs (e.g. [["text", "image"]]) as an explicit # opt-in for combined embeddings — a provider-specific extension of the # OpenAI-faithful EmbeddingInput shape. -GeminiEmbeddingInput = EmbeddingInput | list[list[str]] +GeminiEmbeddingElement = str | ChatCompletionFileObject +GeminiEmbeddingInput = ( + EmbeddingInput | list[GeminiEmbeddingElement] | list[list[str]] | list[list[GeminiEmbeddingElement]] +) class FunctionResponse(TypedDict, total=False): @@ -46,6 +50,12 @@ class FunctionResponsePartType(TypedDict, total=False): file_data: FileDataType +class VideoMetadataType(TypedDict, total=False): + fps: ReadOnly[float] + startOffset: ReadOnly[str] + endOffset: ReadOnly[str] + + class PartType(TypedDict, total=False): text: str inline_data: BlobType @@ -55,6 +65,7 @@ class PartType(TypedDict, total=False): thought: bool thoughtSignature: str media_resolution: Literal["low", "medium", "high"] + video_metadata: ReadOnly[VideoMetadataType] class HttpxFunctionCall(TypedDict, total=False): diff --git a/litellm/types/openai_decisions.py b/litellm/types/openai_decisions.py deleted file mode 100644 index 074c0a803f3..00000000000 --- a/litellm/types/openai_decisions.py +++ /dev/null @@ -1,162 +0,0 @@ -from collections.abc import Mapping, Sequence -from dataclasses import dataclass -from typing import Annotated, Literal, TypeAlias - -from pydantic import ConfigDict, Field, PrivateAttr, StrictBool, StrictStr - -from litellm.types.llms.base import LiteLLMPydanticObjectBase - -ChoiceValue: TypeAlias = StrictStr | StrictBool - - -class DecisionsObjectBase(LiteLLMPydanticObjectBase): - model_config = ConfigDict(extra="allow", frozen=True) - - -class DecisionInputText(DecisionsObjectBase): - type: Literal["input_text"] - text: str - - -class DecisionInputImage(DecisionsObjectBase): - type: Literal["input_image"] - image_url: str - detail: Literal["low", "high", "auto", "original"] | None = None - - -DecisionInputPart: TypeAlias = Annotated[DecisionInputText | DecisionInputImage, Field(discriminator="type")] - - -class DecisionInputMessage(DecisionsObjectBase): - role: Literal["user"] - content: str | Sequence[DecisionInputPart] - type: Literal["message"] | None = None - - -DecisionInput: TypeAlias = str | Sequence[DecisionInputMessage] - - -class DecisionChoice(DecisionsObjectBase): - value: ChoiceValue - description: str | None = None - - -class DecisionLevel(DecisionsObjectBase): - label: str - description: str | None = None - - -class PredicateQuestion(DecisionsObjectBase): - type: Literal["predicate"] - instructions: str - name: str | None = None - - -class ChoiceQuestion(DecisionsObjectBase): - type: Literal["choice"] - instructions: str - choices: Sequence[DecisionChoice] - name: str | None = None - - -class ScoreQuestion(DecisionsObjectBase): - type: Literal["score"] - instructions: str - levels: Sequence[DecisionLevel] - name: str | None = None - - -DecisionQuestion: TypeAlias = Annotated[ - PredicateQuestion | ChoiceQuestion | ScoreQuestion, - Field(discriminator="type"), -] - -DecisionQuestions: TypeAlias = Sequence[DecisionQuestion] - - -class DecisionsRequestBody(DecisionsObjectBase): - input: DecisionInput - questions: DecisionQuestions - safety_identifier: str | None = None - - -@dataclass(frozen=True, slots=True) -class DecisionsRequest: - model: str - body: DecisionsRequestBody - - -class PredicateAnswer(DecisionsObjectBase): - type: Literal["predicate"] - name: str | None = None - probability: float - - -class ChoiceProbability(DecisionsObjectBase): - value: ChoiceValue - probability: float - - -class ChoiceAnswer(DecisionsObjectBase): - type: Literal["choice"] - name: str | None = None - choice: ChoiceValue - probabilities: Sequence[ChoiceProbability] - confidence: float - - -class ScoreProbability(DecisionsObjectBase): - value: int - label: str - probability: float - - -class ScoreAnswer(DecisionsObjectBase): - type: Literal["score"] - name: str | None = None - score: float - probabilities: Sequence[ScoreProbability] - confidence: float - - -class RefusalAnswer(DecisionsObjectBase): - type: Literal["refusal"] - name: str | None = None - - -DecisionAnswer: TypeAlias = Annotated[ - PredicateAnswer | ChoiceAnswer | ScoreAnswer | RefusalAnswer, - Field(discriminator="type"), -] - - -class DecisionInputTokensDetails(DecisionsObjectBase): - cached_tokens: int - cache_write_tokens: int - - -class DecisionOutputTokensDetails(DecisionsObjectBase): - reasoning_tokens: int - - -class DecisionUsage(DecisionsObjectBase): - input_tokens: int - input_tokens_details: DecisionInputTokensDetails - output_tokens: int - output_tokens_details: DecisionOutputTokensDetails - total_tokens: int - - -class DecisionsResponse(DecisionsObjectBase): - model: str - answers: Sequence[DecisionAnswer] - usage: DecisionUsage - - _hidden_params: dict[str, object] = PrivateAttr(default_factory=dict) - - @property - def hidden_params(self) -> dict[str, object]: # mutable-ok: API requires mutation - return self._hidden_params - - def set_hidden_params(self, params: Mapping[str, object]) -> None: - self._hidden_params.update(params) diff --git a/litellm/types/proxy/policy_engine/pipeline_types.py b/litellm/types/proxy/policy_engine/pipeline_types.py index 5278754701d..0a089d2fe9b 100644 --- a/litellm/types/proxy/policy_engine/pipeline_types.py +++ b/litellm/types/proxy/policy_engine/pipeline_types.py @@ -87,7 +87,7 @@ class PipelineStepResult(LiteLLMBaseModel): """Result of executing a single pipeline step.""" guardrail_name: str - outcome: Literal["pass", "fail", "error"] + outcome: Literal["pass", "fail", "error", "skip"] action_taken: str modified_data: dict[str, Any] | None = None error_detail: str | None = None diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 663c6d1a2a7..bbd0d9ab989 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -42320,6 +42320,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "input_cost_per_token": 1.5e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, @@ -42338,6 +42339,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "input_cost_per_token": 7e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, @@ -42363,6 +42365,7 @@ "output_cost_per_token_flex": 3e-07, "source": "https://aws.amazon.com/bedrock/pricing/", "supports_audio_input": false, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -42380,6 +42383,7 @@ "output_cost_per_token_flex": 1e-07, "source": "https://aws.amazon.com/bedrock/pricing/", "supports_audio_input": false, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -48105,6 +48109,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -48123,6 +48128,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -65951,6 +65957,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -65969,6 +65976,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -66219,6 +66227,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -66237,6 +66246,7 @@ "/v1/chat/completions" ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -80458,29 +80468,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "xai.grok-4.7": { - "supports_regex_lookaround": false, - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_flex": 2.5e-07, - "cache_read_input_token_cost_priority": 8.75e-07, - "input_cost_per_token": 2e-06, - "input_cost_per_token_flex": 1e-06, - "input_cost_per_token_priority": 3.5e-06, - "litellm_provider": "bedrock_converse", - "max_input_tokens": 500000, - "max_output_tokens": 500000, - "max_tokens": 500000, - "mode": "chat", - "output_cost_per_token": 6e-06, - "output_cost_per_token_flex": 3e-06, - "output_cost_per_token_priority": 1.05e-05, - "source": "https://b0.p.awsstatic.com/pricing/2.0/meteredUnitMaps/bedrock/USD/current/bedrock.json", - "supports_function_calling": true, - "supports_prompt_caching": false, - "supports_reasoning": true, - "supports_tool_choice": true, - "supports_vision": true - }, "openrouter/anthropic/claude-sonnet-5.5:batch": { "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index e51ae0453ea..d3843b78ff4 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -1039,6 +1039,9 @@ "supports_audio_output": { "type": "boolean" }, + "supports_bedrock_runtime_chat_completions_inline_reasoning": { + "type": "boolean" + }, "supports_bedrock_runtime_chat_completions_response_format": { "type": "boolean" }, diff --git a/tests/e2e/llm_translation/test_messages_count_tokens_bedrock_e2e.py b/tests/e2e/llm_translation/test_messages_count_tokens_bedrock_e2e.py new file mode 100644 index 00000000000..ee5cd54e4e0 --- /dev/null +++ b/tests/e2e/llm_translation/test_messages_count_tokens_bedrock_e2e.py @@ -0,0 +1,70 @@ +"""Live e2e: `/v1/messages/count_tokens` on a Bedrock Claude model that bedrock-runtime +cannot count. + +Claude Opus 4.8 is offered only through cross-region inference, and bedrock-runtime's +CountTokens answers 400 for it. The proxy then has to count through bedrock-mantle's +Anthropic count_tokens, and the answer must sit within a few percent of what `/v1/messages` +bills as `usage.input_tokens`. The local tokenizer fallback undercounts these models by +about 40%, so this is the line that proves the real count is served. Both calls go through +the real Anthropic SDK, the client customers count with +""" + +from __future__ import annotations + +from typing import Final + +import pytest +from anthropic.types import MessageParam +from e2e_config import unique_marker +from e2e_metadata import Domain, Provider, Route, Subject, meta +from lifecycle import ResourceManager +from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import NO_PROXY_CACHE, SdkClients + +pytestmark = pytest.mark.e2e + +CROSS_REGION_ONLY_CLAUDE_BACKEND: Final = "bedrock/global.anthropic.claude-opus-4-8" +COUNT_TOLERANCE: Final = 0.05 + + +def _register(proxy: ProxyClient, resources: ResourceManager, backend: str) -> tuple[str, str]: + model: Final = f"e2e-count-tokens-bedrock-{unique_marker()}" + model_id: Final = proxy.create_model( + model, + LiteLLMParamsBody( + model=backend, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ), + ) + resources.defer(lambda: proxy.delete_model(model_id)) + return model, resources.key() + + +class TestBedrockMessagesCountTokens: + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.COUNT_TOKENS, + providers=(Provider.BEDROCK,), + models=(CROSS_REGION_ONLY_CLAUDE_BACKEND,), + ) + ) + def test_count_matches_billed_input_tokens_for_a_model_bedrock_runtime_cannot_count( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model, key = _register(proxy, resources, CROSS_REGION_ONLY_CLAUDE_BACKEND) + client: Final = sdk.anthropic(key) + prompt: Final = f"{unique_marker()} " + "The quick brown fox jumps over the lazy dog. " * 40 + message: Final[MessageParam] = {"role": "user", "content": prompt} + + counted: Final = client.messages.count_tokens(model=model, messages=[message]) + answered: Final = client.messages.create( + model=model, max_tokens=1, messages=[message], extra_body=NO_PROXY_CACHE + ) + + billed: Final = answered.usage.input_tokens + assert billed, answered.usage + assert abs(counted.input_tokens - billed) <= billed * COUNT_TOLERANCE, (counted, answered.usage) diff --git a/tests/e2e/ui/fixtures/migratedPages.ts b/tests/e2e/ui/fixtures/migratedPages.ts index d44dfc625aa..5fa1f9d2444 100644 --- a/tests/e2e/ui/fixtures/migratedPages.ts +++ b/tests/e2e/ui/fixtures/migratedPages.ts @@ -117,7 +117,7 @@ export const MIGRATED_E2E_PAGES: Readonly> = { linkName: "AI Hub", content: { role: "heading", name: "AI Hub" }, }, - new_usage: { segment: "usage", linkName: "Usage", content: { role: "heading", name: "Usage View" } }, + new_usage: { segment: "usage", linkName: "Usage", content: { role: "heading", name: "Usage" } }, usage: { segment: "old-usage", linkName: "Old Usage", diff --git a/tests/e2e/ui/tests/usage/usagePage.spec.ts b/tests/e2e/ui/tests/usage/usagePage.spec.ts index 3d057cfa2c9..d532042c7af 100644 --- a/tests/e2e/ui/tests/usage/usagePage.spec.ts +++ b/tests/e2e/ui/tests/usage/usagePage.spec.ts @@ -13,9 +13,8 @@ import { /** Covers /ui/usage. The legacy /ui/old-usage view is deprecated and deliberately not covered. */ -/** Stepping up from the title is exact; the page renders several other tables. */ const topKeysCard = (page: PlaywrightPage): Locator => - page.getByText("Top Virtual Keys", { exact: true }).locator("xpath=.."); + page.locator("section").filter({ has: page.getByRole("heading", { name: "Top Virtual Keys", exact: true }) }); async function openUsage(page: PlaywrightPage): Promise { await navigateToPage(page, Page.NewUsage); diff --git a/tests/integration/_support/wire.py b/tests/integration/_support/wire.py index d118da03daf..e6a166e8457 100644 --- a/tests/integration/_support/wire.py +++ b/tests/integration/_support/wire.py @@ -3,13 +3,13 @@ from __future__ import annotations import ssl import threading import time -from collections.abc import Callable, Generator, Mapping +from collections.abc import Callable, Generator, Iterator, Mapping from contextlib import contextmanager from dataclasses import dataclass from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from queue import SimpleQueue from types import MappingProxyType -from typing import Final +from typing import BinaryIO, Final @dataclass(frozen=True, slots=True) @@ -20,6 +20,37 @@ class Request: body: bytes +class AbortedBody(Exception): + """The client closed the connection before the body it announced was complete.""" + + +def exactly(stream: BinaryIO, size: int) -> bytes: + data: Final = stream.read(size) + if len(data) < size: + raise AbortedBody + return data + + +def chunked_body(stream: BinaryIO) -> Iterator[bytes]: + while True: + size_line: Final = stream.readline() + if not size_line: + raise AbortedBody + size: Final = int(size_line.split(b";")[0].strip(), 16) + if size == 0: + while stream.readline().strip(): + pass + return + yield exactly(stream, size) + stream.readline() + + +def read_body(headers: Mapping[str, str], stream: BinaryIO) -> bytes: + if headers.get("transfer-encoding", "").lower() == "chunked": + return b"".join(chunked_body(stream)) + return exactly(stream, int(headers.get("content-length", "0"))) + + @dataclass(frozen=True, slots=True) class Reply: status: int = 200 @@ -78,12 +109,14 @@ def wire_server( connected.put(f"{self.client_address[0]}:{self.client_address[1]}") def respond(self) -> None: - request: Final = Request( - self.command, - self.path, - {name.lower(): value for name, value in self.headers.items()}, - self.rfile.read(int(self.headers.get("content-length", "0"))), - ) + headers: Final = {name.lower(): value for name, value in self.headers.items()} + try: + body: Final = read_body(headers, self.rfile) + except AbortedBody: + self.close_connection = True + disconnected.put(self.path) + return + request: Final = Request(self.command, self.path, headers, body) received.put(request) try: reply = respond(request) diff --git a/tests/integration/authorization/test_cli_sso_login_ui_disabled.py b/tests/integration/authorization/test_cli_sso_login_ui_disabled.py index 126fe27bef2..a5e9b2752ea 100644 --- a/tests/integration/authorization/test_cli_sso_login_ui_disabled.py +++ b/tests/integration/authorization/test_cli_sso_login_ui_disabled.py @@ -23,7 +23,14 @@ import pytest import yaml from pydantic import JsonValue -from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, gateway_from_environment, string_value +from tests.integration._support.client import ( + JSON_OBJECT, + Gateway, + eventually, + gateway_from_environment, + object_value, + string_value, +) from tests.integration._support.database import read_rows from tests.integration._support.process import OwnedProxy, group_members, owned_proxy_process from tests.integration._support.provider import SharedProvider @@ -648,3 +655,52 @@ def test_flag_values_that_do_not_disable_keep_the_gates_open(idp: Idp, tmp_path: for method, path in (("GET", "/sso/saml/login"), ("POST", "/sso/saml/callback")): saml: Final = proxy.client.request(method, path) assert DISABLED_PAGE_TITLE not in saml.text and saml.status_code != 200, f"{path}: {saml.status_code}" + + +def test_cli_session_token_is_denied_once_its_team_budget_is_exhausted( + one_worker: OneWorkerProxy, provider: SharedProvider +) -> None: + proxy: Final = one_worker.owned.gateway + subject: Final = f"cli-sso-team-budget-{uuid.uuid4().hex[:12]}" + budget: Final = 0.0000000005 + with proxy.scenario() as scenario: + team: Final = scenario.team(max_budget=budget, models=[MESSAGE_MODEL]) + scenario.user(user_id=subject, user_email=f"{subject}@example.com", user_role="internal_user") + added: Final = proxy.request( + "POST", "/team/member_add", {"team_id": team, "member": {"user_id": subject, "role": "user"}} + ) + assert added.status_code == 200, added.text + session: Final = _start_lite_login(proxy) + with _browser() as browser: + _sign_in(proxy, one_worker.idp, browser, session, subject=subject) + ready: Final = proxy.client.get( + f"/sso/cli/poll/{session.login_id}", + params={"team_id": team}, + headers={POLL_SECRET_HEADER: session.poll_secret}, + ) + assert ready.status_code == 200, f"{ready.status_code} {ready.text}" + body: Final = JSON_OBJECT.validate_json(ready.content) + assert body["status"] == "ready" and body["user_id"] == subject, ready.text + key: Final = string_value(body["key"]) + assert not key.startswith("sk-"), key + _send_message(proxy, provider, key) + eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team,)), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > budget, + seconds=70, + ) + refused: Final = proxy.request( + "POST", + "/v1/messages", + {"model": MESSAGE_MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "over team budget"}]}, + key=key, + ) + assert refused.status_code == 422, f"{refused.status_code} {refused.text}" + error: Final = object_value(JSON_OBJECT.validate_json(refused.content)["error"]) + assert error["type"] == "budget_exceeded", refused.text + assert error["code"] == "422", refused.text + message: Final = string_value(error["message"]) + assert "Budget has been exceeded!" in message, refused.text + assert f"Team={team}" in message, refused.text + assert "Current cost:" in message and f"Max budget: {budget}" in message, refused.text + assert provider.received() == () diff --git a/tests/integration/authorization/test_model_access_allow_lists.py b/tests/integration/authorization/test_model_access_allow_lists.py new file mode 100644 index 00000000000..cdefc8c8a3e --- /dev/null +++ b/tests/integration/authorization/test_model_access_allow_lists.py @@ -0,0 +1,171 @@ +from __future__ import annotations + +import json +import uuid +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Final, Literal + +import pytest + +from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, object_value, string_value +from tests.integration._support.provider import PROVIDER_URL, SharedProvider +from tests.integration._support.wire import Reply + +_Denial = Literal["key_model_access_denied", "team_model_access_denied"] + + +@dataclass(frozen=True, slots=True) +class _Models: + gpt: str + gpt_mini: str + claude: str + bedrock_claude: str + bedrock_titan: str + + +@dataclass(frozen=True, slots=True) +class _AccessCase: + name: str + allowed: Sequence[str] | None + requested: str + served: bool + + +_CASES: Final = ( + _AccessCase("openai_wildcard_denies_anthropic", ["openai/*"], "claude", False), + _AccessCase("exact_name_allows_itself", ["gpt"], "gpt", True), + _AccessCase("provider_wildcard_allows_bedrock", ["bedrock/*"], "bedrock_claude", True), + _AccessCase("family_wildcard_allows_its_family", ["bedrock/anthropic.*"], "bedrock_claude", True), + _AccessCase("family_wildcard_denies_another_family", ["bedrock/anthropic.*"], "bedrock_titan", False), + _AccessCase("unset_models_allow_everything", None, "gpt", True), + _AccessCase("empty_models_allow_everything", [], "gpt", True), +) + + +def _deployment(scenario: Scenario, gateway: Gateway, name: str) -> str: + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{PROVIDER_URL}/v1", + "api_key": "sk-fixture", + }, + "model_info": {"id": f"access-{uuid.uuid4().hex}"}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return name + + +def _models(scenario: Scenario, gateway: Gateway) -> _Models: + tag: Final = uuid.uuid4().hex[:10] + return _Models( + gpt=_deployment(scenario, gateway, f"openai/gpt-{tag}"), + gpt_mini=_deployment(scenario, gateway, f"openai/gpt-mini-{tag}"), + claude=_deployment(scenario, gateway, f"anthropic/claude-{tag}"), + bedrock_claude=_deployment(scenario, gateway, f"bedrock/anthropic.claude-{tag}"), + bedrock_titan=_deployment(scenario, gateway, f"bedrock/amazon.titan-{tag}"), + ) + + +def _pick(models: _Models, alias: str) -> str: + return string_value(getattr(models, alias)) + + +def _completion() -> Reply: + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "scripted"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + } + ).encode() + ) + + +def _served(gateway: Gateway, provider: SharedProvider, key: str, model: str) -> None: + provider.expect(_completion()) + body: Final = gateway.chat(model, key=key, text=f"access {uuid.uuid4().hex}") + assert object_value(object_value(body["choices"][0])["message"])["content"] == "scripted", body + assert len(provider.received()) == 1 + + +def _denied(gateway: Gateway, provider: SharedProvider, key: str, model: str, denial: _Denial) -> None: + refused: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "hi"}]}, key=key + ) + assert refused.status_code == 403, f"{refused.status_code} {refused.text}" + error: Final = object_value(JSON_OBJECT.validate_json(refused.content)["error"]) + assert error["type"] == denial, refused.text + assert error["param"] == "model", refused.text + assert error["code"] == "403", refused.text + message: Final = string_value(error["message"]) + assert "is not available for this API key" in message, message + assert "not allowed to access model" not in message, message + assert provider.received() == () + + +@pytest.mark.parametrize("case", _CASES, ids=[case.name for case in _CASES]) +def test_a_key_model_allow_list_decides_which_models_it_reaches( + gateway: Gateway, provider: SharedProvider, case: _AccessCase +) -> None: + with gateway.scenario() as scenario: + models: Final = _models(scenario, gateway) + allowed: Final = ( + None + if case.allowed is None + else [_pick(models, item) if hasattr(models, item) else item for item in case.allowed] + ) + key: Final = scenario.key(models=allowed) + requested: Final = _pick(models, case.requested) + if case.served: + _served(gateway, provider, key, requested) + else: + _denied(gateway, provider, key, requested, "key_model_access_denied") + + +def test_widening_a_key_allow_list_to_a_wildcard_takes_effect_on_the_next_request( + gateway: Gateway, provider: SharedProvider +) -> None: + with gateway.scenario() as scenario: + models: Final = _models(scenario, gateway) + key: Final = scenario.key(models=[models.gpt]) + _served(gateway, provider, key, models.gpt) + _denied(gateway, provider, key, models.gpt_mini, "key_model_access_denied") + gateway.post("/key/update", {"key": key, "models": ["openai/*"]}) + _served(gateway, provider, key, models.gpt) + _served(gateway, provider, key, models.gpt_mini) + _denied(gateway, provider, key, models.claude, "key_model_access_denied") + + +def test_a_team_allow_list_denies_its_keys_a_model_outside_it(gateway: Gateway, provider: SharedProvider) -> None: + with gateway.scenario() as scenario: + models: Final = _models(scenario, gateway) + team: Final = scenario.team(models=["openai/*"]) + key: Final = scenario.key(team_id=team) + _served(gateway, provider, key, models.gpt) + _denied(gateway, provider, key, models.claude, "team_model_access_denied") + + +def test_widening_a_team_allow_list_takes_effect_for_its_keys_on_the_next_request( + gateway: Gateway, provider: SharedProvider +) -> None: + with gateway.scenario() as scenario: + models: Final = _models(scenario, gateway) + team: Final = scenario.team(models=[models.gpt]) + key: Final = scenario.key(team_id=team) + _served(gateway, provider, key, models.gpt) + _denied(gateway, provider, key, models.gpt_mini, "team_model_access_denied") + gateway.post("/team/update", {"team_id": team, "models": ["openai/*"]}) + _served(gateway, provider, key, models.gpt) + _served(gateway, provider, key, models.gpt_mini) + _denied(gateway, provider, key, models.claude, "team_model_access_denied") diff --git a/tests/integration/caching/test_embedding_file_block_cache.py b/tests/integration/caching/test_embedding_file_block_cache.py new file mode 100644 index 00000000000..9f753e996e8 --- /dev/null +++ b/tests/integration/caching/test_embedding_file_block_cache.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +import asyncio +import json +import threading +import zlib +from collections.abc import Mapping +from hashlib import sha256 +from typing import Final + +import pytest +from openai import AsyncOpenAI, OpenAI +from openai.types import CreateEmbeddingResponse +from pydantic import JsonValue + +from tests.integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from tests.integration._support.database import read_rows +from tests.integration._support.wire import Reply, Request, Wire, wire_server + +BACKEND: Final = "gemini-embedding-001" +TARGET: Final = f"/models/{BACKEND}:batchEmbedContents" +METADATA: Final[dict[str, JsonValue]] = {"fps": 1.0, "start_offset": "0s", "end_offset": "3s"} +WIRE_METADATA: Final[dict[str, JsonValue]] = {"fps": 1.0, "startOffset": "0s", "endOffset": "3s"} +TEXT_PART: Final[dict[str, JsonValue]] = {"text": "a red bus"} +SETTLE_SECONDS: Final = 15 +SPEND_SQL: Final = ( + 'SELECT request_id, cache_hit, spend FROM "LiteLLM_SpendLogs" WHERE api_key = %s ORDER BY "startTime", request_id' +) + + +def clip(name: str) -> str: + return f"gs://scripted-bucket/cache/{name}.mp4" + + +def block(name: str, metadata: Mapping[str, JsonValue] | None = METADATA) -> dict[str, JsonValue]: + file: Final[dict[str, JsonValue]] = {"file_id": clip(name)} + return {"type": "file", "file": file if metadata is None else {**file, "video_metadata": dict(metadata)}} + + +def gcs_part(name: str, metadata: Mapping[str, JsonValue] | None = WIRE_METADATA) -> dict[str, JsonValue]: + part: Final[dict[str, JsonValue]] = {"file_data": {"mime_type": "video/mp4", "file_uri": clip(name)}} + return part if metadata is None else {**part, "video_metadata": dict(metadata)} + + +def list_value(value: JsonValue) -> list[JsonValue]: + assert isinstance(value, list), value + return value + + +def source_of(part: JsonValue) -> str: + item: Final = object_value(part) + if "text" in item: + return string_value(item["text"]) + return string_value(object_value(item["file_data"])["file_uri"]) + + +def vector(parts: list[JsonValue]) -> list[float]: + return [zlib.crc32("|".join(source_of(part) for part in parts).encode()) / 2**32, 0.5] + + +def request_parts(request: Request) -> list[list[JsonValue]]: + body: Final = object_value(json.loads(request.body)) + return [list_value(object_value(object_value(item)["content"])["parts"]) for item in list_value(body["requests"])] + + +def embed_peer(request: Request) -> Reply: + assert request.method == "POST", request.method + assert request.target.split("?")[0] == TARGET, request.target + embeddings: Final = [{"values": vector(parts)} for parts in request_parts(request)] + return Reply(body=json.dumps({"embeddings": embeddings}).encode()) + + +def wire_parts(wire: Wire) -> list[list[JsonValue]]: + received: Final = wire.drain() + assert len(received) == 1, [item.target for item in received] + return request_parts(received[0]) + + +def embeddings_of(answer: Mapping[str, JsonValue]) -> list[JsonValue]: + data: Final = [object_value(item) for item in list_value(answer["data"])] + assert [item["index"] for item in data] == list(range(len(data))), answer + return [item["embedding"] for item in data] + + +def gemini_model(scenario: Scenario, url: str) -> str: + return scenario.model(model=f"gemini/{BACKEND}", api_key="scripted-gemini-key", api_base=url) + + +def embed_body(model: str, elements: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "input": elements} + + +def spend_rows(key: str) -> list[dict[str, JsonValue]]: + return read_rows(SPEND_SQL, (sha256(key.encode()).hexdigest(),)) + + +def served_from_cache( + gateway: Gateway, wire: Wire, body: Mapping[str, JsonValue], key: str +) -> dict[str, JsonValue] | None: + answer: Final = gateway.post("/v1/embeddings", body, key=key) + return None if wire.drain() else answer + + +def await_hit(gateway: Gateway, wire: Wire, body: Mapping[str, JsonValue], key: str) -> dict[str, JsonValue]: + served: Final = eventually( + lambda: served_from_cache(gateway, wire, body, key), lambda answer: answer is not None, SETTLE_SECONDS + ) + assert served is not None + return served + + +def proxy_root(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def test_block_embedding_is_served_from_the_cache_with_a_zero_spend_row(gateway: Gateway) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + key: Final = scenario.key(models=[model]) + body: Final = embed_body(model, [block("one")]) + first: Final = gateway.request("POST", "/v1/embeddings", body, key=key) + assert first.status_code == 200, first.text + first_id: Final = first.headers["x-litellm-call-id"] + assert wire_parts(wire) == [[gcs_part("one")]] + served: Final = await_hit(gateway, wire, body, key) + assert embeddings_of(served) == embeddings_of(object_value(first.json())) == [vector([gcs_part("one")])] + rows: Final = eventually( + lambda: spend_rows(key), lambda found: any(row["cache_hit"] == "True" for row in found), seconds=70 + ) + by_id: Final = {string_value(row["request_id"]): row for row in rows} + assert by_id[first_id]["cache_hit"] != "True", rows + hits: Final = [row for row in rows if row["cache_hit"] == "True"] + assert hits and all(float(str(row["spend"])) == 0 for row in hits), rows + + +def test_cached_block_is_not_resent_when_a_new_text_joins_it(gateway: Gateway) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + key: Final = scenario.key(models=[model]) + cached: Final = embed_body(model, [block("one")]) + gateway.post("/v1/embeddings", cached, key=key) + assert wire_parts(wire) == [[gcs_part("one")]] + await_hit(gateway, wire, cached, key) + mixed: Final = gateway.post("/v1/embeddings", embed_body(model, [block("one"), "a red bus"]), key=key) + assert wire_parts(wire) == [[TEXT_PART]] + assert embeddings_of(mixed) == [vector([gcs_part("one")]), vector([TEXT_PART])] + + +@pytest.mark.parametrize("names", (("a", "b", "c", "d", "e"), ("same", "same")), ids=("five-distinct", "two-identical")) +def test_block_lists_are_cached_whole(gateway: Gateway, names: tuple[str, ...]) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + key: Final = scenario.key(models=[model]) + body: Final = embed_body(model, [block(name) for name in names]) + first: Final = gateway.post("/v1/embeddings", body, key=key) + assert wire_parts(wire) == [[gcs_part(name)] for name in names] + assert embeddings_of(first) == [vector([gcs_part(name)]) for name in names] + assert embeddings_of(await_hit(gateway, wire, body, key)) == embeddings_of(first) + + +def test_provider_failure_is_not_cached(gateway: Gateway) -> None: + healed: Final = threading.Event() + + def respond(request: Request) -> Reply: + if healed.is_set(): + return embed_peer(request) + return Reply(status=500, body=b'{"error": {"message": "scripted outage"}}') + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + key: Final = scenario.key(models=[model]) + body: Final = embed_body(model, [block("flaky")]) + failed: Final = gateway.request("POST", "/v1/embeddings", body, key=key) + assert failed.status_code >= 500, failed.text + assert wire.drain(), "the failing call never reached the provider" + healed.set() + recovered: Final = gateway.post("/v1/embeddings", body, key=key) + assert wire_parts(wire) == [[gcs_part("flaky")]] + assert embeddings_of(await_hit(gateway, wire, body, key)) == embeddings_of(recovered) + + +def test_video_metadata_is_part_of_the_cache_key(gateway: Gateway) -> None: + two_seconds: Final[dict[str, JsonValue]] = {**METADATA, "end_offset": "2s"} + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + key: Final = scenario.key(models=[model]) + settled: Final = embed_body(model, [block("keyed")]) + gateway.post("/v1/embeddings", settled, key=key) + assert wire_parts(wire) == [[gcs_part("keyed")]] + await_hit(gateway, wire, settled, key) + variants: Final = ( + (block("keyed", two_seconds), gcs_part("keyed", {**WIRE_METADATA, "endOffset": "2s"})), + (block("keyed", None), gcs_part("keyed", None)), + ) + for variant, part in variants: + gateway.post("/v1/embeddings", embed_body(model, [variant]), key=key) + assert wire_parts(wire) == [[part]] + + +def test_openai_sdk_clients_fill_and_hit_the_cache(gateway: Gateway) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + key: Final = scenario.key(models=[model]) + body: Final = embed_body(model, [block("sdk")]) + with OpenAI(base_url=f"{proxy_root(gateway)}/v1", api_key=key, max_retries=0) as client: + filled: Final = client.post("/embeddings", body=body, cast_to=CreateEmbeddingResponse) + assert wire_parts(wire) == [[gcs_part("sdk")]] + await_hit(gateway, wire, body, key) + + async def drive() -> CreateEmbeddingResponse: + async with AsyncOpenAI(base_url=f"{proxy_root(gateway)}/v1", api_key=key, max_retries=0) as client: + return await client.post("/embeddings", body=body, cast_to=CreateEmbeddingResponse) + + served: Final = asyncio.run(drive()) + assert served.data[0].embedding == filled.data[0].embedding == vector([gcs_part("sdk")]) + assert wire.drain() == (), "the cached answer reached the provider again" diff --git a/tests/integration/configuration/test_config_declared_model_behaviours.py b/tests/integration/configuration/test_config_declared_model_behaviours.py new file mode 100644 index 00000000000..6b1f25c35a3 --- /dev/null +++ b/tests/integration/configuration/test_config_declared_model_behaviours.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +import json +import uuid +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import pytest +import yaml +from pydantic import JsonValue, TypeAdapter + +from tests.integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment, object_value, string_value +from tests.integration._support.process import owned_proxy +from tests.integration._support.provider import PROVIDER_URL, SharedProvider +from tests.integration._support.wire import Reply + +_TAGGED_GROUP: Final = "tag-filtered-group" +_TAGGED_IDS: Final = frozenset({"tag-filtered-team-a", "tag-filtered-team-b"}) +_VISION_MODEL: Final = "llava-hf" +_ITEMS: Final = TypeAdapter(list[dict[str, JsonValue]]) + + +def _config(directory: Path) -> Path: + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["model_list"] = [ + *configuration["model_list"], + *( + { + "model_name": _TAGGED_GROUP, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-fixture", + "api_base": f"{PROVIDER_URL}/v1", + "tags": [tag], + }, + "model_info": {"id": identity}, + } + for tag, identity in (("teamA", "tag-filtered-team-a"), ("teamB", "tag-filtered-team-b")) + ), + { + "model_name": _VISION_MODEL, + "litellm_params": { + "model": "openai/llava-hf", + "api_key": "sk-fixture", + "api_base": "http://127.0.0.1:9/v1", + }, + "model_info": {"supports_vision": True}, + }, + ] + configuration.setdefault("router_settings", {})["enable_tag_filtering"] = True + path: Final = directory / "config-declared-models.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@pytest.fixture(scope="module") +def declared(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("config-declared-models") + with ( + gateway_from_environment() as shared, + owned_proxy( + shared, + directory, + {"OPENAI_API_KEY": "sk-fixture", "OPENAI_BASE_URL": f"{PROVIDER_URL}/v1", "GCS_FLUSH_INTERVAL": "1"}, + config=_config(directory), + remove_environment=("GCS_BUCKET_NAME", "OPENAI_API_BASE"), + ) as owned, + ): + yield owned + + +def test_model_info_reports_the_vision_capability_declared_in_the_config(declared: Gateway) -> None: + listing: Final = _ITEMS.validate_python(declared.get("/model/info")["data"]) + vision: Final = [item for item in listing if item["model_name"] == _VISION_MODEL] + assert len(vision) == 1, [item["model_name"] for item in listing] + assert object_value(vision[0]["model_info"])["supports_vision"] is True, vision[0] + + +def test_an_untagged_request_is_served_by_a_group_whose_deployments_are_all_tagged( + declared: Gateway, provider: SharedProvider +) -> None: + provider.expect( + Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "tagged"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } + ).encode() + ) + ) + response: Final = declared.request( + "POST", + "/v1/chat/completions", + {"model": _TAGGED_GROUP, "messages": [{"role": "user", "content": f"untagged {uuid.uuid4().hex}"}]}, + ) + assert response.status_code == 200, response.text + assert response.headers["x-litellm-model-id"] in _TAGGED_IDS, dict(response.headers) + message: Final = object_value(object_value(JSON_OBJECT.validate_json(response.content)["choices"][0])["message"]) + assert message["content"] == "tagged" + assert [request.target for request in provider.received()] == ["/v1/chat/completions"] + + +def test_a_moderation_request_without_a_model_reaches_the_provider_default( + declared: Gateway, provider: SharedProvider +) -> None: + provider.expect( + Reply( + body=json.dumps( + { + "id": f"modr-{uuid.uuid4().hex}", + "model": "omni-moderation-latest", + "results": [ + {"flagged": True, "categories": {"violence": True}, "category_scores": {"violence": 0.9}} + ], + } + ).encode() + ) + ) + phrase: Final = f"I want to harm someone {uuid.uuid4().hex}" + response: Final = declared.request("POST", "/moderations", {"input": phrase}) + assert response.status_code == 200, response.text + body: Final = JSON_OBJECT.validate_json(response.content) + assert body["model"] == "omni-moderation-latest", body + assert object_value(_ITEMS.validate_python(body["results"])[0])["flagged"] is True, body + sent: Final = provider.received() + assert [request.target for request in sent] == ["/v1/moderations"] + payload: Final = JSON_OBJECT.validate_json(sent[0].body) + assert payload == {"input": phrase}, payload + + +def test_key_health_reports_an_unconfigured_key_logging_callback_as_unhealthy(declared: Gateway) -> None: + with declared.scenario() as scenario: + key: Final = scenario.key( + metadata={ + "logging": [ + { + "callback_name": "gcs_bucket", + "callback_type": "success_and_failure", + "callback_vars": { + "gcs_bucket_name": "key-logging-project1", + "gcs_path_service_account": "bad-service-account", + }, + } + ] + } + ) + health: Final = declared.request("POST", "/key/health", {}, key=key) + assert health.status_code == 200, health.text + body: Final = JSON_OBJECT.validate_json(health.content) + assert "key" in body, body + status: Final = object_value(body["logging_callbacks"]) + assert status["callbacks"] == ["gcs_bucket"], status + assert status["status"] == "unhealthy", status + assert "GCS_BUCKET_NAME is not set in the environment" in string_value(status["details"]), status diff --git a/tests/integration/database/test_lens_repository.py b/tests/integration/database/test_lens_repository.py index ff46944b911..6349da7ab2a 100644 --- a/tests/integration/database/test_lens_repository.py +++ b/tests/integration/database/test_lens_repository.py @@ -1,6 +1,7 @@ import asyncio import os -from collections.abc import AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator +from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone from pathlib import Path from types import SimpleNamespace @@ -15,6 +16,7 @@ from fastapi import HTTPException from prisma import Prisma from psycopg import sql from pydantic import TypeAdapter +from typing_extensions import LiteralString from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.db.prisma_client import PrismaWrapper @@ -35,7 +37,7 @@ from litellm.proxy.lens.models import ( TraceIdentity, Worker, ) -from litellm.proxy.lens.repository import LensRepository, WriterDatabase +from litellm.proxy.lens.repository import Database, LensRepository, Row, WriterDatabase from litellm.proxy.lens.state import cancel_job, claim_job, current_job, due_at, end_job, queue_job, replace_job @@ -770,6 +772,156 @@ async def test_delayed_progress_cannot_replace_a_newer_checkpoint(lens_db: Prism await lens_db.execute_raw('DELETE FROM "LiteLLM_Lens" WHERE id=$1', claimed.id) +class ProgressInterleavingDatabase: + def __init__( + self, + database: Database, + read: asyncio.Future[int], + resume: asyncio.Event, + committed: asyncio.Event, + ) -> None: + self.database: Final = database + self.read: Final = read + self.resume: Final = resume + self.committed: Final = committed + + async def query_raw(self, query: LiteralString, *args: object) -> object: + rows: Final = await self.database.query_raw(query, *args) + if query == 'SELECT data FROM "LiteLLM_Lens" WHERE id=$1' and not self.read.done(): + backend: Final = TypeAdapter(tuple[Row, ...]).validate_python( + await self.database.query_raw("SELECT to_jsonb(pg_backend_pid()) AS data") + ) + self.read.set_result(TypeAdapter(int).validate_python(backend[0].data)) + await self.resume.wait() + return rows + + async def execute_raw(self, query: LiteralString, *args: object) -> int: + return await self.database.execute_raw(query, *args) + + @asynccontextmanager + async def transaction(self) -> AsyncGenerator[Database]: + async with self.database.transaction() as database: + yield ProgressInterleavingDatabase(database, self.read, self.resume, self.committed) + self.committed.set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("renewing", (False, True)) +async def test_budget_reservation_survives_competing_progress(lens_db: Prisma, renewing: bool) -> None: + from litellm.proxy.lens.inference import ( + BUDGET_LEASE, + renew_budget_reservation, + reserve_attempt, + wait_for_reservation, + ) + from litellm.proxy.lens.models import BudgetReservation + from tests.unit.proxy.lens.test_state import lens, worker + + now: Final = datetime.now(timezone.utc) + claimed: Final = claim_job(queue_job(lens(), now, uuid4().hex), worker(), now).model_copy( + update={"id": uuid4().hex, "budget_month": now.strftime("%Y-%m")} + ) + job: Final = claimed.jobs[0] + hold: Final = BudgetReservation( + id=uuid4().hex, job_id=job.id, amount=1, month=claimed.budget_month, expires_at=now + BUDGET_LEASE + ) + database: Final = WriterDatabase(PrismaWrapper(lens_db)) + repo: Final = LensRepository(database) + read: Final[asyncio.Future[int]] = asyncio.get_running_loop().create_future() + resume: Final = asyncio.Event() + committed: Final = asyncio.Event() + competing: Final = ProgressInterleavingDatabase(database, read, resume, committed) + await repo.create(claimed.model_copy(update={"reservations": (hold,) if renewing else ()})) + admitted: Final = asyncio.Event() + admitted.set() + operation: Final = asyncio.create_task( + renew_budget_reservation(LensRepository(competing), claimed.id, hold.id, admitted) + if renewing + else wait_for_reservation( + LensRepository(competing), + claimed.id, + hold.id, + lambda current: reserve_attempt(current, job, worker().id, hold, now), + ) + ) + + async def write_progress() -> Lens | None: + await read + return await repo.progress(claimed.id, job, Progress(stage="Reviewing traces concurrently")) + + progress: Final = asyncio.create_task(write_progress()) + try: + async with asyncio.timeout(45): + blocker: Final = await read + async with asyncio.timeout(5): + while not await lens_db.query_raw( + "SELECT pid FROM pg_stat_activity WHERE $1::int=ANY(pg_blocking_pids(pid))", blocker + ): + assert not progress.done(), "Progress committed before the reservation released its lock" + await asyncio.sleep(0.01) + assert not progress.done() + assert not committed.is_set() + resume.set() + await committed.wait() + updated: Final = await progress + assert updated is not None + assert updated.jobs[0].stage == "Reviewing traces concurrently" + assert tuple(reservation.id for reservation in updated.reservations) == (hold.id,) + if not renewing: + await operation + stored: Final = await repo.get(claimed.id) + assert stored is not None + assert tuple(reservation.id for reservation in stored.reservations) == (hold.id,) + assert stored.spent == claimed.spent + assert stored.jobs[0].cost == 0 + assert stored.jobs[0].stage == "Reviewing traces concurrently" + if renewing: + assert stored.reservations[0].expires_at is not None + assert stored.reservations[0].expires_at > now + BUDGET_LEASE + finally: + operation.cancel() + progress.cancel() + await asyncio.gather(operation, progress, return_exceptions=True) + await lens_db.execute_raw('DELETE FROM "LiteLLM_Lens" WHERE id=$1', claimed.id) + + +@pytest.mark.asyncio +async def test_parallel_reservations_release_the_lock_while_waiting_for_budget(lens_db: Prisma) -> None: + from litellm.proxy.lens.inference import reserve_amount, settle_amount, wait_for_reservation + from litellm.proxy.lens.models import BudgetReservation + from tests.unit.proxy.lens.test_state import lens + + now: Final = datetime.now(timezone.utc) + original: Final = lens() + queued: Final = queue_job(original, now, uuid4().hex).model_copy( + update={"id": uuid4().hex, "settings": original.settings.model_copy(update={"monthly_budget": 2})} + ) + repo: Final = LensRepository(WriterDatabase(PrismaWrapper(lens_db))) + holds: Final = tuple( + BudgetReservation(id=uuid4().hex, job_id=queued.jobs[0].id, amount=1, month=queued.budget_month) + for _ in range(8) + ) + await repo.create(queued) + try: + + async def analyze(hold: BudgetReservation) -> None: + await wait_for_reservation(repo, queued.id, hold.id, lambda current: reserve_amount(current, hold)) + stored: Final = await repo.get(queued.id) + assert stored is not None and hold in stored.reservations + assert stored.spent + sum(reservation.amount for reservation in stored.reservations) <= 2 + assert await repo.update_locked(queued.id, lambda current: settle_amount(current, hold.id, 0.125, None)) + + async with asyncio.timeout(15): + await asyncio.gather(*(analyze(hold) for hold in holds)) + stored: Final = await repo.get(queued.id) + assert stored is not None + assert stored.spent == 1 + assert stored.jobs[0].cost == 1 + assert stored.reservations == () + finally: + await lens_db.execute_raw('DELETE FROM "LiteLLM_Lens" WHERE id=$1', queued.id) + + @pytest.mark.asyncio async def test_locked_settlement_charges_every_concurrent_call_exactly_once(lens_db: Prisma) -> None: from litellm.proxy.lens.inference import settle_amount diff --git a/tests/integration/management/test_fallback_management_public_names.py b/tests/integration/management/test_fallback_management_public_names.py new file mode 100644 index 00000000000..b54465c1fd0 --- /dev/null +++ b/tests/integration/management/test_fallback_management_public_names.py @@ -0,0 +1,379 @@ +import uuid +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from functools import partial +from itertools import chain +from pathlib import Path +from typing import Final + +import httpx +import psutil +import pytest +from pydantic import JsonValue + +from tests.integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from tests.integration._support.database import read_rows +from tests.integration._support.process import OwnedProxy, owned_proxy_process +from tests.integration._support.redis_process import owned_redis + +pytestmark: Final = pytest.mark.timeout(300) + +PROVIDER_KEY: Final = "integration-provider-key" +EVICTION_WARNING: Final = "config cache eviction of router_settings failed" +REFUSED: Final = frozenset({401, 403}) + + +@pytest.fixture +def upstream(gateway: Gateway) -> Iterator[httpx.Client]: + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as client: + client.get("/__observations").raise_for_status() + yield client + + +def _observed_requests(upstream: httpx.Client) -> list[JsonValue]: + observed: Final = upstream.get("/__observations") + observed.raise_for_status() + requests: Final = object_value(observed.json())["requests"] + assert isinstance(requests, list) + return requests + + +def _calls_to(observed: list[JsonValue], provider_model: str) -> int: + return sum(object_value(object_value(request)["body"]).get("model") == provider_model for request in observed) + + +def _fallback_body(model: str, fallback_models: list[str], fallback_type: str = "general") -> dict[str, JsonValue]: + return {"model": model, "fallback_models": list(fallback_models), "fallback_type": fallback_type} + + +def _forget_fallback(gateway: Gateway, model: str, fallback_type: str) -> None: + gateway.request("DELETE", f"/fallback/{model}", params={"fallback_type": fallback_type}) + + +def _create_fallback( + gateway: Gateway, scenario: Scenario, model: str, fallback_models: list[str], fallback_type: str = "general" +) -> httpx.Response: + scenario.cleanups.callback(_forget_fallback, gateway, model, fallback_type) + return eventually( + lambda: gateway.request("POST", "/fallback", _fallback_body(model, fallback_models, fallback_type)), + lambda response: response.status_code == 200, + seconds=30, + return_last_on_timeout=True, + ) + + +def _fallback_models(response: httpx.Response) -> list[str]: + body: Final = response.json() + models: Final = body.get("fallback_models") if isinstance(body, dict) else None + return [string_value(entry) for entry in models] if isinstance(models, list) else [] + + +def _available_models(response: httpx.Response) -> list[str]: + body: Final = response.json() + if not isinstance(body, dict): + return [] + detail: Final = body.get("detail") + source: Final = detail if isinstance(detail, dict) else body + models: Final = source.get("available_models") + return [string_value(entry) for entry in models] if isinstance(models, list) else [] + + +def _stored_fallbacks(fallback_key: str = "fallbacks") -> list[JsonValue]: + rows: Final = read_rows('SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', ("router_settings",)) + if not rows: + return [] + settings: Final = object_value(rows[0]["param_value"]) + entries: Final = settings.get(fallback_key) or [] + assert isinstance(entries, list), entries + return entries + + +def _covered_models(entries: list[JsonValue]) -> frozenset[str]: + dict_entries: Final = (object_value(entry) for entry in entries if isinstance(entry, dict)) + return frozenset(chain.from_iterable(dict_entries)) + + +def _models_over_a_fresh_connection(gateway: Gateway, _: int) -> frozenset[str]: + with httpx.Client(base_url=gateway.client.base_url, timeout=15, trust_env=False) as client: + listed: Final = client.get("/v1/models", headers={"Authorization": f"Bearer {gateway.key}"}) + assert listed.status_code == 200, listed.text + data: Final = object_value(listed.json())["data"] + assert isinstance(data, list), listed.text + return frozenset(string_value(object_value(entry)["id"]) for entry in data) + + +def _every_worker_serves(gateway: Gateway, model: str) -> bool: + with ThreadPoolExecutor(max_workers=16) as pool: + rounds: Final = tuple( + tuple(pool.map(partial(_models_over_a_fresh_connection, gateway), range(16))) for _ in range(2) + ) + return all(model in seen for seen in chain.from_iterable(rounds)) + + +def _wait_until_served(gateway: Gateway, model: str) -> None: + eventually(lambda: _every_worker_serves(gateway, model), lambda served: served, seconds=90) + + +def _team_model(gateway: Gateway, scenario: Scenario, team: str, provider_model: str) -> tuple[str, str]: + public: Final = f"ipub-{uuid.uuid4().hex}" + created: Final = gateway.post( + "/model/new", + { + "model_name": public, + "litellm_params": { + "model": f"openai/{provider_model}", + "api_key": PROVIDER_KEY, + "api_base": f"{gateway.upstream_url}/v1", + "num_retries": 0, + }, + "model_info": {"team_id": team}, + }, + ) + info: Final = object_value(created["model_info"]) + scenario.cleanups.callback(scenario.delete_model, string_value(info["id"])) + return public, string_value(created["model_name"]) + + +def _workers(owned: OwnedProxy) -> tuple[psutil.Process, ...]: + return tuple(child for child in psutil.Process(owned.process.pid).children() if _is_worker(child)) + + +def _is_worker(child: psutil.Process) -> bool: + try: + return "spawn_main" in " ".join(child.cmdline()) and child.status() != psutil.STATUS_ZOMBIE + except psutil.Error: + return False + + +def test_post_by_team_public_name_creates_reads_back_and_peer_converges(gateway: Gateway, peer: Gateway) -> None: + with gateway.scenario() as scenario: + team: Final = scenario.team() + primary: Final = scenario.model(model=f"openai/n1p-{uuid.uuid4().hex}", model_info={"team_id": team}) + fallback: Final = scenario.model(model=f"openai/n1f-{uuid.uuid4().hex}", model_info={"team_id": team}) + created: Final = _create_fallback(gateway, scenario, primary, [fallback]) + assert created.status_code == 200, created.text + here: Final = eventually( + lambda: gateway.request("GET", f"/fallback/{primary}", params={"fallback_type": "general"}), + lambda response: response.status_code == 200 and fallback in _fallback_models(response), + seconds=30, + return_last_on_timeout=True, + ) + assert here.status_code == 200 and fallback in _fallback_models(here), here.text + there: Final = eventually( + lambda: peer.request("GET", f"/fallback/{primary}", params={"fallback_type": "general"}), + lambda response: response.status_code == 200 and fallback in _fallback_models(response), + seconds=30, + return_last_on_timeout=True, + ) + assert there.status_code == 200 and fallback in _fallback_models(there), there.text + + +def test_rule_by_team_public_name_fires_on_chat_completions(gateway: Gateway, upstream: httpx.Client) -> None: + with gateway.scenario() as scenario: + team: Final = scenario.team() + primary_provider: Final = f"n2prim-{uuid.uuid4().hex}" + fallback_provider: Final = f"n2fbk-{uuid.uuid4().hex}" + primary: Final = scenario.model(model=f"openai/{primary_provider}", num_retries=0, model_info={"team_id": team}) + fallback: Final = scenario.model(model=f"openai/{fallback_provider}", model_info={"team_id": team}) + team_key: Final = scenario.key(team_id=team) + created: Final = _create_fallback(gateway, scenario, primary, [fallback]) + assert created.status_code == 200, created.text + upstream.post(f"/__scripts/{primary_provider}", json={"statuses": [500]}).raise_for_status() + upstream.get("/__observations").raise_for_status() + answered: Final = eventually( + lambda: gateway.request( + "POST", + "/v1/chat/completions", + {"model": primary, "messages": [{"role": "user", "content": f"fire {uuid.uuid4().hex}"}]}, + key=team_key, + ), + lambda response: response.status_code == 200, + seconds=60, + return_last_on_timeout=True, + ) + assert answered.status_code == 200, answered.text + observed: Final = _observed_requests(upstream) + assert _calls_to(observed, primary_provider) >= 1, observed + assert _calls_to(observed, fallback_provider) >= 1, observed + + +def test_consecutive_creates_within_the_cache_window_keep_every_rule(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + target: Final = scenario.model(model=f"openai/n5t-{uuid.uuid4().hex}") + primaries: Final = tuple(scenario.model(model=f"openai/n5{tag}-{uuid.uuid4().hex}") for tag in "abc") + for model in (target, *primaries): + _wait_until_served(gateway, model) + created: Final = tuple(_create_fallback(gateway, scenario, primary, [target]) for primary in primaries) + assert all(response.status_code == 200 for response in created), [response.text for response in created] + covered: Final = _covered_models(_stored_fallbacks()) + assert frozenset(primaries) <= covered, (primaries, covered) + + +def test_unknown_model_404_lists_team_public_and_internal_names(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + team: Final = scenario.team() + public, internal = _team_model(gateway, scenario, team, f"n6-{uuid.uuid4().hex}") + unknown: Final = f"n6-unknown-{uuid.uuid4().hex}" + refused: Final = eventually( + lambda: gateway.request("POST", "/fallback", _fallback_body(unknown, [public])), + lambda response: response.status_code == 404 and internal in _available_models(response), + seconds=30, + return_last_on_timeout=True, + ) + assert refused.status_code == 404, refused.text + available: Final = _available_models(refused) + assert internal in available, (internal, available) + assert public in available, (public, available) + + +def test_create_by_generated_internal_name_still_works(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + team: Final = scenario.team() + _, primary_internal = _team_model(gateway, scenario, team, f"c1p-{uuid.uuid4().hex}") + _, fallback_internal = _team_model(gateway, scenario, team, f"c1f-{uuid.uuid4().hex}") + created: Final = _create_fallback(gateway, scenario, primary_internal, [fallback_internal]) + assert created.status_code == 200, created.text + here: Final = eventually( + lambda: gateway.request("GET", f"/fallback/{primary_internal}", params={"fallback_type": "general"}), + lambda response: response.status_code == 200 and fallback_internal in _fallback_models(response), + seconds=30, + return_last_on_timeout=True, + ) + assert here.status_code == 200 and fallback_internal in _fallback_models(here), here.text + + +def test_delete_reads_fresh_rules_and_keeps_the_others(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + target: Final = scenario.model(model=f"openai/c2t-{uuid.uuid4().hex}") + keep: Final = scenario.model(model=f"openai/c2k-{uuid.uuid4().hex}") + drop: Final = scenario.model(model=f"openai/c2d-{uuid.uuid4().hex}") + _wait_until_served(gateway, target) + _wait_until_served(gateway, keep) + _wait_until_served(gateway, drop) + assert _create_fallback(gateway, scenario, keep, [target]).status_code == 200 + assert _create_fallback(gateway, scenario, drop, [target]).status_code == 200 + removed: Final = gateway.request("DELETE", f"/fallback/{drop}", params={"fallback_type": "general"}) + assert removed.status_code == 200, removed.text + covered: Final = _covered_models(_stored_fallbacks()) + assert keep in covered, covered + assert drop not in covered, covered + gone: Final = gateway.request("GET", f"/fallback/{drop}", params={"fallback_type": "general"}) + assert gone.status_code == 404, gone.text + + +def test_self_fallback_and_unknown_fallback_models_are_rejected(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/c3-{uuid.uuid4().hex}") + _wait_until_served(gateway, model) + itself: Final = gateway.request("POST", "/fallback", _fallback_body(model, [model])) + assert itself.status_code == 400, itself.text + unknown_target: Final = gateway.request( + "POST", "/fallback", _fallback_body(model, [f"c3-missing-{uuid.uuid4().hex}"]) + ) + assert unknown_target.status_code == 400, unknown_target.text + + +def test_malformed_requests_are_rejected_and_the_proxy_stays_up(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/c4-{uuid.uuid4().hex}") + fallback: Final = scenario.model(model=f"openai/c4f-{uuid.uuid4().hex}") + _wait_until_served(gateway, model) + _wait_until_served(gateway, fallback) + malformed: Final = ( + {"model": model, "fallback_models": [], "fallback_type": "general"}, + {"model": model, "fallback_type": "general"}, + {"model": model, "fallback_models": [fallback], "fallback_type": "nonsense"}, + {"model": 123, "fallback_models": [fallback], "fallback_type": "general"}, + ) + for body in malformed: + assert gateway.request("POST", "/fallback", body).status_code in (400, 422), body + healthy: Final = _create_fallback(gateway, scenario, model, [fallback]) + assert healthy.status_code == 200, healthy.text + listed: Final = gateway.request("GET", "/v1/models") + assert listed.status_code == 200, listed.text + + +def test_team_key_is_forbidden_on_create_and_delete(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + team: Final = scenario.team() + model: Final = scenario.model(model=f"openai/c5-{uuid.uuid4().hex}", model_info={"team_id": team}) + fallback: Final = scenario.model(model=f"openai/c5f-{uuid.uuid4().hex}", model_info={"team_id": team}) + team_key: Final = scenario.key(team_id=team) + creating: Final = gateway.request("POST", "/fallback", _fallback_body(model, [fallback]), key=team_key) + assert creating.status_code in REFUSED, creating.text + assert model not in _covered_models(_stored_fallbacks()) + admitted: Final = _create_fallback(gateway, scenario, model, [fallback]) + assert admitted.status_code == 200, admitted.text + deleting: Final = gateway.request( + "DELETE", f"/fallback/{model}", params={"fallback_type": "general"}, key=team_key + ) + assert deleting.status_code in REFUSED, deleting.text + assert model in _covered_models(_stored_fallbacks()) + + +def test_context_window_and_content_policy_types_persist(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + target: Final = scenario.model(model=f"openai/c6t-{uuid.uuid4().hex}") + windowed: Final = scenario.model(model=f"openai/c6w-{uuid.uuid4().hex}") + policy: Final = scenario.model(model=f"openai/c6p-{uuid.uuid4().hex}") + _wait_until_served(gateway, target) + _wait_until_served(gateway, windowed) + _wait_until_served(gateway, policy) + window_create: Final = _create_fallback(gateway, scenario, windowed, [target], "context_window") + assert window_create.status_code == 200, window_create.text + window_read: Final = eventually( + lambda: gateway.request("GET", f"/fallback/{windowed}", params={"fallback_type": "context_window"}), + lambda response: response.status_code == 200 and target in _fallback_models(response), + seconds=30, + return_last_on_timeout=True, + ) + assert window_read.status_code == 200 and target in _fallback_models(window_read), window_read.text + policy_create: Final = _create_fallback(gateway, scenario, policy, [target], "content_policy") + assert policy_create.status_code == 200, policy_create.text + covered: Final = _covered_models(_stored_fallbacks("content_policy_fallbacks")) + assert policy in covered, covered + + +def test_fallback_write_survives_a_redis_outage(gateway: Gateway, tmp_path: Path) -> None: + with owned_redis(tmp_path) as cache: + overrides: Final = {"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port)} + with owned_proxy_process(gateway, tmp_path, overrides, workers=2, database_setup=()) as owned: + with owned.gateway.scenario() as scenario: + target: Final = scenario.model(model=f"openai/x1t-{uuid.uuid4().hex}") + first: Final = scenario.model(model=f"openai/x1a-{uuid.uuid4().hex}") + second: Final = scenario.model(model=f"openai/x1b-{uuid.uuid4().hex}") + third: Final = scenario.model(model=f"openai/x1c-{uuid.uuid4().hex}") + for model in (target, first, second, third): + _wait_until_served(owned.gateway, model) + assert _create_fallback(owned.gateway, scenario, first, [target]).status_code == 200 + cache.stop() + degraded: Final = tuple( + _create_fallback(owned.gateway, scenario, model, [target]) for model in (second, third) + ) + assert all(response.status_code == 200 for response in degraded), [r.text for r in degraded] + covered: Final = _covered_models(_stored_fallbacks()) + assert {first, second, third} <= covered, covered + assert EVICTION_WARNING in owned.log.read_text(), owned.log.read_text()[-3000:] + cache.start() + + +def test_fallback_write_survives_a_worker_kill(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy_process(gateway, tmp_path, {}, workers=2, database_setup=()) as owned: + with owned.gateway.scenario() as scenario: + target: Final = scenario.model(model=f"openai/x2t-{uuid.uuid4().hex}") + before: Final = scenario.model(model=f"openai/x2a-{uuid.uuid4().hex}") + after: Final = scenario.model(model=f"openai/x2b-{uuid.uuid4().hex}") + last: Final = scenario.model(model=f"openai/x2c-{uuid.uuid4().hex}") + for model in (target, before, after, last): + _wait_until_served(owned.gateway, model) + assert _create_fallback(owned.gateway, scenario, before, [target]).status_code == 200 + victim: Final = eventually(lambda: _workers(owned), lambda workers: len(workers) == 2)[0] + victim.kill() + fresh: Final = scenario.cleanups.enter_context( + httpx.Client(base_url=owned.gateway.client.base_url, timeout=15, trust_env=False) + ) + survivor: Final = Gateway(fresh, owned.gateway.key, owned.gateway.upstream_url) + degraded: Final = tuple(_create_fallback(survivor, scenario, model, [target]) for model in (after, last)) + assert all(response.status_code == 200 for response in degraded), [r.text for r in degraded] + covered: Final = _covered_models(_stored_fallbacks()) + assert {before, after, last} <= covered, covered diff --git a/tests/integration/observability/_openinference_support.py b/tests/integration/observability/_openinference_support.py new file mode 100644 index 00000000000..ee6dee07c46 --- /dev/null +++ b/tests/integration/observability/_openinference_support.py @@ -0,0 +1,1171 @@ +from __future__ import annotations + +import base64 +import json +import uuid +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import AbstractContextManager, ExitStack, contextmanager, nullcontext +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from typing import Final + +import httpx +import yaml +from integration._support.client import Gateway, eventually, object_value, string_value +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest +from opentelemetry.proto.common.v1.common_pb2 import AnyValue +from pydantic import JsonValue, TypeAdapter + +JSON_OBJECT: Final[TypeAdapter[dict[str, JsonValue]]] = TypeAdapter(dict[str, JsonValue]) +JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +JSON_MESSAGES: Final[TypeAdapter[list[dict[str, JsonValue]]]] = TypeAdapter(list[dict[str, JsonValue]]) + +CHAT_TOOLS: Final = [ + { + "type": "function", + "function": { + "name": "lookup_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + } +] + +RESPONSES_TOOLS: Final = [ + { + "type": "function", + "name": "lookup_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + } +] + +ANTHROPIC_TOOLS: Final = [ + { + "name": "lookup_weather", + "description": "Get weather", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + } +] + +_TOOL_PREFIX: Final = "llm.output_messages.{message}.message.tool_calls.{tool}.tool_call." + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + owned: OwnedProxy + model: str + provider: Wire + destination: Wire + + +def _json_object(body: bytes) -> dict[str, JsonValue]: + return JSON_OBJECT.validate_json(body) + + +def _json_object_value(value: JsonValue) -> dict[str, JsonValue]: + return JSON_OBJECT.validate_python(value) + + +def _json_messages(value: str) -> list[dict[str, JsonValue]]: + return JSON_MESSAGES.validate_json(value) + + +def _assert_chat_request( + request: Request, + *, + messages: JsonValue, + tools: JsonValue = CHAT_TOOLS, + tool_choice: JsonValue | None = None, + include_tools: bool = True, + stream: bool | None = None, + stream_options: JsonValue | None = None, + n: int | None = None, +) -> dict[str, JsonValue]: + request_tools: Final = ( + { + "tools": tools, + "tool_choice": ( + tool_choice if tool_choice is not None else {"type": "function", "function": {"name": "lookup_weather"}} + ), + } + if include_tools + else {} + ) + expected: Final[dict[str, JsonValue]] = { + "model": "gpt-4o-mini", + "messages": messages, + **request_tools, + **({"stream": stream} if stream is not None else {}), + **({"stream_options": stream_options} if stream_options is not None else {}), + **({"n": n} if n is not None else {}), + } + observed: Final = _json_object(request.body) + assert observed == expected, (observed, expected) + return observed + + +def _chat_request_marker(request: Request) -> str: + body: Final = _json_object(request.body) + messages: Final = JSON_MESSAGES.validate_python(body["messages"]) + marker: Final = messages[0].get("content") + assert isinstance(marker, str), body + return marker + + +def _assert_responses_request( + request: Request, + *, + marker: str, + input_value: str = "weather in Paris?", + stream: bool = False, +) -> None: + expected: Final[dict[str, JsonValue]] = { + "model": "gpt-4o-mini", + "input": input_value, + "tools": RESPONSES_TOOLS, + "tool_choice": {"type": "function", "name": "lookup_weather"}, + "metadata": {"trace_marker": marker}, + **({"stream": True} if stream else {}), + } + observed: Final = _json_object(request.body) + assert observed == expected, request + + +def _assert_messages_request( + request: Request, + *, + marker: str, + prompt: str = "weather in Paris?", + stream: bool = False, +) -> None: + expected: Final[dict[str, JsonValue]] = { + "model": "claude-opus-5-5", + "max_tokens": 64, + "messages": [{"role": "user", "content": prompt}], + "tools": ANTHROPIC_TOOLS, + "tool_choice": {"type": "auto"}, + "metadata": {}, + "stream": stream, + } + observed: Final = _json_object(request.body) + assert observed == expected, request + + +def _stream_values(reply: Reply) -> tuple[dict[str, JsonValue], ...]: + return tuple( + _json_object(chunk.split(b"data: ", 1)[1].splitlines()[0]) + for chunk in reply.chunks or () + if b"data: " in chunk and b"[DONE]" not in chunk + ) + + +def _sse_json_values(body: bytes) -> tuple[dict[str, JsonValue], ...]: + return tuple( + _json_object(line.removeprefix(b"data: ").strip()) + for line in body.splitlines() + if line.startswith(b"data: ") and line != b"data: [DONE]" + ) + + +def _chat_caller_response(reply: Reply, model: str) -> dict[str, JsonValue]: + body: Final = _json_object(reply.body) + choices: Final = JSON_MESSAGES.validate_python(body["choices"]) + return { + **body, + "model": model, + "choices": [ + { + **choice, + "message": _chat_caller_message(_json_object_value(choice["message"])), + "provider_specific_fields": {}, + } + for choice in choices + ], + } + + +def _chat_caller_message(message: dict[str, JsonValue]) -> dict[str, JsonValue]: + raw_tool_calls: Final = message.get("tool_calls") + tool_calls: Final = JSON_MESSAGES.validate_python(raw_tool_calls) if isinstance(raw_tool_calls, list) else () + return { + **{key: value for key, value in message.items() if key != "tool_calls"}, + **({"tool_calls": [_chat_caller_tool_call(call) for call in tool_calls]} if tool_calls else {}), + "provider_specific_fields": {"refusal": None}, + } + + +def _chat_caller_tool_call(call: dict[str, JsonValue]) -> dict[str, JsonValue]: + function: Final = _json_object_value(call["function"]) + arguments: Final = function.get("arguments") + return { + **call, + "function": { + **function, + **( + {"arguments": json.dumps(arguments)} + if "arguments" in function and not isinstance(arguments, str) + else {} + ), + }, + } + + +def _chat_output_tool_call(call: dict[str, JsonValue]) -> dict[str, JsonValue]: + function: Final = _json_object_value(call["function"]) + return { + **call, + "id": call["id"] if isinstance(call.get("id"), str) else None, + "function": { + **function, + **({"name": None} if "name" not in function else {}), + }, + } + + +def _chat_output_value(reply: Reply) -> str: + body: Final = _json_object(reply.body) + choices: Final = JSON_MESSAGES.validate_python(body["choices"]) + assert len(choices) == 1, body + message: Final = _chat_caller_message(_json_object_value(choices[0]["message"])) + tool_calls: Final = JSON_MESSAGES.validate_python(message["tool_calls"]) if "tool_calls" in message else () + return json.dumps( + [ + { + **{ + key: value + for key, value in message.items() + if key not in {"provider_specific_fields", "tool_calls"} + }, + **({"tool_calls": [_chat_output_tool_call(call) for call in tool_calls]} if tool_calls else {}), + } + ] + ) + + +def _chat_caller_stream(reply: Reply, model: str) -> tuple[dict[str, JsonValue], ...]: + return tuple(chain.from_iterable(_chat_caller_stream_events(event, model) for event in _stream_values(reply))) + + +def _chat_cache_hit_caller_stream( + marker: str, model: str, calls: Sequence[dict[str, JsonValue]] +) -> tuple[dict[str, JsonValue], ...]: + return ( + { + "id": marker, + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [ + { + "index": 0, + "delta": { + "role": "assistant", + "tool_calls": [{"index": index, **call} for index, call in enumerate(calls)], + }, + } + ], + }, + { + "id": marker, + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}], + }, + ) + + +def _chat_caller_stream_event(event: dict[str, JsonValue], model: str) -> dict[str, JsonValue]: + choices: Final = JSON_MESSAGES.validate_python(event["choices"]) + return { + "id": event["id"], + "object": event["object"], + "created": 1, + "model": model, + "choices": [_chat_caller_stream_choice(choice) for choice in choices], + } + + +def _chat_caller_stream_events(event: dict[str, JsonValue], model: str) -> tuple[dict[str, JsonValue], ...]: + return ( + _chat_caller_stream_event(event, model), + *((_chat_caller_stream_usage_event(event, model),) if "usage" in event else ()), + ) + + +def _chat_caller_stream_usage_event(event: dict[str, JsonValue], model: str) -> dict[str, JsonValue]: + return { + "id": event["id"], + "object": event["object"], + "created": 1, + "model": model, + "choices": [{"index": 0, "delta": {}}], + "usage": {**_json_object_value(event["usage"]), "cost": 4.05e-6}, + } + + +def _chat_caller_stream_choice(choice: dict[str, JsonValue]) -> dict[str, JsonValue]: + delta: Final = _json_object_value(choice["delta"]) + tool_calls: Final = JSON_MESSAGES.validate_python(delta["tool_calls"]) if "tool_calls" in delta else None + finish_reason: Final = choice.get("finish_reason") + return { + **{ + key: value + for key, value in choice.items() + if key not in {"delta", "finish_reason", "index"} and value is not None + }, + "index": choice["index"], + "delta": { + **{key: value for key, value in delta.items() if key != "tool_calls" and value is not None}, + **( + {"tool_calls": [_chat_caller_stream_tool_call(call) for call in tool_calls]} + if tool_calls is not None + else {} + ), + }, + **({"finish_reason": finish_reason} if finish_reason is not None else {}), + } + + +def _chat_caller_stream_tool_call(call: dict[str, JsonValue]) -> dict[str, JsonValue]: + function: Final = _json_object_value(call["function"]) + return { + **{key: value for key, value in call.items() if key not in {"function", "type"}}, + "type": "function", + "function": {**function, **({"arguments": ""} if "arguments" not in function else {})}, + } + + +def _normalize_chat_caller_stream(events: Sequence[JsonValue]) -> tuple[dict[str, JsonValue], ...]: + return tuple({**_json_object_value(event), "created": 1} for event in events) + + +def _responses_caller_body(body: dict[str, JsonValue], model: str) -> dict[str, JsonValue]: + usage: Final = _json_object_value(body["usage"]) + output: Final = JSON_MESSAGES.validate_python(body["output"]) + return { + **body, + "id": "", + "model": model, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": None, + "parallel_tool_calls": None, + "temperature": None, + "tool_choice": None, + "tools": None, + "top_p": None, + "max_output_tokens": None, + "previous_response_id": None, + "reasoning": None, + "text": None, + "truncation": None, + "user": None, + "store": None, + "output": [{**item, **({"namespace": None} if item.get("type") == "function_call" else {})} for item in output], + "usage": { + "input_tokens_details": None, + "output_tokens_details": None, + "cost": None, + **usage, + }, + } + + +def _responses_stream_caller_body(body: dict[str, JsonValue]) -> dict[str, JsonValue]: + response_id: Final = body.get("id") + assert isinstance(response_id, str), body + usage: Final = _json_object_value(body["usage"]) if body.get("status") == "completed" else {} + return { + **body, + "id": "", + **({"usage": {**usage, "cost": 4.05e-06}} if usage else {}), + } + + +def _responses_caller_response(reply: Reply, model: str) -> dict[str, JsonValue]: + return _responses_caller_body(_json_object(reply.body), model) + + +def _normalize_responses_caller_body(body: dict[str, JsonValue]) -> dict[str, JsonValue]: + response_id: Final = body.get("id") + assert isinstance(response_id, str) and response_id.startswith("resp_"), body + return {**body, "id": ""} + + +def _responses_caller_stream(reply: Reply, model: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + { + **event, + "model": model, + **( + {"response": _responses_stream_caller_body(_json_object_value(event["response"]))} + if "response" in event + else {} + ), + } + for event in _stream_values(reply) + ) + + +def _normalize_responses_caller_stream( + events: Sequence[dict[str, JsonValue]], +) -> tuple[dict[str, JsonValue], ...]: + return tuple( + { + **event, + **( + {"response": _normalize_responses_caller_body(_json_object_value(event["response"]))} + if "response" in event + else {} + ), + } + for event in events + ) + + +def _messages_caller_response(reply: Reply, model: str) -> dict[str, JsonValue]: + return {**_json_object(reply.body), "model": model} + + +def _messages_caller_raw_stream(reply: Reply, model: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + { + **event, + **( + {"message": {**_json_object_value(event["message"]), "model": model}} + if event.get("type") == "message_start" + else {} + ), + } + for event in _stream_values(reply) + ) + + +def _messages_caller_stream_response(reply: Reply, model: str) -> dict[str, JsonValue]: + body: Final = _messages_caller_response(reply, model) + content: Final = JSON_MESSAGES.validate_python(body["content"]) + return { + **body, + "content": [{**block, **({"caller": None} if block.get("type") == "tool_use" else {})} for block in content], + } + + +def _messages_caller_stream( + reply: Reply, model: str, *, final_message: dict[str, JsonValue] +) -> tuple[dict[str, JsonValue], ...]: + events: Final = tuple( + chain.from_iterable( + _messages_caller_stream_event(event, model, final_message) for event in _stream_values(reply) + ) + ) + return (*events, {"type": "message_stop", "message": final_message}) + + +def _messages_caller_stream_event( + event: dict[str, JsonValue], model: str, final_message: dict[str, JsonValue] +) -> tuple[dict[str, JsonValue], ...]: + if event.get("type") == "message_stop": + return () + if event.get("type") == "content_block_stop": + index: Final = event["index"] + assert isinstance(index, int), event + content: Final = JSON_MESSAGES.validate_python(final_message["content"]) + return ({**event, "content_block": content[index]},) + caller_event: Final = { + **event, + **({"message": {**_json_object_value(event["message"]), "model": model}} if "message" in event else {}), + } + if event.get("type") != "content_block_delta": + return (caller_event,) + delta: Final = _json_object_value(event["delta"]) + if delta.get("type") != "input_json_delta": + return (caller_event,) + partial_json: Final = delta["partial_json"] + assert isinstance(partial_json, str), event + return ( + caller_event, + { + "type": "input_json", + "partial_json": partial_json, + "snapshot": JSON_OBJECT.validate_json(partial_json), + }, + ) + + +def _span_attributes(request: Request) -> Iterator[dict[str, str]]: + if request.headers.get("content-type") != "application/x-protobuf": + return + batch: Final = ExportTraceServiceRequest.FromString(request.body) + for resource_spans in batch.resource_spans: + for scope_spans in resource_spans.scope_spans: + for span in scope_spans.spans: + yield {attribute.key: _attribute_text(attribute.value) for attribute in span.attributes} + + +def _attribute_text(value: object) -> str: + assert isinstance(value, AnyValue) + match value.WhichOneof("value"): + case "string_value": + return value.string_value + case "int_value": + return str(value.int_value) + case "double_value": + return str(value.double_value) + case "bool_value": + return str(value.bool_value) + case _: + return "" + + +def _spans(requests: tuple[Request, ...]) -> Iterator[dict[str, str]]: + for request in requests: + yield from _span_attributes(request) + + +def _canonical_response_id(value: str) -> str: + try: + return base64.b64decode(value.removeprefix("resp_").encode()).decode() + except (ValueError, UnicodeDecodeError): + return value + + +def _matching_llm_spans(requests: tuple[Request, ...], response_id: str) -> Iterator[dict[str, str]]: + wanted: Final = _canonical_response_id(response_id) + for attributes in _spans(requests): + if attributes.get("openinference.span.kind") != "LLM": + continue + observed: Final = attributes.get("gen_ai.response.id", "") + if _canonical_response_id(observed) == wanted or response_id in attributes.values(): + yield attributes + + +def _matching_marker_spans(requests: tuple[Request, ...], marker: str) -> Iterator[dict[str, str]]: + for attributes in _spans(requests): + if attributes.get("openinference.span.kind") == "LLM" and marker in attributes.values(): + yield attributes + + +def _single_span(spans: tuple[dict[str, str], ...]) -> dict[str, str]: + assert len(spans) == 1, spans + return spans[0] + + +def _matching_span(destination: Wire, response_id: str) -> dict[str, str]: + spans: Final = eventually( + lambda: tuple(_matching_llm_spans(destination.drain(), response_id)), + bool, + seconds=30, + ) + return _single_span(spans) + + +def _matching_marker_span(destination: Wire, marker: str) -> dict[str, str]: + spans: Final = eventually( + lambda: tuple(_matching_marker_spans(destination.drain(), marker)), + bool, + seconds=30, + ) + return _single_span(spans) + + +def _matching_output_value_span(destination: Wire, marker: str) -> dict[str, str]: + spans: Final = eventually( + lambda: tuple( + attributes + for attributes in _spans(destination.drain()) + if attributes.get("openinference.span.kind") == "LLM" and marker in attributes.get("output.value", "") + ), + bool, + seconds=30, + ) + return _single_span(spans) + + +def _matching_genai_marker_span(destination: Wire, marker: str) -> dict[str, str]: + spans: Final = eventually( + lambda: tuple( + attributes + for attributes in _spans(destination.drain()) + if attributes.get("gen_ai.operation.name") == "chat" and marker in attributes.values() + ), + bool, + seconds=30, + ) + return _single_span(spans) + + +def _matching_any_marker_span(destination: Wire, marker: str) -> dict[str, str]: + def matches(requests: tuple[Request, ...]) -> tuple[dict[str, str], ...]: + return tuple(attributes for attributes in _spans(requests) if marker in attributes.values()) + + spans: Final = eventually( + lambda: matches(destination.drain()), + bool, + seconds=30, + ) + return _single_span(spans) + + +def _collect_marker_spans( + destination: Wire, markers: tuple[str, ...], *, timeout_seconds: float = 30 +) -> tuple[dict[str, str], ...]: + expected: Final = frozenset(markers) + + def matches(requests: tuple[Request, ...]) -> tuple[dict[str, str], ...]: + return tuple( + attributes + for attributes in _spans(requests) + if attributes.get("openinference.span.kind") == "LLM" + and any(marker in attributes.values() for marker in expected) + ) + + def complete(spans: tuple[dict[str, str], ...]) -> bool: + return all(any(marker in attributes.values() for attributes in spans) for marker in expected) + + def collect(previous: tuple[dict[str, str], ...]) -> tuple[dict[str, str], ...]: + current: Final = eventually(lambda: matches(destination.drain()), bool, seconds=timeout_seconds) + combined: Final = (*previous, *current) + return combined if complete(combined) else collect(combined) + + return collect(()) + + +def _llm_spans_through_markers(destination: Wire, markers: tuple[str, ...]) -> tuple[dict[str, str], ...]: + expected: Final = frozenset(markers) + + def llm_spans(requests: tuple[Request, ...]) -> tuple[dict[str, str], ...]: + return tuple( + attributes for attributes in _spans(requests) if attributes.get("openinference.span.kind") == "LLM" + ) + + def complete(spans: tuple[dict[str, str], ...]) -> bool: + return all(any(marker in attributes.values() for attributes in spans) for marker in expected) + + def collect(previous: tuple[dict[str, str], ...]) -> tuple[dict[str, str], ...]: + current: Final = eventually(lambda: llm_spans(destination.drain()), bool, seconds=30) + combined: Final = (*previous, *current) + return combined if complete(combined) else collect(combined) + + return collect(()) + + +def _write_config( + directory: Path, + *, + callbacks: Sequence[str] = ("arize",), + callback_settings: Mapping[str, JsonValue] | None = None, + litellm_settings: Mapping[str, JsonValue] | None = None, + general_settings: Mapping[str, JsonValue] | None = None, +) -> Path: + loaded: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = { + **loaded, + "litellm_settings": { + **loaded["litellm_settings"], + "callbacks": list(callbacks), + **(litellm_settings or {}), + }, + **({"callback_settings": dict(callback_settings)} if callback_settings is not None else {}), + "general_settings": { + **loaded["general_settings"], + **(general_settings or {}), + }, + } + path: Final = directory / f"openinference-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _environment(destination: Wire, extra: Mapping[str, str] | None = None) -> dict[str, str]: + return { + "LITELLM_OTEL_V2": "1", + "OTEL_BSP_SCHEDULE_DELAY": "100", + "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT": "span_only", + "LITELLM_OTEL_BAGGAGE_METADATA_KEYS": "requester_metadata.trace_marker", + "ARIZE_HTTP_ENDPOINT": destination.url + "/v1/traces", + "ARIZE_SPACE_ID": "integration-space", + "ARIZE_API_KEY": "integration-arize-key", + **(extra or {}), + } + + +def _collector(_request: Request) -> Reply: + return Reply(body=b"", content_type="application/x-protobuf") + + +def _owned_sink_handler(sink: Callable[[Request], Reply]) -> Callable[[Request], Reply]: + def handle(request: Request) -> Reply: + if ( + request.method != "POST" + or not request.body + or request.headers.get("content-type", "").split(";", maxsplit=1)[0].strip().lower() + != "application/x-protobuf" + ): + return Reply(body=b"", content_type="application/x-protobuf") + try: + return sink(request) + except (AssertionError, IndexError, KeyError, TypeError, ValueError) as error: + raise AssertionError(f"{request.method} {request.target}: {error!r}") from error + + return handle + + +def _provider_handler(upstream: Callable[[Request], Reply]) -> Callable[[Request], Reply]: + def handle(request: Request) -> Reply: + if ( + request.method != "POST" + or not request.body + or request.headers.get("content-type", "").split(";", maxsplit=1)[0].strip().lower() != "application/json" + ): + return Reply(body=b"{}", content_type="application/json") + try: + return upstream(request) + except (AssertionError, IndexError, KeyError, TypeError, ValueError) as error: + raise AssertionError(f"{request.method} {request.target}: {error!r}") from error + + return handle + + +@contextmanager +def _rig( + gateway: Gateway, + directory: Path, + upstream: Callable[[Request], Reply], + *, + callbacks: Sequence[str] = ("arize",), + callback_settings: Mapping[str, JsonValue] | None = None, + litellm_settings: Mapping[str, JsonValue] | None = None, + general_settings: Mapping[str, JsonValue] | None = None, + environment: Mapping[str, str] | None = None, + remove_environment: tuple[str, ...] = (), + disabled_environment: tuple[str, ...] = (), + model_name: str = "openai/gpt-4o-mini", + api_base_suffix: str = "/v1", + destination_handler: Callable[[Request], Reply] | None = None, + destination_wire: Wire | None = None, + fresh_client_connections: bool = False, + workers: int = 2, +) -> Iterator[Rig]: + destination_context: Final[AbstractContextManager[Wire]] = ( + nullcontext(destination_wire) + if destination_wire is not None + else wire_server(_owned_sink_handler(destination_handler or _collector)) + ) + with wire_server(_provider_handler(upstream)) as provider, destination_context as destination: + resolved_settings: Final = { + key: ( + {**value, "endpoint": destination.url + "/v1/traces"} + if key == "otel" and isinstance(value, dict) and value.get("endpoint") == "unused" + else value + ) + for key, value in (callback_settings or {}).items() + } + config: Final = _write_config( + directory, + callbacks=callbacks, + callback_settings=resolved_settings, + litellm_settings=litellm_settings, + general_settings=general_settings, + ) + preset_environment: Final = { + **({"PHOENIX_COLLECTOR_ENDPOINT": destination.url + "/v1/traces"} if "arize_phoenix" in callbacks else {}), + **( + { + "WANDB_HOST": destination.url, + "WANDB_API_KEY": "integration-weave-key", + "WANDB_PROJECT_ID": "integration/project", + } + if "weave_otel" in callbacks + else {} + ), + **( + { + "LANGFUSE_OTEL_HOST": destination.url, + "LANGFUSE_PUBLIC_KEY": "integration-public", + "LANGFUSE_SECRET_KEY": "integration-secret", + } + if "langfuse_otel" in callbacks + else {} + ), + **( + { + "LEVOAI_API_KEY": "integration-levo-key", + "LEVOAI_ORG_ID": "integration-org", + "LEVOAI_WORKSPACE_ID": "integration-workspace", + "LEVOAI_COLLECTOR_URL": destination.url + "/v1/traces", + } + if "levo" in callbacks + else {} + ), + **( + { + "SIGNOZ_INGESTION_ENDPOINT": destination.url + "/v1/traces", + "SIGNOZ_INGESTION_KEY": "integration-signoz-key", + } + if "signoz" in callbacks + else {} + ), + **( + {"OTEL_EXPORTER_OTLP_ENDPOINT": destination.url} + if any(callback_name in callbacks for callback_name in ("langtrace", "newrelic", "agentops")) + else {} + ), + } + overrides: Final = { + key: value + for key, value in _environment(destination, {**preset_environment, **(environment or {})}).items() + if key not in disabled_environment + } + with ( + owned_proxy_process( + gateway, + directory, + overrides, + config=config, + remove_environment=remove_environment, + workers=workers, + ) as owned, + ExitStack() as resources, + ): + proxy_client: Final = ( + resources.enter_context( + httpx.Client( + base_url=str(owned.gateway.client.base_url), + timeout=15, + trust_env=False, + limits=httpx.Limits(max_keepalive_connections=0), + ) + ) + if fresh_client_connections + else owned.gateway.client + ) + proxy: Final = Gateway( + client=proxy_client, + key=owned.gateway.key, + upstream_url=owned.gateway.upstream_url, + ) + scenario: Final = resources.enter_context(proxy.scenario()) + model: Final = scenario.model( + model=model_name, + api_base=provider.url + api_base_suffix, + ) + yield Rig(proxy, owned, model, provider, destination) + + +def _chat_tool_call(identity: str, city: str = "Paris") -> dict[str, JsonValue]: + return { + "id": "call_" + identity, + "type": "function", + "function": { + "name": "lookup_weather", + "arguments": json.dumps({"city": city}), + }, + } + + +def _chat_response(identity: str, calls: Sequence[dict[str, JsonValue]] | None = None) -> Reply: + tool_calls: Final = list(calls if calls is not None else (_chat_tool_call(identity),)) + return Reply( + body=json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": {"role": "assistant", "content": None, "tool_calls": tool_calls}, + } + ], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + ).encode() + ) + + +def _chat_plain_response(identity: str, content: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": content}, + } + ], + "usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9}, + } + ).encode() + ) + + +def _chat_stream_response(identity: str, calls: Sequence[dict[str, JsonValue]], *, include_usage: bool = True) -> Reply: + def first_chunk(index: int, call: dict[str, JsonValue]) -> dict[str, JsonValue]: + fields: Final = object_value(JSON_VALUE.validate_python(call)) + return { + "index": index, + "id": string_value(fields["id"]), + "type": "function", + "function": {"name": "lookup_weather"}, + } + + def arguments_chunk(index: int, call: dict[str, JsonValue]) -> dict[str, JsonValue]: + fields: Final = object_value(JSON_VALUE.validate_python(call)) + function: Final = object_value(fields["function"]) + return { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": index, + "function": {"arguments": string_value(function["arguments"])}, + } + ] + }, + "finish_reason": None, + } + ], + } + + chunks: Final = [ + { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "delta": { + "role": "assistant", + "tool_calls": [first_chunk(index, call) for index, call in enumerate(calls)], + }, + "finish_reason": None, + } + ], + }, + *(arguments_chunk(index, call) for index, call in enumerate(calls)), + { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}], + **({"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}} if include_usage else {}), + }, + ] + return Reply( + content_type="text/event-stream", + chunks=tuple(b"data: " + json.dumps(chunk).encode() + b"\n\n" for chunk in chunks) + (b"data: [DONE]\n\n",), + ) + + +def _responses_response(identity: str, calls: Sequence[dict[str, JsonValue]] | None = None) -> Reply: + def response_item(call: dict[str, JsonValue]) -> dict[str, JsonValue]: + fields: Final = object_value(JSON_VALUE.validate_python(call)) + function: Final = object_value(fields["function"]) + call_id: Final = string_value(fields["id"]) + return { + "type": "function_call", + "id": "fc_" + call_id, + "call_id": call_id, + "name": string_value(function["name"]), + "arguments": string_value(function["arguments"]), + "status": "completed", + } + + response: Final = { + "id": identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [response_item(call) for call in calls if calls is not None] + if calls is not None + else [response_item(_chat_tool_call(identity))], + "usage": {"input_tokens": 11, "output_tokens": 4, "total_tokens": 15}, + } + return Reply(body=json.dumps(response).encode()) + + +def _responses_stream_response(identity: str, calls: Sequence[dict[str, JsonValue]]) -> Reply: + response: Final = _json_object(_responses_response(identity, calls).body) + output: Final = response["output"] + items: Final = tuple(object_value(item) for item in output) if isinstance(output, list) else () + events: Final = ( + { + "type": "response.created", + "sequence_number": 0, + "response": {**response, "status": "in_progress", "output": []}, + }, + *( + { + "type": "response.output_item.added", + "sequence_number": index + 1, + "output_index": index, + "item": item, + } + for index, item in enumerate(items) + ), + *( + { + "type": "response.function_call_arguments.delta", + "sequence_number": index + len(items) + 1, + "item_id": str(item["id"]), + "output_index": index, + "delta": str(item["arguments"]), + } + for index, item in enumerate(items) + ), + { + "type": "response.completed", + "sequence_number": len(items) * 2 + 1, + "response": response, + }, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events), + ) + + +def _anthropic_response(identity: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": identity, + "type": "message", + "role": "assistant", + "model": "claude-opus-5-5", + "content": [ + { + "type": "tool_use", + "id": "call_" + identity, + "name": "lookup_weather", + "input": {"city": "Paris"}, + } + ], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 4}, + } + ).encode() + ) + + +def _anthropic_stream_response(identity: str) -> Reply: + message: Final = { + "id": identity, + "type": "message", + "role": "assistant", + "model": "claude-opus-5-5", + "content": [ + { + "type": "tool_use", + "id": "call_" + identity, + "name": "lookup_weather", + "input": {"city": "Paris"}, + } + ], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 4}, + } + events: Final = ( + {"type": "message_start", "message": {**message, "content": [], "stop_reason": None}}, + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "tool_use", + "id": "call_" + identity, + "name": "lookup_weather", + "input": {}, + }, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"city": "Paris"}'}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None}, + "usage": {"output_tokens": 4}, + }, + {"type": "message_stop"}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events), + ) + + +def _tool_call_attributes( + attributes: Mapping[str, str], + *, + message_index: int = 0, + tool_index: int = 0, +) -> tuple[str, str, dict[str, JsonValue]]: + prefix: Final = _TOOL_PREFIX.format(message=message_index, tool=tool_index) + return ( + attributes[prefix + "id"], + attributes[prefix + "function.name"], + JSON_OBJECT.validate_json(attributes[prefix + "function.arguments"]), + ) + + +def _assert_tool_span( + attributes: Mapping[str, str], + *, + marker: str, + output: JsonValue, + calls: Sequence[tuple[str, str, Mapping[str, JsonValue]]], + metadata: Mapping[str, JsonValue] | None = None, + baggage: Mapping[str, str] | None = None, +) -> None: + for index, (call_id, name, arguments) in enumerate(calls): + observed: Final = _tool_call_attributes(attributes, tool_index=index) + assert observed == (call_id, name, dict(arguments)), f"tool call {index} for {marker}: {observed!r}" + observed_output: Final = JSON_MESSAGES.validate_json(attributes["output.value"]) + assert observed_output == output, (marker, observed_output) + if metadata is None: + assert "metadata" not in attributes, attributes + else: + assert json.loads(attributes["metadata"]) == dict(metadata), attributes + for key, value in (baggage or {}).items(): + assert attributes.get("litellm.metadata." + key) == value, attributes + + +def _response_tool_calls(identity: str, cities: Sequence[str] = ("Paris",)) -> list[dict[str, JsonValue]]: + return [_chat_tool_call(identity + "-" + city, city) for city in cities] diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py b/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py new file mode 100644 index 00000000000..9a86bb28577 --- /dev/null +++ b/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py @@ -0,0 +1,403 @@ +from __future__ import annotations + +import threading +import uuid +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +from _openinference_support import ( + CHAT_TOOLS, + RESPONSES_TOOLS, + _anthropic_response, + _anthropic_stream_response, + _assert_chat_request, + _assert_messages_request, + _assert_responses_request, + _chat_caller_response, + _chat_caller_stream, + _chat_request_marker, + _chat_response, + _chat_stream_response, + _chat_tool_call, + _collect_marker_spans, + _json_object, + _llm_spans_through_markers, + _matching_marker_span, + _messages_caller_raw_stream, + _messages_caller_response, + _normalize_chat_caller_stream, + _normalize_responses_caller_body, + _normalize_responses_caller_stream, + _owned_sink_handler, + _responses_caller_response, + _responses_caller_stream, + _responses_response, + _responses_stream_response, + _rig, + _spans, + _sse_json_values, +) +from integration._support.client import Gateway, eventually +from integration._support.wire import Reply, Request, wire_server + + +def _call( + proxy: Gateway, + model: str, + marker: str, + *, + surface: str = "chat", + stream: bool = False, + prompt: str | None = None, +) -> httpx.Response: + match surface: + case "chat": + return proxy.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": prompt or "weather in Paris?"}], + "tools": CHAT_TOOLS, + "tool_choice": {"type": "function", "function": {"name": "lookup_weather"}}, + "metadata": {"trace_marker": marker}, + **({"stream": True} if stream else {}), + **({"stream_options": {"include_usage": True}} if stream and surface == "chat" else {}), + "cache": {"no-cache": True}, + }, + ) + case "responses": + return proxy.request( + "POST", + "/v1/responses", + { + "model": model, + "input": prompt or "weather in Paris?", + "tools": RESPONSES_TOOLS, + "tool_choice": {"type": "function", "name": "lookup_weather"}, + "metadata": {"trace_marker": marker}, + **({"stream": True} if stream else {}), + "cache": {"no-cache": True}, + }, + ) + case "messages": + return proxy.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": prompt or "weather in Paris?"}], + "tools": [ + { + "name": "lookup_weather", + "description": "Get weather", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + } + ], + "tool_choice": {"type": "auto"}, + "metadata": {"trace_marker": marker}, + **({"stream": True} if stream else {}), + "cache": {"no-cache": True}, + }, + ) + case _: + raise AssertionError(f"Unknown endpoint: {surface}") + + +def _assert_response(response: httpx.Response, marker: str, surface: str, stream: bool, model: str) -> None: + assert response.status_code == 200, response.text + if stream: + observed_stream: Final = _sse_json_values(response.content) + if surface == "chat": + chat_reply: Final = _chat_stream_response(marker, (_chat_tool_call(marker),)) + assert _normalize_chat_caller_stream(observed_stream) == _chat_caller_stream(chat_reply, model), ( + response.text + ) + return + if surface == "responses": + responses_reply: Final = _responses_stream_response(marker, (_chat_tool_call(marker),)) + assert _normalize_responses_caller_stream(observed_stream) == _responses_caller_stream( + responses_reply, model + ), response.text + return + if surface == "messages": + messages_reply: Final = _anthropic_stream_response(marker) + assert observed_stream == _messages_caller_raw_stream(messages_reply, model), response.text + return + raise AssertionError(f"Unknown endpoint: {surface}") + + observed: Final = _json_object(response.content) + if surface == "chat": + chat_reply: Final = _chat_response(marker) + assert observed == _chat_caller_response(chat_reply, model), response.text + return + if surface == "responses": + responses_reply: Final = _responses_response(marker) + assert _normalize_responses_caller_body(observed) == _responses_caller_response(responses_reply, model), ( + response.text + ) + return + if surface == "messages": + messages_reply: Final = _anthropic_response(marker) + assert observed == _messages_caller_response(messages_reply, model), response.text + return + raise AssertionError(f"Unknown endpoint: {surface}") + + +def _call_without_worker_error( + proxy: Gateway, model: str, marker: str, *, prompt: str | None = None +) -> httpx.Response | None: + try: + return _call(proxy, model, marker, prompt=prompt) + except httpx.HTTPError: + return None + + +def test_arize_otel_v2_f1_sink_outage_and_recovery(gateway: Gateway, tmp_path: Path) -> None: + outage_markers: Final = tuple("f1-outage-" + uuid.uuid4().hex for _ in range(30)) + recovery_markers: Final = tuple("f1-recovery-" + uuid.uuid4().hex for _ in range(30)) + surfaces: Final = ("chat", "responses", "messages") + surface_streams: Final = tuple( + (surfaces[index % len(surfaces)], index % 2 == 0) for index in range(len(outage_markers)) + ) + calls: Final = tuple( + (marker, surface, stream) for marker, (surface, stream) in zip(outage_markers, surface_streams, strict=True) + ) + recovery_calls: Final = tuple( + (marker, surface, stream) for marker, (surface, stream) in zip(recovery_markers, surface_streams, strict=True) + ) + + def upstream(request: Request) -> Reply: + body: Final = _json_object(request.body) + if request.target.endswith("/messages"): + messages: Final = body.get("messages") + assert isinstance(messages, list) and isinstance(messages[0], dict), body + marker: Final = messages[0].get("content") + assert isinstance(marker, str), body + _assert_messages_request( + request, + marker=marker, + prompt=marker, + stream=True if body.get("stream") is True else False, + ) + return _anthropic_stream_response(marker) if body.get("stream") is True else _anthropic_response(marker) + if request.target.endswith("/responses"): + marker: Final = body.get("input") + assert isinstance(marker, str), body + _assert_responses_request( + request, + marker=marker, + input_value=marker, + stream=body.get("stream") is True, + ) + return ( + _responses_stream_response(marker, (_chat_tool_call(marker),)) + if body.get("stream") is True + else _responses_response(marker) + ) + marker: Final = _chat_request_marker(request) + _assert_chat_request( + request, + messages=[{"role": "user", "content": marker}], + stream=True if body.get("stream") is True else None, + stream_options={"include_usage": True} if body.get("stream") is True else None, + ) + return ( + _chat_stream_response(marker, (_chat_tool_call(marker),)) + if body.get("stream") is True + else _chat_response(marker) + ) + + def sink(_request: Request) -> Reply: + return Reply(body=b"", content_type="application/x-protobuf") + + def run_burst( + burst: tuple[tuple[str, str, bool], ...], + proxy: Gateway, + chat_model: str, + messages_model: str, + ) -> None: + with ThreadPoolExecutor(max_workers=len(burst)) as executor: + futures: Final = tuple( + executor.submit( + _call, + proxy, + messages_model if surface == "messages" else chat_model, + marker, + surface=surface, + stream=stream, + prompt=marker, + ) + for marker, surface, stream in burst + ) + responses: Final = tuple( + (marker, surface, stream, future.result(timeout=60)) + for (marker, surface, stream), future in zip(burst, futures, strict=True) + ) + for marker, surface, stream, response in responses: + model_name: Final = messages_model if surface == "messages" else chat_model + _assert_response(response, marker, surface, stream, model_name) + + with ExitStack() as servers: + initial_stack: Final = servers.enter_context(ExitStack()) + stopped_destination: Final = initial_stack.enter_context(wire_server(_owned_sink_handler(sink))) + sink_port: Final = urlsplit(stopped_destination.url).port + assert sink_port is not None, stopped_destination.url + initial_stack.close() + with _rig( + gateway, + tmp_path, + upstream, + destination_wire=stopped_destination, + workers=1, + ) as rig: + with httpx.Client(trust_env=False) as client, pytest.raises(httpx.ConnectError): + client.get(stopped_destination.url + "/health", timeout=2) + with rig.proxy.scenario() as scenario: + messages_model: Final = scenario.model( + model="anthropic/claude-opus-5-5", + api_base=rig.provider.url, + ) + run_burst(calls, rig.proxy, rig.model, messages_model) + with httpx.Client(trust_env=False) as client, pytest.raises(httpx.ConnectError): + client.get(stopped_destination.url + "/health", timeout=2) + recovered_stack: Final = servers.enter_context(ExitStack()) + recovered_destination: Final = recovered_stack.enter_context( + wire_server(_owned_sink_handler(sink), port=sink_port) + ) + run_burst(recovery_calls, rig.proxy, rig.model, messages_model) + sentinel: Final = f"f1-sentinel-{uuid.uuid4().hex}" + sentinel_response: Final = _call( + rig.proxy, + rig.model, + sentinel, + surface="chat", + stream=False, + prompt=sentinel, + ) + _assert_response(sentinel_response, sentinel, "chat", False, rig.model) + spans: Final = _llm_spans_through_markers(recovered_destination, (*recovery_markers, sentinel)) + assert all( + sum(span.get("litellm.metadata.trace_marker") == marker for span in spans) == 1 + for marker in recovery_markers + ), spans + assert all( + sum(span.get("litellm.metadata.trace_marker") == marker for span in spans) <= 1 + for marker in outage_markers + ), spans + assert sum(span.get("litellm.metadata.trace_marker") == sentinel for span in spans) == 1, spans + + +def test_arize_otel_v2_f2_slow_sink_does_not_deadlock(gateway: Gateway, tmp_path: Path) -> None: + release: Final = threading.Event() + blocked: Final = threading.Event() + completed: Final = threading.Event() + marker: Final = "f2-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + return _chat_response(marker) + + def sink(request: Request) -> Reply: + if any(attributes.get("litellm.metadata.trace_marker") == marker for attributes in _spans((request,))): + blocked.set() + assert release.wait(timeout=5), "slow sink was not released" + completed.set() + return Reply(body=b"", content_type="application/x-protobuf") + + with _rig(gateway, tmp_path, upstream, destination_handler=sink) as rig: + timer: Final = threading.Timer(2, release.set) + try: + response: Final = _call(rig.proxy, rig.model, marker) + _assert_response(response, marker, "chat", False, rig.model) + assert eventually(lambda: blocked.is_set(), bool, seconds=10) + timer.start() + assert eventually(lambda: completed.is_set(), bool, seconds=10) + timer.join(timeout=5) + assert not timer.is_alive(), "Slow sink timer did not finish" + finally: + release.set() + timer.cancel() + if timer.ident is not None: + timer.join(timeout=5) + _matching_marker_span(rig.destination, marker) + + +def test_arize_otel_v2_f3_one_proxy_worker_can_die(gateway: Gateway, tmp_path: Path) -> None: + markers: Final = tuple("f3-" + uuid.uuid4().hex for _ in range(8)) + release: Final = threading.Event() + + def upstream(request: Request) -> Reply: + request_marker: Final = _chat_request_marker(request) + _assert_chat_request( + request, + messages=[{"role": "user", "content": request_marker}], + ) + assert release.wait(timeout=20), "F3 upstream barrier was not released" + return _chat_response(request_marker) + + def sink(_request: Request) -> Reply: + return Reply(body=b"", content_type="application/x-protobuf") + + with _rig( + gateway, + tmp_path, + upstream, + destination_handler=sink, + fresh_client_connections=True, + ) as rig: + children: Final = psutil.Process(rig.owned.process.pid).children(recursive=True) + workers: Final = tuple( + child for child in children if child.is_running() and "resource_tracker" not in " ".join(child.cmdline()) + ) + assert len(workers) >= 2, tuple((worker.pid, worker.name()) for worker in workers) + with ThreadPoolExecutor(max_workers=len(markers)) as executor: + try: + futures: Final = tuple( + executor.submit( + _call_without_worker_error, + rig.proxy, + rig.model, + marker, + prompt=marker, + ) + for marker in markers + ) + observed: Final = eventually( + lambda: rig.provider.received.qsize(), + lambda count: count >= 2, + seconds=10, + ) + assert observed >= 2, observed + workers[0].kill() + assert eventually(lambda: not workers[0].is_running(), bool, seconds=10), workers[0] + release.set() + in_flight: Final = tuple( + (marker, future.result(timeout=60)) for marker, future in zip(markers, futures, strict=True) + ) + finally: + release.set() + survivor: Final = "f3-survivor-" + uuid.uuid4().hex + survivor_response: Final = _call(rig.proxy, rig.model, survivor, prompt=survivor) + _assert_response(survivor_response, survivor, "chat", False, rig.model) + served_responses: Final = tuple( + (marker, response) for marker, response in in_flight if response is not None and response.status_code == 200 + ) + for marker, response in served_responses: + _assert_response(response, marker, "chat", False, rig.model) + served: Final = tuple(marker for marker, _response in served_responses) + (survivor,) + collected: Final = _collect_marker_spans(rig.destination, served) + assert len(collected) == len(served), collected + exported: Final = tuple(span["litellm.metadata.trace_marker"] for span in collected) + assert len(exported) == len(served), exported + assert frozenset(exported) == frozenset(served), exported diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_family.py b/tests/integration/observability/test_arize_otel_v2_openinference_family.py new file mode 100644 index 00000000000..242b7bfe8c1 --- /dev/null +++ b/tests/integration/observability/test_arize_otel_v2_openinference_family.py @@ -0,0 +1,368 @@ +from __future__ import annotations + +import uuid +from collections.abc import Mapping +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +from _openinference_support import ( + CHAT_TOOLS, + _assert_chat_request, + _chat_caller_response, + _chat_response, + _json_messages, + _json_object, + _matching_genai_marker_span, + _matching_marker_span, + _rig, +) +from integration._support.client import Gateway +from integration._support.wire import Reply, Request +from pydantic import JsonValue + +_GENAI_B3_KEYS_WITHOUT_BAGGAGE: Final = frozenset( + { + "gen_ai.input.messages", + "gen_ai.operation.name", + "gen_ai.output.messages", + "gen_ai.provider.name", + "gen_ai.request.model", + "gen_ai.response.finish_reasons", + "gen_ai.response.id", + "gen_ai.response.model", + "gen_ai.system", + "gen_ai.tool.0.description", + "gen_ai.tool.0.name", + "gen_ai.tool.0.parameters", + "gen_ai.usage.completion_tokens", + "gen_ai.usage.input_tokens", + "gen_ai.usage.output_tokens", + "gen_ai.usage.prompt_tokens", + "gen_ai.usage.total_tokens", + "litellm.api_key.hash", + "litellm.call_id", + "litellm.call_type", + "litellm.cost.discount_amount", + "litellm.cost.discount_percent", + "litellm.cost.input", + "litellm.cost.margin_fixed_amount", + "litellm.cost.margin_percent", + "litellm.cost.margin_total_amount", + "litellm.cost.original", + "litellm.cost.output", + "litellm.cost.tool_usage", + "litellm.cost.total", + "litellm.provider.model", + "litellm.request.route", + "litellm.request.tools.declared", + "llm.request.functions.0.description", + "llm.request.functions.0.name", + "llm.request.functions.0.parameters", + "server.address", + "server.port", + } +) +_GENAI_B3_KEYS_WITH_BAGGAGE: Final = _GENAI_B3_KEYS_WITHOUT_BAGGAGE | frozenset({"litellm.metadata.trace_marker"}) +_LANGFUSE_B3_KEYS: Final = frozenset( + { + "gen_ai.input.messages", + "gen_ai.operation.name", + "gen_ai.output.messages", + "gen_ai.provider.name", + "gen_ai.request.model", + "gen_ai.response.finish_reasons", + "gen_ai.response.id", + "gen_ai.response.model", + "gen_ai.system", + "gen_ai.tool.0.description", + "gen_ai.tool.0.name", + "gen_ai.tool.0.parameters", + "gen_ai.usage.completion_tokens", + "gen_ai.usage.input_tokens", + "gen_ai.usage.output_tokens", + "gen_ai.usage.prompt_tokens", + "gen_ai.usage.total_tokens", + "langfuse.observation.cost_details", + "langfuse.observation.id", + "langfuse.observation.input", + "langfuse.observation.metadata.provider", + "langfuse.observation.model.name", + "langfuse.observation.output", + "langfuse.observation.type", + "langfuse.observation.usage_details", + "litellm.api_key.hash", + "litellm.call_id", + "litellm.call_type", + "litellm.cost.discount_amount", + "litellm.cost.discount_percent", + "litellm.cost.input", + "litellm.cost.margin_fixed_amount", + "litellm.cost.margin_percent", + "litellm.cost.margin_total_amount", + "litellm.cost.original", + "litellm.cost.output", + "litellm.cost.tool_usage", + "litellm.cost.total", + "litellm.provider.model", + "litellm.request.route", + "litellm.request.tools.declared", + "llm.request.functions.0.description", + "llm.request.functions.0.name", + "llm.request.functions.0.parameters", + "server.address", + "server.port", + } +) +_B3_ATTRIBUTE_KEYS: Final[Mapping[str, frozenset[str]]] = MappingProxyType( + { + "langfuse_otel": _LANGFUSE_B3_KEYS, + "langtrace": _GENAI_B3_KEYS_WITHOUT_BAGGAGE, + "signoz": _GENAI_B3_KEYS_WITH_BAGGAGE, + "newrelic": _GENAI_B3_KEYS_WITH_BAGGAGE, + "levo": _GENAI_B3_KEYS_WITH_BAGGAGE, + "agentops": _GENAI_B3_KEYS_WITH_BAGGAGE, + "otel": _GENAI_B3_KEYS_WITH_BAGGAGE, + } +) +_B3_BAGGAGE_CALLBACKS: Final = frozenset({"signoz", "newrelic", "levo", "agentops", "otel"}) +_B4_ATTRIBUTE_KEYS: Final = frozenset( + { + "input.value", + "litellm.trace_id", + "llm.cost.total", + "llm.input_messages.0.message.content", + "llm.input_messages.0.message.role", + "llm.invocation_parameters", + "llm.is_streaming", + "llm.model_name", + "llm.output_messages.0.message.content", + "llm.output_messages.0.message.role", + "llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments", + "llm.output_messages.0.message.tool_calls.0.tool_call.function.name", + "llm.output_messages.0.message.tool_calls.0.tool_call.id", + "llm.provider", + "llm.request.type", + "llm.response.cost", + "llm.response.id", + "llm.response.model", + "llm.token_count.completion", + "llm.token_count.prompt", + "llm.token_count.total", + "llm.tools.0.description", + "llm.tools.0.name", + "llm.tools.0.parameters", + "metadata", + "openinference.span.kind", + "output.value", + "user.id", + } +) + + +def _request(proxy: Gateway, model: str, marker: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "weather in Paris?"}], + "tools": CHAT_TOOLS, + "tool_choice": {"type": "function", "function": {"name": "lookup_weather"}}, + "metadata": {"trace_marker": marker}, + "cache": {"no-cache": True}, + }, + ) + + +def _assert_openinference(attributes: dict[str, str], marker: str) -> None: + assert attributes["llm.output_messages.0.message.tool_calls.0.tool_call.id"] == f"call_{marker}", attributes + assert attributes["llm.output_messages.0.message.tool_calls.0.tool_call.function.name"] == "lookup_weather", ( + attributes + ) + assert ( + attributes["llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments"] == '{"city": "Paris"}' + ), attributes + assert _json_object(attributes["metadata"].encode()) == {"trace_marker": marker}, attributes + assert attributes["litellm.metadata.trace_marker"] == marker, attributes + assert _json_messages(attributes["output.value"]) == [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": f"call_{marker}", + "type": "function", + "function": {"name": "lookup_weather", "arguments": '{"city": "Paris"}'}, + } + ], + } + ], attributes + + +def test_arize_otel_v2_b1_phoenix_openinference(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "b1-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + return _chat_response(marker) + + with _rig( + gateway, + tmp_path, + upstream, + callbacks=("arize_phoenix",), + callback_settings={"otel": {"exporter": "http/protobuf", "endpoint": "unused"}}, + environment={"PHOENIX_PROJECT_NAME": "integration"}, + ) as rig: + response: Final = _request(rig.proxy, rig.model, marker) + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(_chat_response(marker), rig.model), response.text + _assert_openinference(_matching_marker_span(rig.destination, marker), marker) + + +def test_arize_otel_v2_b2_weave_openinference(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "b2-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + return _chat_response(marker) + + with _rig( + gateway, + tmp_path, + upstream, + callbacks=("weave_otel",), + callback_settings={"otel": {"exporter": "http/protobuf", "endpoint": "unused"}}, + environment={"WANDB_BASE_URL": "http://127.0.0.1"}, + ) as rig: + response: Final = _request(rig.proxy, rig.model, marker) + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(_chat_response(marker), rig.model), response.text + _assert_openinference(_matching_marker_span(rig.destination, marker), marker) + + +@pytest.mark.parametrize( + "callback", + ("langfuse_otel", "langtrace", "signoz", "newrelic", "levo", "agentops", "otel"), +) +def test_arize_otel_v2_b3_non_openinference_callback_family(callback: str, gateway: Gateway, tmp_path: Path) -> None: + marker: Final = f"b3-{callback}-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + body: Final = _json_object(request.body) + assert body == { + "messages": [{"role": "user", "content": "weather in Paris?"}], + "model": "gpt-4o-mini", + "tool_choice": {"type": "function", "function": {"name": "lookup_weather"}}, + "tools": CHAT_TOOLS, + }, body + return _chat_response(marker) + + callback_settings: Final[dict[str, JsonValue]] = { + "otel": {"exporter": "http/protobuf", "endpoint": "unused", "mapper_names": ["genai"]} + } + with _rig( + gateway, + tmp_path, + upstream, + callbacks=(callback,), + callback_settings=callback_settings, + environment={ + "LANGFUSE_HOST": "http://127.0.0.1", + **( + { + "HTTPS_PROXY": "http://127.0.0.1:0", + "https_proxy": "http://127.0.0.1:0", + "NO_PROXY": "", + "no_proxy": "", + } + if callback == "agentops" + else {} + ), + }, + remove_environment=( + ("NEW_RELIC_LICENSE_KEY",) + if callback == "newrelic" + else ("AGENTOPS_API_KEY",) + if callback == "agentops" + else () + ), + ) as rig: + response: Final = _request(rig.proxy, rig.model, marker) + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(_chat_response(marker), rig.model), response.text + attributes: Final = _matching_genai_marker_span(rig.destination, marker) + assert frozenset(attributes) == _B3_ATTRIBUTE_KEYS[callback], attributes + assert "metadata" not in attributes, attributes + assert not any(".tool_calls." in key for key in attributes), attributes + baggage: Final = tuple( + sorted((key, value) for key, value in attributes.items() if key.startswith("litellm.metadata.")) + ) + expected_baggage: Final = ( + (("litellm.metadata.trace_marker", marker),) if callback in _B3_BAGGAGE_CALLBACKS else () + ) + assert baggage == expected_baggage, attributes + + +def test_arize_otel_v2_b3_langfuse_carries_metadata_baggage(gateway: Gateway, tmp_path: Path) -> None: + pytest.skip("BUG: LIT-9128 Langfuse OTel v2 preset omits request-metadata baggage") + marker: Final = "b3-langfuse-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + body: Final = _json_object(request.body) + assert body == { + "messages": [{"role": "user", "content": "weather in Paris?"}], + "model": "gpt-4o-mini", + "tool_choice": {"type": "function", "function": {"name": "lookup_weather"}}, + "tools": CHAT_TOOLS, + }, body + return _chat_response(marker) + + callback_settings: Final[dict[str, JsonValue]] = { + "otel": {"exporter": "http/protobuf", "endpoint": "unused", "mapper_names": ["genai"]} + } + with _rig( + gateway, + tmp_path, + upstream, + callbacks=("langfuse_otel",), + callback_settings=callback_settings, + environment={"LANGFUSE_HOST": "http://127.0.0.1"}, + ) as rig: + response: Final = _request(rig.proxy, rig.model, marker) + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(_chat_response(marker), rig.model), response.text + attributes: Final = _matching_genai_marker_span(rig.destination, marker) + assert attributes["litellm.metadata.trace_marker"] == marker, attributes + + +def test_arize_otel_v2_b4_legacy_otel_is_unchanged(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "b4-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + body: Final = _json_object(request.body) + assert body == { + "messages": [{"role": "user", "content": "weather in Paris?"}], + "model": "gpt-4o-mini", + "tool_choice": {"type": "function", "function": {"name": "lookup_weather"}}, + "tools": CHAT_TOOLS, + }, body + return _chat_response(marker) + + with _rig( + gateway, + tmp_path, + upstream, + callbacks=("arize",), + callback_settings={"otel": {"exporter": "http/protobuf", "endpoint": "unused"}}, + remove_environment=("LITELLM_OTEL_V2",), + disabled_environment=("LITELLM_OTEL_V2",), + ) as rig: + response: Final = _request(rig.proxy, rig.model, marker) + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(_chat_response(marker), rig.model), response.text + attributes: Final = _matching_marker_span(rig.destination, marker) + assert frozenset(attributes) == _B4_ATTRIBUTE_KEYS, attributes diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_modes.py b/tests/integration/observability/test_arize_otel_v2_openinference_modes.py new file mode 100644 index 00000000000..7d0f3f31787 --- /dev/null +++ b/tests/integration/observability/test_arize_otel_v2_openinference_modes.py @@ -0,0 +1,484 @@ +from __future__ import annotations + +import uuid +from collections.abc import Callable +from pathlib import Path +from typing import Final + +import httpx +import pytest +from _openinference_support import ( + CHAT_TOOLS, + _assert_chat_request, + _chat_caller_response, + _chat_plain_response, + _chat_request_marker, + _chat_response, + _json_messages, + _json_object, + _llm_spans_through_markers, + _matching_marker_span, + _matching_output_value_span, + _matching_span, + _response_tool_calls, + _rig, +) +from integration._support.client import Gateway +from integration._support.wire import Reply, Request + +_DEFAULT_METADATA: Final = { + "requester_ip_address": "127.0.0.1", + "user_api_key_user_id": "default_user_id", +} +_DEFAULT_METADATA_BAGGAGE: Final = frozenset({("litellm.metadata.user_api_key_user_id", "default_user_id")}) + + +def _request( + proxy: Gateway, + model: str, + marker: str, + *, + prompt: str = "weather in Paris?", + headers: dict[str, str] | None = None, + key: str | None = None, +) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "tools": CHAT_TOOLS, + "tool_choice": {"type": "function", "function": {"name": "lookup_weather"}}, + "metadata": {"trace_marker": marker}, + "cache": {"no-cache": True}, + }, + key=key, + headers=headers, + ) + + +def _assert_success_body(response: httpx.Response, marker: str, model: str) -> None: + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(_chat_response(marker), model), response.text + + +def _upstream(marker: str) -> Callable[[Request], Reply]: + def reply(request: Request) -> Reply: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + return _chat_response(marker) + + return reply + + +def _assert_output_tool_call(attributes: dict[str, str], marker: str) -> None: + assert attributes["llm.output_messages.0.message.tool_calls.0.tool_call.id"] == f"call_{marker}", attributes + assert attributes["llm.output_messages.0.message.tool_calls.0.tool_call.function.name"] == "lookup_weather", ( + attributes + ) + assert ( + attributes["llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments"] == '{"city": "Paris"}' + ), attributes + assert _json_messages(attributes["output.value"]) == [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": f"call_{marker}", + "type": "function", + "function": {"name": "lookup_weather", "arguments": '{"city": "Paris"}'}, + } + ], + } + ], attributes + + +def _assert_default_allowlist_attributes(attributes: dict[str, str], marker: str) -> None: + observed_baggage: Final = frozenset( + (key, value) for key, value in attributes.items() if key.startswith("litellm.metadata.") + ) + assert observed_baggage == _DEFAULT_METADATA_BAGGAGE, attributes + metadata: Final = _json_object(attributes["metadata"].encode()) + assert metadata == _DEFAULT_METADATA, attributes + assert "trace_marker" not in metadata, attributes + assert "litellm.metadata.trace_marker" not in attributes, attributes + _assert_output_tool_call(attributes, marker) + + +def test_arize_otel_v2_c1_absent_allowlist(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "c1-" + uuid.uuid4().hex + with _rig( + gateway, + tmp_path, + _upstream(marker), + remove_environment=("LITELLM_OTEL_BAGGAGE_METADATA_KEYS",), + disabled_environment=("LITELLM_OTEL_BAGGAGE_METADATA_KEYS",), + ) as rig: + response: Final = _request(rig.proxy, rig.model, marker) + _assert_success_body(response, marker, rig.model) + attributes: Final = _matching_span(rig.destination, marker) + _assert_default_allowlist_attributes(attributes, marker) + + +def test_arize_otel_v2_c2_empty_allowlist(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "c2-" + uuid.uuid4().hex + with _rig(gateway, tmp_path, _upstream(marker), environment={"LITELLM_OTEL_BAGGAGE_METADATA_KEYS": ""}) as rig: + response: Final = _request(rig.proxy, rig.model, marker) + _assert_success_body(response, marker, rig.model) + attributes: Final = _matching_span(rig.destination, marker) + assert "metadata" not in attributes, attributes + assert not any(key.startswith("litellm.metadata.") for key in attributes), attributes + _assert_output_tool_call(attributes, marker) + + +def test_arize_otel_v2_c3_absent_allowlisted_key(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "c3-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + return _chat_response(marker) + + with _rig(gateway, tmp_path, upstream) as rig: + response: Final = rig.proxy.request( + "POST", + "/v1/chat/completions", + { + "model": rig.model, + "messages": [{"role": "user", "content": "weather in Paris?"}], + "tools": CHAT_TOOLS, + "tool_choice": {"type": "function", "function": {"name": "lookup_weather"}}, + "metadata": {"other": "value"}, + "cache": {"no-cache": True}, + }, + ) + _assert_success_body(response, marker, rig.model) + attributes: Final = _matching_span(rig.destination, marker) + assert "metadata" not in attributes, attributes + assert "litellm.metadata.trace_marker" not in attributes, attributes + _assert_output_tool_call(attributes, marker) + + +def test_arize_otel_v2_c4_promotes_marker_and_alias(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "c4-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + return _chat_response(marker) + + with ( + _rig( + gateway, + tmp_path, + upstream, + environment={"LITELLM_OTEL_BAGGAGE_METADATA_KEYS": "requester_metadata.trace_marker,user_api_key_alias"}, + ) as rig, + rig.proxy.scenario() as scenario, + ): + key: Final = scenario.key(key_alias="alias-c4") + response: Final = rig.proxy.request( + "POST", + "/v1/chat/completions", + { + "model": rig.model, + "messages": [{"role": "user", "content": "weather in Paris?"}], + "tools": CHAT_TOOLS, + "tool_choice": {"type": "function", "function": {"name": "lookup_weather"}}, + "metadata": {"trace_marker": marker}, + "cache": {"no-cache": True}, + }, + key=key, + ) + _assert_success_body(response, marker, rig.model) + attributes: Final = _matching_marker_span(rig.destination, marker) + assert _json_object(attributes["metadata"].encode()) == { + "trace_marker": marker, + "user_api_key_alias": "alias-c4", + }, attributes + assert attributes["litellm.metadata.trace_marker"] == marker, attributes + assert attributes["litellm.metadata.user_api_key_alias"] == "alias-c4", attributes + + +def test_arize_otel_v2_c5_yaml_allowlist_does_not_reach_preset_so_default_applies( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "c5-" + uuid.uuid4().hex + with _rig( + gateway, + tmp_path, + _upstream(marker), + callback_settings={"otel": {"baggage_metadata_keys": ["requester_metadata.trace_marker"]}}, + remove_environment=("LITELLM_OTEL_BAGGAGE_METADATA_KEYS",), + disabled_environment=("LITELLM_OTEL_BAGGAGE_METADATA_KEYS",), + ) as rig: + response: Final = _request(rig.proxy, rig.model, marker) + _assert_success_body(response, marker, rig.model) + attributes: Final = _matching_marker_span(rig.destination, marker) + _assert_default_allowlist_attributes(attributes, marker) + + +def test_arize_otel_v2_c5_yaml_allowlist_reaches_preset(gateway: Gateway, tmp_path: Path) -> None: + pytest.skip("BUG: LIT-9124 arize preset ignores callback_settings.otel.baggage_metadata_keys from config.yaml") + marker: Final = "c5-allowlist-" + uuid.uuid4().hex + with _rig( + gateway, + tmp_path, + _upstream(marker), + callback_settings={"otel": {"baggage_metadata_keys": ["requester_metadata.trace_marker"]}}, + remove_environment=("LITELLM_OTEL_BAGGAGE_METADATA_KEYS",), + disabled_environment=("LITELLM_OTEL_BAGGAGE_METADATA_KEYS",), + ) as rig: + response: Final = _request(rig.proxy, rig.model, marker) + _assert_success_body(response, marker, rig.model) + attributes: Final = _matching_marker_span(rig.destination, marker) + assert _json_object(attributes["metadata"].encode()) == {"trace_marker": marker}, attributes + assert attributes["litellm.metadata.trace_marker"] == marker, attributes + + +def test_arize_otel_v2_c6_content_capture_disabled(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "c6-" + uuid.uuid4().hex + with _rig( + gateway, + tmp_path, + _upstream(marker), + environment={ + "LITELLM_OTEL_BAGGAGE_METADATA_KEYS": "requester_metadata.trace_marker", + }, + remove_environment=("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT",), + disabled_environment=("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT",), + ) as rig: + response: Final = _request(rig.proxy, rig.model, marker) + _assert_success_body(response, marker, rig.model) + attributes: Final = _matching_marker_span(rig.destination, marker) + assert _json_object(attributes["metadata"].encode()) == {"trace_marker": marker}, attributes + assert attributes["litellm.metadata.trace_marker"] == marker, attributes + assert "gen_ai.input.messages" not in attributes, attributes + assert "gen_ai.output.messages" not in attributes, attributes + assert not any( + key.startswith("llm.input_messages.") or key.startswith("llm.output_messages.") for key in attributes + ), attributes + assert "input.value" not in attributes, attributes + assert "output.value" not in attributes, attributes + assert not any(".tool_calls." in key for key in attributes), attributes + + +def test_arize_otel_v2_c7_key_and_team_logging_callbacks(gateway: Gateway, tmp_path: Path) -> None: + key_marker: Final = "c7-key-" + uuid.uuid4().hex + team_marker: Final = "c7-team-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + marker: Final = _chat_request_marker(request) + assert marker in (key_marker, team_marker), request + _assert_chat_request(request, messages=[{"role": "user", "content": marker}]) + return _chat_response(marker) + + logging_metadata: Final = {"logging": [{"callback_name": "arize", "callback_type": "success"}]} + with _rig(gateway, tmp_path, upstream) as rig, rig.proxy.scenario() as scenario: + key: Final = scenario.key(key_alias="key-c7", metadata=logging_metadata) + team: Final = scenario.team(metadata=logging_metadata) + team_key: Final = scenario.key(team_id=team) + key_response: Final = _request(rig.proxy, rig.model, key_marker, prompt=key_marker, key=key) + _assert_success_body(key_response, key_marker, rig.model) + key_attributes: Final = _matching_marker_span(rig.destination, key_marker) + team_response: Final = _request(rig.proxy, rig.model, team_marker, prompt=team_marker, key=team_key) + _assert_success_body(team_response, team_marker, rig.model) + team_attributes: Final = _matching_marker_span(rig.destination, team_marker) + assert _json_object(key_attributes["metadata"].encode()) == {"trace_marker": key_marker}, key_attributes + assert _json_object(team_attributes["metadata"].encode()) == {"trace_marker": team_marker}, team_attributes + assert key_attributes["litellm.metadata.trace_marker"] == key_marker, key_attributes + assert team_attributes["litellm.metadata.trace_marker"] == team_marker, team_attributes + _assert_output_tool_call(key_attributes, key_marker) + _assert_output_tool_call(team_attributes, team_marker) + + +def test_arize_otel_v2_c8_request_callback_disable(gateway: Gateway, tmp_path: Path) -> None: + control_marker: Final = "c8-control-" + uuid.uuid4().hex + disabled_marker: Final = "c8-disabled-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + marker: Final = _chat_request_marker(request) + assert marker in (control_marker, disabled_marker), marker + return _chat_response(marker) + + with _rig( + gateway, + tmp_path, + upstream, + litellm_settings={"allow_dynamic_callback_disabling": True}, + ) as rig: + control_response: Final = _request(rig.proxy, rig.model, control_marker, prompt=control_marker) + _assert_success_body(control_response, control_marker, rig.model) + control_attributes: Final = _matching_marker_span(rig.destination, control_marker) + _assert_output_tool_call(control_attributes, control_marker) + + disabled_response: Final = _request( + rig.proxy, + rig.model, + disabled_marker, + prompt=disabled_marker, + headers={"x-litellm-disable-callbacks": "arize"}, + ) + _assert_success_body(disabled_response, disabled_marker, rig.model) + + +def test_arize_otel_v2_c8_disabled_callback_exports_no_span(gateway: Gateway, tmp_path: Path) -> None: + pytest.skip("BUG: LIT-9049 x-litellm-disable-callbacks: arize still exports the OTel v2 span") + disabled_marker: Final = "c8-disabled-" + uuid.uuid4().hex + sentinel: Final = "c8-sentinel-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + marker: Final = _chat_request_marker(request) + assert marker in (disabled_marker, sentinel), request + _assert_chat_request(request, messages=[{"role": "user", "content": marker}]) + return _chat_response(marker) + + with _rig( + gateway, + tmp_path, + upstream, + litellm_settings={"allow_dynamic_callback_disabling": True}, + workers=1, + ) as rig: + disabled_response: Final = _request( + rig.proxy, + rig.model, + disabled_marker, + prompt=disabled_marker, + headers={"x-litellm-disable-callbacks": "arize"}, + ) + _assert_success_body(disabled_response, disabled_marker, rig.model) + sentinel_response: Final = _request(rig.proxy, rig.model, sentinel, prompt=sentinel) + _assert_success_body(sentinel_response, sentinel, rig.model) + spans: Final = _llm_spans_through_markers(rig.destination, (sentinel,)) + assert not any(attributes.get("litellm.metadata.trace_marker") == disabled_marker for attributes in spans), ( + spans + ) + + +@pytest.mark.parametrize("failure_status", (401, 500)) +def test_arize_otel_v2_c9_upstream_failures_are_recorded(failure_status: int, gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "c9-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + assert request.method == "POST", request.method + assert request.body, f"{request.method} {request.target}" + body: Final = _json_object(request.body) + assert body == { + "messages": [{"role": "user", "content": "weather in Paris?"}], + "model": "gpt-4o-mini", + }, body + return Reply( + status=failure_status, + body=b'{"error":{"message":"upstream failure"}}', + content_type="application/json", + ) + + with _rig(gateway, tmp_path, upstream) as rig: + response: Final = rig.proxy.request( + "POST", + "/v1/chat/completions", + { + "model": rig.model, + "messages": [{"role": "user", "content": "weather in Paris?"}], + "metadata": {"trace_marker": marker, "failure_status": failure_status}, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == failure_status, response.text + error_body: Final = _json_object(response.content) + error_name: Final = {401: "AuthenticationError", 500: "InternalServerError"}[failure_status] + error_type: Final = {401: "authentication_error", 500: "internal_server_error"}[failure_status] + provider_message: Final = f"litellm.{error_name}: {error_name}: OpenAIException - upstream failure" + caller_message: Final = ( + f"{provider_message}\n\nLiteLLM: model group '{rig.model}' failed with the error above. " + "No fallback was attempted." + ) + assert error_body == { + "error": { + "message": caller_message, + "type": error_type, + "param": None, + "code": str(failure_status), + } + }, error_body + attributes: Final = _matching_marker_span(rig.destination, marker) + metadata: Final = attributes.get("metadata") + assert metadata is not None, attributes + assert _json_object(metadata.encode()) == {"trace_marker": marker}, attributes + assert attributes["litellm.metadata.trace_marker"] == marker, attributes + assert attributes["error.message"] == provider_message, attributes + assert attributes["error.type"] == error_name, attributes + assert not any(".tool_calls." in key for key in attributes), attributes + + +def test_arize_otel_v2_c10_attribute_limit_keeps_tool_prefix(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "c10-" + uuid.uuid4().hex + calls: Final = _response_tool_calls(marker, ("Paris", "Berlin", "Rome", "Tokyo", "Oslo", "Lima", "Accra", "Delhi")) + + def upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + return _chat_response(marker, calls) + + with _rig(gateway, tmp_path, upstream, environment={"OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT": "57"}) as rig: + response: Final = _request(rig.proxy, rig.model, marker) + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(_chat_response(marker, calls), rig.model), ( + response.text + ) + attributes: Final = _matching_output_value_span(rig.destination, marker) + indexes: Final = tuple( + int(key.split(".tool_calls.")[1].split(".")[0]) + for key in attributes + if ".tool_calls." in key and key.endswith(".tool_call.id") + ) + assert indexes == (0,), attributes + assert attributes["llm.output_messages.0.message.role"] == "assistant", attributes + fields: Final = ("id", "function.name", "function.arguments") + expected_tool_call_keys: Final = frozenset().union( + *( + frozenset(f"llm.output_messages.0.message.tool_calls.{index}.tool_call.{field}" for field in fields) + for index in indexes + ) + ) + observed_tool_call_keys: Final = frozenset(key for key in attributes if ".tool_calls." in key) + assert observed_tool_call_keys == expected_tool_call_keys, attributes + assert _json_messages(attributes["output.value"])[0]["tool_calls"] == calls, attributes + assert _json_object(attributes["metadata"].encode()) == {"trace_marker": marker}, attributes + assert attributes["litellm.metadata.trace_marker"] == marker, attributes + + history: Final = [{"role": "user", "content": f"history-{index}"} for index in range(40)] + history_marker: Final = marker + "-history" + + def history_upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=history, include_tools=False) + return _chat_plain_response(history_marker, "history retained") + + with _rig(gateway, tmp_path, history_upstream) as history_rig: + history_response: Final = history_rig.proxy.request( + "POST", + "/v1/chat/completions", + { + "model": history_rig.model, + "messages": history, + "metadata": {"trace_marker": history_marker}, + "cache": {"no-cache": True}, + }, + ) + assert history_response.status_code == 200, history_response.text + assert _json_object(history_response.content) == _chat_caller_response( + _chat_plain_response(history_marker, "history retained"), history_rig.model + ), history_response.text + history_attributes: Final = _matching_marker_span(history_rig.destination, history_marker) + assert _json_object(history_attributes["metadata"].encode()) == {"trace_marker": history_marker}, ( + history_attributes + ) + assert history_attributes["litellm.metadata.trace_marker"] == history_marker, history_attributes + assert tuple(history_attributes[f"llm.input_messages.{index}.message.role"] for index in range(40)) == tuple( + str(message["role"]) for message in history + ) + assert tuple(history_attributes[f"llm.input_messages.{index}.message.content"] for index in range(40)) == tuple( + str(message["content"]) for message in history + ) + assert _json_messages(history_attributes["output.value"]) == [ + {"role": "assistant", "content": "history retained"} + ], history_attributes diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_sad_edge.py b/tests/integration/observability/test_arize_otel_v2_openinference_sad_edge.py new file mode 100644 index 00000000000..a05b3f0f045 --- /dev/null +++ b/tests/integration/observability/test_arize_otel_v2_openinference_sad_edge.py @@ -0,0 +1,547 @@ +from __future__ import annotations + +import json +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import httpx +import pytest +from _openinference_support import ( + CHAT_TOOLS, + _assert_chat_request, + _chat_caller_response, + _chat_output_value, + _chat_request_marker, + _collect_marker_spans, + _json_messages, + _json_object, + _json_object_value, + _matching_marker_span, + _rig, + _span_attributes, + _spans, +) +from integration._support.client import Gateway +from integration._support.wire import Reply, Request +from pydantic import JsonValue + + +def _call( + proxy: Gateway, + model: str, + marker: str, + metadata: JsonValue | None = None, + key: str | None = None, + *, + prompt: str = "weather in Paris?", +) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "tools": CHAT_TOOLS, + "tool_choice": {"type": "function", "function": {"name": "lookup_weather"}}, + "metadata": metadata if metadata is not None else {"trace_marker": marker}, + "cache": {"no-cache": True}, + }, + key=key, + ) + + +def _success(marker: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": marker, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_" + marker, + "type": "function", + "function": {"name": "lookup_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}, + } + ).encode() + ) + + +def _chat_message(body: bytes) -> dict[str, JsonValue]: + response: Final = _json_object(body) + choices: Final = response["choices"] + assert isinstance(choices, list) and len(choices) == 1 and isinstance(choices[0], dict), response + message: Final = choices[0]["message"] + assert isinstance(message, dict), response + return message + + +@pytest.mark.parametrize( + ("value", "expected_trace"), + ( + (7, "7"), + (["one", 2], None), + ("", None), + ("x" * 5000, "x" * 5000), + ({"enabled": True}, None), + ), +) +def test_arize_otel_v2_d1_metadata_value_shapes( + value: JsonValue, expected_trace: str | None, gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "d1-" + uuid.uuid4().hex + metadata: Final = {"trace_marker": value} + + def upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + return _success(marker) + + with _rig(gateway, tmp_path, upstream) as rig: + response: Final = _call(rig.proxy, rig.model, marker, metadata) + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(_success(marker), rig.model), response.text + attributes: Final = _matching_marker_span(rig.destination, marker) + if expected_trace is None: + assert "metadata" not in attributes, attributes + assert "litellm.metadata.trace_marker" not in attributes, attributes + else: + assert json.loads(attributes["metadata"]) == {"trace_marker": expected_trace}, attributes + assert attributes["litellm.metadata.trace_marker"] == expected_trace, attributes + + +def test_arize_otel_v2_d2_duplicate_json_metadata_keys(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "d2-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + return _success(marker) + + with _rig(gateway, tmp_path, upstream) as rig: + request_body: Final = ( + '{"model":"' + + rig.model + + '","messages":' + + json.dumps([{"role": "user", "content": "weather in Paris?"}]) + + ',"tools":' + + json.dumps(CHAT_TOOLS) + + "," + + '"tool_choice":{"type":"function","function":{"name":"lookup_weather"}},' + + '"metadata":{"trace_marker":"' + + marker + + '","trace_marker":"' + + marker + + '"},"cache":{"no-cache":true}}' + ) + response: Final = rig.proxy.client.post( + "/v1/chat/completions", + content=request_body, + headers={ + "authorization": f"Bearer {rig.proxy.key}", + "content-type": "application/json", + }, + ) + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(_success(marker), rig.model), response.text + attributes: Final = _matching_marker_span(rig.destination, marker) + assert _json_object(attributes["metadata"].encode()) == {"trace_marker": marker}, attributes + assert attributes["litellm.metadata.trace_marker"] == marker, attributes + + +def test_arize_otel_v2_d4_unauthenticated_request_has_no_span(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "d4-" + uuid.uuid4().hex + with _rig(gateway, tmp_path, lambda _request: _success(marker)) as rig: + with httpx.Client(base_url=str(rig.proxy.client.base_url), trust_env=False) as client: + response: Final = client.post( + "/v1/chat/completions", + json={"model": rig.model, "messages": [{"role": "user", "content": marker}]}, + ) + assert response.status_code == 401, response.text + body: Final = _json_object(response.content) + assert body == { + "error": { + "message": "Authentication Error, No api key passed in.", + "type": "auth_error", + "param": "None", + "code": "401", + } + }, body + assert rig.provider.received.qsize() == 0 + spans: Final = tuple( + attributes + for attributes in _spans(rig.destination.drain()) + if attributes.get("openinference.span.kind") == "LLM" + ) + assert spans == (), "unauthenticated request exported an LLM span" + + +def test_arize_otel_v2_d5_unknown_model_leaves_proxy_ready(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "d5-" + uuid.uuid4().hex + unknown_model: Final = "unknown-model-" + uuid.uuid4().hex + with _rig(gateway, tmp_path, lambda _request: _success(marker)) as rig: + response: Final = rig.proxy.request( + "POST", + "/v1/chat/completions", + {"model": unknown_model, "messages": [{"role": "user", "content": marker}]}, + ) + assert response.status_code == 400, response.text + body: Final = _json_object(response.content) + error_message: Final = ( + f"/chat/completions: Invalid model name passed in model={unknown_model}. " + "Call `/v1/models` to view available models for your key." + ) + assert body == { + "error": { + "message": error_message, + "type": "invalid_request_error", + "param": None, + "code": "400", + "provider_specific_fields": {"error": error_message}, + } + }, body + assert rig.provider.received.qsize() == 0 + readiness: Final = rig.proxy.client.get("/health/readiness") + assert readiness.status_code == 200, readiness.text + assert _json_object(readiness.content) == {"status": "healthy", "db": "connected"}, readiness.text + spans: Final = tuple( + attributes + for attributes in _spans(rig.destination.drain()) + if attributes.get("openinference.span.kind") == "LLM" + ) + assert spans == (), "unknown model exported an LLM span" + + +def test_arize_otel_v2_d6_sink_rejections_do_not_change_caller_response(gateway: Gateway, tmp_path: Path) -> None: + markers: Final = tuple(f"d6-{status}-" + uuid.uuid4().hex for status in (403, 404)) + unrelated_marker: Final = "d6-unrelated-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + response_marker: Final = _chat_request_marker(request) + _assert_chat_request(request, messages=[{"role": "user", "content": response_marker}]) + return _success(response_marker) + + def sink(request: Request) -> Reply: + exported_markers: Final = tuple( + attributes.get("litellm.metadata.trace_marker") + for attributes in _span_attributes(request) + if "litellm.metadata.trace_marker" in attributes + ) + if markers[0] in exported_markers: + return Reply(status=403, body=b"rejected") + if markers[1] in exported_markers: + return Reply(status=404, body=b"rejected") + return Reply(body=b"", content_type="application/x-protobuf") + + with ( + _rig( + gateway, + tmp_path, + upstream, + destination_handler=sink, + ) as rig, + rig.proxy.scenario() as scenario, + ): + unrelated_key: Final = scenario.key(key_alias="unrelated-d6") + first: Final = _call(rig.proxy, rig.model, markers[0], prompt=markers[0]) + assert first.status_code == 200, first.text + assert _json_object(first.content) == _chat_caller_response(_success(markers[0]), rig.model), first.text + assert _matching_marker_span(rig.destination, markers[0])["litellm.metadata.trace_marker"] == markers[0] + second: Final = _call(rig.proxy, rig.model, markers[1], prompt=markers[1]) + assert second.status_code == 200, second.text + assert _json_object(second.content) == _chat_caller_response(_success(markers[1]), rig.model), second.text + assert _matching_marker_span(rig.destination, markers[1])["litellm.metadata.trace_marker"] == markers[1] + unrelated: Final = _call( + rig.proxy, + rig.model, + unrelated_marker, + key=unrelated_key, + prompt=unrelated_marker, + ) + assert unrelated.status_code == 200, unrelated.text + assert _json_object(unrelated.content) == _chat_caller_response(_success(unrelated_marker), rig.model), ( + unrelated.text + ) + assert ( + _matching_marker_span(rig.destination, unrelated_marker)["litellm.metadata.trace_marker"] + == unrelated_marker + ) + + +def test_arize_otel_v2_d7_missing_space_id_is_stable(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "d7-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + return _success(marker) + + with _rig( + gateway, + tmp_path, + upstream, + remove_environment=("ARIZE_SPACE_ID",), + disabled_environment=("ARIZE_SPACE_ID",), + ) as rig: + response: Final = _call(rig.proxy, rig.model, marker) + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(_success(marker), rig.model), response.text + assert _matching_marker_span(rig.destination, marker)["litellm.metadata.trace_marker"] == marker + + +def test_arize_otel_v2_e1_uncached_request_exports_one_llm_span(gateway: Gateway, tmp_path: Path) -> None: + markers: Final = tuple("e1-" + uuid.uuid4().hex for _ in range(3)) + + def upstream(request: Request) -> Reply: + marker: Final = _chat_request_marker(request) + assert marker in markers, request + _assert_chat_request(request, messages=[{"role": "user", "content": marker}]) + return _success(marker) + + with _rig(gateway, tmp_path, upstream) as rig: + responses: Final = tuple(_call(rig.proxy, rig.model, marker, prompt=marker) for marker in markers) + assert all(response.status_code == 200 for response in responses), tuple( + response.text for response in responses + ) + assert tuple(_json_object(response.content) for response in responses) == tuple( + _chat_caller_response(_success(marker), rig.model) for marker in markers + ), responses + spans: Final = _collect_marker_spans(rig.destination, markers) + assert len(spans) == len(markers), spans + spans_by_id: Final = {span["gen_ai.response.id"]: span for span in spans} + assert frozenset(spans_by_id) == frozenset(markers), spans + assert all( + spans_by_id[marker]["gen_ai.response.id"] in response.text + for marker, response in zip(markers, responses, strict=True) + ), spans + + +def test_arize_otel_v2_e2_concurrent_unique_markers(gateway: Gateway, tmp_path: Path) -> None: + markers: Final = tuple("e2-" + uuid.uuid4().hex for _ in range(20)) + + def upstream(request: Request) -> Reply: + marker: Final = _chat_request_marker(request) + assert marker.startswith("e2-"), request + _assert_chat_request(request, messages=[{"role": "user", "content": marker}]) + return _success(marker) + + with _rig(gateway, tmp_path, upstream) as rig: + with ThreadPoolExecutor(max_workers=20) as executor: + responses: Final = tuple( + executor.map( + lambda marker: _call(rig.proxy, rig.model, marker, prompt=marker), + markers, + ) + ) + assert all(response.status_code == 200 for response in responses), tuple( + response.text for response in responses + ) + assert tuple(_json_object(response.content) for response in responses) == tuple( + _chat_caller_response(_success(marker), rig.model) for marker in markers + ), responses + spans: Final = _collect_marker_spans(rig.destination, markers) + assert len(spans) == len(markers), spans + assert tuple( + _json_object(next(span for span in spans if marker in span.values())["metadata"].encode()) + for marker in markers + ) == tuple({"trace_marker": marker} for marker in markers), spans + assert ( + tuple( + next(span for span in spans if marker in span.values())["litellm.metadata.trace_marker"] + for marker in markers + ) + == markers + ), spans + + +@pytest.mark.parametrize( + "shape", + ("object-arguments", "missing-name", "non-dict-call", "null-tool-calls", "integer-id"), +) +def test_arize_otel_v2_d3_malformed_tool_calls_are_normalized(shape: str, gateway: Gateway, tmp_path: Path) -> None: + marker: Final = f"d3-{shape}-" + uuid.uuid4().hex + + def malformed_reply() -> Reply: + call: Final = { + "id": 17 if shape == "integer-id" else f"call_{marker}", + "type": "function", + "function": { + **({} if shape == "missing-name" else {"name": "lookup_weather"}), + "arguments": {"city": "Paris"} if shape == "object-arguments" else '{"city": "Paris"}', + }, + } + tool_calls: Final[JsonValue] = ( + None if shape == "null-tool-calls" else ["not-a-call"] if shape == "non-dict-call" else [call] + ) + message: Final = { + "role": "assistant", + "content": None, + "tool_calls": tool_calls, + } + return Reply( + body=json.dumps( + { + "id": marker, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "finish_reason": "tool_calls", "message": message}], + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}, + } + ).encode() + ) + + def upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + return malformed_reply() + + with _rig(gateway, tmp_path, upstream) as rig: + proxy_log: Final = rig.owned.log + response: Final = _call(rig.proxy, rig.model, marker) + if shape == "non-dict-call": + assert response.status_code == 400, response.text + error_body: Final = _json_object(response.content) + assert frozenset(error_body) == frozenset({"error"}), error_body + error: Final = _json_object_value(error_body["error"]) + assert frozenset(error) == frozenset({"type", "code", "param", "message"}), error + assert error["type"] == "invalid_request_error", error + assert error["code"] == "400", error + assert error["param"] is None, error + message: Final = error["message"] + assert isinstance(message, str), error + assert "AttributeError: 'str' object has no attribute 'get'" in message, error + spans: Final = tuple(_spans(rig.destination.drain())) + assert all(not any(".tool_calls." in key for key in attributes) for attributes in spans), spans + readiness: Final = rig.proxy.client.get("/health/readiness") + assert readiness.status_code == 200, readiness.text + assert _json_object(readiness.content) == {"status": "healthy", "db": "connected"}, readiness.text + else: + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(malformed_reply(), rig.model), response.text + attributes: Final = _matching_marker_span(rig.destination, marker) + indexed_prefix: Final = "llm.output_messages.0.message.tool_calls.0.tool_call." + expected_fields: Final = ( + ("id", "function.name", "function.arguments"), + ("id", "function.arguments"), + (), + ("function.name", "function.arguments"), + )[("object-arguments", "missing-name", "null-tool-calls", "integer-id").index(shape)] + expected_keys: Final = frozenset(indexed_prefix + field for field in expected_fields) + observed_keys: Final = frozenset(key for key in attributes if ".tool_calls." in key) + assert observed_keys == expected_keys, attributes + if shape == "object-arguments": + assert attributes["llm.output_messages.0.message.tool_calls.0.tool_call.id"] == f"call_{marker}", ( + attributes + ) + assert ( + attributes["llm.output_messages.0.message.tool_calls.0.tool_call.function.name"] == "lookup_weather" + ), attributes + assert ( + attributes["llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments"] + == '{"city": "Paris"}' + ), attributes + elif shape == "missing-name": + assert attributes["llm.output_messages.0.message.tool_calls.0.tool_call.id"] == f"call_{marker}", ( + attributes + ) + assert ( + attributes["llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments"] + == '{"city": "Paris"}' + ), attributes + elif shape == "integer-id": + assert ( + attributes["llm.output_messages.0.message.tool_calls.0.tool_call.function.name"] == "lookup_weather" + ), attributes + assert ( + attributes["llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments"] + == '{"city": "Paris"}' + ), attributes + assert _json_messages(attributes["output.value"]) == _json_messages( + _chat_output_value(malformed_reply()) + ), attributes + assert "Exception while exporting Span batch" not in proxy_log.read_text(), proxy_log.read_text() + + +def test_arize_otel_v2_d3_non_dict_tool_call_is_not_a_caller_error(gateway: Gateway, tmp_path: Path) -> None: + pytest.skip("BUG: LIT-9125 non-dict tool_calls entry returns HTTP 400 with a server traceback") + marker: Final = "d3-non-dict-call-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + return Reply( + body=json.dumps( + { + "id": marker, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": {"role": "assistant", "content": None, "tool_calls": ["not-a-call"]}, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}, + } + ).encode() + ) + + with _rig(gateway, tmp_path, upstream) as rig: + response: Final = _call(rig.proxy, rig.model, marker) + assert response.status_code not in range(400, 500), response.text + assert "Traceback" not in response.text, response.text + assert "AttributeError" not in response.text, response.text + readiness: Final = rig.proxy.client.get("/health/readiness") + assert readiness.status_code == 200, readiness.text + assert _json_object(readiness.content) == {"status": "healthy", "db": "connected"}, readiness.text + + +@pytest.mark.parametrize("shape", ("empty", "null", "missing")) +def test_arize_otel_v2_e3_empty_or_missing_tool_calls_never_indexed( + shape: str, gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = f"e3-{shape}-" + uuid.uuid4().hex + message: Final = ( + {"role": "assistant", "content": None, "tool_calls": []} + if shape == "empty" + else {"role": "assistant", "content": None, "tool_calls": None} + if shape == "null" + else {"role": "assistant", "content": None} + ) + expected_response: Final[dict[str, JsonValue]] = { + "id": marker, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "finish_reason": "stop", "message": message}], + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}, + } + + expected_reply: Final = Reply(body=json.dumps(expected_response).encode()) + + def upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + return expected_reply + + with _rig(gateway, tmp_path, upstream) as rig: + response: Final = _call(rig.proxy, rig.model, marker) + assert response.status_code == 200, response.text + response_body: Final = _json_object(response.content) + assert response_body == _chat_caller_response(expected_reply, rig.model), response.text + attributes: Final = _matching_marker_span(rig.destination, marker) + assert not any(".tool_calls." in key for key in attributes), attributes + assert "tool_calls" not in _json_messages(attributes["output.value"])[0], attributes diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_spans.py b/tests/integration/observability/test_arize_otel_v2_openinference_spans.py new file mode 100644 index 00000000000..d5b4dbeacea --- /dev/null +++ b/tests/integration/observability/test_arize_otel_v2_openinference_spans.py @@ -0,0 +1,966 @@ +from __future__ import annotations + +import asyncio +import json +import uuid +from itertools import chain +from pathlib import Path +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from _openinference_support import ( + CHAT_TOOLS, + RESPONSES_TOOLS, + Rig, + _anthropic_stream_response, + _assert_chat_request, + _assert_messages_request, + _assert_responses_request, + _assert_tool_span, + _chat_cache_hit_caller_stream, + _chat_caller_response, + _chat_caller_stream, + _chat_plain_response, + _chat_request_marker, + _chat_response, + _chat_stream_response, + _chat_tool_call, + _json_messages, + _json_object, + _llm_spans_through_markers, + _matching_marker_span, + _messages_caller_response, + _messages_caller_stream, + _messages_caller_stream_response, + _normalize_chat_caller_stream, + _normalize_responses_caller_body, + _normalize_responses_caller_stream, + _response_tool_calls, + _responses_caller_response, + _responses_caller_stream, + _responses_response, + _responses_stream_response, + _rig, +) +from integration._support.client import Gateway, object_value, string_value +from integration._support.wire import Reply, Request +from pydantic import JsonValue + + +def _chat_upstream(request: Request) -> None: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + + +def _responses_upstream(request: Request, marker: str, *, stream: bool = False) -> None: + _assert_responses_request(request, marker=marker, stream=stream) + + +def _messages_upstream(request: Request, marker: str, *, stream: bool = False) -> None: + _assert_messages_request(request, marker=marker, stream=stream) + + +def _messages_response(identity: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": identity, + "type": "message", + "role": "assistant", + "model": "claude-opus-5-5", + "content": [ + { + "type": "tool_use", + "id": "call_" + identity, + "name": "lookup_weather", + "input": {"city": "Paris"}, + } + ], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 4}, + } + ).encode() + ) + + +def _assert_tool_span_for_marker(attributes: dict[str, str], marker: str, *, content: str | None = None) -> None: + _assert_tool_span( + attributes, + marker=marker, + output=[ + { + "role": "assistant", + "content": content, + "tool_calls": [ + { + "id": "call_" + marker, + "type": "function", + "function": {"name": "lookup_weather", "arguments": '{"city": "Paris"}'}, + } + ], + } + ], + calls=[("call_" + marker, "lookup_weather", {"city": "Paris"})], + metadata={"trace_marker": marker}, + baggage={"trace_marker": marker}, + ) + + +def _chat_request( + proxy: Gateway, model: str, marker: str, *, stream: bool = False, no_cache: bool = True +) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "weather in Paris?"}], + "tools": CHAT_TOOLS, + "tool_choice": {"type": "function", "function": {"name": "lookup_weather"}}, + "metadata": {"trace_marker": marker}, + **({"stream": True} if stream else {}), + **({"cache": {"no-cache": True}} if no_cache else {}), + }, + ) + + +def _responses_request(proxy: Gateway, model: str, marker: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/responses", + { + "model": model, + "input": "weather in Paris?", + "tools": RESPONSES_TOOLS, + "tool_choice": {"type": "function", "name": "lookup_weather"}, + "metadata": {"trace_marker": marker}, + "cache": {"no-cache": True}, + }, + ) + + +def _messages_request(proxy: Gateway, model: str, marker: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": "weather in Paris?"}], + "tools": [ + { + "name": "lookup_weather", + "description": "Get weather", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + } + ], + "tool_choice": {"type": "auto"}, + "metadata": {"trace_marker": marker}, + "cache": {"no-cache": True}, + }, + ) + + +def _openai_client(proxy: Gateway) -> openai.OpenAI: + return openai.OpenAI(base_url=str(proxy.client.base_url) + "/v1", api_key=proxy.key, max_retries=0) + + +def _async_openai_client(proxy: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI(base_url=str(proxy.client.base_url) + "/v1", api_key=proxy.key, max_retries=0) + + +def test_arize_otel_v2_a1_chat_sync_sdk(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "a1-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _chat_upstream(request) + body: Final = _json_object(request.body) + assert body == { + "messages": [{"role": "user", "content": "weather in Paris?"}], + "model": "gpt-4o-mini", + "tool_choice": {"type": "function", "function": {"name": "lookup_weather"}}, + "tools": CHAT_TOOLS, + }, body + return _chat_response(marker) + + with _rig(gateway, tmp_path, upstream) as rig: + client: Final = _openai_client(rig.proxy) + response: Final = client.chat.completions.create( + model=rig.model, + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=CHAT_TOOLS, + tool_choice={"type": "function", "function": {"name": "lookup_weather"}}, + extra_body={"metadata": {"trace_marker": marker}, "cache": {"no-cache": True}}, + ) + assert response.id == marker, response + assert response.model_dump(mode="json", exclude_unset=True) == _chat_caller_response( + _chat_response(marker), rig.model + ), response + calls: Final = response.choices[0].message.tool_calls + assert calls is not None and len(calls) == 1, response + assert (calls[0].id, calls[0].function.name, calls[0].function.arguments) == ( + f"call_{marker}", + "lookup_weather", + '{"city": "Paris"}', + ), response + _assert_tool_span_for_marker(_matching_marker_span(rig.destination, marker), marker) + + +def test_arize_otel_v2_a2_chat_async_sdk(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "a2-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _chat_upstream(request) + return _chat_response(marker) + + async def call() -> None: + with _rig(gateway, tmp_path, upstream) as rig: + client: Final = _async_openai_client(rig.proxy) + response: Final = await client.chat.completions.create( + model=rig.model, + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=CHAT_TOOLS, + tool_choice={"type": "function", "function": {"name": "lookup_weather"}}, + extra_body={"metadata": {"trace_marker": marker}, "cache": {"no-cache": True}}, + ) + assert response.id == marker, response + assert response.model_dump(mode="json", exclude_unset=True) == _chat_caller_response( + _chat_response(marker), rig.model + ), response + calls: Final = response.choices[0].message.tool_calls + assert calls is not None and len(calls) == 1, response + assert (calls[0].id, calls[0].function.name, calls[0].function.arguments) == ( + f"call_{marker}", + "lookup_weather", + '{"city": "Paris"}', + ), response + _assert_tool_span_for_marker(_matching_marker_span(rig.destination, marker), marker) + + asyncio.run(call()) + + +def test_arize_otel_v2_a5_responses_sync_sdk(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "a5-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _responses_upstream(request, marker) + return _responses_response(marker) + + with _rig(gateway, tmp_path, upstream) as rig: + client: Final = _openai_client(rig.proxy) + response: Final = client.responses.create( + model=rig.model, + input="weather in Paris?", + tools=RESPONSES_TOOLS, + tool_choice={"type": "function", "name": "lookup_weather"}, + extra_body={"metadata": {"trace_marker": marker}, "cache": {"no-cache": True}}, + ) + assert response.id.startswith("resp_"), response + assert _normalize_responses_caller_body(response.model_dump(mode="json", exclude_unset=True)) == ( + _responses_caller_response(_responses_response(marker), rig.model) + ), response + assert (response.status, response.model) == ("completed", rig.model), response + assert ( + response.output[0].type, + response.output[0].call_id, + response.output[0].name, + response.output[0].arguments, + ) == ("function_call", f"call_{marker}", "lookup_weather", '{"city": "Paris"}'), response + _assert_tool_span_for_marker(_matching_marker_span(rig.destination, marker), marker) + + +def test_arize_otel_v2_a6_responses_async_streaming_sdk(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "a6-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _responses_upstream(request, marker, stream=True) + return _responses_stream_response(marker, (_chat_tool_call(marker),)) + + async def call() -> None: + with _rig(gateway, tmp_path, upstream) as rig: + client: Final = _async_openai_client(rig.proxy) + stream: Final = await client.responses.create( + model=rig.model, + input="weather in Paris?", + tools=RESPONSES_TOOLS, + tool_choice={"type": "function", "name": "lookup_weather"}, + stream=True, + extra_body={"metadata": {"trace_marker": marker}, "cache": {"no-cache": True}}, + ) + events: Final = tuple([event async for event in stream]) + assert events[-1].type == "response.completed", events + assert events[-1].response.id.startswith("resp_"), events[-1] + assert _normalize_responses_caller_stream( + tuple(event.model_dump(mode="json", exclude_unset=True) for event in events) + ) == _responses_caller_stream(_responses_stream_response(marker, (_chat_tool_call(marker),)), rig.model), ( + events + ) + call: Final = events[-1].response.output[0] + assert (call.type, call.call_id, call.name, call.arguments) == ( + "function_call", + f"call_{marker}", + "lookup_weather", + '{"city": "Paris"}', + ), events[-1] + _assert_tool_span_for_marker(_matching_marker_span(rig.destination, marker), marker) + + asyncio.run(call()) + + +def test_arize_otel_v2_a7_messages_sync_sdk(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "a7-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _messages_upstream(request, marker) + return _messages_response(marker) + + with _rig(gateway, tmp_path, upstream, model_name="anthropic/claude-opus-5-5", api_base_suffix="") as rig: + client: Final = anthropic.Anthropic( + base_url=str(rig.proxy.client.base_url), api_key=rig.proxy.key, max_retries=0 + ) + response: Final = client.messages.create( + model=rig.model, + max_tokens=64, + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=[ + { + "name": "lookup_weather", + "description": "Get weather", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}}, + } + ], + tool_choice={"type": "auto"}, + metadata={"trace_marker": marker}, + ) + assert response.id == marker, response + assert response.model_dump(mode="json", exclude_unset=True) == _messages_caller_response( + _messages_response(marker), rig.model + ), response + call: Final = response.content[0] + assert (call.type, call.id, call.name, call.input) == ( + "tool_use", + f"call_{marker}", + "lookup_weather", + {"city": "Paris"}, + ), response + _assert_tool_span_for_marker(_matching_marker_span(rig.destination, marker), marker) + + +def test_arize_otel_v2_a3_chat_streaming(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "a3-" + uuid.uuid4().hex + stream_reply: Final = _chat_stream_response(marker, (_chat_tool_call(marker),), include_usage=False) + + def upstream(request: Request) -> Reply: + _assert_chat_request( + request, + messages=[{"role": "user", "content": "weather in Paris?"}], + stream=True, + stream_options={"include_usage": False}, + ) + return stream_reply + + with _rig(gateway, tmp_path, upstream, general_settings={"always_include_stream_usage": False}) as rig: + client: Final = _openai_client(rig.proxy) + stream: Final = client.chat.completions.create( + model=rig.model, + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=CHAT_TOOLS, + tool_choice={"type": "function", "function": {"name": "lookup_weather"}}, + stream=True, + stream_options={"include_usage": False}, + extra_body={"metadata": {"trace_marker": marker}, "cache": {"no-cache": True}}, + ) + chunks: Final = tuple(stream) + assert chunks[0].id == marker and chunks[-1].id == marker, chunks + assert _normalize_chat_caller_stream( + tuple(chunk.model_dump(mode="json", exclude_unset=True) for chunk in chunks) + ) == _chat_caller_stream(stream_reply, rig.model), chunks + assert chunks[-1].choices[0].finish_reason == "tool_calls", chunks + tool_call_deltas: Final = tuple( + chain.from_iterable(chunk.choices[0].delta.tool_calls or () for chunk in chunks) + ) + assert len(tool_call_deltas) == 2, chunks + assert ( + tool_call_deltas[0].id, + tool_call_deltas[0].function.name, + tool_call_deltas[1].function.arguments, + ) == (f"call_{marker}", "lookup_weather", '{"city": "Paris"}'), chunks + _assert_tool_span_for_marker(_matching_marker_span(rig.destination, marker), marker) + + +def test_arize_otel_v2_a4_chat_async_streaming(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "a4-" + uuid.uuid4().hex + stream_reply: Final = _chat_stream_response(marker, (_chat_tool_call(marker),), include_usage=False) + + def upstream(request: Request) -> Reply: + _assert_chat_request( + request, + messages=[{"role": "user", "content": "weather in Paris?"}], + stream=True, + stream_options={"include_usage": False}, + ) + return stream_reply + + async def call() -> None: + with _rig(gateway, tmp_path, upstream, general_settings={"always_include_stream_usage": False}) as rig: + client: Final = _async_openai_client(rig.proxy) + stream: Final = await client.chat.completions.create( + model=rig.model, + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=CHAT_TOOLS, + tool_choice={"type": "function", "function": {"name": "lookup_weather"}}, + stream=True, + stream_options={"include_usage": False}, + extra_body={"metadata": {"trace_marker": marker}, "cache": {"no-cache": True}}, + ) + chunks: Final = tuple([chunk async for chunk in stream]) + assert chunks[0].id == marker and chunks[-1].id == marker, chunks + assert _normalize_chat_caller_stream( + tuple(chunk.model_dump(mode="json", exclude_unset=True) for chunk in chunks) + ) == _chat_caller_stream(stream_reply, rig.model), chunks + assert chunks[-1].choices[0].finish_reason == "tool_calls", chunks + tool_call_deltas: Final = tuple( + chain.from_iterable(chunk.choices[0].delta.tool_calls or () for chunk in chunks) + ) + assert len(tool_call_deltas) == 2, chunks + assert ( + tool_call_deltas[0].id, + tool_call_deltas[0].function.name, + tool_call_deltas[1].function.arguments, + ) == (f"call_{marker}", "lookup_weather", '{"city": "Paris"}'), chunks + _assert_tool_span_for_marker(_matching_marker_span(rig.destination, marker), marker) + + asyncio.run(call()) + + +def test_arize_otel_v2_a8_messages_streaming(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "a8-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _messages_upstream(request, marker, stream=True) + return _anthropic_stream_response(marker) + + async def call() -> None: + with _rig(gateway, tmp_path, upstream, model_name="anthropic/claude-opus-5-5", api_base_suffix="") as rig: + client: Final = anthropic.AsyncAnthropic( + base_url=str(rig.proxy.client.base_url), api_key=rig.proxy.key, max_retries=0 + ) + async with client.messages.stream( + model=rig.model, + max_tokens=64, + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=[ + { + "name": "lookup_weather", + "description": "Get weather", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}}, + } + ], + tool_choice={"type": "auto"}, + metadata={"trace_marker": marker}, + ) as stream: + events: Final = tuple([event async for event in stream]) + response: Final = await stream.get_final_message() + assert events[-1].type == "message_stop", events + assert response.id == marker, response + assert response.model_dump(mode="json", exclude_unset=True) == _messages_caller_stream_response( + _messages_response(marker), rig.model + ), response + assert tuple(event.model_dump(mode="json", exclude_unset=True) for event in events) == ( + _messages_caller_stream( + _anthropic_stream_response(marker), + rig.model, + final_message=_messages_caller_stream_response(_messages_response(marker), rig.model), + ) + ), events + call: Final = response.content[0] + assert (call.type, call.id, call.name, call.input) == ( + "tool_use", + f"call_{marker}", + "lookup_weather", + {"city": "Paris"}, + ), response + _assert_tool_span_for_marker(_matching_marker_span(rig.destination, marker), marker, content="") + + asyncio.run(call()) + + +def test_arize_otel_v2_llm_span_carries_openinference_tool_calls_and_metadata(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "a9-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _chat_upstream(request) + return _chat_response(marker) + + with _rig(gateway, tmp_path, upstream) as rig: + response: Final = _chat_request(rig.proxy, rig.model, marker) + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(_chat_response(marker), rig.model), response.text + attributes: Final = _matching_marker_span(rig.destination, marker) + _assert_tool_span_for_marker(attributes, marker) + + +def test_arize_otel_v2_responses_span_carries_openinference_tool_calls_and_metadata( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "a10-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _responses_upstream(request, marker) + return _responses_response(marker) + + with _rig(gateway, tmp_path, upstream) as rig: + response: Final = _responses_request(rig.proxy, rig.model, marker) + assert response.status_code == 200, response.text + response_body: Final = _json_object(response.content) + assert _normalize_responses_caller_body(response_body) == _responses_caller_response( + _responses_response(marker), rig.model + ), response.text + _assert_tool_span_for_marker(_matching_marker_span(rig.destination, marker), marker) + + +def test_arize_otel_v2_a11_parallel_output_tool_calls(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "a11-" + uuid.uuid4().hex + calls: Final = _response_tool_calls(marker, ("Paris", "Berlin")) + + def upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=[{"role": "user", "content": "weather in Paris?"}]) + return _chat_response(marker, calls) + + with _rig(gateway, tmp_path, upstream) as rig: + response: Final = _chat_request(rig.proxy, rig.model, marker) + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(_chat_response(marker, calls), rig.model), ( + response.text + ) + attributes: Final = _matching_marker_span(rig.destination, marker) + expected_calls: Final = ( + ("call_" + marker + "-Paris", "lookup_weather", {"city": "Paris"}), + ("call_" + marker + "-Berlin", "lookup_weather", {"city": "Berlin"}), + ) + _assert_tool_span( + attributes, + marker=marker, + output=[{"role": "assistant", "content": None, "tool_calls": calls}], + calls=expected_calls, + metadata={"trace_marker": marker}, + baggage={"trace_marker": marker}, + ) + + +def test_arize_otel_v2_a12_plain_text_has_metadata_without_tool_calls(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "a12-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + _assert_chat_request( + request, + messages=[{"role": "user", "content": "weather in Paris?"}], + include_tools=False, + ) + return _chat_plain_response(marker, "The weather is clear") + + with _rig(gateway, tmp_path, upstream) as rig: + response: Final = rig.proxy.request( + "POST", + "/v1/chat/completions", + { + "model": rig.model, + "messages": [{"role": "user", "content": "weather in Paris?"}], + "metadata": {"trace_marker": marker}, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response( + _chat_plain_response(marker, "The weather is clear"), rig.model + ), response.text + attributes: Final = _matching_marker_span(rig.destination, marker) + assert _json_object(attributes["metadata"].encode()) == {"trace_marker": marker}, attributes + assert attributes["litellm.metadata.trace_marker"] == marker, attributes + assert not any(".tool_calls." in key for key in attributes), attributes + assert _json_messages(attributes["output.value"]) == [ + {"role": "assistant", "content": "The weather is clear"} + ], attributes + + +def test_arize_otel_v2_a13_multiturn_input_tool_calls(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "a13-" + uuid.uuid4().hex + messages: Final = [ + {"role": "user", "content": "weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call-prior", + "type": "function", + "function": {"name": "lookup_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call-prior", "content": '{"temperature": 20}'}, + ] + upstream_messages: Final = [ + {key: value for key, value in message.items() if value is not None} for message in messages + ] + + def upstream(request: Request) -> Reply: + _assert_chat_request(request, messages=upstream_messages) + return _chat_response(marker) + + with _rig(gateway, tmp_path, upstream) as rig: + response: Final = rig.proxy.request( + "POST", + "/v1/chat/completions", + { + "model": rig.model, + "messages": messages, + "tools": CHAT_TOOLS, + "tool_choice": {"type": "function", "function": {"name": "lookup_weather"}}, + "metadata": {"trace_marker": marker}, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(_chat_response(marker), rig.model), response.text + attributes: Final = _matching_marker_span(rig.destination, marker) + _assert_tool_span_for_marker(attributes, marker) + assert not any(key.startswith("llm.input_messages.") and ".tool_calls." in key for key in attributes), ( + attributes + ) + assert tuple(attributes[f"llm.input_messages.{index}.message.role"] for index in range(3)) == ( + "user", + "assistant", + "tool", + ), attributes + assert tuple(attributes[f"llm.input_messages.{index}.message.content"] for index in (0, 2)) == ( + "weather in Paris?", + '{"temperature": 20}', + ), attributes + expected_input_value: Final = [ + messages[0], + messages[1], + {"role": "tool", "content": '{"temperature": 20}'}, + ] + assert _json_messages(attributes["input.value"]) == expected_input_value, attributes + + +def test_arize_otel_v2_a14_two_choices_each_with_tool_calls(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "a14-" + uuid.uuid4().hex + first_call: Final = _response_tool_calls(marker, ("Paris",))[0] + second_call: Final = _response_tool_calls(marker, ("Berlin",))[0] + + def choice(index: int, call: dict[str, JsonValue]) -> dict[str, JsonValue]: + return { + "index": index, + "finish_reason": "tool_calls", + "message": {"role": "assistant", "content": None, "tool_calls": [call]}, + } + + expected_response: Final = Reply( + body=json.dumps( + { + "id": marker, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [choice(0, first_call), choice(1, second_call)], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + ).encode() + ) + + def upstream(request: Request) -> Reply: + _assert_chat_request( + request, + messages=[{"role": "user", "content": "weather in Paris and Berlin?"}], + n=2, + ) + return expected_response + + with _rig(gateway, tmp_path, upstream) as rig: + response: Final = rig.proxy.request( + "POST", + "/v1/chat/completions", + { + "model": rig.model, + "messages": [{"role": "user", "content": "weather in Paris and Berlin?"}], + "tools": CHAT_TOOLS, + "tool_choice": {"type": "function", "function": {"name": "lookup_weather"}}, + "n": 2, + "metadata": {"trace_marker": marker}, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + assert _json_object(response.content) == _chat_caller_response(expected_response, rig.model), response.text + attributes: Final = _matching_marker_span(rig.destination, marker) + assert _json_messages(attributes["output.value"]) == [ + {"role": "assistant", "content": None, "tool_calls": [first_call]}, + {"role": "assistant", "content": None, "tool_calls": [second_call]}, + ], attributes + assert attributes["llm.output_messages.0.message.tool_calls.0.tool_call.id"] == str(first_call["id"]), ( + attributes + ) + assert attributes["llm.output_messages.1.message.tool_calls.0.tool_call.id"] == str(second_call["id"]), ( + attributes + ) + assert attributes["llm.output_messages.0.message.tool_calls.0.tool_call.function.name"] == "lookup_weather", ( + attributes + ) + assert ( + attributes["llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments"] == '{"city": "Paris"}' + ), attributes + assert attributes["llm.output_messages.1.message.tool_calls.0.tool_call.function.name"] == "lookup_weather", ( + attributes + ) + assert ( + attributes["llm.output_messages.1.message.tool_calls.0.tool_call.function.arguments"] + == '{"city": "Berlin"}' + ), attributes + assert _json_object(attributes["metadata"].encode()) == {"trace_marker": marker}, attributes + assert attributes["litellm.metadata.trace_marker"] == marker, attributes + + +def _cache_call( + rig: Rig, surface: str, marker: str, *, cache_hit: bool = False +) -> tuple[tuple[str, str, str], str, httpx.Headers]: + match surface: + case "chat": + client: Final = _openai_client(rig.proxy) + raw: Final = client.chat.completions.with_raw_response.create( + model=rig.model, + messages=[{"role": "user", "content": marker}], + tools=CHAT_TOOLS, + tool_choice={"type": "function", "function": {"name": "lookup_weather"}}, + extra_body={"metadata": {"trace_marker": marker}, "cache": {"no-cache": False}}, + ) + response: Final = raw.parse() + assert response.model_dump(mode="json", exclude_unset=True) == _chat_caller_response( + _chat_response(marker), rig.model + ), response + assert response.choices[0].message.tool_calls is not None, response + call: Final = response.choices[0].message.tool_calls[0] + return (call.id, call.function.name, call.function.arguments), response.id, raw.headers + case "chat-stream": + client: Final = _openai_client(rig.proxy) + with client.chat.completions.with_streaming_response.create( + model=rig.model, + messages=[{"role": "user", "content": marker}], + tools=CHAT_TOOLS, + tool_choice={"type": "function", "function": {"name": "lookup_weather"}}, + stream=True, + stream_options={"include_usage": False}, + extra_body={"metadata": {"trace_marker": marker}, "cache": {"no-cache": False}}, + ) as raw: + chunks: Final = tuple(raw.parse()) + expected_chunks: Final = ( + _chat_cache_hit_caller_stream(marker, rig.model, (_chat_tool_call(marker),)) + if cache_hit + else _chat_caller_stream( + _chat_stream_response(marker, (_chat_tool_call(marker),), include_usage=False), rig.model + ) + ) + assert ( + _normalize_chat_caller_stream( + tuple(chunk.model_dump(mode="json", exclude_unset=True) for chunk in chunks) + ) + == expected_chunks + ), chunks + calls: Final = tuple(chain.from_iterable(chunk.choices[0].delta.tool_calls or () for chunk in chunks)) + if cache_hit: + assert len(calls) == 1, chunks + call: Final = calls[0] + assert ( + call.id is not None and call.function.name is not None and call.function.arguments is not None + ), chunks + return ( + (call.id, call.function.name, call.function.arguments), + chunks[0].id, + raw.headers, + ) + assert len(calls) == 2, chunks + call: Final = calls[0] + arguments_call: Final = calls[1] + assert ( + call.id is not None + and call.function.name is not None + and arguments_call.function.arguments is not None + ), chunks + return ( + (call.id, call.function.name, arguments_call.function.arguments), + chunks[0].id, + raw.headers, + ) + case "responses": + client: Final = _openai_client(rig.proxy) + raw: Final = client.responses.with_raw_response.create( + model=rig.model, + input=marker, + tools=RESPONSES_TOOLS, + tool_choice={"type": "function", "name": "lookup_weather"}, + extra_body={"metadata": {"trace_marker": marker}, "cache": {"no-cache": False}}, + ) + response: Final = raw.parse() + assert _normalize_responses_caller_body(response.model_dump(mode="json", exclude_unset=True)) == ( + _responses_caller_response(_responses_response(marker), rig.model) + ), response + call: Final = response.output[0] + assert call.type == "function_call", response + return (call.call_id, call.name, call.arguments), response.id, raw.headers + case "messages": + client: Final = anthropic.Anthropic( + base_url=str(rig.proxy.client.base_url), api_key=rig.proxy.key, max_retries=0 + ) + raw: Final = client.messages.with_raw_response.create( + model=rig.model, + max_tokens=64, + messages=[{"role": "user", "content": marker}], + tools=[ + { + "name": "lookup_weather", + "description": "Get weather", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}}, + } + ], + tool_choice={"type": "auto"}, + extra_body={"cache": {"no-cache": False}, "metadata": {"trace_marker": marker}}, + ) + response: Final = raw.parse() + assert response.model_dump(mode="json", exclude_unset=True) == _messages_caller_response( + _messages_response(marker), rig.model + ), response + call: Final = response.content[0] + assert call.type == "tool_use", response + return (call.id, call.name, json.dumps(call.input)), response.id, raw.headers + case _: + raise AssertionError(f"Unknown cache surface: {surface}") + + +@pytest.mark.parametrize("surface", ("chat", "chat-stream", "responses", "messages")) +def test_arize_otel_v2_a_cache(surface: str, gateway: Gateway, tmp_path: Path) -> None: + marker: Final = f"a-cache-{surface}-" + uuid.uuid4().hex + sentinel: Final = f"{marker}-sentinel" + + def upstream(request: Request) -> Reply: + body: Final = _json_object(request.body) + request_marker: Final = ( + _chat_request_marker(request) + if surface in ("chat", "chat-stream", "messages") + else string_value(object_value(body["metadata"])["trace_marker"]) + ) + assert request_marker in (marker, sentinel), request + if surface in ("chat", "chat-stream"): + _assert_chat_request( + request, + messages=[{"role": "user", "content": request_marker}], + stream=True if surface == "chat-stream" else None, + stream_options={"include_usage": False} if surface == "chat-stream" else None, + ) + return ( + _chat_stream_response(request_marker, (_chat_tool_call(request_marker),), include_usage=False) + if surface == "chat-stream" + else _chat_response(request_marker) + ) + if surface == "responses": + assert body.get("metadata") == {"trace_marker": request_marker}, body + _assert_responses_request(request, marker=request_marker, input_value=request_marker) + return _responses_response(request_marker) + _assert_messages_request(request, marker=request_marker, prompt=request_marker) + return _messages_response(request_marker) + + with _rig( + gateway, + tmp_path, + upstream, + general_settings={"always_include_stream_usage": False} if surface == "chat-stream" else None, + model_name="anthropic/claude-opus-5-5" if surface == "messages" else "gpt-4o-mini", + api_base_suffix="" if surface == "messages" else "/v1", + workers=1, + ) as rig: + expected_arguments: Final = '{"city": "Paris"}' + first: Final = _cache_call(rig, surface, marker) + assert first[0] == (f"call_{marker}", "lookup_weather", expected_arguments), first + assert first[1].startswith("resp_") if surface == "responses" else first[1] == marker, first + assert not first[2].get("x-litellm-cache-key"), first[2] + first_span: Final = _matching_marker_span(rig.destination, marker) + _assert_tool_span_for_marker(first_span, marker) + second: Final = _cache_call(rig, surface, marker, cache_hit=True) + assert second[0] == first[0], second + assert second[1].startswith("resp_") if surface == "responses" else second[1] == first[1], second + forwarded: Final = tuple( + request for request in rig.provider.drain() if request.method == "POST" and marker.encode() in request.body + ) + assert len(forwarded) == 1, forwarded + if surface == "messages": + assert not second[2].get("x-litellm-cache-key"), second[2] + else: + assert second[2].get("x-litellm-cache-key"), second[2] + forwarded_body: Final = _json_object(forwarded[0].body) + if surface in ("chat", "chat-stream"): + assert "metadata" not in forwarded_body, forwarded[0] + elif surface == "responses": + assert forwarded_body["metadata"] == {"trace_marker": marker}, forwarded[0] + else: + assert forwarded_body["metadata"] == {}, forwarded[0] + sentinel_response: Final = _cache_call(rig, surface, sentinel) + assert not sentinel_response[2].get("x-litellm-cache-key"), sentinel_response[2] + sentinel_forwarded: Final = tuple( + request + for request in rig.provider.drain() + if request.method == "POST" and sentinel.encode() in request.body + ) + assert len(sentinel_forwarded) == 1, sentinel_forwarded + spans: Final = _llm_spans_through_markers(rig.destination, (sentinel,)) + sentinel_spans: Final = tuple( + attributes for attributes in spans if attributes.get("litellm.metadata.trace_marker") == sentinel + ) + assert len(sentinel_spans) == 1, spans + _assert_tool_span_for_marker(sentinel_spans[0], sentinel) + assert not any(attributes.get("litellm.metadata.trace_marker") == marker for attributes in spans), spans + + +def test_arize_otel_v2_cache_hit_exports_llm_span(gateway: Gateway, tmp_path: Path) -> None: + pytest.skip("BUG: LIT-9127 response-cache hits emit no OTel v2 LLM span") + marker: Final = "a-cache-hit-" + uuid.uuid4().hex + sentinel: Final = f"{marker}-sentinel" + + def upstream(request: Request) -> Reply: + request_marker: Final = _chat_request_marker(request) + assert request_marker in (marker, sentinel), request + _assert_chat_request(request, messages=[{"role": "user", "content": request_marker}]) + return _chat_response(request_marker) + + with _rig(gateway, tmp_path, upstream, workers=1) as rig: + first: Final = _cache_call(rig, "chat", marker) + assert first[0] == (f"call_{marker}", "lookup_weather", '{"city": "Paris"}'), first + assert first[1] == marker, first + assert not first[2].get("x-litellm-cache-key"), first[2] + second: Final = _cache_call(rig, "chat", marker, cache_hit=True) + assert second[0] == first[0], second + assert second[1] == first[1], second + assert second[2].get("x-litellm-cache-key"), second[2] + forwarded: Final = tuple( + request for request in rig.provider.drain() if request.method == "POST" and marker.encode() in request.body + ) + assert len(forwarded) == 1, forwarded + sentinel_response: Final = _cache_call(rig, "chat", sentinel) + assert not sentinel_response[2].get("x-litellm-cache-key"), sentinel_response[2] + sentinel_forwarded: Final = tuple( + request + for request in rig.provider.drain() + if request.method == "POST" and sentinel.encode() in request.body + ) + assert len(sentinel_forwarded) == 1, sentinel_forwarded + spans: Final = _llm_spans_through_markers(rig.destination, (sentinel,)) + marker_spans: Final = tuple( + attributes for attributes in spans if attributes.get("litellm.metadata.trace_marker") == marker + ) + assert len(marker_spans) == 2, spans + _assert_tool_span_for_marker(marker_spans[1], marker) diff --git a/tests/integration/observability/test_guardrail_attachment.py b/tests/integration/observability/test_guardrail_attachment.py new file mode 100644 index 00000000000..f291b3877f3 --- /dev/null +++ b/tests/integration/observability/test_guardrail_attachment.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import json +import shutil +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml + +from tests.integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment, object_value, string_value +from tests.integration._support.process import owned_proxy +from tests.integration._support.wire import Reply, Request, Wire, wire_server + +_ATTACHABLE: Final = "attachable-during-guard" +_WORDS: Final = "custom-words-during-guard" +_HEADER: Final = "x-litellm-applied-guardrails" + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + policy: Wire + upstream: Wire + model: str + + +def _completion(request: Request) -> Reply: + assert request.target == "/v1/chat/completions", request.target + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "guarded"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } + ).encode() + ) + + +def _allow(request: Request) -> Reply: + assert request.target == "/beta/litellm_basic_guardrail_api", request.target + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +def _config(directory: Path, policy: Wire) -> Path: + shutil.copy(Path("litellm/proxy/example_config_yaml/custom_guardrail.py"), directory / "custom_guardrail.py") + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["guardrails"] = [ + { + "guardrail_name": _ATTACHABLE, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "during_call", + "api_base": policy.url, + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": _WORDS, + "litellm_params": {"guardrail": "custom_guardrail.myCustomGuardrail", "mode": "during_call"}, + }, + ] + path: Final = directory / "guardrail-attachment.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("guardrail-attachment") + with ( + wire_server(_allow) as policy, + wire_server(_completion) as upstream, + gateway_from_environment() as shared, + owned_proxy(shared, directory, {}, config=_config(directory, policy)) as owned, + owned.scenario() as scenario, + ): + yield _Rig(owned, policy, upstream, scenario.model(api_base=f"{upstream.url}/v1", api_key="sk-fixture")) + + +def _prompt(text: str) -> str: + return f"{text} {uuid.uuid4().hex}" + + +def _ask(rig: _Rig, key: str | None, prompt: str, guardrails: list[str] | None = None) -> httpx.Response: + body: Final = {"model": rig.model, "messages": [{"role": "user", "content": prompt}]} + return rig.gateway.request( + "POST", "/v1/chat/completions", body if guardrails is None else {**body, "guardrails": guardrails}, key=key + ) + + +def _served_without_guardrail(rig: _Rig, response: httpx.Response, prompt: str) -> None: + assert response.status_code == 200, response.text + assert _HEADER not in response.headers, dict(response.headers) + assert rig.policy.drain() == () + forwarded: Final = rig.upstream.drain() + assert len(forwarded) == 1 and prompt in forwarded[0].body.decode(), forwarded + + +def _served_with_attachable(rig: _Rig, response: httpx.Response, prompt: str) -> None: + assert response.status_code == 200, response.text + assert response.headers[_HEADER] == _ATTACHABLE, dict(response.headers) + inspected: Final = rig.policy.drain() + assert len(inspected) == 1 and prompt in inspected[0].body.decode(), inspected + forwarded: Final = rig.upstream.drain() + assert len(forwarded) == 1 and prompt in forwarded[0].body.decode(), forwarded + + +def test_a_request_with_an_empty_guardrail_list_is_served_without_the_applied_header(rig: _Rig) -> None: + prompt: Final = _prompt("no guardrails") + _served_without_guardrail(rig, _ask(rig, None, prompt, []), prompt) + + +def test_a_key_carrying_a_guardrail_applies_it_and_a_plain_key_does_not(rig: _Rig) -> None: + with rig.gateway.scenario() as scenario: + plain: Final = scenario.key() + guarded: Final = scenario.key(guardrails=[_ATTACHABLE]) + plain_prompt: Final = _prompt("plain key") + _served_without_guardrail(rig, _ask(rig, plain, plain_prompt), plain_prompt) + guarded_prompt: Final = _prompt("guarded key") + _served_with_attachable(rig, _ask(rig, guarded, guarded_prompt), guarded_prompt) + + +def test_a_team_carrying_a_guardrail_applies_it_to_its_keys_only(rig: _Rig) -> None: + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(guardrails=[_ATTACHABLE]) + outside: Final = scenario.key() + member: Final = scenario.key(team_id=team) + outside_prompt: Final = _prompt("outside team") + _served_without_guardrail(rig, _ask(rig, outside, outside_prompt), outside_prompt) + member_prompt: Final = _prompt("team key") + _served_with_attachable(rig, _ask(rig, member, member_prompt), member_prompt) + + +def test_a_during_call_custom_guardrail_rejects_a_request_naming_the_banned_word(rig: _Rig) -> None: + unguarded_prompt: Final = _prompt("what is litellm") + _served_without_guardrail(rig, _ask(rig, None, unguarded_prompt), unguarded_prompt) + refused: Final = _ask(rig, None, _prompt("what is litellm"), [_WORDS]) + rig.upstream.drain() + assert refused.status_code >= 400, refused.text + error: Final = object_value(JSON_OBJECT.validate_json(refused.content)["error"]) + assert "Guardrail failed words - `litellm` detected" in string_value(error["message"]), refused.text + assert rig.policy.drain() == () diff --git a/tests/integration/observability/test_guardrail_stream_scope.py b/tests/integration/observability/test_guardrail_stream_scope.py new file mode 100644 index 00000000000..8b832f07aa4 --- /dev/null +++ b/tests/integration/observability/test_guardrail_stream_scope.py @@ -0,0 +1,856 @@ +from __future__ import annotations + +import base64 +import json +import os +import uuid +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, JsonValue, Scenario, eventually +from integration._support.database import read_rows, write_rows +from integration._support.process import owned_proxy_process +from integration._support.upstream import ( + _aws_event_frame, # pyright: ignore[reportPrivateUsage] # project Bedrock event-stream encoder +) +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import TypeAdapter + +import litellm + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +EndpointScope: TypeAlias = Literal["streaming", "non_streaming"] +ModelKind: TypeAlias = Literal["router", "direct"] +StreamingAction: TypeAlias = Literal["converse-stream", "invoke-with-response-stream"] +NonStreamingAction: TypeAlias = Literal["converse", "invoke"] +BedrockAction: TypeAlias = StreamingAction | NonStreamingAction +BEDROCK_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" +BEDROCK_FALSE_POSITIVE_MODEL_ID: Final = "anthropic.claude-converse-stream-test-v1:0" +BEDROCK_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +GUARDRAIL_PATH: Final = "/beta/litellm_basic_guardrail_api" +SERVER_STREAMING_CLASSIFICATION_KEY: Final = "litellm_server_streaming_classification" +STREAMING_ACTIONS: Final[tuple[StreamingAction, ...]] = ( + "converse-stream", + "invoke-with-response-stream", +) +NON_STREAMING_ACTIONS: Final[tuple[NonStreamingAction, ...]] = ("converse", "invoke") +SCOPES: Final[tuple[EndpointScope, ...]] = ("streaming", "non_streaming") +HOSTILE_CLASSIFICATION_CASES: Final = ( + pytest.param("is_streaming_request", True, id="boolean-marker"), + pytest.param("is_streaming_request", "litellm-server-streaming", id="server-marker-string"), + pytest.param("litellm_server_streaming_classification", True, id="classification-field"), +) +WORKTREE: Final = Path(__file__).resolve().parents[3] +LITELLM_PATH: Final = Path(litellm.__file__).resolve() +assert LITELLM_PATH.is_relative_to(WORKTREE), (LITELLM_PATH, WORKTREE) +print(f"stream_scope repro litellm import: {LITELLM_PATH}") # noqa: T201 # required worktree evidence + + +def _json(value: object) -> bytes: + return json.dumps(value, separators=(",", ":")).encode() + + +def _strings(value: JsonValue) -> tuple[str, ...]: + if isinstance(value, str): + return (value,) + if isinstance(value, list): + return tuple(chain.from_iterable(_strings(item) for item in value)) + if isinstance(value, dict): + return tuple(chain.from_iterable(_strings(item) for item in value.values())) + return () + + +def _key_names(value: JsonValue) -> tuple[str, ...]: + if isinstance(value, dict): + return tuple(value) + tuple(chain.from_iterable(_key_names(item) for item in value.values())) + if isinstance(value, list): + return tuple(chain.from_iterable(_key_names(item) for item in value)) + return () + + +def _marker(body: JsonValue) -> str: + return next(value for value in _strings(body) if value.startswith("scope-")) + + +def _bedrock_converse_stream(marker: str) -> bytes: + return b"".join( + _aws_event_frame(event_type, payload, marker, marker) + for event_type, payload in ( + ("messageStart", {"role": "assistant"}), + ( + "contentBlockDelta", + {"delta": {"text": f"scripted Bedrock reply {marker}"}, "contentBlockIndex": 0}, + ), + ("messageStop", {"stopReason": "end_turn"}), + ("metadata", {"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}}), + ) + ) + + +def _invoke_chunk(payload: Mapping[str, JsonValue], marker: str) -> bytes: + encoded: Final = base64.b64encode(_json(payload)).decode() + return _aws_event_frame("chunk", {"bytes": encoded}, marker, marker) + + +def _bedrock_invoke_stream(marker: str) -> bytes: + events: Final = ( + { + "type": "message_start", + "message": { + "id": f"msg-{marker}", + "type": "message", + "role": "assistant", + "model": BEDROCK_MODEL_ID, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 0}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": f"scripted Bedrock reply {marker}"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"input_tokens": 11, "output_tokens": 4}, + }, + {"type": "message_stop"}, + ) + return b"".join(_invoke_chunk(event, marker) for event in events) + + +def _provider(request: Request) -> Reply: + if not request.body: + return Reply(status=400, body=_json({"error": "empty request body"})) + body: Final = JSON_OBJECT.validate_json(request.body) + marker: Final = _marker(body) + target: Final = request.target.split("?", 1)[0] + if target.startswith("/passthrough"): + if body.get("stream") is True: + streamed_response: Final = _json({"received": body}) + return Reply( + content_type="text/event-stream", + chunks=(b"data: " + streamed_response + b"\n\n", b"data: [DONE]\n\n"), + ) + return Reply(body=_json({"received": body})) + if target.endswith("/converse-stream"): + return Reply(body=_bedrock_converse_stream(marker), content_type=BEDROCK_EVENT_STREAM) + if target.endswith("/invoke-with-response-stream"): + return Reply(body=_bedrock_invoke_stream(marker), content_type=BEDROCK_EVENT_STREAM) + if target.endswith("/converse") or target.endswith("/invoke"): + return Reply(body=_json({"output": f"scripted Bedrock reply {marker}"})) + if target == "/v1/chat/completions": + if body.get("stream") is True: + chunk: Final = { + "id": f"chatcmpl-{marker}", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": f"scripted chat reply {marker}"}, + "finish_reason": None, + } + ], + } + final_chunk: Final = { + **chunk, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + } + return Reply( + content_type="text/event-stream", + chunks=( + b"data: " + _json(chunk) + b"\n\n", + b"data: " + _json(final_chunk) + b"\n\n", + b"data: [DONE]\n\n", + ), + ) + return Reply( + body=_json( + { + "id": f"chatcmpl-{marker}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": f"scripted chat reply {marker}"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + ) + ) + return Reply(status=404, body=_json({"error": f"unexpected upstream path: {target}"})) + + +def _sink(request: Request) -> Reply: + assert request.target.endswith(GUARDRAIL_PATH), request.target + assert b"scope-" in request.body, request.body.decode() + return Reply(body=_json({"action": "NONE"})) + + +def _rail(name: str, sink: Wire, scope: EndpointScope) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "stream_scope": scope, + "api_base": f"{sink.url}/{name}", + "api_key": "synthetic-guardrail-key", + }, + } + + +def _chat_proxy_config(provider_url: str, guardrails: list[dict[str, JsonValue]]) -> dict[str, object]: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + return { + **config, + "guardrails": guardrails, + "model_list": [ + { + "model_name": "scope-invalid-config-chat", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{provider_url}/v1", + "api_key": "synthetic-provider-key", + }, + } + ], + } + + +@dataclass(frozen=True, slots=True) +class ReproRig: + candidate: Gateway + direct_candidate: Gateway + scenario: Scenario + models: Mapping[str, str] + rails: Mapping[str, str] + provider: Wire + sink: Wire + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[ReproRig]: + with httpx.Client( + base_url=os.environ["INTEGRATION_PROXY_URL"], + timeout=30, + trust_env=False, + ) as root_client: + root_gateway: Final = Gateway( + root_client, + os.environ.get("INTEGRATION_MASTER_KEY", "sk-integration-master"), + os.environ["INTEGRATION_UPSTREAM_URL"], + ) + directory: Final = tmp_path_factory.mktemp("guardrail-stream-scope-repro") + with wire_server(_provider) as provider, wire_server(_sink) as sink: + rails: Final = MappingProxyType( + { + "bedrock_streaming": "bedrock_streaming", + "bedrock_non_streaming": "bedrock_non_streaming", + "chat_streaming": "chat_streaming", + "passthrough_streaming": "passthrough_streaming", + "passthrough_non_streaming": "passthrough_non_streaming", + "passthrough_spoof_streaming": "passthrough_spoof_streaming", + } + ) + models: Final = MappingProxyType( + { + "chat": "scope-chat", + "bedrock_router": "scope-bedrock-router", + "bedrock_false_positive_router": "scope-bedrock-converse-stream-model", + } + ) + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["guardrails"] = [ + _rail(rails["bedrock_streaming"], sink, "streaming"), + _rail(rails["bedrock_non_streaming"], sink, "non_streaming"), + _rail(rails["chat_streaming"], sink, "streaming"), + _rail(rails["passthrough_streaming"], sink, "streaming"), + _rail(rails["passthrough_non_streaming"], sink, "non_streaming"), + _rail(rails["passthrough_spoof_streaming"], sink, "streaming"), + ] + config["model_list"] = [ + { + "model_name": models["chat"], + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{provider.url}/v1", + "api_key": "synthetic-provider-key", + }, + }, + { + "model_name": models["bedrock_router"], + "litellm_params": { + "model": f"bedrock/{BEDROCK_MODEL_ID}", + "api_base": provider.url, + "aws_access_key_id": "AKIASYNTHETICSTREAMSCOPE", + "aws_secret_access_key": "synthetic-bedrock-secret", + "aws_region_name": "us-east-1", + }, + }, + { + "model_name": models["bedrock_false_positive_router"], + "litellm_params": { + "model": f"bedrock/{BEDROCK_FALSE_POSITIVE_MODEL_ID}", + "api_base": provider.url, + "aws_access_key_id": "AKIASYNTHETICSTREAMSCOPE", + "aws_secret_access_key": "synthetic-bedrock-secret", + "aws_region_name": "us-east-1", + }, + }, + ] + config["environment_variables"] = { + "AWS_BEDROCK_RUNTIME_ENDPOINT": provider.url, + "AWS_ACCESS_KEY_ID": "AKIASYNTHETICSTREAMSCOPE", + "AWS_SECRET_ACCESS_KEY": "synthetic-bedrock-secret", + "AWS_REGION": "us-east-1", + "AWS_REGION_NAME": "us-east-1", + } + config["general_settings"]["pass_through_endpoints"] = [ + { + "path": "/pt-forward", + "target": f"{provider.url}/passthrough", + "include_subpath": True, + }, + { + "path": "/pt-spoof", + "target": f"{provider.url}/passthrough", + "include_subpath": True, + "guardrails": {rails["passthrough_spoof_streaming"]: None}, + }, + { + "path": "/pt-scope", + "target": f"{provider.url}/passthrough", + "include_subpath": True, + "guardrails": { + rails["passthrough_streaming"]: None, + rails["passthrough_non_streaming"]: None, + }, + }, + ] + config_path: Final = directory / "stream-scope-repro.yaml" + config_path.write_text(yaml.safe_dump(config)) + with owned_proxy_process(root_gateway, directory, {}, config=config_path, workers=1) as owned: + with owned.gateway.scenario() as scenario: + direct_config: Final = { + **config, + "model_list": [ + *config["model_list"], + { + "model_name": f"scope-unused-{uuid.uuid4().hex}*", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{provider.url}/v1", + "api_key": "synthetic-provider-key", + }, + }, + ], + } + direct_config_path: Final = directory / "stream-scope-direct.yaml" + direct_config_path.write_text(yaml.safe_dump(direct_config)) + with owned_proxy_process( + root_gateway, + directory, + {}, + config=direct_config_path, + workers=1, + ) as direct: + yield ReproRig( + owned.gateway, + direct.gateway, + scenario, + models, + rails, + provider, + sink, + ) + + +def _bedrock_request( + action: BedrockAction, + model_path: str, + marker: str, +) -> tuple[str, dict[str, JsonValue]]: + if action in ("converse", "converse-stream"): + body: Final = { + "messages": [{"role": "user", "content": [{"text": marker}]}], + "inferenceConfig": {"maxTokens": 16}, + } + else: + body = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 16, + "messages": [{"role": "user", "content": [{"type": "text", "text": marker}]}], + } + return f"/bedrock/model/{model_path}/{action}", body + + +def _matching_requests(wire: Wire, marker: str) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if marker.encode() in request.body) + + +def _rail_scans(rows: Sequence[Request], rail_name: str, marker: str) -> tuple[Request, ...]: + return tuple( + request for request in rows if request.target.startswith(f"/{rail_name}/") and marker.encode() in request.body + ) + + +def _chat_request_with_scans( + gateway: Gateway, + sink: Wire, + model: str, + marker: str, + streamed: bool, + rail_name: str, +) -> tuple[httpx.Response, tuple[Request, ...]]: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": marker}], + "stream": streamed, + "guardrails": [rail_name], + }, + ) + return response, _rail_scans(sink.drain(), rail_name, marker) + + +def _spend_row_for_call(call_id: str, content: bytes) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + "SELECT request_id, litellm_call_id, spend, prompt_tokens, completion_tokens, metadata " + 'FROM "LiteLLM_SpendLogs" WHERE request_id=%s OR litellm_call_id=%s', + (call_id, call_id), + ), + lambda values: len(values) >= 1, + seconds=70, + ) + assert len(rows) == 1, (call_id, rows, content) + assert call_id in (rows[0]["request_id"], rows[0]["litellm_call_id"]), content + return rows[0] + + +@pytest.mark.parametrize("model_kind", ("router", "direct")) +@pytest.mark.parametrize("action", STREAMING_ACTIONS) +@pytest.mark.parametrize("scope", SCOPES) +def test_bedrock_streaming_actions_run_streaming_scoped_rails( + rig: ReproRig, + model_kind: ModelKind, + action: StreamingAction, + scope: EndpointScope, +) -> None: + marker: Final = f"scope-bedrock-stream-{uuid.uuid4().hex}" + model_path: Final = rig.models["bedrock_router"] if model_kind == "router" else BEDROCK_MODEL_ID + path, body = _bedrock_request(action, model_path, marker) + expected_body: Final = ( + _bedrock_converse_stream(marker) if action == "converse-stream" else _bedrock_invoke_stream(marker) + ) + call_id: Final = f"stream-scope-{uuid.uuid4().hex}" + key: Final = rig.scenario.key(guardrails=[rig.rails[f"bedrock_{scope}"]]) + candidate: Final = rig.candidate if model_kind == "router" else rig.direct_candidate + response: Final = candidate.request( + "POST", + path, + body, + key=key, + headers={"x-litellm-call-id": call_id}, + ) + assert response.status_code == 200, response.content + assert response.headers.get("content-type") == BEDROCK_EVENT_STREAM, dict(response.headers) + assert response.content == expected_body, response.content + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, provider_rows, response.content) + provider_body: Final = JSON_OBJECT.validate_json(provider_rows[0].body) + key_names: Final = _key_names(provider_body) + assert "is_streaming_request" not in key_names, provider_body + assert not tuple(name for name in key_names if name.startswith("litellm_")), provider_body + _spend_row_for_call(call_id, response.content) + sink_rows: Final = _rail_scans(rig.sink.drain(), rig.rails[f"bedrock_{scope}"], marker) + expected_scans: Final = int(scope == "streaming") + assert len(sink_rows) == expected_scans, (marker, model_kind, action, scope, sink_rows, response.content) + + +@pytest.mark.parametrize("model_kind", ("router", "direct")) +@pytest.mark.parametrize("action", NON_STREAMING_ACTIONS) +@pytest.mark.parametrize("scope", SCOPES) +def test_bedrock_non_streaming_actions_run_non_streaming_scoped_rails( + rig: ReproRig, + model_kind: ModelKind, + action: NonStreamingAction, + scope: EndpointScope, +) -> None: + marker: Final = f"scope-bedrock-nonstream-{uuid.uuid4().hex}" + model_path: Final = rig.models["bedrock_router"] if model_kind == "router" else BEDROCK_MODEL_ID + path, body = _bedrock_request(action, model_path, marker) + key: Final = rig.scenario.key(guardrails=[rig.rails[f"bedrock_{scope}"]]) + candidate: Final = rig.candidate if model_kind == "router" else rig.direct_candidate + response: Final = candidate.request("POST", path, body, key=key) + assert response.status_code == 200, response.text + assert marker in response.text, response.text + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, provider_rows, response.text) + provider_body: Final = JSON_OBJECT.validate_json(provider_rows[0].body) + key_names: Final = _key_names(provider_body) + assert "is_streaming_request" not in key_names, provider_body + assert not tuple(name for name in key_names if name.startswith("litellm_")), provider_body + sink_rows: Final = _rail_scans(rig.sink.drain(), rig.rails[f"bedrock_{scope}"], marker) + expected_scans: Final = int(scope == "non_streaming") + assert len(sink_rows) == expected_scans, (marker, model_kind, action, scope, sink_rows, response.text) + + +@pytest.mark.parametrize("model_kind", ("router", "direct")) +def test_bedrock_model_id_streaming_action_text_on_converse_is_non_streaming( + rig: ReproRig, + model_kind: ModelKind, +) -> None: + marker: Final = f"scope-bedrock-converse-model-{uuid.uuid4().hex}" + model_path: Final = ( + rig.models["bedrock_false_positive_router"] if model_kind == "router" else BEDROCK_FALSE_POSITIVE_MODEL_ID + ) + path, body = _bedrock_request("converse", model_path, marker) + call_id: Final = f"stream-scope-bedrock-{uuid.uuid4().hex}" + key: Final = rig.scenario.key( + guardrails=[ + rig.rails["bedrock_streaming"], + rig.rails["bedrock_non_streaming"], + ] + ) + candidate: Final = rig.candidate if model_kind == "router" else rig.direct_candidate + response: Final = candidate.request( + "POST", + path, + body, + key=key, + headers={"x-litellm-call-id": call_id}, + ) + assert response.status_code == 200, response.text + assert marker in response.text, response.text + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, model_kind, provider_rows, response.text) + sink_rows: Final = rig.sink.drain() + streaming_rows: Final = _rail_scans(sink_rows, rig.rails["bedrock_streaming"], marker) + non_streaming_rows: Final = _rail_scans(sink_rows, rig.rails["bedrock_non_streaming"], marker) + assert streaming_rows == (), (marker, model_kind, streaming_rows, response.text) + assert len(non_streaming_rows) == 1, (marker, model_kind, non_streaming_rows, response.text) + + +@pytest.mark.parametrize("streamed", (False, True), ids=("stream-absent", "stream-true")) +def test_configured_passthrough_forwards_caller_is_streaming_request_field( + rig: ReproRig, + streamed: bool, +) -> None: + marker: Final = f"scope-passthrough-forward-{uuid.uuid4().hex}" + caller_value: Final = f"caller-{uuid.uuid4().hex}" + body: Final = { + "marker": marker, + "is_streaming_request": caller_value, + **({"stream": True} if streamed else {}), + } + response: Final = rig.candidate.request("POST", "/pt-forward", body) + assert response.status_code == 200, response.text + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, provider_rows, response.text) + upstream_body: Final = JSON_OBJECT.validate_json(provider_rows[0].body) + assert upstream_body.get("is_streaming_request") == caller_value, ( + marker, + caller_value, + upstream_body, + response.text, + ) + assert upstream_body == body, (marker, body, upstream_body, response.text) + if streamed: + assert response.headers.get("content-type", "").lower().startswith("text/event-stream"), dict(response.headers) + event_body: Final = response.text.removeprefix("data: ").split("\n", maxsplit=1)[0] + response_body: Final = JSON_OBJECT.validate_json(event_body) + else: + response_body = JSON_OBJECT.validate_json(response.content) + assert response_body == {"received": body}, response.text + + +@pytest.mark.parametrize(("hostile_field", "hostile_value"), HOSTILE_CLASSIFICATION_CASES) +def test_client_cannot_spoof_server_stream_classification( + rig: ReproRig, + hostile_field: str, + hostile_value: JsonValue, +) -> None: + chat_marker: Final = f"scope-chat-spoof-{uuid.uuid4().hex}" + chat_body: Final = { + "model": rig.models["chat"], + "messages": [{"role": "user", "content": chat_marker}], + "stream": False, + hostile_field: hostile_value, + } + chat_key: Final = rig.scenario.key(guardrails=[rig.rails["chat_streaming"]]) + chat_response: Final = rig.candidate.request( + "POST", + "/v1/chat/completions", + chat_body, + key=chat_key, + ) + chat_provider_rows: Final = _matching_requests(rig.provider, chat_marker) + chat_upstream_body: Final = chat_provider_rows[0].body.decode() if chat_provider_rows else "" + print( # noqa: T201 # required chat hostile-body observation + f"chat hostile {hostile_field}={hostile_value!r}: " + f"status={chat_response.status_code}, response={chat_response.text!r}, upstream={chat_upstream_body}" + ) + assert chat_response.status_code == 200, chat_response.text + assert len(chat_provider_rows) == 1, (chat_marker, chat_provider_rows, chat_response.text) + assert JSON_OBJECT.validate_json(chat_provider_rows[0].body) == { + "messages": [{"role": "user", "content": chat_marker}], + "model": "gpt-4o-mini", + hostile_field: hostile_value, + }, (hostile_field, hostile_value, chat_upstream_body) + assert _rail_scans(rig.sink.drain(), rig.rails["chat_streaming"], chat_marker) == (), ( + chat_marker, + hostile_field, + hostile_value, + chat_response.text, + ) + + +@pytest.mark.parametrize(("hostile_field", "hostile_value"), HOSTILE_CLASSIFICATION_CASES) +def test_configured_passthrough_cannot_spoof_server_stream_classification( + rig: ReproRig, + hostile_field: str, + hostile_value: JsonValue, +) -> None: + passthrough_marker: Final = f"scope-passthrough-spoof-{uuid.uuid4().hex}" + passthrough_body: Final = { + "marker": passthrough_marker, + "stream": False, + hostile_field: hostile_value, + } + passthrough_response: Final = rig.candidate.request("POST", "/pt-spoof", passthrough_body) + assert passthrough_response.status_code == 200, passthrough_response.text + passthrough_provider_rows: Final = _matching_requests(rig.provider, passthrough_marker) + assert len(passthrough_provider_rows) == 1, ( + passthrough_marker, + passthrough_provider_rows, + passthrough_response.text, + ) + assert _rail_scans(rig.sink.drain(), rig.rails["passthrough_spoof_streaming"], passthrough_marker) == (), ( + passthrough_marker, + hostile_field, + hostile_value, + passthrough_response.text, + ) + + +@pytest.mark.parametrize( + ("streamed", "request_fields"), + ((True, {"stream": True}), (False, {"stream": False}), (False, {})), + ids=("stream-true", "stream-false", "stream-absent"), +) +def test_passthrough_scope_follows_proxy_stream_decision( + rig: ReproRig, + streamed: bool, + request_fields: dict[str, JsonValue], +) -> None: + marker: Final = f"scope-passthrough-scope-{uuid.uuid4().hex}" + body: Final = {"marker": marker, **request_fields} + response: Final = rig.candidate.request("POST", "/pt-scope", body) + assert response.status_code == 200, response.text + if streamed: + expected_frame: Final = f"data: {_json({'received': body}).decode()}\n\ndata: [DONE]\n\n" + assert response.headers.get("content-type", "").lower().startswith("text/event-stream"), dict(response.headers) + assert response.text == expected_frame, response.text + assert response.headers.get("transfer-encoding", "").lower() == "chunked", dict(response.headers) + else: + response_body: Final = JSON_OBJECT.validate_json(response.content) + assert response_body == {"received": body}, response.text + assert "content-length" in response.headers, dict(response.headers) + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, provider_rows, response.text) + upstream_body: Final = JSON_OBJECT.validate_json(provider_rows[0].body) + assert upstream_body == body, (marker, body, upstream_body, response.text) + assert SERVER_STREAMING_CLASSIFICATION_KEY not in upstream_body, upstream_body + sink_rows: Final = rig.sink.drain() + streaming_rows: Final = _rail_scans(sink_rows, rig.rails["passthrough_streaming"], marker) + non_streaming_rows: Final = _rail_scans(sink_rows, rig.rails["passthrough_non_streaming"], marker) + assert len(streaming_rows) == int(streamed), (marker, streamed, streaming_rows, response.text) + assert len(non_streaming_rows) == int(not streamed), (marker, streamed, non_streaming_rows, response.text) + + +def test_invalid_yaml_stream_scope_keeps_rail_running_on_both_shapes(rig: ReproRig, tmp_path: Path) -> None: + name: Final = f"scope-invalid-yaml-{uuid.uuid4().hex}" + invalid_rail: Final = { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "stream_scope": "sometimes", + "api_base": f"{rig.sink.url}/{name}", + "api_key": "synthetic-guardrail-key", + }, + } + config: Final = _chat_proxy_config(rig.provider.url, [invalid_rail]) + config_path: Final = tmp_path / "invalid-stream-scope.yaml" + config_path.write_text(yaml.safe_dump(config)) + with owned_proxy_process(rig.candidate, tmp_path, {}, config=config_path, workers=1) as owned: + markers: Final = tuple(f"scope-invalid-yaml-{int(streamed)}-{uuid.uuid4().hex}" for streamed in (False, True)) + observations: Final = tuple( + _chat_request_with_scans( + owned.gateway, + rig.sink, + "scope-invalid-config-chat", + marker, + streamed, + name, + ) + for streamed, marker in zip((False, True), markers) + ) + assert tuple(response.status_code for response, _ in observations) == (200, 200), tuple( + response.text for response, _ in observations + ) + assert tuple(len(sink_rows) for _, sink_rows in observations) == (1, 1), (markers, observations) + listed: Final = owned.gateway.request("GET", "/guardrails/list") + assert listed.status_code == 200, listed.text + list_payload: Final = JSON_OBJECT.validate_json(listed.content) + listed_guardrails: Final = list_payload.get("guardrails") + assert isinstance(listed_guardrails, list), list_payload + listed_rail: Final = next( + (row for row in listed_guardrails if isinstance(row, dict) and row.get("guardrail_name") == name), + None, + ) + assert isinstance(listed_rail, dict), list_payload + listed_params: Final = listed_rail.get("litellm_params") + assert isinstance(listed_params, dict), listed_rail + assert listed_params.get("stream_scope") is None, listed_rail + + +def test_persisted_invalid_stream_scope_row_stays_readable_and_enforced( + rig: ReproRig, + tmp_path: Path, +) -> None: + guardrail_id: Final = str(uuid.uuid4()) + guardrail_name: Final = f"scope-invalid-persisted-{uuid.uuid4().hex}" + params: Final = { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": f"{rig.sink.url}/{guardrail_name}", + "api_key": "synthetic-guardrail-key", + "stream_scope": "sometimes", + } + database_url: Final = os.environ.get("INTEGRATION_PROXY_DATABASE_URL") or os.environ["DATABASE_URL"] + write_rows( + 'INSERT INTO "LiteLLM_GuardrailsTable" ' + '("guardrail_id", "guardrail_name", "litellm_params", "created_at", "updated_at") ' + "VALUES (%s, %s, %s::jsonb, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)", + (guardrail_id, guardrail_name, json.dumps(params)), + database_url=database_url, + ) + try: + config: Final = _chat_proxy_config(rig.provider.url, []) + config_path: Final = tmp_path / "persisted-invalid-stream-scope.yaml" + config_path.write_text(yaml.safe_dump(config)) + with owned_proxy_process(rig.candidate, tmp_path, {}, config=config_path, workers=1) as owned: + info: Final = eventually( + lambda: owned.gateway.request("GET", f"/guardrails/{guardrail_id}/info"), + lambda response: response.status_code != 404, + seconds=20, + ) + listed: Final = owned.gateway.request("GET", "/v2/guardrails/list") + list_payload: Final = JSON_OBJECT.validate_json(listed.content) if listed.status_code == 200 else {} + listed_guardrails: Final = list_payload.get("guardrails") + includes_row: Final = isinstance(listed_guardrails, list) and any( + isinstance(row, dict) and row.get("guardrail_id") == guardrail_id for row in listed_guardrails + ) + markers: Final = ( + f"scope-invalid-persisted-0-{uuid.uuid4().hex}", + f"scope-invalid-persisted-1-{uuid.uuid4().hex}", + ) + observations: Final = tuple( + _chat_request_with_scans( + owned.gateway, + rig.sink, + "scope-invalid-config-chat", + marker, + streamed, + guardrail_name, + ) + for streamed, marker in zip((False, True), markers) + ) + assert ( + info.status_code == 200 + and listed.status_code == 200 + and includes_row + and tuple(response.status_code for response, _ in observations) == (200, 200) + and tuple(len(sink_rows) for _, sink_rows in observations) == (1, 1) + ), { + "info": (info.status_code, info.text), + "list": (listed.status_code, listed.text), + "includes_row": includes_row, + "responses": tuple((response.status_code, response.text) for response, _ in observations), + "scan_counts": tuple(len(sink_rows) for _, sink_rows in observations), + } + finally: + write_rows( + 'DELETE FROM "LiteLLM_GuardrailsTable" WHERE guardrail_id=%s', + (guardrail_id,), + database_url=database_url, + ) + + +def test_management_rejects_invalid_stream_scope(rig: ReproRig) -> None: + name: Final = f"scope-invalid-management-{uuid.uuid4().hex}" + invalid_params: Final = { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": f"{rig.sink.url}/{name}", + "api_key": "synthetic-guardrail-key", + "stream_scope": "sometimes", + } + created_invalid: Final = rig.candidate.request( + "POST", + "/guardrails", + {"guardrail": {"guardrail_name": name, "litellm_params": invalid_params}}, + ) + assert created_invalid.status_code == 422, created_invalid.text + + valid_params: Final = {**invalid_params, "stream_scope": "both"} + created: Final = rig.candidate.request( + "POST", + "/guardrails", + {"guardrail": {"guardrail_name": name, "litellm_params": valid_params}}, + ) + assert created.status_code == 200, created.text + guardrail_id: Final = str(created.json()["guardrail_id"]) + try: + put_response: Final = rig.candidate.request( + "PUT", + f"/guardrails/{guardrail_id}", + {"guardrail": {"guardrail_name": name, "litellm_params": invalid_params}}, + ) + patch_response: Final = rig.candidate.request( + "PATCH", + f"/guardrails/{guardrail_id}", + {"litellm_params": {"stream_scope": "sometimes"}}, + ) + assert put_response.status_code == 422, put_response.text + assert patch_response.status_code == 422, patch_response.text + finally: + deleted: Final = rig.candidate.request("DELETE", f"/guardrails/{guardrail_id}") + assert deleted.status_code == 200, deleted.text diff --git a/tests/integration/observability/test_guardrail_stream_scope_chaos.py b/tests/integration/observability/test_guardrail_stream_scope_chaos.py new file mode 100644 index 00000000000..f8a9d0a4bba --- /dev/null +++ b/tests/integration/observability/test_guardrail_stream_scope_chaos.py @@ -0,0 +1,1079 @@ +from __future__ import annotations + +import json +import os +import shutil +import signal +import socket +import subprocess +import uuid +from collections.abc import Iterator, Mapping, Sequence +from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from threading import Barrier, Event +from types import MappingProxyType +from typing import Final, Literal, TypeAlias, cast + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, JsonValue, eventually +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import TypeAdapter + +ChaosEndpoint: TypeAlias = Literal["chat", "messages", "responses"] +CHAOS_ENDPOINTS: Final[tuple[ChaosEndpoint, ...]] = ("chat", "messages", "responses") +CHAOS_MODELS: Final = MappingProxyType( + {"chat": "chaos-chat", "messages": "chaos-messages", "responses": "chaos-responses"} +) +GUARDRAIL_PATH: Final = "/beta/litellm_basic_guardrail_api" +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +POSTGRES_IMAGE: Final = "postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5" + + +def _json(value: object) -> bytes: + return json.dumps(value, separators=(",", ":")).encode() + + +def _texts(value: JsonValue) -> tuple[str, ...]: + if isinstance(value, str): + return (value,) + if isinstance(value, list): + return tuple(chain.from_iterable(_texts(item) for item in value)) + if isinstance(value, dict): + return tuple(chain.from_iterable(_texts(item) for item in value.values())) + return () + + +def _marker(body: Mapping[str, JsonValue]) -> str: + return next((text for text in _texts(dict(body)) if text.startswith("audit-")), "audit-chaos") + + +def _sse(events: Sequence[Mapping[str, JsonValue]]) -> tuple[bytes, ...]: + return tuple(f"data: {json.dumps(event, separators=(',', ':'))}\n\n".encode() for event in events) + ( + b"data: [DONE]\n\n", + ) + + +def _messages_stream(message: Mapping[str, JsonValue]) -> tuple[bytes, ...]: + content: Final = cast(list[JsonValue], message["content"]) + text: Final = cast(dict[str, JsonValue], content[0])["text"] + assert isinstance(text, str) + return ( + f"event: message_start\ndata: {json.dumps({**message, 'content': [], 'stop_reason': None, 'usage': {'input_tokens': 2, 'output_tokens': 0}})}\n\n".encode(), + b'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}\n\n', + f"event: content_block_delta\ndata: {json.dumps({'type': 'content_block_delta', 'index': 0, 'delta': {'type': 'text_delta', 'text': text}})}\n\n".encode(), + b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n', + f"event: message_delta\ndata: {json.dumps({'type': 'message_delta', 'delta': {'stop_reason': 'end_turn', 'stop_sequence': None}, 'usage': {'output_tokens': 2}})}\n\n".encode(), + b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + ) + + +def _responses_stream( + response: Mapping[str, JsonValue], + output: Mapping[str, JsonValue], + marker: str, +) -> tuple[bytes, ...]: + events: Final[tuple[dict[str, JsonValue], ...]] = ( + {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, + {"type": "response.in_progress", "response": {**response, "status": "in_progress", "output": []}}, + {"type": "response.output_item.added", "item": dict(output), "output_index": 0}, + { + "type": "response.content_part.added", + "item_id": f"msg-{marker}", + "output_index": 0, + "content_index": 0, + "part": {"type": "output_text", "text": "", "annotations": []}, + }, + { + "type": "response.output_text.delta", + "item_id": f"msg-{marker}", + "output_index": 0, + "content_index": 0, + "delta": marker, + }, + { + "type": "response.output_text.done", + "item_id": f"msg-{marker}", + "output_index": 0, + "content_index": 0, + "text": marker, + }, + { + "type": "response.content_part.done", + "item_id": f"msg-{marker}", + "output_index": 0, + "content_index": 0, + "part": cast(list[JsonValue], output["content"])[0], + }, + {"type": "response.output_item.done", "item": dict(output), "output_index": 0}, + {"type": "response.completed", "response": dict(response)}, + ) + return tuple( + f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': index}, separators=(',', ':'))}\n\n".encode() + for index, event in enumerate(events) + ) + + +def _provider(request: Request) -> Reply: + if request.method == "GET" and request.target.partition("?")[0] == "/v1/models": + return Reply(body=_json({"object": "list", "data": [{"id": "gpt-4o-mini", "object": "model"}]})) + if not request.body: + return Reply(status=400, body=_json({"error": "request body is required"})) + body: Final = JSON_OBJECT.validate_json(request.body) + marker: Final = _marker(body) + streamed: Final = bool(body.get("stream")) + if request.target == "/v1/chat/completions": + if streamed: + return Reply( + content_type="text/event-stream", + chunks=_sse( + ( + { + "id": f"chatcmpl-{marker}", + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {"content": marker}, "finish_reason": None}], + }, + ) + ), + ) + return Reply( + body=_json( + { + "id": f"chatcmpl-{marker}", + "object": "chat.completion", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": marker}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 2, "completion_tokens": 2, "total_tokens": 4}, + } + ) + ) + if request.target == "/v1/messages": + message: Final = { + "id": f"msg-{marker}", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": marker}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 2}, + } + return ( + Reply( + content_type="text/event-stream", + chunks=_messages_stream(message), + ) + if streamed + else Reply(body=_json(message)) + ) + if request.target == "/v1/responses": + response: Final = { + "id": f"resp-{marker}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": f"msg-{marker}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": marker, "annotations": []}], + } + ], + "usage": {"input_tokens": 2, "output_tokens": 2, "total_tokens": 4}, + } + return ( + Reply( + content_type="text/event-stream", + chunks=_responses_stream(response, cast(dict[str, JsonValue], response["output"][0]), marker), + ) + if streamed + else Reply(body=_json(response)) + ) + return Reply(status=404, body=_json({"error": f"unexpected provider target {request.target}"})) + + +def _sink(request: Request) -> Reply: + assert request.target.endswith(GUARDRAIL_PATH), request.target + body: Final = JSON_OBJECT.validate_json(request.body) + assert body.get("litellm_call_id") is not None or any("audit-" in text for text in _texts(body)), ( + request.body.decode() + ) + return Reply(body=_json({"action": "NONE"})) + + +def _rail( + name: str, + sink_url: str, + scope: Literal["streaming", "non_streaming"], + *, + default_on: bool = False, +) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": default_on, + "stream_scope": scope, + "api_base": f"{sink_url}/{name}", + "api_key": "synthetic-chaos-key", + }, + } + + +def _config(provider_url: str, rails: Sequence[dict[str, JsonValue]]) -> dict[str, JsonValue]: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + return cast( + dict[str, JsonValue], + { + **base, + "guardrails": list(rails), + "model_list": [ + { + "model_name": CHAOS_MODELS["chat"], + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{provider_url}/v1", + "api_key": "synthetic-provider-key", + "num_retries": 0, + }, + }, + { + "model_name": CHAOS_MODELS["messages"], + "litellm_params": { + "model": "anthropic/claude-sonnet-4-5-20250929", + "api_base": provider_url, + "api_key": "synthetic-provider-key", + "num_retries": 0, + }, + }, + { + "model_name": CHAOS_MODELS["responses"], + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{provider_url}/v1", + "api_key": "synthetic-provider-key", + "num_retries": 0, + }, + }, + ], + "environment_variables": { + **base.get("environment_variables", {}), + "OPENAI_API_BASE": provider_url, + "OPENAI_API_KEY": "synthetic-provider-key", + "ANTHROPIC_API_BASE": provider_url, + "ANTHROPIC_API_KEY": "synthetic-provider-key", + }, + }, + ) + + +@dataclass(frozen=True, slots=True) +class ChaosRig: + gateway: Gateway + provider: Wire + directory: Path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[ChaosRig]: + with ( + httpx.Client( + base_url=os.environ["INTEGRATION_PROXY_URL"], + timeout=30, + trust_env=False, + ) as root_client, + wire_server(_provider) as provider, + ): + root_gateway: Final = Gateway( + root_client, + os.environ.get("INTEGRATION_MASTER_KEY", "sk-integration-master"), + os.environ["INTEGRATION_UPSTREAM_URL"], + ) + yield ChaosRig(root_gateway, provider, tmp_path_factory.mktemp("stream-scope-chaos")) + + +@dataclass(frozen=True, slots=True) +class CallPlan: + marker: str + call_id: str + endpoint: ChaosEndpoint + streamed: bool + + +def _plans(prefix: str, count: int) -> tuple[CallPlan, ...]: + return tuple( + CallPlan( + f"audit-{prefix}-{index}-{uuid.uuid4().hex}", + f"{prefix}-{uuid.uuid4().hex}", + CHAOS_ENDPOINTS[index % len(CHAOS_ENDPOINTS)], + index % 2 == 1, + ) + for index in range(count) + ) + + +def _request(gateway: Gateway, plan: CallPlan, rails: Sequence[str]) -> httpx.Response: + guardrail_field: Final[dict[str, JsonValue]] = {"guardrails": list(rails)} if rails else {} + match plan.endpoint: + case "chat": + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": CHAOS_MODELS["chat"], + "messages": [{"role": "user", "content": plan.marker}], + **guardrail_field, + **({"stream": True} if plan.streamed else {}), + }, + headers={"x-litellm-call-id": plan.call_id}, + ) + case "messages": + return gateway.request( + "POST", + "/v1/messages", + { + "model": CHAOS_MODELS["messages"], + "max_tokens": 32, + "messages": [{"role": "user", "content": plan.marker}], + **guardrail_field, + **({"stream": True} if plan.streamed else {}), + }, + headers={"x-litellm-call-id": plan.call_id}, + ) + case "responses": + return gateway.request( + "POST", + "/v1/responses", + { + "model": CHAOS_MODELS["responses"], + "input": plan.marker, + **guardrail_field, + **({"stream": True} if plan.streamed else {}), + }, + headers={"x-litellm-call-id": plan.call_id}, + ) + raise AssertionError(plan.endpoint) + + +def _rows_for_marker(rows: Sequence[Request], marker: str) -> tuple[Request, ...]: + return tuple(request for request in rows if marker.encode() in request.body) + + +def _rail_scans(rows: Sequence[Request], rail_name: str, marker: str) -> tuple[Request, ...]: + return tuple( + request for request in rows if request.target.startswith(f"/{rail_name}/") and marker.encode() in request.body + ) + + +def _spend_rows(call_id: str, database_url: str | None = None) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id, litellm_call_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s OR litellm_call_id=%s', + (call_id, call_id), + database_url=database_url, + ) + + +def _one_spend_row( + call_id: str, + database_url: str | None = None, + *, + seconds: float = 70, +) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: _spend_rows(call_id, database_url), + lambda values: len(values) >= 1, + seconds=seconds, + ) + assert len(rows) == 1, (call_id, rows) + return rows[0] + + +def _expected_in_scope(plan: CallPlan, sink_a_name: str, sink_b_name: str) -> tuple[str, str]: + return (sink_a_name, sink_b_name) if plan.streamed else (sink_b_name, sink_a_name) + + +def _assert_successful_calls( + plans: Sequence[CallPlan], + responses: Sequence[httpx.Response], + provider_rows: Sequence[Request], + sink_a_rows: Sequence[Request], + sink_b_rows: Sequence[Request], + sink_a_name: str, + sink_b_name: str, +) -> None: + for plan, response in zip(plans, responses): + if response.status_code != 200: + continue + assert plan.marker in response.text, (plan, response.text) + provider_match: Final = _rows_for_marker(provider_rows, plan.marker) + assert len(provider_match) == 1, (plan, provider_match) + in_sink, out_sink = _expected_in_scope(plan, sink_a_name, sink_b_name) + in_rows: Final = _rail_scans( + sink_a_rows if in_sink == sink_a_name else sink_b_rows, + in_sink, + plan.marker, + ) + out_rows: Final = _rail_scans( + sink_a_rows if out_sink == sink_a_name else sink_b_rows, + out_sink, + plan.marker, + ) + assert len(in_rows) == 1 and len(out_rows) == 0, (plan, in_rows, out_rows) + spend: Final = _one_spend_row(plan.call_id) + assert plan.call_id in (spend.get("request_id"), spend.get("litellm_call_id")), (plan, spend) + + +def _call_wave(gateway: Gateway, plans: Sequence[CallPlan], rails: Sequence[str]) -> tuple[httpx.Response, ...]: + with ThreadPoolExecutor(max_workers=20) as pool: + futures: Final[tuple[Future[httpx.Response], ...]] = tuple( + pool.submit(_request, gateway, plan, rails) for plan in plans + ) + return tuple(future.result() for future in futures) + + +@contextmanager +def _owned_proxy( + rig: ChaosRig, + directory: Path, + rails: Sequence[dict[str, JsonValue]], + *, + workers: int = 1, +) -> Iterator[OwnedProxy]: + config_path: Final = directory / f"chaos-{uuid.uuid4().hex}.yaml" + config_path.write_text(yaml.safe_dump(_config(rig.provider.url, rails))) + with owned_proxy_process(rig.gateway, directory, {}, config=config_path, workers=workers) as owned: + yield owned + + +def _free_port() -> int: + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + return int(reserve.getsockname()[1]) + + +def _assert_down_wave( + plans: Sequence[CallPlan], + responses: Sequence[httpx.Response], + provider_rows: Sequence[Request], + sink_a_rows: Sequence[Request], + sink_b_rows: Sequence[Request], + sink_a_name: str, + sink_b_name: str, +) -> None: + for plan, response in zip(plans, responses): + if plan.streamed: + assert response.status_code >= 500 and response.content, (plan, response.status_code, response.text) + assert _rows_for_marker(provider_rows, plan.marker) == (), (plan, provider_rows) + assert _rail_scans(sink_b_rows, sink_b_name, plan.marker) == (), (plan, sink_b_name) + continue + assert response.status_code == 200 and plan.marker in response.text, (plan, response.status_code, response.text) + assert len(_rows_for_marker(provider_rows, plan.marker)) == 1, (plan, provider_rows) + assert _rail_scans(sink_a_rows, sink_a_name, plan.marker) == (), (plan, sink_a_name) + assert len(_rail_scans(sink_b_rows, sink_b_name, plan.marker)) == 1, (plan, sink_b_name) + spend: Final = _one_spend_row(plan.call_id) + assert plan.call_id in (spend.get("request_id"), spend.get("litellm_call_id")), (plan, spend) + + +def test_h1_sink_outage_keeps_scope_isolated_through_recovery(rig: ChaosRig, tmp_path: Path) -> None: + port_a: Final = _free_port() + started: Final = Event() + unavailable: Final = Event() + release: Final = Event() + + def gated_sink(request: Request) -> Reply: + started.set() + assert release.wait(timeout=45), "sink outage gate was not released" + if unavailable.is_set(): + return Reply(status=503, body=_json({"error": "synthetic sink outage"})) + return _sink(request) + + with wire_server(_sink) as sink_b, ExitStack() as sink_a_stack: + sink_a: Final = sink_a_stack.enter_context(wire_server(gated_sink, port=port_a)) + name_a: Final = f"h1-stream-{uuid.uuid4().hex}" + name_b: Final = f"h1-non-stream-{uuid.uuid4().hex}" + rails: Final = ( + _rail(name_a, sink_a.url, "streaming"), + _rail(name_b, sink_b.url, "non_streaming"), + ) + with _owned_proxy(rig, tmp_path, rails) as owned, ThreadPoolExecutor(max_workers=30) as pool: + outage_plans: Final = _plans("h1-burst", 30) + futures: Final[tuple[Future[httpx.Response], ...]] = tuple( + pool.submit(_request, owned.gateway, plan, (name_a, name_b)) for plan in outage_plans + ) + try: + assert started.wait(timeout=30), "streaming rail did not reach sink A" + unavailable.set() + finally: + release.set() + sink_a_stack.close() + outage_responses: Final = tuple(future.result(timeout=70) for future in futures) + outage_provider: Final = rig.provider.drain() + outage_a: Final = sink_a.drain() + outage_b: Final = sink_b.drain() + _assert_down_wave(outage_plans, outage_responses, outage_provider, outage_a, outage_b, name_a, name_b) + + with wire_server(_sink, port=port_a) as recovered_sink_a: + recovery_plans: Final = _plans("h1-recovery", 20) + recovery_responses: Final = _call_wave(owned.gateway, recovery_plans, (name_a, name_b)) + recovery_provider: Final = rig.provider.drain() + recovery_a: Final = recovered_sink_a.drain() + recovery_b: Final = sink_b.drain() + assert tuple(response.status_code for response in recovery_responses) == (200,) * 20, recovery_responses + _assert_successful_calls( + recovery_plans, + recovery_responses, + recovery_provider, + recovery_a, + recovery_b, + name_a, + name_b, + ) + + +def test_h2_stream_sink_gate_does_not_block_out_of_scope_calls(rig: ChaosRig, tmp_path: Path) -> None: + started: Final = Event() + release: Final = Event() + blocked_marker: Final = f"audit-h2-stream-{uuid.uuid4().hex}" + + def gated_sink(request: Request) -> Reply: + if blocked_marker.encode() in request.body: + started.set() + assert release.wait(timeout=45), "stream sink gate was not released" + return _sink(request) + + with wire_server(gated_sink) as sink_a, wire_server(_sink) as sink_b: + name_a: Final = f"h2-stream-{uuid.uuid4().hex}" + name_b: Final = f"h2-non-stream-{uuid.uuid4().hex}" + rails: Final = (_rail(name_a, sink_a.url, "streaming"), _rail(name_b, sink_b.url, "non_streaming")) + with _owned_proxy(rig, tmp_path, rails) as owned, ThreadPoolExecutor(max_workers=2) as pool: + streaming_plan: Final = CallPlan(blocked_marker, f"h2-stream-{uuid.uuid4().hex}", "chat", True) + non_streaming_plan: Final = CallPlan( + f"audit-h2-non-stream-{uuid.uuid4().hex}", + f"h2-non-stream-{uuid.uuid4().hex}", + "messages", + False, + ) + streaming_future: Final = pool.submit(_request, owned.gateway, streaming_plan, (name_a, name_b)) + assert started.wait(timeout=30), "stream request did not reach the gated sink" + try: + non_streaming_future: Final = pool.submit( + _request, + owned.gateway, + non_streaming_plan, + (name_a, name_b), + ) + non_streaming_response: Final = non_streaming_future.result(timeout=15) + assert non_streaming_response.status_code == 200, non_streaming_response.text + assert non_streaming_plan.marker in non_streaming_response.text, non_streaming_response.text + finally: + release.set() + streaming_response: Final = streaming_future.result(timeout=30) + assert streaming_response.status_code == 200, streaming_response.text + assert streaming_plan.marker in streaming_response.text, streaming_response.text + provider_rows: Final = rig.provider.drain() + sink_a_rows: Final = sink_a.drain() + sink_b_rows: Final = sink_b.drain() + _assert_successful_calls( + (streaming_plan, non_streaming_plan), + (streaming_response, non_streaming_response), + provider_rows, + sink_a_rows, + sink_b_rows, + name_a, + name_b, + ) + + +def _worker_processes(owned: OwnedProxy) -> tuple[psutil.Process, ...]: + return tuple( + child for child in psutil.Process(owned.process.pid).children() if "spawn_main" in _process_command(child) + ) + + +def _process_command(process: psutil.Process) -> str: + try: + return " ".join(process.cmdline()) + except psutil.Error: + return "" + + +def _safe_response(future: Future[httpx.Response]) -> httpx.Response | None: + try: + return future.result(timeout=70) + except (httpx.HTTPError, TimeoutError): + return None + + +def test_h3_worker_kill_mid_burst_keeps_remaining_worker_serving(rig: ChaosRig, tmp_path: Path) -> None: + plans: Final = _plans("h3-burst", 30) + gate_markers: Final = frozenset(plan.marker for plan in plans if plan.streamed) + started: Final = Event() + release: Final = Event() + + def gated_sink(request: Request) -> Reply: + if any(marker.encode() in request.body for marker in gate_markers): + started.set() + assert release.wait(timeout=60), "worker-kill sink gate was not released" + return _sink(request) + + with wire_server(gated_sink) as sink_a, wire_server(_sink) as sink_b: + name_a: Final = f"h3-stream-{uuid.uuid4().hex}" + name_b: Final = f"h3-non-stream-{uuid.uuid4().hex}" + rails: Final = ( + _rail(name_a, sink_a.url, "streaming", default_on=True), + _rail(name_b, sink_b.url, "non_streaming", default_on=True), + ) + with _owned_proxy(rig, tmp_path, rails, workers=2) as owned, ThreadPoolExecutor(max_workers=30) as pool: + workers: Final = _worker_processes(owned) + assert len(workers) == 2, tuple(worker.pid for worker in workers) + futures: Final[tuple[Future[httpx.Response], ...]] = tuple( + pool.submit(_request, owned.gateway, plan, ()) for plan in plans + ) + assert started.wait(timeout=30), "stream requests did not reach the owned sink" + victim: Final = workers[0] + survivor: Final = workers[1] + survivor_plan: Final = CallPlan( + f"audit-h3-survivor-{uuid.uuid4().hex}", + f"h3-survivor-{uuid.uuid4().hex}", + "messages", + False, + ) + try: + victim.send_signal(signal.SIGKILL) + survivor_response: Final = _request(owned.gateway, survivor_plan, ()) + assert survivor_response.status_code == 200 and survivor_plan.marker in survivor_response.text, ( + survivor_response.status_code, + survivor_response.text, + ) + assert survivor.is_running(), survivor.pid + finally: + release.set() + responses: Final = tuple(_safe_response(future) for future in futures) + assert owned.gateway.request("GET", "/health/liveliness").status_code == 200 + successful: Final = tuple( + (plan, response) + for plan, response in zip(plans, responses) + if response is not None and response.status_code == 200 + ) + provider_rows: Final = rig.provider.drain() + sink_a_rows: Final = sink_a.drain() + sink_b_rows: Final = sink_b.drain() + successful_plans: Final = (*tuple(plan for plan, _ in successful), survivor_plan) + successful_responses: Final = (*tuple(response for _, response in successful), survivor_response) + _assert_successful_calls( + successful_plans, + successful_responses, + provider_rows, + sink_a_rows, + sink_b_rows, + name_a, + name_b, + ) + + +def _create_stored_rail( + gateway: Gateway, + name: str, + sink_url: str, + stream_scope: Literal["streaming"] | None = "streaming", +) -> str: + response: Final = gateway.request( + "POST", + "/guardrails", + { + "guardrail": { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": f"{sink_url}/{name}", + "api_key": "synthetic-chaos-key", + **({"stream_scope": stream_scope} if stream_scope is not None else {}), + }, + } + }, + ) + assert response.status_code == 200, response.text + identity: Final = JSON_OBJECT.validate_json(response.content).get("guardrail_id") + assert isinstance(identity, str), response.text + return identity + + +def _assert_one_scope_wave( + plans: Sequence[CallPlan], + responses: Sequence[httpx.Response], + provider_rows: Sequence[Request], + sink_rows: Sequence[Request], + sink_name: str, + *, + database_url: str | None = None, +) -> None: + for plan, response in zip(plans, responses): + assert response.status_code == 200 and plan.marker in response.text, (plan, response.status_code, response.text) + provider_match: Final = _rows_for_marker(provider_rows, plan.marker) + sink_match: Final = _rail_scans(sink_rows, sink_name, plan.marker) + assert len(provider_match) == 1, (plan, provider_match) + assert len(sink_match) == int(plan.streamed), (plan, sink_match) + spend: Final = _one_spend_row(plan.call_id, database_url) + assert plan.call_id in (spend.get("request_id"), spend.get("litellm_call_id")), (plan, spend) + + +def test_h4_stored_scope_survives_owned_proxy_restart(rig: ChaosRig, tmp_path: Path) -> None: + with wire_server(_sink) as sink: + name: Final = f"h4-stored-{uuid.uuid4().hex}" + identity: Final = _create_stored_rail(rig.gateway, name, sink.url) + try: + with ExitStack() as first_stack: + first: Final = first_stack.enter_context(_owned_proxy(rig, tmp_path, ())) + first_plans: Final = _plans("h4-before", 20) + first_responses: Final = _call_wave(first.gateway, first_plans, (name,)) + first_provider: Final = rig.provider.drain() + first_sink: Final = sink.drain() + assert tuple(response.status_code for response in first_responses) == (200,) * 20, first_responses + _assert_one_scope_wave(first_plans, first_responses, first_provider, first_sink, name) + first_stack.close() + with _owned_proxy(rig, tmp_path, ()) as restarted: + recovery_plans: Final = _plans("h4-after", 20) + recovery_responses: Final = _call_wave(restarted.gateway, recovery_plans, (name,)) + recovery_provider: Final = rig.provider.drain() + recovery_sink: Final = sink.drain() + assert tuple(response.status_code for response in recovery_responses) == (200,) * 20, ( + recovery_responses, + ) + _assert_one_scope_wave( + recovery_plans, + recovery_responses, + recovery_provider, + recovery_sink, + name, + ) + finally: + deleted: Final = rig.gateway.request("DELETE", f"/guardrails/{identity}") + assert deleted.status_code == 200, deleted.text + + +@contextmanager +def _owned_postgres(directory: Path) -> Iterator[PostgresCluster]: + docker: Final = shutil.which("docker") + assert docker is not None, "Docker CLI is required for the H5 PostgreSQL outage test" + docker_info: Final = subprocess.run( + [docker, "info", "--format", "{{.ServerVersion}}"], + capture_output=True, + text=True, + check=False, + ) + assert docker_info.returncode == 0, docker_info.stderr + directory.mkdir(parents=True, exist_ok=True) + port: Final = _free_port() + container_name: Final = f"litellm-stream-scope-h5-{uuid.uuid4().hex}" + password: Final = uuid.uuid4().hex + created: Final = subprocess.run( + [ + docker, + "create", + "--name", + container_name, + "--env", + "POSTGRES_USER=postgres", + "--env", + "POSTGRES_PASSWORD", + "--env", + "POSTGRES_DB=postgres", + "--publish", + f"127.0.0.1:{port}:5432/tcp", + POSTGRES_IMAGE, + ], + env=os.environ | {"POSTGRES_PASSWORD": password}, + capture_output=True, + text=True, + check=False, + ) + assert created.returncode == 0, created.stderr + cluster: Final = PostgresCluster( + f"postgresql://postgres:{password}@127.0.0.1:{port}/postgres?sslmode=disable", + container_name, + docker, + directory / "postgres.log", + ) + try: + started: Final = _start_postgres(cluster) + assert started.returncode == 0, started.stderr + assert eventually(lambda: _postgres_is_ready(cluster), bool, seconds=70) + yield cluster + finally: + logs: Final = subprocess.run( + [docker, "logs", container_name], + capture_output=True, + text=True, + check=False, + ) + cluster.log_path.write_text(logs.stdout + logs.stderr) + removed: Final = subprocess.run( + [docker, "rm", "-f", container_name], + capture_output=True, + text=True, + check=False, + ) + assert removed.returncode == 0, removed.stderr + remaining: Final = subprocess.run( + [docker, "ps", "--all", "--quiet", "--filter", f"name={container_name}"], + capture_output=True, + text=True, + check=False, + ) + assert remaining.returncode == 0, remaining.stderr + assert not remaining.stdout.strip(), remaining.stdout + + +@dataclass(frozen=True, slots=True) +class PostgresCluster: + database_url: str + container_name: str + docker: str + log_path: Path + + +def _start_postgres(cluster: PostgresCluster) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [cluster.docker, "start", cluster.container_name], + capture_output=True, + text=True, + check=False, + ) + + +def _postgres_is_ready(cluster: PostgresCluster) -> bool: + readiness: Final = subprocess.run( + [ + cluster.docker, + "exec", + cluster.container_name, + "pg_isready", + "-h", + "127.0.0.1", + "-p", + "5432", + "-d", + "postgres", + "-U", + "postgres", + ], + capture_output=True, + text=True, + check=False, + ) + return readiness.returncode == 0 + + +def _postgres_is_running(cluster: PostgresCluster) -> bool: + state: Final = subprocess.run( + [cluster.docker, "inspect", "--format", "{{.State.Running}}", cluster.container_name], + capture_output=True, + text=True, + check=False, + ) + return state.returncode == 0 and state.stdout.strip() == "true" + + +@pytest.mark.timeout(180) +def test_h5_stored_scope_survives_owned_postgres_outage_and_recovers_once( + rig: ChaosRig, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + postgres_directory: Final = Path(os.environ["INTEGRATION_RESULTS_DIR"]) / f"owned-postgres-{uuid.uuid4().hex}" + with _owned_postgres(postgres_directory) as database, wire_server(_sink) as sink: + monkeypatch.setenv("INTEGRATION_PROXY_DATABASE_URL", database.database_url) + name: Final = f"h5-stored-{uuid.uuid4().hex}" + with _owned_proxy(rig, tmp_path, ()) as registrar: + _create_stored_rail(registrar.gateway, name, sink.url) + with _owned_proxy(rig, tmp_path, ()) as owned: + try: + preflight_plans: Final = ( + CallPlan( + f"audit-h5-preflight-{uuid.uuid4().hex}", + f"h5-preflight-{uuid.uuid4().hex}", + "chat", + True, + ), + ) + preflight_responses: Final = _call_wave(owned.gateway, preflight_plans, (name,)) + preflight_provider: Final = rig.provider.drain() + preflight_sink: Final = sink.drain() + assert tuple(response.status_code for response in preflight_responses) == (200,), (preflight_responses,) + _assert_one_scope_wave( + preflight_plans, + preflight_responses, + preflight_provider, + preflight_sink, + name, + database_url=database.database_url, + ) + outage_result: Final = subprocess.run( + [database.docker, "stop", database.container_name], + capture_output=True, + text=True, + check=False, + ) + assert outage_result.returncode == 0, outage_result.stderr + outage_plans: Final = _plans("h5-outage", 20) + outage_responses: Final = _call_wave(owned.gateway, outage_plans, (name,)) + outage_provider: Final = rig.provider.drain() + outage_sink: Final = sink.drain() + for plan, response in zip(outage_plans, outage_responses): + assert response.status_code == 200 and plan.marker in response.text, ( + plan, + response.status_code, + response.text, + ) + assert len(_rows_for_marker(outage_provider, plan.marker)) == 1, (plan, outage_provider) + sink_match: Final = _rail_scans(outage_sink, name, plan.marker) + assert len(sink_match) == int(plan.streamed), (plan, sink_match) + + recovered_database: Final = _start_postgres(database) + assert recovered_database.returncode == 0, recovered_database.stderr + postgres_ready: Final = eventually(lambda: _postgres_is_ready(database), bool, seconds=70) + assert postgres_ready + readiness: Final = eventually( + lambda: owned.gateway.client.get("/health/readiness"), + lambda response: response.status_code == 200 and response.json().get("db") == "connected", + seconds=70, + ) + assert readiness.json().get("db") == "connected", readiness.text + recovery_plans: Final = _plans("h5-recovery", 20) + recovery_responses: Final = _call_wave(owned.gateway, recovery_plans, (name,)) + recovery_provider: Final = rig.provider.drain() + recovery_sink: Final = sink.drain() + assert tuple(response.status_code for response in recovery_responses) == (200,) * 20, recovery_responses + _assert_one_scope_wave( + recovery_plans, + recovery_responses, + recovery_provider, + recovery_sink, + name, + database_url=database.database_url, + ) + finally: + if not _postgres_is_running(database): + restarted: Final = _start_postgres(database) + assert restarted.returncode == 0, restarted.stderr + assert eventually(lambda: _postgres_is_ready(database), bool, seconds=70) + + +@pytest.mark.timeout(360) +def test_h6_spend_rows_for_requests_served_during_postgres_restart_land_once( + rig: ChaosRig, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + pytest.skip("BUG: LIT-9050 spend rows for requests served during a Postgres restart are dropped") + postgres_directory: Final = Path(os.environ["INTEGRATION_RESULTS_DIR"]) / f"owned-postgres-{uuid.uuid4().hex}" + with _owned_postgres(postgres_directory) as database, wire_server(_sink) as sink: + monkeypatch.setenv("INTEGRATION_PROXY_DATABASE_URL", database.database_url) + name: Final = f"h6-stored-{uuid.uuid4().hex}" + plans: Final = _plans("h6-postgres-restart", 20) + outage_markers: Final = frozenset(plan.marker for plan in plans[:10]) + arrivals: Final = Barrier(len(plans) + 1) + during_outage: Final = Event() + after_restart: Final = Event() + + def _gated_provider(request: Request) -> Reply: + if request.method == "GET" and request.target.partition("?")[0] == "/v1/models": + return _provider(request) + body: Final = JSON_OBJECT.validate_json(request.body) + marker: Final = _marker(body) + arrivals.wait(timeout=70) + gate: Final = during_outage if marker in outage_markers else after_restart + assert gate.wait(timeout=70), marker + return _provider(request) + + with wire_server(_gated_provider) as provider: + h6_rig: Final = ChaosRig(rig.gateway, provider, rig.directory) + with _owned_proxy(h6_rig, tmp_path, ()) as registrar: + _create_stored_rail(registrar.gateway, name, sink.url, stream_scope=None) + with _owned_proxy(h6_rig, tmp_path, ()) as owned, ThreadPoolExecutor(max_workers=len(plans)) as pool: + futures: Final = tuple( + pool.submit(_request, owned.gateway, plan, (name,)) for plan in plans + ) + try: + arrivals.wait(timeout=70) + stopped: Final = subprocess.run( + [database.docker, "stop", database.container_name], + capture_output=True, + text=True, + check=False, + ) + assert stopped.returncode == 0, stopped.stderr + during_outage.set() + outage_responses: Final = tuple(future.result(timeout=70) for future in futures[:10]) + for plan, response in zip(plans[:10], outage_responses): + assert response.status_code == 200 and plan.marker in response.text, ( + plan, + response.status_code, + response.text, + ) + + restarted: Final = _start_postgres(database) + assert restarted.returncode == 0, restarted.stderr + postgres_ready: Final = eventually(lambda: _postgres_is_ready(database), bool, seconds=70) + assert postgres_ready + readiness: Final = eventually( + lambda: owned.gateway.client.get("/health/readiness"), + lambda response: response.status_code == 200 + and JSON_OBJECT.validate_python(cast(object, response.json())).get("db") == "connected", + seconds=70, + ) + readiness_body: Final = JSON_OBJECT.validate_python(cast(object, readiness.json())) + assert readiness_body.get("db") == "connected", readiness.text + after_restart.set() + responses: Final = tuple(future.result(timeout=70) for future in futures) + finally: + during_outage.set() + after_restart.set() + if not _postgres_is_running(database): + recovered: Final = _start_postgres(database) + assert recovered.returncode == 0, recovered.stderr + assert eventually(lambda: _postgres_is_ready(database), bool, seconds=70) + + provider_rows: Final = provider.drain() + sink_rows: Final = sink.drain() + for plan, response in zip(plans, responses): + assert response.status_code == 200 and plan.marker in response.text, ( + plan, + response.status_code, + response.text, + ) + assert len(_rows_for_marker(provider_rows, plan.marker)) == 1, (plan, provider_rows) + assert len(_rail_scans(sink_rows, name, plan.marker)) == 1, (plan, sink_rows) + + spend_rows: Final = eventually( + lambda: tuple(_spend_rows(plan.call_id, database.database_url) for plan in plans), + lambda values: all(len(rows) == 1 for rows in values), + seconds=70, + return_last_on_timeout=True, + ) + counts: Final = tuple(len(rows) for rows in spend_rows) + missing_ids: Final = tuple( + plan.call_id for plan, rows in zip(plans, spend_rows) if not rows + ) + duplicate_ids: Final = tuple( + plan.call_id for plan, rows in zip(plans, spend_rows) if len(rows) > 1 + ) + assert counts == (1,) * len(plans), { + "missing_ids": missing_ids, + "duplicate_ids": duplicate_ids, + "counts": counts, + } diff --git a/tests/integration/observability/test_guardrail_stream_scope_matrix.py b/tests/integration/observability/test_guardrail_stream_scope_matrix.py new file mode 100644 index 00000000000..b5c66b202cc --- /dev/null +++ b/tests/integration/observability/test_guardrail_stream_scope_matrix.py @@ -0,0 +1,2076 @@ +from __future__ import annotations + +import asyncio +import contextlib +import json +import os +import socket +import socketserver +import ssl +import threading +import uuid +from asyncio import run, wait_for +from collections.abc import Generator, Iterable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, Literal, TypeAlias, cast + +import anthropic +import httpx +import openai +import pytest +import websockets +import yaml +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from integration._support.client import Gateway, JsonValue, Scenario, eventually +from integration._support.database import read_rows +from integration._support.mcp import McpCaller, echo_tool, register_mcp, scripted_peer +from integration._support.process import owned_proxy_process +from integration._support.tls import server_context, write_self_signed_cert +from integration._support.upstream import delete_scenario, register_scenario +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.cost_calculation.cost_tracking_case import RealtimeResponse +from pydantic import TypeAdapter +from websockets.asyncio.server import ServerConnection, serve + +Endpoint: TypeAlias = Literal["chat", "messages", "responses"] +Mode: TypeAlias = Literal["pre_call", "during_call", "post_call", "logging_only"] +Scope: TypeAlias = Literal["streaming", "non_streaming"] +MatrixScope: TypeAlias = Scope | None +LoggingPhase: TypeAlias = Literal["request", "response"] +Endpoints: Final[tuple[Endpoint, ...]] = ("chat", "messages", "responses") +Modes: Final[tuple[Mode, ...]] = ("pre_call", "during_call", "post_call", "logging_only") +Scopes: Final[tuple[MatrixScope, ...]] = ("streaming", "non_streaming", None) +PROVIDER_TEXT: Final = "provider stream_scope control" +GUARDRAIL_PATH: Final = "/beta/litellm_basic_guardrail_api" +F1_BLOCKED_WORD: Final = "streamscope-mcp-blocked-word" +F3_STREAM_BLOCKED_WORD: Final = "matrix-f3-stream-block" +F3_NON_STREAM_BLOCKED_WORD: Final = "matrix-f3-non-stream-block" +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +LOGGING_PHASE: Final = TypeAdapter(LoggingPhase) + + +def _json(value: object) -> bytes: + return json.dumps(value, separators=(",", ":")).encode() + + +def _strings(value: JsonValue) -> tuple[str, ...]: + if isinstance(value, str): + return (value,) + if isinstance(value, list): + return tuple(chain.from_iterable(_strings(item) for item in value)) + if isinstance(value, dict): + return tuple(chain.from_iterable(_strings(item) for item in value.values())) + return () + + +def _marker(body: Mapping[str, JsonValue]) -> str: + return next((text for text in _strings(dict(body)) if text.startswith("audit-")), "audit-provider") + + +def _sse(events: Iterable[Mapping[str, JsonValue]]) -> tuple[bytes, ...]: + return tuple(f"data: {json.dumps(event, separators=(',', ':'))}\n\n".encode() for event in events) + ( + b"data: [DONE]\n\n", + ) + + +def _messages_stream(message: Mapping[str, JsonValue]) -> tuple[bytes, ...]: + content: Final = cast(list[JsonValue], message["content"]) + text: Final = cast(dict[str, JsonValue], content[0])["text"] + return ( + f"event: message_start\ndata: {json.dumps({**message, 'content': [], 'stop_reason': None, 'usage': {'input_tokens': 11, 'output_tokens': 0}})}\n\n".encode(), + b'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}\n\n', + f"event: content_block_delta\ndata: {json.dumps({'type': 'content_block_delta', 'index': 0, 'delta': {'type': 'text_delta', 'text': text}})}\n\n".encode(), + b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n', + f"event: message_delta\ndata: {json.dumps({'type': 'message_delta', 'delta': {'stop_reason': 'end_turn', 'stop_sequence': None}, 'usage': {'output_tokens': 4}})}\n\n".encode(), + b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + ) + + +def _responses_stream( + response: Mapping[str, JsonValue], output: Mapping[str, JsonValue], marker: str +) -> tuple[bytes, ...]: + events: Final[tuple[dict[str, JsonValue], ...]] = ( + {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, + {"type": "response.in_progress", "response": {**response, "status": "in_progress", "output": []}}, + {"type": "response.output_item.added", "item": output, "output_index": 0}, + { + "type": "response.content_part.added", + "item_id": f"msg-{marker}", + "output_index": 0, + "content_index": 0, + "part": {"type": "output_text", "text": "", "annotations": []}, + }, + { + "type": "response.output_text.delta", + "item_id": f"msg-{marker}", + "output_index": 0, + "content_index": 0, + "delta": f"{PROVIDER_TEXT} {marker}", + }, + { + "type": "response.output_text.done", + "item_id": f"msg-{marker}", + "output_index": 0, + "content_index": 0, + "text": f"{PROVIDER_TEXT} {marker}", + }, + { + "type": "response.content_part.done", + "item_id": f"msg-{marker}", + "output_index": 0, + "content_index": 0, + "part": cast(list[JsonValue], output["content"])[0], + }, + {"type": "response.output_item.done", "item": output, "output_index": 0}, + {"type": "response.completed", "response": response}, + ) + return tuple( + f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': index}, separators=(',', ':'))}\n\n".encode() + for index, event in enumerate(events) + ) + + +def _provider_reply(target: str, body: Mapping[str, JsonValue]) -> Reply: + marker: Final = _marker(body) + streamed: Final = bool(body.get("stream")) + if target == "/v1/chat/completions": + if streamed: + common: Final = {"id": f"chatcmpl-{marker}", "object": "chat.completion.chunk", "created": 1} + return Reply( + content_type="text/event-stream", + chunks=_sse( + ( + { + **common, + "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}], + }, + { + **common, + "choices": [ + {"index": 0, "delta": {"content": f"{PROVIDER_TEXT} {marker}"}, "finish_reason": None} + ], + }, + {**common, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + ) + ), + ) + return Reply( + body=_json( + { + "id": f"chatcmpl-{marker}", + "object": "chat.completion", + "created": 1, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": f"{PROVIDER_TEXT} {marker}"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + ) + ) + if target == "/v1/messages": + message: Final = { + "id": f"msg-{marker}", + "type": "message", + "role": "assistant", + "model": "synthetic-anthropic-model", + "content": [{"type": "text", "text": f"{PROVIDER_TEXT} {marker}"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 4}, + } + return ( + Reply(content_type="text/event-stream", chunks=_messages_stream(message)) + if streamed + else Reply(body=_json(message)) + ) + if target == "/v1/responses": + output: Final = { + "type": "message", + "id": f"msg-{marker}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": f"{PROVIDER_TEXT} {marker}", "annotations": []}], + } + response: Final = { + "id": f"resp-{marker}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [output], + "usage": {"input_tokens": 11, "output_tokens": 4, "total_tokens": 15}, + } + return ( + Reply(content_type="text/event-stream", chunks=_responses_stream(response, output, marker)) + if streamed + else Reply(body=_json(response)) + ) + if target.endswith(":generateContent"): + return Reply( + body=_json( + {"candidates": [{"content": {"role": "model", "parts": [{"text": f"{PROVIDER_TEXT} {marker}"}]}}]} + ) + ) + if target.endswith(":streamGenerateContent"): + return Reply( + content_type="text/event-stream", + chunks=( + f"data: {_json({'candidates': [{'content': {'role': 'model', 'parts': [{'text': f'{PROVIDER_TEXT} {marker}'}]}}]}).decode()}\n\n".encode(), + ), + ) + return Reply(status=404, body=_json({"error": {"message": f"unexpected target {target}"}})) + + +def _provider(request: Request) -> Reply: + if not request.body: + return Reply(status=400, body=_json({"error": "empty request body"})) + body: Final = JSON_OBJECT.validate_json(request.body) + if request.target.startswith("/passthrough"): + marker: Final = _marker(body) + if bool(body.get("stream")): + return Reply( + content_type="text/event-stream", + chunks=_sse(({"text": f"{PROVIDER_TEXT} {marker}"},)), + ) + return Reply(body=_json({"received": body})) + if "audit-g3-401-" in request.body.decode(): + return Reply(status=401, body=_json({"error": {"message": "synthetic provider unauthorized"}})) + target: Final = request.target.split("?", 1)[0] + provider_target: Final = "/v1/messages" if target.endswith("/anthropic/v1/messages") else target + return _provider_reply(provider_target, body) + + +def _sink(request: Request) -> Reply: + assert request.target.endswith(GUARDRAIL_PATH), request.target + JSON_OBJECT.validate_json(request.body) + if "/g1-" in request.target: + return Reply(status=500, body=_json({"error": "synthetic sink failure"})) + if "/g2-" in request.target or "/f2-block-" in request.target: + return Reply(body=_json({"action": "BLOCKED", "blocked_reason": "synthetic policy block"})) + return Reply(body=_json({"action": "NONE"})) + + +def _rail( + name: str, + sink: Wire, + *, + mode: str | list[str], + scope: str | Mapping[str, str] | None = None, + default_on: bool = False, +) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": mode, + "default_on": default_on, + **({"stream_scope": dict(scope)} if isinstance(scope, Mapping) else {}), + **({"stream_scope": scope} if isinstance(scope, str) else {}), + "api_base": f"{sink.url}/{name}", + "api_key": "synthetic-guardrail-key", + }, + } + + +def _realtime_filter_rail(name: str, scope: Scope, keyword: str) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "realtime_input_transcription", + "default_on": True, + "stream_scope": scope, + "blocked_words": [{"keyword": keyword, "action": "BLOCK"}], + }, + } + + +def _mcp_filter_rail(name: str, scope: MatrixScope) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_mcp_call", + "default_on": False, + **({"stream_scope": scope} if scope is not None else {}), + "blocked_words": [{"keyword": F1_BLOCKED_WORD, "action": "BLOCK"}], + }, + } + + +def _mode_scope_pairs() -> tuple[tuple[Mode, MatrixScope], ...]: + return tuple(chain.from_iterable(tuple((mode, scope) for scope in Scopes) for mode in Modes)) + + +def _rail_names() -> Mapping[str, str]: + return MappingProxyType( + { + **{f"a_{mode}_{scope or 'unset'}": f"a_{mode}_{scope or 'unset'}" for mode, scope in _mode_scope_pairs()}, + "c1_default": "c1_default", + "c2_key": "c2_key", + "c3_team": "c3_team", + "c4_modes": "c4_modes", + "c5_omitted": "c5_omitted", + "c6_both": "c6_both", + "c6_unset": "c6_unset", + "e_stream": "e_stream", + "e_non_stream": "e_non_stream", + "f1_stream": "f1_stream", + "f1_non_stream": "f1_non_stream", + "f1_unset": "f1_unset", + "f2_stream": "f2_stream", + "f2_non_stream": "f2-block-non-stream", + "f3_stream": "f3_stream", + "f3_non_stream": "f3_non_stream", + "g1_failure": "g1-failure", + "g2_block": "g2-block", + "g3_provider": "g3-provider", + "z_logging_only_barrier": "z_logging_only_barrier", + } + ) + + +def _configured_rails(names: Mapping[str, str], sink: Wire) -> tuple[dict[str, JsonValue], ...]: + return tuple( + _rail(names[f"a_{mode}_{scope or 'unset'}"], sink, mode=mode, scope=scope) + for mode, scope in _mode_scope_pairs() + ) + ( + _rail(names["c1_default"], sink, mode="post_call", scope="streaming"), + _rail(names["c2_key"], sink, mode="post_call", scope="streaming"), + _rail(names["c3_team"], sink, mode="post_call", scope="streaming"), + _rail( + names["c4_modes"], + sink, + mode=["pre_call", "post_call"], + scope={"pre_call": "non_streaming", "post_call": "streaming"}, + ), + _rail( + names["c5_omitted"], + sink, + mode=["pre_call", "post_call"], + scope={"post_call": "streaming"}, + ), + _rail(names["c6_both"], sink, mode="post_call", scope="both"), + _rail(names["c6_unset"], sink, mode="post_call"), + _rail(names["e_stream"], sink, mode="pre_call", scope="streaming"), + _rail(names["e_non_stream"], sink, mode="pre_call", scope="non_streaming"), + _mcp_filter_rail(names["f1_stream"], "streaming"), + _mcp_filter_rail(names["f1_non_stream"], "non_streaming"), + _mcp_filter_rail(names["f1_unset"], None), + _rail(names["f2_stream"], sink, mode="post_call", scope="streaming"), + _rail(names["f2_non_stream"], sink, mode="post_call", scope="non_streaming"), + _realtime_filter_rail(names["f3_stream"], "streaming", F3_STREAM_BLOCKED_WORD), + _realtime_filter_rail(names["f3_non_stream"], "non_streaming", F3_NON_STREAM_BLOCKED_WORD), + _rail(names["g1_failure"], sink, mode="pre_call", scope="streaming"), + _rail(names["g2_block"], sink, mode="pre_call", scope="streaming"), + _rail(names["g3_provider"], sink, mode="pre_call", scope="both"), + _rail(names["z_logging_only_barrier"], sink, mode="logging_only"), + ) + + +def _realtime_transcription_response(transcript: str) -> RealtimeResponse: + return RealtimeResponse( + content_type="application/x-realtime", + events=( + { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": "evt_$REQUEST_ID", + "item_id": "item_$REQUEST_ID", + "content_index": 0, + "transcript": transcript, + }, + { + "type": "response.done", + "event_id": "evt_$REQUEST_ID", + "response": { + "id": "resp_$REQUEST_ID", + "object": "realtime.response", + "status": "completed", + "output": [], + "usage": { + "total_tokens": 0, + "input_tokens": 0, + "output_tokens": 0, + "input_token_details": { + "text_tokens": 0, + "audio_tokens": 0, + "cached_tokens": 0, + "cached_tokens_details": {"text_tokens": 0, "audio_tokens": 0}, + }, + "output_token_details": {"text_tokens": 0, "audio_tokens": 0}, + }, + }, + }, + ), + ) + + +async def _collect_realtime_events(websocket: websockets.ClientConnection) -> tuple[dict[str, JsonValue], ...]: + event: Final = JSON_OBJECT.validate_json(await websocket.recv()) + if event.get("type") == "response.done": + return (event,) + return (event, *await _collect_realtime_events(websocket)) + + +async def _realtime_transcription_events( + url: str, + key: str, + model: str, +) -> tuple[dict[str, JsonValue], ...]: + websocket_url: Final = ( + f"{url.replace('http://', 'ws://').replace('https://', 'wss://').rstrip('/')}/v1/realtime?model={model}" + ) + async with websockets.connect(websocket_url, additional_headers={"Authorization": f"Bearer {key}"}) as websocket: + session: Final = JSON_OBJECT.validate_json(await websocket.recv()) + assert session.get("type") == "session.created", session + await websocket.send(json.dumps({"type": "response.create"})) + return await _collect_realtime_events(websocket) + + +@dataclass(frozen=True, slots=True) +class MatrixRig: + candidate: Gateway + scenario: Scenario + models: Mapping[str, str] + rails: Mapping[str, str] + key: str + provider: Wire + sink: Wire + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[MatrixRig]: + with httpx.Client( + base_url=os.environ["INTEGRATION_PROXY_URL"], + timeout=30, + trust_env=False, + ) as root_client: + root_gateway: Final = Gateway( + root_client, + os.environ.get("INTEGRATION_MASTER_KEY", "sk-integration-master"), + os.environ["INTEGRATION_UPSTREAM_URL"], + ) + directory: Final = tmp_path_factory.mktemp("guardrail-stream-scope-matrix") + with wire_server(_provider) as provider, wire_server(_sink) as sink: + names: Final = _rail_names() + config_base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + pass_through_paths: Final = ( + { + "path": "/pt", + "target": f"{provider.url}/passthrough", + "include_subpath": True, + "guardrails": {names["e_stream"]: None, names["e_non_stream"]: None}, + }, + { + "path": "/anthropic/v1/messages", + "target": f"{provider.url}/v1/messages", + "guardrails": {names["e_stream"]: None, names["e_non_stream"]: None}, + }, + { + "path": "/gemini/v1beta/models", + "target": "", + "include_subpath": True, + "guardrails": {names["e_stream"]: None, names["e_non_stream"]: None}, + }, + ) + config: Final = { + **config_base, + "guardrails": _configured_rails(names, sink), + "model_list": [ + { + "model_name": "matrix-pipeline-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{provider.url}/v1", + "api_key": "synthetic-provider-key", + }, + } + ], + "environment_variables": { + **config_base.get("environment_variables", {}), + "ANTHROPIC_API_BASE": provider.url, + "ANTHROPIC_API_KEY": "synthetic-anthropic-key", + "GEMINI_API_BASE": provider.url, + "GEMINI_API_KEY": "synthetic-gemini-key", + }, + "general_settings": { + **config_base["general_settings"], + "pass_through_endpoints": pass_through_paths, + }, + "policies": { + "matrix-f2": { + "guardrails": {"add": [names["f2_stream"], names["f2_non_stream"]]}, + "pipeline": { + "mode": "post_call", + "steps": [ + {"guardrail": names["f2_stream"], "on_pass": "next", "on_fail": "block"}, + {"guardrail": names["f2_non_stream"], "on_pass": "next", "on_fail": "block"}, + ], + }, + } + }, + "policy_attachments": [{"policy": "matrix-f2", "models": ["matrix-pipeline-model"]}], + } + config_path: Final = directory / "stream-scope-matrix.yaml" + config_path.write_text(yaml.safe_dump(config)) + with owned_proxy_process(root_gateway, directory, {}, config=config_path, workers=1) as owned: + with owned.gateway.scenario() as scenario: + models: Final = MappingProxyType( + { + "chat": scenario.model( + model="openai/gpt-4o-mini", + api_base=f"{provider.url}/v1", + api_key="synthetic-provider-key", + ), + "messages": scenario.model( + model="anthropic/claude-sonnet-4-5-20250929", + api_base=provider.url, + api_key="synthetic-provider-key", + ), + "responses": scenario.model( + model="openai/gpt-4o-mini", + api_base=f"{provider.url}/v1", + api_key="synthetic-provider-key", + ), + } + ) + key: Final = scenario.key(guardrails=[names["c2_key"]]) + yield MatrixRig(owned.gateway, scenario, models, names, key, provider, sink) + + +def _cell_body( + rig: MatrixRig, + endpoint: Endpoint, + marker: str, + streamed: bool, + rail: str | None, +) -> tuple[str, dict[str, JsonValue]]: + selected: Final = [] if rail is None else [rail] + if endpoint == "chat": + body: Final = { + "model": rig.models["chat"], + "messages": [{"role": "user", "content": marker}], + "guardrails": selected, + **({"stream": True} if streamed else {}), + } + return "/v1/chat/completions", body + if endpoint == "messages": + body = { + "model": rig.models["messages"], + "max_tokens": 32, + "messages": [{"role": "user", "content": marker}], + "guardrails": selected, + **({"stream": True} if streamed else {}), + } + return "/v1/messages", body + body = { + "model": rig.models["responses"], + "input": marker, + "guardrails": selected, + **({"stream": True} if streamed else {}), + } + return "/v1/responses", body + + +def _matching_requests(wire: Wire, marker: str) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if marker.encode() in request.body) + + +def _rail_scans(rows: Sequence[Request], rail_name: str, marker: str) -> tuple[Request, ...]: + return tuple( + request for request in rows if request.target.startswith(f"/{rail_name}/") and marker.encode() in request.body + ) + + +def _logging_phases(rows: Sequence[Request]) -> tuple[str, ...]: + return tuple( + sorted(LOGGING_PHASE.validate_python(JSON_OBJECT.validate_json(request.body)["input_type"]) for request in rows) + ) + + +@dataclass(slots=True) +class _SinkRowsAccumulator: + sink: Wire + rows: tuple[Request, ...] = () + + def drain(self) -> tuple[Request, ...]: + self.rows = (*self.rows, *self.sink.drain()) + return self.rows + + +def _logging_only_scans( + sink: Wire, + marker: str, + rail: str, + barrier: str, +) -> tuple[Request, ...]: + accumulator: Final = _SinkRowsAccumulator(sink) + eventually( + accumulator.drain, + lambda rows: bool(_rail_scans(rows, barrier, marker)), + seconds=70, + ) + return _rail_scans(accumulator.rows, rail, marker) + + +def _scope_scan_count(scope: MatrixScope, streamed: bool) -> int: + if scope is None or scope == "both": + return 1 + return int((scope == "streaming") == streamed) + + +def _expected_logging_phases(scope: MatrixScope, streamed: bool) -> tuple[str, ...]: + return ("request", "response") if _scope_scan_count(scope, streamed) else () + + +def _spend_row_for_call(call_id: str, content: bytes) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + "SELECT request_id, litellm_call_id, spend, prompt_tokens, completion_tokens, metadata " + 'FROM "LiteLLM_SpendLogs" WHERE request_id=%s OR litellm_call_id=%s', + (call_id, call_id), + ), + lambda values: len(values) >= 1, + seconds=70, + ) + assert len(rows) == 1, (call_id, rows, content) + assert call_id in (rows[0]["request_id"], rows[0]["litellm_call_id"]), content + return rows[0] + + +def _cell_rows( + rig: MatrixRig, + marker: str, + response: httpx.Response, + mode: Mode, + streamed: bool, + scope: MatrixScope, + call_id: str, + rail: str, +) -> tuple[tuple[Request, ...], tuple[Request, ...]]: + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, provider_rows, response.text) + _spend_row_for_call(call_id, response.content) + sink_rows: Final = ( + _logging_only_scans( + rig.sink, + marker, + rail, + rig.rails["z_logging_only_barrier"], + ) + if mode == "logging_only" + else _rail_scans(rig.sink.drain(), rail, marker) + ) + return provider_rows, sink_rows + + +def _run_matrix_cell( + rig: MatrixRig, + endpoint: Endpoint, + streamed: bool, + mode: Mode, + scope: MatrixScope, +) -> None: + marker: Final = f"audit-a-{endpoint}-{int(streamed)}-{mode}-{scope or 'unset'}-{uuid.uuid4().hex}" + call_id: Final = f"matrix-a-{uuid.uuid4().hex}" + rail: Final = rig.rails[f"a_{mode}_{scope or 'unset'}"] + path, body = _cell_body(rig, endpoint, marker, streamed, rail) + response: Final = rig.candidate.request( + "POST", + path, + body, + headers={"x-litellm-call-id": call_id}, + ) + assert response.status_code == 200, response.text + assert PROVIDER_TEXT in response.text and marker in response.text, response.text + provider_rows, sink_rows = _cell_rows(rig, marker, response, mode, streamed, scope, call_id, rail) + assert len(provider_rows) == 1, (marker, provider_rows, response.text) + if mode == "logging_only": + assert _logging_phases(sink_rows) == _expected_logging_phases(scope, streamed), ( + marker, + sink_rows, + response.text, + ) + else: + expected_scans: Final = _scope_scan_count(scope, streamed) + assert len(sink_rows) == expected_scans, (marker, sink_rows, response.text) + + +@pytest.mark.parametrize("endpoint", Endpoints) +@pytest.mark.parametrize("streamed", (False, True), ids=("S0", "S1")) +@pytest.mark.parametrize("mode", Modes) +@pytest.mark.parametrize("scope", Scopes, ids=("streaming", "non_streaming", "unset")) +def test_a_stream_scope_matrix( + rig: MatrixRig, + endpoint: Endpoint, + streamed: bool, + mode: Mode, + scope: MatrixScope, +) -> None: + _run_matrix_cell(rig, endpoint, streamed, mode, scope) + + +SDK_CASES: Final[tuple[str, ...]] = ( + "openai-chat-sync", + "openai-chat-async", + "openai-responses-sync", + "openai-responses-async", + "anthropic-messages-sync", + "anthropic-messages-async", +) + + +def _openai_sdk_base_url(rig: MatrixRig) -> str: + return f"{str(rig.candidate.client.base_url).rstrip('/')}/v1" + + +def _anthropic_sdk_base_url(rig: MatrixRig) -> str: + return str(rig.candidate.client.base_url).rstrip("/") + + +def _verify_sdk_call(rig: MatrixRig, marker: str, call_id: str, text: str, streamed: bool) -> None: + assert marker in text and PROVIDER_TEXT in text, text + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, provider_rows) + _spend_row_for_call(call_id, text.encode()) + sink_rows: Final = _rail_scans(rig.sink.drain(), rig.rails["c2_key"], marker) + assert len(sink_rows) == int(streamed), (marker, streamed, sink_rows) + + +def _run_sync_sdk(rig: MatrixRig, sdk: str, marker: str, call_id: str, streamed: bool) -> str: + if sdk == "openai-chat-sync": + with openai.OpenAI( + api_key=rig.key, + base_url=_openai_sdk_base_url(rig), + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=30), + ) as client: + response: Final = client.chat.completions.create( + model=rig.models["chat"], + messages=[{"role": "user", "content": marker}], + stream=streamed, + extra_headers={"x-litellm-call-id": call_id}, + ) + if streamed: + return "".join( + chunk.choices[0].delta.content or "" for chunk in response if chunk.choices[0].delta.content + ) + return cast(str, response.choices[0].message.content) + if sdk == "openai-responses-sync": + with openai.OpenAI( + api_key=rig.key, + base_url=_openai_sdk_base_url(rig), + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=30), + ) as client: + response = client.responses.create( + model=rig.models["responses"], + input=marker, + stream=streamed, + extra_headers={"x-litellm-call-id": call_id}, + ) + if streamed: + return "".join(event.delta for event in response if event.type == "response.output_text.delta") + return response.output_text + if sdk == "anthropic-messages-sync": + with anthropic.Anthropic( + api_key=rig.key, + base_url=_anthropic_sdk_base_url(rig), + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=30), + ) as client: + response = client.messages.create( + model=rig.models["messages"], + max_tokens=32, + messages=[{"role": "user", "content": marker}], + stream=streamed, + extra_headers={"x-litellm-call-id": call_id}, + ) + if streamed: + return "".join( + event.delta.text + for event in response + if event.type == "content_block_delta" and event.delta.type == "text_delta" + ) + return "".join(block.text for block in response.content if block.type == "text") + raise AssertionError(sdk) + + +async def _run_async_sdk(rig: MatrixRig, sdk: str, marker: str, call_id: str, streamed: bool) -> str: + if sdk == "openai-chat-async": + async with openai.AsyncOpenAI( + api_key=rig.key, + base_url=_openai_sdk_base_url(rig), + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False, timeout=30), + ) as client: + response = await client.chat.completions.create( + model=rig.models["chat"], + messages=[{"role": "user", "content": marker}], + stream=streamed, + extra_headers={"x-litellm-call-id": call_id}, + ) + if streamed: + chunks: Final = [ + chunk.choices[0].delta.content or "" async for chunk in response if chunk.choices[0].delta.content + ] + return "".join(chunks) + return cast(str, response.choices[0].message.content) + if sdk == "openai-responses-async": + async with openai.AsyncOpenAI( + api_key=rig.key, + base_url=_openai_sdk_base_url(rig), + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False, timeout=30), + ) as client: + response = await client.responses.create( + model=rig.models["responses"], + input=marker, + stream=streamed, + extra_headers={"x-litellm-call-id": call_id}, + ) + if streamed: + events: Final = [event.delta async for event in response if event.type == "response.output_text.delta"] + return "".join(events) + return response.output_text + if sdk == "anthropic-messages-async": + async with anthropic.AsyncAnthropic( + api_key=rig.key, + base_url=_anthropic_sdk_base_url(rig), + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False, timeout=30), + ) as client: + response = await client.messages.create( + model=rig.models["messages"], + max_tokens=32, + messages=[{"role": "user", "content": marker}], + stream=streamed, + extra_headers={"x-litellm-call-id": call_id}, + ) + if streamed: + events: Final = [ + event.delta.text + async for event in response + if event.type == "content_block_delta" and event.delta.type == "text_delta" + ] + return "".join(events) + return "".join(block.text for block in response.content if block.type == "text") + raise AssertionError(sdk) + + +@pytest.mark.parametrize("sdk", SDK_CASES) +@pytest.mark.parametrize("streamed", (False, True), ids=("S0", "S1")) +def test_b_streaming_scope_classifies_sdk_streams(rig: MatrixRig, sdk: str, streamed: bool) -> None: + marker: Final = f"audit-b-{sdk}-{int(streamed)}-{uuid.uuid4().hex}" + call_id: Final = f"matrix-b-{uuid.uuid4().hex}" + if sdk.endswith("-async"): + text: Final = run(_run_async_sdk(rig, sdk, marker, call_id, streamed)) + else: + text = _run_sync_sdk(rig, sdk, marker, call_id, streamed) + _verify_sdk_call(rig, marker, call_id, text, streamed) + + +def _yaml_proxy_config( + rig: MatrixRig, + guardrail_name: str, + parameters: Mapping[str, JsonValue], +) -> dict[str, JsonValue]: + config_base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + return cast( + dict[str, JsonValue], + { + **config_base, + "guardrails": [{"guardrail_name": guardrail_name, "litellm_params": dict(parameters)}], + "model_list": [ + { + "model_name": "scope-yaml-chat", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{rig.provider.url}/v1", + "api_key": "synthetic-provider-key", + }, + } + ], + "environment_variables": { + **config_base.get("environment_variables", {}), + "OPENAI_API_BASE": rig.provider.url, + "OPENAI_API_KEY": "synthetic-provider-key", + }, + }, + ) + + +def _management_params( + name: str, + rig: MatrixRig, + *, + mode: str | list[str] = "pre_call", + scope: JsonValue = "streaming", + default_on: bool = False, +) -> dict[str, JsonValue]: + return { + "guardrail": "generic_guardrail_api", + "mode": mode, + "default_on": default_on, + "api_base": f"{rig.sink.url}/{name}", + "api_key": "synthetic-guardrail-key", + "stream_scope": scope, + } + + +def _raw_chat( + gateway: Gateway, + marker: str, + streamed: bool, + model: str, + *, + rails: Sequence[str] = (), + key: str | None = None, + call_id: str | None = None, +) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": marker}], + "guardrails": list(rails), + **({"stream": True} if streamed else {}), + }, + key=key, + headers={} if call_id is None else {"x-litellm-call-id": call_id}, + ) + + +def _assert_raw_call( + rig: MatrixRig, + marker: str, + response: httpx.Response, + streamed: bool, + expected_scans: int, + rail: str, + call_id: str | None = None, +) -> tuple[Request, ...]: + assert response.status_code == 200, response.text + assert PROVIDER_TEXT in response.text and marker in response.text, response.text + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, provider_rows) + if call_id is not None: + _spend_row_for_call(call_id, response.content) + sink_rows: Final = _rail_scans(rig.sink.drain(), rail, marker) + assert len(sink_rows) == expected_scans, (marker, streamed, expected_scans, sink_rows) + return sink_rows + + +def test_c1_default_on_rail_respects_stream_scope(rig: MatrixRig, tmp_path: Path) -> None: + name: Final = f"c1-default-{uuid.uuid4().hex}" + parameters: Final = { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": True, + "stream_scope": "streaming", + "api_base": f"{rig.sink.url}/{name}", + "api_key": "synthetic-guardrail-key", + } + config_path: Final = tmp_path / "default-on-stream-scope.yaml" + config_path.write_text(yaml.safe_dump(_yaml_proxy_config(rig, name, parameters))) + with owned_proxy_process(rig.candidate, tmp_path, {}, config=config_path, workers=1) as owned: + for streamed in (False, True): + marker: Final = f"audit-c1-{int(streamed)}-{uuid.uuid4().hex}" + response: Final = _raw_chat(owned.gateway, marker, streamed, "scope-yaml-chat") + _assert_raw_call(rig, marker, response, streamed, int(streamed), name) + + +def test_c2_key_attached_rail_respects_stream_scope(rig: MatrixRig) -> None: + marker0: Final = f"audit-c2-0-{uuid.uuid4().hex}" + response0: Final = _raw_chat(rig.candidate, marker0, False, rig.models["chat"], key=rig.key) + _assert_raw_call(rig, marker0, response0, False, 0, rig.rails["c2_key"]) + marker1: Final = f"audit-c2-1-{uuid.uuid4().hex}" + response1: Final = _raw_chat(rig.candidate, marker1, True, rig.models["chat"], key=rig.key) + _assert_raw_call(rig, marker1, response1, True, 1, rig.rails["c2_key"]) + + +def test_c3_team_attached_rail_respects_stream_scope(rig: MatrixRig) -> None: + team: Final = rig.scenario.team(guardrails=[rig.rails["c3_team"]]) + key: Final = rig.scenario.key(team_id=team) + marker0: Final = f"audit-c3-0-{uuid.uuid4().hex}" + response0: Final = _raw_chat(rig.candidate, marker0, False, rig.models["chat"], key=key) + _assert_raw_call(rig, marker0, response0, False, 0, rig.rails["c3_team"]) + marker1: Final = f"audit-c3-1-{uuid.uuid4().hex}" + response1: Final = _raw_chat(rig.candidate, marker1, True, rig.models["chat"], key=key) + _assert_raw_call(rig, marker1, response1, True, 1, rig.rails["c3_team"]) + + +@pytest.mark.parametrize("streamed", (False, True), ids=("S0", "S1")) +def test_c4_per_mode_scope_map_selects_each_mode(rig: MatrixRig, streamed: bool) -> None: + marker: Final = f"audit-c4-{int(streamed)}-{uuid.uuid4().hex}" + response: Final = _raw_chat( + rig.candidate, + marker, + streamed, + rig.models["chat"], + rails=(rig.rails["c4_modes"],), + ) + assert response.status_code == 200, response.text + assert PROVIDER_TEXT in response.text and marker in response.text, response.text + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, provider_rows) + sink_rows: Final = _rail_scans(rig.sink.drain(), rig.rails["c4_modes"], marker) + request_bodies: Final = tuple(JSON_OBJECT.validate_json(row.body) for row in sink_rows) + request_texts: Final = tuple(chain.from_iterable(_strings(body.get("texts", [])) for body in request_bodies)) + assert len(sink_rows) == 1, (marker, streamed, sink_rows) + assert any(marker in text for text in request_texts), (marker, request_texts) + assert (any(PROVIDER_TEXT in text for text in request_texts)) == streamed, ( + marker, + streamed, + request_texts, + ) + + +@pytest.mark.parametrize("streamed", (False, True), ids=("S0", "S1")) +def test_c5_mode_omitted_from_scope_map_means_both(rig: MatrixRig, streamed: bool) -> None: + marker: Final = f"audit-c5-{int(streamed)}-{uuid.uuid4().hex}" + response: Final = _raw_chat( + rig.candidate, + marker, + streamed, + rig.models["chat"], + rails=(rig.rails["c5_omitted"],), + ) + sink_rows: Final = _assert_raw_call( + rig, + marker, + response, + streamed, + 1 + int(streamed), + rig.rails["c5_omitted"], + ) + request_bodies: Final = tuple(JSON_OBJECT.validate_json(row.body) for row in sink_rows) + request_texts: Final = tuple(chain.from_iterable(_strings(body.get("texts", [])) for body in request_bodies)) + assert sum(marker in text and PROVIDER_TEXT not in text for text in request_texts) == 1, (marker, request_texts) + assert sum(PROVIDER_TEXT in text for text in request_texts) == int(streamed), (marker, request_texts) + + +@pytest.mark.parametrize("streamed", (False, True), ids=("S0", "S1")) +def test_c6_explicit_both_matches_unset_scope(rig: MatrixRig, streamed: bool) -> None: + marker: Final = f"audit-c6-{int(streamed)}-{uuid.uuid4().hex}" + response: Final = _raw_chat( + rig.candidate, + marker, + streamed, + rig.models["chat"], + rails=(rig.rails["c6_both"], rig.rails["c6_unset"]), + ) + assert response.status_code == 200, response.text + assert PROVIDER_TEXT in response.text and marker in response.text, response.text + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, provider_rows) + sink_rows: Final = rig.sink.drain() + both_rows: Final = _rail_scans(sink_rows, rig.rails["c6_both"], marker) + unset_rows: Final = _rail_scans(sink_rows, rig.rails["c6_unset"], marker) + assert (len(both_rows), len(unset_rows)) == (1, 1), (marker, streamed, both_rows, unset_rows) + + +@pytest.mark.parametrize( + ("case", "scope"), + ( + ("missing", None), + ("null", None), + ("empty", ""), + ("invalid-scalar", "sometimes"), + ("invalid-map", {"unknown_mode": "streaming"}), + ), + ids=("missing", "null", "empty", "invalid-scalar", "invalid-map"), +) +def test_c7_yaml_unset_and_invalid_scope_values_are_tolerated( + rig: MatrixRig, + tmp_path: Path, + case: str, + scope: JsonValue | None, +) -> None: + name: Final = f"c7-yaml-{case}-{uuid.uuid4().hex}" + parameters: Final = { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": True, + "api_base": f"{rig.sink.url}/{name}", + "api_key": "synthetic-guardrail-key", + **({} if case == "missing" else {"stream_scope": scope}), + } + config_path: Final = tmp_path / f"{case}-stream-scope.yaml" + config_path.write_text(yaml.safe_dump(_yaml_proxy_config(rig, name, parameters))) + with owned_proxy_process(rig.candidate, tmp_path, {}, config=config_path, workers=1) as owned: + for streamed in (False, True): + marker: Final = f"audit-c7-{case}-{int(streamed)}-{uuid.uuid4().hex}" + response: Final = _raw_chat(owned.gateway, marker, streamed, "scope-yaml-chat") + _assert_raw_call(rig, marker, response, streamed, 1, name) + if case == "empty": + assert "Ignoring invalid stored stream_scope value of type str" in owned.log.read_text() + + +def test_c8_identical_requests_have_one_scan_and_spend_each(rig: MatrixRig) -> None: + marker: Final = f"audit-c8-identical-{uuid.uuid4().hex}" + call_ids: Final = tuple(f"matrix-c8-{uuid.uuid4().hex}" for _ in range(3)) + responses: Final = tuple( + _raw_chat( + rig.candidate, + marker, + True, + rig.models["chat"], + rails=(rig.rails["a_post_call_streaming"],), + call_id=call_id, + ) + for call_id in call_ids + ) + assert tuple(response.status_code for response in responses) == (200, 200, 200), tuple( + response.text for response in responses + ) + for call_id, response in zip(call_ids, responses): + _spend_row_for_call(call_id, response.content) + sink_rows: Final = _rail_scans(rig.sink.drain(), rig.rails["a_post_call_streaming"], marker) + assert len(sink_rows) == 3, (marker, sink_rows) + guardrail_call_ids: Final = tuple( + cast(str, JSON_OBJECT.validate_json(row.body).get("litellm_call_id")) for row in sink_rows + ) + assert set(guardrail_call_ids) == set(call_ids), (call_ids, guardrail_call_ids) + + +def _create_management_rail(rig: MatrixRig, name: str, params: Mapping[str, JsonValue]) -> str: + created: Final = rig.candidate.request( + "POST", + "/guardrails", + {"guardrail": {"guardrail_name": name, "litellm_params": dict(params)}}, + ) + assert created.status_code == 200, created.text + payload: Final = JSON_OBJECT.validate_json(created.content) + identity: Final = payload.get("guardrail_id") + assert isinstance(identity, str), payload + return identity + + +def _management_observation( + rig: MatrixRig, + name: str, + streamed: bool, +) -> tuple[int, int, int]: + marker: Final = f"audit-d-management-{int(streamed)}-{uuid.uuid4().hex}" + call_id: Final = f"matrix-d1-{uuid.uuid4().hex}" + response: Final = _raw_chat( + rig.candidate, + marker, + streamed, + rig.models["chat"], + rails=(name,), + call_id=call_id, + ) + provider_rows: Final = _matching_requests(rig.provider, marker) + sink_rows: Final = _rail_scans(rig.sink.drain(), name, marker) + if response.status_code == 200 and len(provider_rows) == 1: + assert marker in response.text, response.text + _spend_row_for_call(call_id, response.content) + return response.status_code, len(provider_rows), len(sink_rows) + + +def _eventually_management_scope( + rig: MatrixRig, + name: str, + expected: tuple[int, int], +) -> tuple[tuple[int, int, int], tuple[int, int, int]]: + return eventually( + lambda: ( + _management_observation(rig, name, False), + _management_observation(rig, name, True), + ), + lambda values: ( + (values[0][0], values[0][2]) == (200, expected[0]) + and (values[1][0], values[1][2]) == (200, expected[1]) + and values[0][1] == values[1][1] == 1 + ), + seconds=30, + ) + + +def test_d1_management_create_read_and_runtime_scope(rig: MatrixRig) -> None: + name: Final = f"d1-management-{uuid.uuid4().hex}" + params: Final = _management_params(name, rig) + identity: Final = _create_management_rail(rig, name, params) + try: + info: Final = rig.candidate.request("GET", f"/guardrails/{identity}/info") + listing: Final = rig.candidate.request("GET", "/v2/guardrails/list") + info_payload: Final = JSON_OBJECT.validate_json(info.content) + list_payload: Final = JSON_OBJECT.validate_json(listing.content) + list_rows: Final = list_payload.get("guardrails") + listed: Final = ( + tuple(row for row in list_rows if isinstance(row, dict) and row.get("guardrail_id") == identity) + if isinstance(list_rows, list) + else () + ) + info_params: Final = info_payload.get("litellm_params") + list_params: Final = listed[0].get("litellm_params") if listed else None + assert ( + info.status_code == 200 + and listing.status_code == 200 + and isinstance(info_params, dict) + and info_params.get("stream_scope") == "streaming" + and len(listed) == 1 + and isinstance(list_params, dict) + and list_params.get("stream_scope") == "streaming" + ), (info.status_code, info.text, listing.status_code, listing.text) + observed: Final = _eventually_management_scope(rig, name, (0, 1)) + assert observed[0][2] == 0 and observed[1][2] == 1, observed + finally: + deleted: Final = rig.candidate.request("DELETE", f"/guardrails/{identity}") + assert deleted.status_code == 200, deleted.text + + +def test_d2_patch_then_put_updates_runtime_scope(rig: MatrixRig) -> None: + name: Final = f"d2-management-{uuid.uuid4().hex}" + identity: Final = _create_management_rail(rig, name, _management_params(name, rig)) + try: + patched: Final = rig.candidate.request( + "PATCH", + f"/guardrails/{identity}", + {"litellm_params": {"stream_scope": "non_streaming"}}, + ) + assert patched.status_code == 200, patched.text + after_patch: Final = _eventually_management_scope(rig, name, (1, 0)) + assert after_patch[0][2] == 1 and after_patch[1][2] == 0, after_patch + params: Final = _management_params(name, rig) + updated: Final = rig.candidate.request( + "PUT", + f"/guardrails/{identity}", + {"guardrail": {"guardrail_name": name, "litellm_params": params}}, + ) + assert updated.status_code == 200, updated.text + after_put: Final = _eventually_management_scope(rig, name, (0, 1)) + assert after_put[0][2] == 0 and after_put[1][2] == 1, after_put + finally: + deleted: Final = rig.candidate.request("DELETE", f"/guardrails/{identity}") + assert deleted.status_code == 200, deleted.text + + +HOSTILE_SCOPE_VALUES: Final[tuple[JsonValue, ...]] = ( + 1, + [], + "", + "x" * 5000, + {"unknown_mode": "streaming"}, + {"pre_call": 1}, +) + + +@pytest.mark.parametrize("operation", ("POST", "PUT", "PATCH")) +@pytest.mark.parametrize( + "scope", HOSTILE_SCOPE_VALUES, ids=("integer", "list", "empty", "long", "unknown-key", "wrong-value") +) +@pytest.mark.parametrize("repeat", (1, 2), ids=("first", "second")) +def test_d3_management_rejects_invalid_scope_values( + rig: MatrixRig, + operation: str, + scope: JsonValue, + repeat: int, +) -> None: + name: Final = f"d3-management-{operation.lower()}-{repeat}-{uuid.uuid4().hex}" + identity: Final = _create_management_rail(rig, name, _management_params(name, rig)) if operation != "POST" else None + payload: Final = ( + {"litellm_params": {"stream_scope": scope}} + if operation == "PATCH" + else { + "guardrail": { + "guardrail_name": name, + "litellm_params": _management_params(name, rig, scope=scope), + } + } + ) + path: Final = "/guardrails" if operation == "POST" else f"/guardrails/{identity}" + try: + response: Final = rig.candidate.request(operation, path, payload) + if operation == "POST" and response.status_code == 200: + created_payload: Final = JSON_OBJECT.validate_json(response.content) + created_identity: Final = created_payload.get("guardrail_id") + if isinstance(created_identity, str): + rig.candidate.request("DELETE", f"/guardrails/{created_identity}") + assert response.status_code == 422 and "stream_scope" in response.text, ( + operation, + scope, + repeat, + response.status_code, + response.text, + ) + if operation == "POST": + stored: Final = read_rows( + 'SELECT guardrail_id FROM "LiteLLM_GuardrailsTable" WHERE guardrail_name=%s', + (name,), + ) + assert stored == [], (operation, scope, stored) + else: + existing: Final = rig.candidate.request("GET", f"/guardrails/{identity}/info") + existing_payload: Final = JSON_OBJECT.validate_json(existing.content) + existing_params: Final = existing_payload.get("litellm_params") + assert ( + existing.status_code == 200 + and isinstance(existing_params, dict) + and existing_params.get("stream_scope") == "streaming" + ), (operation, scope, existing.status_code, existing.text) + finally: + if isinstance(identity, str): + deleted: Final = rig.candidate.request("DELETE", f"/guardrails/{identity}") + assert deleted.status_code == 200, deleted.text + + +def test_d3_management_normalizes_uppercase_scope(rig: MatrixRig) -> None: + name: Final = f"d3-uppercase-{uuid.uuid4().hex}" + identity: Final = _create_management_rail(rig, name, _management_params(name, rig, scope="STREAMING")) + try: + info: Final = rig.candidate.request("GET", f"/guardrails/{identity}/info") + payload: Final = JSON_OBJECT.validate_json(info.content) + parameters: Final = payload.get("litellm_params") + assert ( + info.status_code == 200 and isinstance(parameters, dict) and parameters.get("stream_scope") == "streaming" + ), (info.status_code, info.text) + finally: + deleted: Final = rig.candidate.request("DELETE", f"/guardrails/{identity}") + assert deleted.status_code == 200, deleted.text + + +def test_d4_unauthenticated_management_create_is_rejected(rig: MatrixRig) -> None: + name: Final = f"d4-unauthenticated-{uuid.uuid4().hex}" + response: Final = rig.candidate.client.request( + "POST", + "/guardrails", + json={"guardrail": {"guardrail_name": name, "litellm_params": _management_params(name, rig)}}, + ) + assert response.status_code == 401, response.text + + +def _scoped_rows( + rows: Sequence[Request], + marker: str, + streaming_name: str, + non_streaming_name: str, +) -> tuple[tuple[Request, ...], tuple[Request, ...]]: + return ( + _rail_scans(rows, streaming_name, marker), + _rail_scans(rows, non_streaming_name, marker), + ) + + +HOSTILE_CLASSIFICATION_VALUES: Final[tuple[JsonValue, ...]] = ( + True, + 1, + [], + "", + "litellm-server-streaming", + "x" * 5000, +) + + +@pytest.mark.parametrize( + "value", + HOSTILE_CLASSIFICATION_VALUES, + ids=("true", "integer", "list", "empty", "marker", "long"), +) +def test_ea_is_streaming_request_body_does_not_change_chat_classification( + rig: MatrixRig, + value: JsonValue, +) -> None: + marker: Final = f"audit-ea-{uuid.uuid4().hex}" + response: Final = rig.candidate.request( + "POST", + "/v1/chat/completions", + { + "model": rig.models["chat"], + "messages": [{"role": "user", "content": marker}], + "guardrails": [rig.rails["e_stream"], rig.rails["e_non_stream"]], + "is_streaming_request": value, + }, + ) + assert response.status_code == 200, response.text + assert PROVIDER_TEXT in response.text and marker in response.text, response.text + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, provider_rows) + provider_body: Final = JSON_OBJECT.validate_json(provider_rows[0].body) + assert provider_body.get("is_streaming_request") == value, provider_rows[0].body.decode() + sink_rows: Final = rig.sink.drain() + streaming_rows, non_streaming_rows = _scoped_rows( + sink_rows, + marker, + rig.rails["e_stream"], + rig.rails["e_non_stream"], + ) + assert (len(streaming_rows), len(non_streaming_rows)) == (0, 1), (marker, streaming_rows, non_streaming_rows) + + +@pytest.mark.parametrize("endpoint", ("chat", "passthrough"), ids=("chat", "configured-pass-through")) +@pytest.mark.parametrize("value", (True, "litellm-server-streaming"), ids=("true", "marker")) +def test_eb_namespaced_caller_body_field_cannot_flip_classification( + rig: MatrixRig, + endpoint: str, + value: JsonValue, +) -> None: + marker: Final = f"audit-eb-{endpoint}-{uuid.uuid4().hex}" + body: Final = ( + { + "model": rig.models["chat"], + "messages": [{"role": "user", "content": marker}], + "guardrails": [rig.rails["e_stream"], rig.rails["e_non_stream"]], + "litellm_server_streaming_classification": value, + } + if endpoint == "chat" + else { + "marker": marker, + "litellm_server_streaming_classification": value, + } + ) + response: Final = rig.candidate.request( + "POST", + "/v1/chat/completions" if endpoint == "chat" else "/pt", + body, + ) + assert response.status_code == 200 and marker in response.text, response.text + actual_streamed: Final = response.headers.get("content-type", "").lower().startswith("text/event-stream") + assert not actual_streamed, (endpoint, response.headers, response.text) + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, provider_rows) + provider_body: Final = JSON_OBJECT.validate_json(provider_rows[0].body) + if endpoint == "passthrough" and value == "litellm-server-streaming": + assert "litellm_server_streaming_classification" not in provider_body, provider_rows[0].body.decode() + else: + assert provider_body.get("litellm_server_streaming_classification") == value, provider_rows[0].body.decode() + sink_rows: Final = rig.sink.drain() + streaming_rows, non_streaming_rows = _scoped_rows( + sink_rows, + marker, + rig.rails["e_stream"], + rig.rails["e_non_stream"], + ) + assert (len(streaming_rows), len(non_streaming_rows)) == (0, 1), (marker, streaming_rows, non_streaming_rows) + + +HOSTILE_STREAM_VALUES: Final[tuple[JsonValue, ...]] = ("true", 1, [], "") + + +@pytest.mark.parametrize("value", HOSTILE_STREAM_VALUES, ids=("string", "integer", "list", "empty")) +def test_ec_hostile_stream_values_follow_observed_response_shape(rig: MatrixRig, value: JsonValue) -> None: + marker: Final = f"audit-ec-{uuid.uuid4().hex}" + response: Final = rig.candidate.request( + "POST", + "/v1/chat/completions", + { + "model": rig.models["chat"], + "messages": [{"role": "user", "content": marker}], + "guardrails": [rig.rails["e_stream"], rig.rails["e_non_stream"]], + "stream": value, + }, + ) + assert response.status_code == 200, response.text + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, provider_rows) + observed_stream: Final = response.headers.get("content-type", "").startswith("text/event-stream") + assert marker in response.text, response.text + sink_rows: Final = rig.sink.drain() + streaming_rows, non_streaming_rows = _scoped_rows( + sink_rows, + marker, + rig.rails["e_stream"], + rig.rails["e_non_stream"], + ) + assert (len(streaming_rows), len(non_streaming_rows)) == ( + int(observed_stream), + int(not observed_stream), + ), (marker, value, response.headers, streaming_rows, non_streaming_rows) + + +@pytest.mark.parametrize("streamed", (False, True), ids=("stream-absent", "stream-true")) +def test_ed_configured_passthrough_forwards_caller_flag_and_uses_route_body_stream( + rig: MatrixRig, + streamed: bool, +) -> None: + marker: Final = f"audit-ed-{int(streamed)}-{uuid.uuid4().hex}" + body: Final = { + "marker": marker, + "is_streaming_request": True, + **({"stream": True} if streamed else {}), + } + response: Final = rig.candidate.request("POST", "/pt", body) + assert response.status_code == 200 and marker in response.text, response.text + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, provider_rows) + provider_body: Final = JSON_OBJECT.validate_json(provider_rows[0].body) + assert provider_body == body, (body, provider_body) + + +@pytest.mark.parametrize( + ("route", "body", "streamed"), + ( + ( + "/anthropic/v1/messages", + { + "model": "synthetic-anthropic-model", + "max_tokens": 32, + "messages": [{"role": "user", "content": "MARKER"}], + }, + False, + ), + ( + "/anthropic/v1/messages", + { + "model": "synthetic-anthropic-model", + "max_tokens": 32, + "messages": [{"role": "user", "content": "MARKER"}], + "stream": True, + }, + True, + ), + ( + "/gemini/v1beta/models/audit-model:generateContent", + {"contents": [{"parts": [{"text": "MARKER"}]}]}, + False, + ), + ( + "/gemini/v1beta/models/audit-model:streamGenerateContent?alt=sse", + {"contents": [{"parts": [{"text": "MARKER"}]}]}, + True, + ), + ), + ids=("anthropic-S0", "anthropic-S1", "gemini-generate", "gemini-stream"), +) +def test_eh_provider_passthrough_routes_classify_effective_streaming( + rig: MatrixRig, + route: str, + body: dict[str, JsonValue], + streamed: bool, +) -> None: + marker: Final = f"audit-eh-{uuid.uuid4().hex}" + request_body: Final = JSON_OBJECT.validate_python(json.loads(json.dumps(body).replace("MARKER", marker))) + team: Final = ( + rig.scenario.team( + metadata={ + "allowed_passthrough_routes": [ + "/gemini/v1beta/models/audit-model:generateContent", + "/gemini/v1beta/models/audit-model:streamGenerateContent", + ] + } + ) + if route.startswith("/gemini/") + else None + ) + key: Final = ( + rig.scenario.key(team_id=team, guardrails=[rig.rails["e_stream"], rig.rails["e_non_stream"]]) + if team is not None + else rig.scenario.key(guardrails=[rig.rails["e_stream"], rig.rails["e_non_stream"]]) + ) + headers: Final = {"x-goog-api-key": key} if route.startswith("/gemini/") else None + response: Final = rig.candidate.request( + "POST", + route, + request_body, + headers=headers, + ) + assert response.status_code == 200 and marker in response.text, response.text + actual_streamed: Final = response.headers.get("content-type", "").lower().startswith("text/event-stream") + assert actual_streamed is streamed, (route, streamed, response.headers, response.text) + provider_rows: Final = _matching_requests(rig.provider, marker) + assert len(provider_rows) == 1, (marker, provider_rows) + sink_rows: Final = rig.sink.drain() + streaming_rows, non_streaming_rows = _scoped_rows( + sink_rows, + marker, + rig.rails["e_stream"], + rig.rails["e_non_stream"], + ) + assert (len(streaming_rows), len(non_streaming_rows)) == ((1, 0) if streamed else (0, 1)), ( + marker, + route, + streaming_rows, + non_streaming_rows, + ) + + +def test_ei_unauthenticated_scoped_request_has_no_upstream_or_guardrail_call(rig: MatrixRig) -> None: + marker: Final = f"audit-ei-{uuid.uuid4().hex}" + response: Final = rig.candidate.client.request( + "POST", + "/v1/chat/completions", + json={ + "model": rig.models["chat"], + "messages": [{"role": "user", "content": marker}], + "stream": True, + "guardrails": [rig.rails["e_stream"]], + }, + ) + assert response.status_code == 401, response.text + assert _matching_requests(rig.provider, marker) == () + assert _rail_scans(rig.sink.drain(), rig.rails["e_stream"], marker) == () + + +@pytest.mark.parametrize( + ("rail_key", "blocked"), + (("f1_stream", False), ("f1_non_stream", True), ("f1_unset", True)), + ids=("streaming", "non_streaming", "unset"), +) +def test_f1_mcp_content_filter_respects_stream_scope( + rig: MatrixRig, + rail_key: Literal["f1_stream", "f1_non_stream", "f1_unset"], + blocked: bool, +) -> None: + with scripted_peer(echo_tool("echo")) as peer: + alias: Final = f"streamscope{uuid.uuid4().hex}" + identity: Final = register_mcp(rig.scenario, peer, alias) + key: Final = rig.scenario.key( + object_permission={"mcp_servers": [identity]}, + guardrails=[rig.rails[rail_key]], + ) + marker: Final = f"audit-f1-{uuid.uuid4().hex}" + caller: Final = McpCaller(rig.candidate, key, "rest", alias) + outcome: Final = caller.call( + f"{alias}-echo", + {"text": f"{F1_BLOCKED_WORD} {marker}"}, + identity, + ) + peer_requests: Final = peer.drain() + marker_requests: Final = tuple(request for request in peer_requests if marker in json.dumps(request)) + if blocked: + assert outcome.error is not None, outcome.raw + assert "block" in outcome.raw.casefold() or F1_BLOCKED_WORD in outcome.raw.casefold(), outcome.raw + assert marker_requests == (), (marker, marker_requests) + return + assert outcome.ok and marker in (outcome.text or ""), outcome.raw + assert len(marker_requests) == 1, (marker, marker_requests) + + +@pytest.mark.parametrize("streamed", (False, True), ids=("S0", "S1")) +def test_f2_mismatched_policy_step_skips_matching_sibling_still_enforces( + rig: MatrixRig, + streamed: bool, +) -> None: + marker: Final = f"audit-f2-{int(streamed)}-{uuid.uuid4().hex}" + response: Final = _raw_chat( + rig.candidate, + marker, + streamed, + "matrix-pipeline-model", + ) + provider_rows: Final = _matching_requests(rig.provider, marker) + sink_rows: Final = rig.sink.drain() + streaming_rows, non_streaming_rows = _scoped_rows( + sink_rows, + marker, + rig.rails["f2_stream"], + rig.rails["f2_non_stream"], + ) + assert len(provider_rows) == 1, (marker, provider_rows) + if streamed: + assert response.status_code == 200 and marker in response.text, response.text + assert (len(streaming_rows), len(non_streaming_rows)) == (1, 0), (marker, streaming_rows, non_streaming_rows) + return + assert response.status_code == 400 and "synthetic policy block" in response.text, response.text + assert (len(streaming_rows), len(non_streaming_rows)) == (0, 1), (marker, streaming_rows, non_streaming_rows) + + +def test_f3_realtime_transcription_uses_streaming_scope(rig: MatrixRig) -> None: + streaming_scenario_id: Final = f"matrix-f3-stream-{uuid.uuid4().hex}" + streaming_upstream: Final = register_scenario( + streaming_scenario_id, + _realtime_transcription_response(F3_STREAM_BLOCKED_WORD), + ) + rig.scenario.cleanups.callback(delete_scenario, streaming_upstream) + streaming_model: Final = rig.scenario.model( + model="openai/gpt-realtime-2", + api_key=streaming_scenario_id, + api_base=rig.candidate.upstream_url, + ) + non_streaming_scenario_id: Final = f"matrix-f3-non-stream-{uuid.uuid4().hex}" + non_streaming_upstream: Final = register_scenario( + non_streaming_scenario_id, + _realtime_transcription_response(F3_NON_STREAM_BLOCKED_WORD), + ) + rig.scenario.cleanups.callback(delete_scenario, non_streaming_upstream) + non_streaming_model: Final = rig.scenario.model( + model="openai/gpt-realtime-2", + api_key=non_streaming_scenario_id, + api_base=rig.candidate.upstream_url, + ) + key: Final = rig.scenario.key() + proxy_url: Final = str(rig.candidate.client.base_url).rstrip("/") + + streaming_events: Final = run(_realtime_transcription_events(proxy_url, key, streaming_model)) + streaming_transcriptions: Final = tuple( + event + for event in streaming_events + if event.get("type") == "conversation.item.input_audio_transcription.completed" + ) + streaming_errors: Final = tuple(event for event in streaming_events if event.get("type") == "error") + assert len(streaming_transcriptions) == 1, streaming_events + assert streaming_transcriptions[0].get("transcript") == F3_STREAM_BLOCKED_WORD, streaming_events + assert len(streaming_errors) == 1, streaming_events + streaming_error: Final = streaming_errors[0].get("error") + assert isinstance(streaming_error, dict) and streaming_error.get("type") == "guardrail_violation", streaming_events + + non_streaming_events: Final = run(_realtime_transcription_events(proxy_url, key, non_streaming_model)) + non_streaming_transcriptions: Final = tuple( + event + for event in non_streaming_events + if event.get("type") == "conversation.item.input_audio_transcription.completed" + ) + non_streaming_errors: Final = tuple(event for event in non_streaming_events if event.get("type") == "error") + completions: Final = tuple(event for event in non_streaming_events if event.get("type") == "response.done") + assert len(non_streaming_transcriptions) == 1, non_streaming_events + assert non_streaming_transcriptions[0].get("transcript") == F3_NON_STREAM_BLOCKED_WORD, non_streaming_events + assert non_streaming_errors == (), non_streaming_events + assert len(completions) == 1, non_streaming_events + + +def test_g1_sink_failure_fails_closed_only_when_rail_is_in_scope(rig: MatrixRig) -> None: + marker_s0: Final = f"audit-g1-S0-{uuid.uuid4().hex}" + response_s0: Final = _raw_chat( + rig.candidate, + marker_s0, + False, + rig.models["chat"], + rails=(rig.rails["g1_failure"],), + ) + assert response_s0.status_code == 200 and marker_s0 in response_s0.text, response_s0.text + provider_rows_s0: Final = _matching_requests(rig.provider, marker_s0) + sink_rows_s0: Final = _rail_scans(rig.sink.drain(), rig.rails["g1_failure"], marker_s0) + assert len(provider_rows_s0) == 1 and sink_rows_s0 == (), (marker_s0, provider_rows_s0, sink_rows_s0) + marker_s1: Final = f"audit-g1-S1-{uuid.uuid4().hex}" + response_s1: Final = _raw_chat( + rig.candidate, + marker_s1, + True, + rig.models["chat"], + rails=(rig.rails["g1_failure"],), + ) + assert response_s1.status_code == 500 and "Generic Guardrail API failed" in response_s1.text, response_s1.text + assert _matching_requests(rig.provider, marker_s1) == () + sink_rows: Final = _rail_scans(rig.sink.drain(), rig.rails["g1_failure"], marker_s1) + assert len(sink_rows) == 1, (marker_s1, sink_rows) + + +def test_g2_blocked_verdict_blocks_only_when_rail_is_in_scope(rig: MatrixRig) -> None: + marker_s0: Final = f"audit-g2-S0-{uuid.uuid4().hex}" + response_s0: Final = _raw_chat( + rig.candidate, + marker_s0, + False, + rig.models["chat"], + rails=(rig.rails["g2_block"],), + ) + assert response_s0.status_code == 200 and marker_s0 in response_s0.text, response_s0.text + assert len(_matching_requests(rig.provider, marker_s0)) == 1 + assert _rail_scans(rig.sink.drain(), rig.rails["g2_block"], marker_s0) == () + marker_s1: Final = f"audit-g2-S1-{uuid.uuid4().hex}" + response_s1: Final = _raw_chat( + rig.candidate, + marker_s1, + True, + rig.models["chat"], + rails=(rig.rails["g2_block"],), + ) + assert response_s1.status_code == 400 and "synthetic policy block" in response_s1.text, response_s1.text + assert _matching_requests(rig.provider, marker_s1) == () + sink_rows: Final = _rail_scans(rig.sink.drain(), rig.rails["g2_block"], marker_s1) + assert len(sink_rows) == 1, (marker_s1, sink_rows) + + +def test_g3_provider_errors_reach_caller_and_proxy_remains_usable(rig: MatrixRig) -> None: + marker: Final = f"audit-g3-401-{uuid.uuid4().hex}" + unauthorized: Final = _raw_chat( + rig.candidate, + marker, + False, + rig.models["chat"], + rails=(rig.rails["g3_provider"],), + ) + assert unauthorized.status_code == 401 and "synthetic provider unauthorized" in unauthorized.text, unauthorized.text + assert len(_matching_requests(rig.provider, marker)) == 1 + assert len(_rail_scans(rig.sink.drain(), rig.rails["g3_provider"], marker)) == 1 + unknown_marker: Final = f"audit-g3-unknown-{uuid.uuid4().hex}" + unknown: Final = _raw_chat( + rig.candidate, + unknown_marker, + False, + "audit-unknown-model", + rails=(rig.rails["g3_provider"],), + ) + assert unknown.status_code in (400, 404) and "model" in unknown.text.lower(), unknown.text + assert _matching_requests(rig.provider, unknown_marker) == () + healthy_marker: Final = f"audit-g3-healthy-{uuid.uuid4().hex}" + healthy: Final = _raw_chat(rig.candidate, healthy_marker, False, rig.models["chat"], key=rig.key) + assert healthy.status_code == 200 and healthy_marker in healthy.text, healthy.text + assert len(_matching_requests(rig.provider, healthy_marker)) == 1 + + +VERTEX_LIVE_HOST: Final = "us-central1-aiplatform.googleapis.com" +VERTEX_LIVE_PATH: Final = "/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" +VERTEX_LIVE_AUTHORITY: Final = f"{VERTEX_LIVE_HOST}:443" + + +@dataclass(frozen=True, slots=True) +class _VertexLivePeer: + port: int + paths: SimpleQueue[str] + authorizations: SimpleQueue[str] + frames: SimpleQueue[dict[str, JsonValue]] + + +@dataclass(frozen=True, slots=True) +class _VertexConnectTunnel: + url: str + authorities: SimpleQueue[str] + + +async def _answer_vertex_live_peer( + connection: ServerConnection, + paths: SimpleQueue[str], + authorizations: SimpleQueue[str], + frames: SimpleQueue[dict[str, JsonValue]], +) -> None: + request: Final = connection.request + assert request is not None + paths.put(request.path) + authorizations.put(request.headers.get("Authorization", "")) + async for raw_frame in connection: + frame: Final = JSON_OBJECT.validate_python(json.loads(raw_frame)) + frames.put(frame) + if frames.qsize() == 1: + await connection.send(json.dumps({"setupComplete": {}})) + + +async def _serve_vertex_live_peer( + tls: ssl.SSLContext, + paths: SimpleQueue[str], + authorizations: SimpleQueue[str], + frames: SimpleQueue[dict[str, JsonValue]], + ports: SimpleQueue[int], + stop: asyncio.Event, +) -> None: + async with serve( + lambda connection: _answer_vertex_live_peer(connection, paths, authorizations, frames), + "127.0.0.1", + 0, + ssl=tls, + ) as server: + ports.put(next(iter(server.sockets)).getsockname()[1]) + await stop.wait() + + +@contextmanager +def _vertex_live_peer(cert: tuple[Path, Path]) -> Iterator[_VertexLivePeer]: + loop: Final = asyncio.new_event_loop() + stop: Final = asyncio.Event() + paths: Final = SimpleQueue[str]() + authorizations: Final = SimpleQueue[str]() + frames: Final = SimpleQueue[dict[str, JsonValue]]() + ports: Final = SimpleQueue[int]() + thread: Final = threading.Thread( + target=loop.run_until_complete, + args=(_serve_vertex_live_peer(server_context(*cert), paths, authorizations, frames, ports, stop),), + daemon=True, + ) + thread.start() + try: + yield _VertexLivePeer(ports.get(timeout=10), paths, authorizations, frames) + finally: + loop.call_soon_threadsafe(stop.set) + thread.join(timeout=10) + loop.close() + + +def _pipe_vertex_socket(source: socket.socket, sink: socket.socket) -> None: + with contextlib.suppress(OSError): + for chunk in iter(lambda: source.recv(65536), b""): + sink.sendall(chunk) + with contextlib.suppress(OSError): + sink.shutdown(socket.SHUT_WR) + + +@contextmanager +def _vertex_connect_tunnel(peer: _VertexLivePeer) -> Generator[_VertexConnectTunnel, None, None]: + authorities: Final = SimpleQueue[str]() + + class ConnectHandler(socketserver.StreamRequestHandler): + rbufsize = 0 + request: socket.socket + + def handle(self) -> None: + request_line: Final = self.rfile.readline().decode().split() + authority: Final = request_line[1] if len(request_line) > 1 else "" + while self.rfile.readline() not in (b"\r\n", b""): + pass + authorities.put(authority) + if authority != VERTEX_LIVE_AUTHORITY: + self.wfile.write(b"HTTP/1.1 403 Forbidden\r\ncontent-length: 0\r\n\r\n") + return + self.wfile.write(b"HTTP/1.1 200 Connection established\r\n\r\n") + self.request.settimeout(10) + with socket.create_connection(("127.0.0.1", peer.port), timeout=10) as upstream: + outbound: Final = threading.Thread(target=_pipe_vertex_socket, args=(self.request, upstream)) + outbound.start() + _pipe_vertex_socket(upstream, self.request) + outbound.join(timeout=12) + + with socketserver.ThreadingTCPServer(("127.0.0.1", 0), ConnectHandler) as server: + thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}) + thread.start() + try: + yield _VertexConnectTunnel(f"http://127.0.0.1:{server.server_address[1]}", authorities) + finally: + server.shutdown() + thread.join(timeout=6) + + +async def _exchange_vertex_live_frame( + url: str, key: str, vertex_project: str, vertex_location: str, frame: str +) -> str | bytes: + websocket_url: Final = ( + f"{url.replace('http://', 'ws://', 1).rstrip('/')}/vertex_ai/live" + f"?vertex_project={vertex_project}&vertex_location={vertex_location}" + ) + async with websockets.connect( + websocket_url, + additional_headers={"Authorization": f"Bearer {key}"}, + open_timeout=10, + ) as websocket: + await websocket.send(frame) + reply: Final = await wait_for(websocket.recv(), timeout=10) + await websocket.close() + return reply + + +def test_ej_websocket_passthrough_runs_only_streaming_scoped_rails(rig: MatrixRig, tmp_path: Path) -> None: + vertex_project: Final = f"matrix-ej-{uuid.uuid4().hex}" + vertex_location: Final = "us-central1" + streaming_name: Final = f"ej-streaming-{uuid.uuid4().hex}" + non_streaming_name: Final = f"ej-non-streaming-{uuid.uuid4().hex}" + + def _token_response(_request: Request) -> Reply: + return Reply( + body=_json( + { + "access_token": "synthetic-vertex-access-token", + "expires_in": 3600, + "token_type": "Bearer", + } + ) + ) + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + private_key_pem: Final = private_key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ).decode("utf-8") + credentials_path: Final = tmp_path / "vertex-service-account.json" + + with ( + wire_server(_token_response) as token_double, + wire_server(lambda _request: Reply(body=_json({"action": "NONE"}))) as sink, + ): + credentials_path.write_text( + json.dumps( + { + "type": "service_account", + "project_id": vertex_project, + "private_key_id": uuid.uuid4().hex, + "private_key": private_key_pem, + "client_email": f"integration-test@{vertex_project}.iam.gserviceaccount.com", + "client_id": "123456789012345678901", + "token_uri": f"{token_double.url}/token", + "auth_uri": "https://accounts.google.com/o/oauth2/auth", + "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", + "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/integration-test", + } + ) + ) + rails: Final = ( + _rail(streaming_name, sink, mode="pre_call", scope="streaming", default_on=True), + _rail(non_streaming_name, sink, mode="pre_call", scope="non_streaming", default_on=True), + ) + config_data: Final = cast(object, yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config_base: Final = JSON_OBJECT.validate_python(config_data) + config: Final = cast(dict[str, JsonValue], {**config_base, "guardrails": list(rails)}) + config_path: Final = tmp_path / "ej-vertex-live-stream-scope.yaml" + config_path.write_text(yaml.safe_dump(config)) + certificate: Final = write_self_signed_cert(tmp_path, (VERTEX_LIVE_HOST,)) + with _vertex_live_peer(certificate) as peer, _vertex_connect_tunnel(peer) as tunnel: + with owned_proxy_process( + rig.candidate, + tmp_path, + { + "DEFAULT_VERTEXAI_PROJECT": vertex_project, + "DEFAULT_VERTEXAI_LOCATION": vertex_location, + "DEFAULT_GOOGLE_APPLICATION_CREDENTIALS": str(credentials_path), + "HTTPS_PROXY": tunnel.url, + "https_proxy": tunnel.url, + "NO_PROXY": "127.0.0.1,localhost", + "no_proxy": "127.0.0.1,localhost", + "SSL_CERT_FILE": str(certificate[0]), + }, + config=config_path, + workers=1, + ) as owned: + with owned.gateway.scenario() as scenario: + key: Final = scenario.key() + marker: Final = f"vertex-live-{uuid.uuid4().hex}" + frame: Final = json.dumps( + { + "clientContent": { + "turns": [{"role": "user", "parts": [{"text": marker}]}], + "turnComplete": True, + } + } + ) + client_reply_raw: Final = run( + _exchange_vertex_live_frame( + str(owned.gateway.client.base_url), + key, + vertex_project, + vertex_location, + frame, + ) + ) + proxy_log: Final = owned.log + + tunnel_authorities: Final = tuple( + tunnel.authorities.get_nowait() for _ in range(tunnel.authorities.qsize()) + ) + peer_path: Final = peer.paths.get_nowait() + peer_authorization: Final = peer.authorizations.get_nowait() + peer_frames: Final = tuple(peer.frames.get_nowait() for _ in range(peer.frames.qsize())) + client_reply: Final = JSON_OBJECT.validate_python(json.loads(client_reply_raw)) + expected_frame: Final = JSON_OBJECT.validate_python(json.loads(frame)) + token_requests: Final = token_double.drain() + sink_rows: Final = sink.drain() + streaming_rows: Final = tuple(row for row in sink_rows if row.target.startswith(f"/{streaming_name}/")) + non_streaming_rows: Final = tuple( + row for row in sink_rows if row.target.startswith(f"/{non_streaming_name}/") + ) + token_request_count: Final = len(token_requests) + streaming_count: Final = len(streaming_rows) + non_streaming_count: Final = len(non_streaming_rows) + assert tunnel_authorities == (VERTEX_LIVE_AUTHORITY,), ( + tunnel_authorities, + proxy_log, + ) + assert peer_path == VERTEX_LIVE_PATH, ( + peer_path, + proxy_log, + ) + assert peer_authorization == "Bearer synthetic-vertex-access-token", (peer_authorization, proxy_log) + assert peer_frames == (expected_frame,), (marker, peer_frames, proxy_log) + assert client_reply == {"setupComplete": {}}, (client_reply, proxy_log) + assert token_request_count == 1, (token_request_count, proxy_log) + assert streaming_count == 1, (streaming_count, non_streaming_count, proxy_log) + assert non_streaming_count == 0, (streaming_count, non_streaming_count, proxy_log) diff --git a/tests/integration/observability/test_prometheus_request_metrics.py b/tests/integration/observability/test_prometheus_request_metrics.py new file mode 100644 index 00000000000..042e4350640 --- /dev/null +++ b/tests/integration/observability/test_prometheus_request_metrics.py @@ -0,0 +1,274 @@ +from __future__ import annotations + +import json +import uuid +from hashlib import sha256 +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS +from litellm.types.integrations.prometheus import LATENCY_BUCKETS + +from tests.integration._support.client import Gateway, eventually, gateway_from_environment, object_value +from tests.integration._support.process import owned_proxy +from tests.integration._support.prometheus_series import Sample, label_values, scrape +from tests.integration._support.wire import Reply, Request, Wire, wire_server + +_GOOD: Final = "prometheus-good-endpoint" +_LIMITED: Final = "prometheus-rate-limited-endpoint" +_FAILING: Final = "prometheus-failing-endpoint" +_LATENCY: Final = "prometheus-latency-endpoint" +_END_USER: Final = f"prometheus-end-user-{uuid.uuid4().hex}" + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + upstream: Wire + + +def _respond(request: Request) -> Reply: + if json.loads(request.body)["model"] == "429": + return Reply( + status=429, + body=json.dumps({"error": {"message": "rate limited", "type": "rate_limit_error", "code": "429"}}).encode(), + ) + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "metered"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } + ).encode() + ) + + +def _config(directory: Path, upstream: Wire) -> Path: + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + params: Final = {"api_key": "sk-fixture", "api_base": f"{upstream.url}/v1"} + configuration["model_list"] = [ + *configuration["model_list"], + *( + { + "model_name": name, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + **params, + }, + } + for name in (_GOOD, _LATENCY) + ), + {"model_name": _LIMITED, "litellm_params": {"model": "openai/429", **params}}, + {"model_name": _FAILING, "litellm_params": {"model": "openai/429", **params}}, + ] + configuration["litellm_settings"]["callbacks"] = ["prometheus"] + configuration["litellm_settings"]["disable_end_user_cost_tracking_prometheus_only"] = True + configuration.setdefault("router_settings", {})["num_retries"] = 0 + path: Final = directory / "prometheus-request-metrics.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("prometheus-request-metrics") + with ( + wire_server(_respond) as upstream, + gateway_from_environment() as shared, + owned_proxy(shared, directory, {}, config=_config(directory, upstream)) as owned, + ): + yield _Rig(owned, upstream) + + +def _ask(rig: _Rig, model: str, key: str | None = None, **extra: list[str] | str) -> httpx.Response: + return rig.gateway.request( + "POST", + "/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"metrics {uuid.uuid4().hex}"}], **extra}, + key=key, + ) + + +def _series(samples: Sequence[Sample], name: str, **labels: str) -> tuple[Sample, ...]: + return tuple( + sample + for sample in samples + if sample.name == name and all(sample.labels.get(label) == value for label, value in labels.items()) + ) + + +def _until(rig: _Rig, name: str, **labels: str) -> tuple[Sample, ...]: + return eventually(lambda: _series(scrape(rig.gateway), name, **labels), bool, seconds=30) + + +def test_a_rate_limited_call_counts_as_a_failed_and_a_429_total_request(rig: _Rig) -> None: + response: Final = _ask(rig, _FAILING) + assert response.status_code == 429, response.text + assert len(rig.upstream.drain()) == 1 + failed: Final = _until( + rig, + "litellm_proxy_failed_requests_metric_total", + api_key_alias="None", + exception_class="Openai.RateLimitError", + exception_status="429", + hashed_api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, + requested_model=_FAILING, + route="/chat/completions", + ) + assert [sample.value for sample in failed] == [1.0] + totals: Final = _until( + rig, + "litellm_proxy_total_requests_metric_total", + hashed_api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, + requested_model=_FAILING, + status_code="429", + ) + assert [sample.value for sample in totals] == [1.0] + + +def test_a_good_call_exports_latency_histograms_on_the_shared_buckets_without_the_end_user(rig: _Rig) -> None: + response: Final = _ask(rig, _LATENCY, user=_END_USER, tags=["teamB"]) + assert response.status_code == 200, response.text + assert len(rig.upstream.drain()) == 1 + master: Final = { + "api_key_alias": "None", + "hashed_api_key": LITELLM_PROXY_MASTER_KEY_ALIAS, + "requested_model": _LATENCY, + } + _until(rig, "litellm_request_total_latency_metric_bucket", le="0.005", **master) + _until(rig, "litellm_llm_api_latency_metric_bucket", le="0.005", **master) + for name in ("litellm_request_total_latency_metric_count", "litellm_llm_api_latency_metric_count"): + assert [sample.value for sample in _until(rig, name, **master)] == [1.0], name + samples: Final = scrape(rig.gateway) + assert _END_USER not in label_values(samples) + expected: Final = {str(bucket).replace("inf", "+Inf") for bucket in LATENCY_BUCKETS} + for name in ( + "litellm_request_total_latency_metric_bucket", + "litellm_llm_api_latency_metric_bucket", + "litellm_overhead_latency_metric_bucket", + ): + assert {sample.labels["le"] for sample in _series(samples, name)} == expected, name + + +def test_client_side_fallbacks_count_one_success_and_one_failure(rig: _Rig) -> None: + recovered: Final = _ask(rig, _LIMITED, fallbacks=[_GOOD]) + assert recovered.status_code == 200, recovered.text + missing: Final = f"unknown-model-{uuid.uuid4().hex[:8]}" + failed: Final = _ask(rig, _LIMITED, fallbacks=[missing]) + assert failed.status_code >= 400, failed.text + rig.upstream.drain() + shared: Final = { + "api_key_alias": "None", + "exception_class": "Openai.RateLimitError", + "exception_status": "429", + "hashed_api_key": LITELLM_PROXY_MASTER_KEY_ALIAS, + "requested_model": _LIMITED, + } + succeeded: Final = _until(rig, "litellm_deployment_successful_fallbacks_total", fallback_model=_GOOD, **shared) + assert [sample.value for sample in succeeded] == [1.0] + lost: Final = _until(rig, "litellm_deployment_failed_fallbacks_total", fallback_model=missing, **shared) + assert [sample.value for sample in lost] == [1.0] + + +@dataclass(frozen=True, slots=True) +class _Budget: + remaining: float + total: float + hours: float + + +def _budget(samples: Sequence[Sample], scope: str, label: str, identity: str) -> _Budget | None: + remaining: Final = _series(samples, f"litellm_remaining_{scope}_budget_metric", **{label: identity}) + total: Final = _series(samples, f"litellm_{scope}_max_budget_metric", **{label: identity}) + hours: Final = _series(samples, f"litellm_{scope}_budget_remaining_hours_metric", **{label: identity}) + if len(remaining) != 1 or len(total) != 1 or len(hours) != 1: + return None + return _Budget(remaining[0].value, total[0].value, hours[0].value) + + +def _reconciled(rig: _Rig, scope: str, label: str, identity: str, info: str, field: str, lookup: str) -> _Budget: + def read() -> tuple[_Budget | None, float]: + record: Final = object_value(rig.gateway.get(info, {field: lookup})[_INFO_FIELDS[info]]) + return _budget(scrape(rig.gateway), scope, label, identity), float(str(record["max_budget"])) - float( + str(record["spend"]) + ) + + budget, remaining = eventually( + read, + lambda state: ( + state[0] is not None and state[0].remaining < 10.0 and abs(state[1] - state[0].remaining) <= 0.001 + ), + seconds=60, + ) + assert budget is not None + assert abs(remaining - budget.remaining) <= 0.001 + return budget + + +_INFO_FIELDS: Final = {"/team/info": "team_info", "/key/info": "info", "/user/info": "user_info"} + + +def test_a_team_call_exports_remaining_max_and_hours_gauges_matching_team_info(rig: _Rig) -> None: + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(max_budget=10, budget_duration="7d") + key: Final = scenario.key(team_id=team) + assert _ask(rig, _GOOD, key).status_code == 200 + assert len(rig.upstream.drain()) == 1 + budget: Final = _reconciled(rig, "team", "team", team, "/team/info", "team_id", team) + assert budget.total == 10.0 + assert 0 < budget.hours <= 168 + + +def test_a_key_call_exports_remaining_max_and_hours_gauges_matching_key_info(rig: _Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(max_budget=10, budget_duration="7d") + assert _ask(rig, _GOOD, key).status_code == 200 + assert len(rig.upstream.drain()) == 1 + hashed: Final = sha256(key.encode()).hexdigest() + budget: Final = _reconciled(rig, "api_key", "hashed_api_key", hashed, "/key/info", "key", key) + assert budget.total == 10.0 + assert 0 <= budget.hours <= 168 + + +def test_a_user_call_exports_remaining_max_and_hours_gauges_matching_user_info(rig: _Rig) -> None: + with rig.gateway.scenario() as scenario: + user: Final = f"prometheus-user-{uuid.uuid4().hex}" + scenario.user(user_id=user, max_budget=10, budget_duration="7d") + key: Final = scenario.key(user_id=user) + assert _ask(rig, _GOOD, key).status_code == 200 + assert len(rig.upstream.drain()) == 1 + budget: Final = _reconciled(rig, "user", "user", user, "/user/info", "user_id", user) + assert budget.total == 10.0 + assert 0 <= budget.hours <= 168 + + +def test_a_user_email_labels_the_spend_and_failed_request_series_of_its_keys(rig: _Rig) -> None: + with rig.gateway.scenario() as scenario: + email: Final = f"prometheus-{uuid.uuid4().hex}@example.com" + user: Final = f"prometheus-email-{uuid.uuid4().hex}" + scenario.user(user_id=user, user_email=email) + key: Final = scenario.key(user_id=user) + assert _ask(rig, _GOOD, key).status_code == 200 + assert len(rig.upstream.drain()) == 1 + spend: Final = _until(rig, "litellm_spend_metric_total", user_email=email) + assert [(sample.labels["user"], sample.value) for sample in spend] == [(user, pytest.approx(0.005))] + assert email in label_values(scrape(rig.gateway)) + assert _ask(rig, _FAILING, key).status_code == 429 + assert len(rig.upstream.drain()) == 1 + failed: Final = _until( + rig, "litellm_proxy_failed_requests_metric_total", user_email=email, requested_model=_FAILING + ) + assert [(sample.labels["user"], sample.value) for sample in failed] == [(user, 1.0)] diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py new file mode 100644 index 00000000000..f9d4f7f91bb --- /dev/null +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -0,0 +1,1072 @@ +import asyncio +import json +import re +import signal +import socket +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping, Sequence +from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import ExitStack +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final + +import anthropic +import httpx +import psutil +import pytest +import yaml +from anthropic.types import MessageCountTokensToolParam, MessageParam +from integration._support.bedrock_runtime_peer import answer, marker_of, target_of +from integration._support.bedrock_runtime_peer import respond as runtime_generation +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_OPUS: Final = "global.anthropic.claude-opus-4-8" +_OPUS_BASE: Final = "anthropic.claude-opus-4-8" +_SONNET: Final = "anthropic.claude-sonnet-4-6" +_NOVA: Final = "amazon.nova-lite-v1:0" +_REGION: Final = "us-east-1" +_OWNED_OPUS: Final = "mantle-opus" +_OWNED_SONNET: Final = "mantle-sonnet" +_OWNED_NOVA: Final = "mantle-nova" +_OWNED_MODELS: Final = MappingProxyType({_OWNED_OPUS: _OPUS, _OWNED_SONNET: _SONNET, _OWNED_NOVA: _NOVA}) +_ACCESS_KEY: Final = "AKIAINTEGRATIONMANTLE" +_SECRET_KEY: Final = "integration-mantle-secret" +_SIGV4_SCOPE: Final = f"/{_REGION}/bedrock/aws4_request" +_MANTLE_TARGET: Final = "/anthropic/v1/messages/count_tokens" +_MANTLE_VERSION: Final = "2023-06-01" +_MANTLE_COUNT: Final = 4242 +_RUNTIME_COUNT: Final = 1345 +_UNSUPPORTED: Final = "The provided model doesn't support counting tokens." +_REJECTION: Final = "scripted mantle rejection" +_COUNT_TARGET: Final = re.compile(r"^/model/(.+)/count-tokens$") +_INVOKE_TARGET: Final = re.compile(r"^/model/(.+)/invoke$") +_SCRIPTED_STATUS: Final = re.compile(r"status=(\d{3})") +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) + +_SDK_MESSAGES: Final[list[MessageParam]] = [{"role": "user", "content": "Count this message"}] +_MESSAGES: Final = _JSON_LIST.validate_json(json.dumps(_SDK_MESSAGES)) +_SYSTEM: Final = "You are a terse assistant that answers in one sentence" +_SYSTEM_BLOCKS: Final[list[JsonValue]] = [ + {"type": "text", "text": "You are a terse assistant"}, + {"type": "text", "text": "Answer in one sentence"}, +] +_SDK_TOOLS: Final[list[MessageCountTokensToolParam]] = [ + { + "name": "get_weather", + "description": "Look up the current weather for a city", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string", "description": "City to look up"}}, + "required": ["city"], + }, + } +] +_TOOLS: Final = _JSON_LIST.validate_json(json.dumps(_SDK_TOOLS)) +_GEMINI_BODY: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": "Count this"}]}]} +_GEMINI_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this"}] + + +def _mantle_body(**fields: JsonValue) -> dict[str, JsonValue]: + return {"model": _OPUS_BASE, "messages": _MESSAGES, **fields} + + +_MANTLE_BARE: Final = _mantle_body() +_MANTLE_FULL: Final = _mantle_body(system=_SYSTEM, tools=_TOOLS) + + +def _json_reply(status: int, payload: Mapping[str, JsonValue]) -> Reply: + return Reply(status=status, body=json.dumps(payload).encode()) + + +def _runtime_count(request: Request, model: str) -> Reply: + scripted: Final = _SCRIPTED_STATUS.search(request.body.decode(errors="replace")) + if scripted is not None: + return _json_reply(int(scripted.group(1)), {"message": f"scripted {scripted.group(1)}"}) + if "sonnet" in model: + return _json_reply(200, {"inputTokens": _RUNTIME_COUNT}) + return _json_reply(400, {"message": _UNSUPPORTED}) + + +def _invoke_reply(marker: str) -> Reply: + return _json_reply( + 200, + { + "id": f"msg_{marker}", + "type": "message", + "role": "assistant", + "model": _OPUS_BASE, + "content": [{"type": "text", "text": answer(marker)}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, + }, + ) + + +def _runtime(request: Request) -> Reply: + target: Final = target_of(request) + counted: Final = _COUNT_TARGET.match(target) + if counted is not None: + return _runtime_count(request, counted.group(1)) + if _INVOKE_TARGET.match(target): + return _invoke_reply(marker_of(request)) + return runtime_generation(request) + + +def _mantle_counted(_request: Request) -> Reply: + return _json_reply(200, {"input_tokens": _MANTLE_COUNT}) + + +def _rejected(status: int) -> Reply: + return _json_reply(status, {"type": "error", "error": {"type": "invalid_request_error", "message": _REJECTION}}) + + +def _rejecting(status: int) -> Callable[[Request], Reply]: + def count(_request: Request) -> Reply: + return _rejected(status) + + return count + + +def _anthropic_message(message: JsonValue) -> bool: + return isinstance(message, dict) and message.get("role") in ("user", "assistant") + + +def _anthropic_tool(tool: JsonValue) -> bool: + return isinstance(tool, dict) and isinstance(tool.get("name"), str) and isinstance(tool.get("input_schema"), dict) + + +def _strict(request: Request) -> Reply: + body: Final = _JSON_OBJECT.validate_json(request.body) + messages: Final = body.get("messages") + tools: Final = body.get("tools", []) + accepted: Final = ( + isinstance(messages, list) + and all(map(_anthropic_message, messages)) + and isinstance(body.get("system", ""), (str, list)) + and isinstance(tools, list) + and all(map(_anthropic_tool, tools)) + ) + return _mantle_counted(request) if accepted else _rejected(400) + + +def _mantle(count: Callable[[Request], Reply] = _mantle_counted) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == _MANTLE_TARGET: + return count(request) + return _json_reply(404, {"error": f"unscripted mantle target {request.target}"}) + + return respond + + +def _mantle_environment(port: int) -> Mapping[str, str]: + return {"BEDROCK_MANTLE_API_BASE": f"http://127.0.0.1:{port}"} + + +_INHERITED_BEARER: Final = ("AWS_BEARER_TOKEN_BEDROCK",) + + +def _reserved_port() -> int: + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + return reserve.getsockname()[1] + + +def _closed_port_url() -> str: + return f"http://127.0.0.1:{_reserved_port()}" + + +@pytest.fixture(scope="module") +def mantle_port() -> int: + return _reserved_port() + + +@pytest.fixture(scope="module") +def counting_proxy(tmp_path_factory: pytest.TempPathFactory, mantle_port: int) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("mantle-count"), + _mantle_environment(mantle_port), + workers=1, + remove_environment=_INHERITED_BEARER, + ) as owned, + ): + yield owned + + +def _litellm_params(model: str, api_base: str) -> dict[str, JsonValue]: + return { + "model": f"bedrock/{model}", + "api_base": api_base, + "aws_access_key_id": _ACCESS_KEY, + "aws_secret_access_key": _SECRET_KEY, + "aws_region_name": _REGION, + } + + +def _owned_config(path: Path, runtime_url: str, settings: Mapping[str, JsonValue]) -> Path: + config: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + litellm_settings: Final = object_value(config["litellm_settings"]) + path.write_text( + yaml.safe_dump( + { + **config, + "model_list": [ + {"model_name": name, "litellm_params": _litellm_params(model, runtime_url)} + for name, model in _OWNED_MODELS.items() + ], + "litellm_settings": {**litellm_settings, **settings}, + } + ) + ) + return path + + +def _deployment(scenario: Scenario, api_base: str, model: str = _OPUS) -> str: + return scenario.model(model_info=None, api_key=None, **_litellm_params(model, api_base)) + + +def _bare(model: str) -> dict[str, JsonValue]: + return {"model": model, "messages": _MESSAGES} + + +def _full(model: str) -> dict[str, JsonValue]: + return {**_bare(model), "system": _SYSTEM, "tools": _TOOLS} + + +def _count(gateway: Gateway, body: Mapping[str, JsonValue]) -> httpx.Response: + return gateway.request("POST", "/v1/messages/count_tokens", body) + + +def _payload(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text + return _JSON_OBJECT.validate_json(response.content) + + +def _local_count(gateway: Gateway, body: Mapping[str, JsonValue]) -> int: + response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "false"}) + payload: Final = _payload(response) + total: Final = payload["total_tokens"] + assert payload["tokenizer_type"] not in ("bedrock_api", "bedrock_mantle_api"), response.text + assert isinstance(total, int) and total > 0 and total not in (_MANTLE_COUNT, _RUNTIME_COUNT), response.text + return total + + +def _assert_sigv4(request: Request) -> None: + authorization: Final = request.headers.get("authorization", "") + assert authorization.startswith(f"AWS4-HMAC-SHA256 Credential={_ACCESS_KEY}/"), request.headers + assert _SIGV4_SCOPE in authorization, authorization + + +def _mantle_requests(requests: Sequence[Request]) -> tuple[dict[str, JsonValue], ...]: + for request in requests: + assert (request.method, request.target) == ("POST", _MANTLE_TARGET), request.target + assert request.headers["anthropic-version"] == _MANTLE_VERSION, request.headers + assert request.headers["content-type"] == "application/json", request.headers + _assert_sigv4(request) + return tuple(_JSON_OBJECT.validate_json(request.body) for request in requests) + + +def _mantle_bodies(wire: Wire) -> tuple[dict[str, JsonValue], ...]: + return _mantle_requests(wire.drain()) + + +def _runtime_count_targets(wire: Wire) -> tuple[str, ...]: + counts: Final = tuple(request for request in wire.drain() if _COUNT_TARGET.match(target_of(request))) + for request in counts: + assert request.method == "POST", request.method + _assert_sigv4(request) + return tuple(target_of(request) for request in counts) + + +def _clients(stack: ExitStack, base_url: str, count: int, timeout: float = 30) -> tuple[httpx.Client, ...]: + return tuple( + stack.enter_context(httpx.Client(base_url=base_url, timeout=timeout, trust_env=False)) for _ in range(count) + ) + + +def _counted_on(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue]: + response: Final = client.post( + "/v1/messages/count_tokens", json=dict(body), headers={"Authorization": f"Bearer {key}"} + ) + return response.status_code, _JSON_OBJECT.validate_json(response.content).get("input_tokens") + + +def _generated_then_counted( + client: httpx.Client, key: str, model: str, body: Mapping[str, JsonValue] +) -> tuple[int, int, JsonValue]: + generated: Final = client.post( + "/v1/messages", + json={ + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"Generate before counting marker-{uuid.uuid4().hex}"}], + }, + headers={"Authorization": f"Bearer {key}"}, + ) + return generated.status_code, *_counted_on(client, key, body) + + +def _counted_or_dropped(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue] | None: + try: + return _counted_on(client, key, body) + except httpx.TransportError: + return None + + +def _probed_then_counted_or_dropped( + client: httpx.Client, key: str, body: Mapping[str, JsonValue], probed: SimpleQueue[int] +) -> tuple[int, tuple[int, JsonValue] | None]: + port: Final = _local_port(client) + probed.put(port) + return port, _counted_or_dropped(client, key, body) + + +def _local_port(client: httpx.Client) -> int: + with client.stream("GET", "/health/liveliness") as response: + port: Final = int(response.extensions["network_stream"].get_extra_info("client_addr")[1]) + response.read() + assert response.status_code == 200, response.text + return port + + +def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]: + return frozenset( + connection.raddr.port + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.raddr and connection.laddr.port == proxy_port + ) + + +def _holding(held: SimpleQueue[str], release: threading.Event, seconds: float) -> Callable[[Request], Reply]: + def hold(request: Request) -> Reply: + held.put(request.target) + assert release.wait(timeout=seconds), "Held count was never released" + return _mantle_counted(request) + + return hold + + +def _anthropic_client(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +def _async_anthropic_client(gateway: Gateway) -> anthropic.AsyncAnthropic: + return anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +@pytest.mark.parametrize("system", [_SYSTEM, _SYSTEM_BLOCKS], ids=["string", "blocks"]) +def test_messages_count_tokens_counts_through_mantle_when_the_runtime_cannot( + counting_proxy: OwnedProxy, mantle_port: int, system: JsonValue +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) + assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) + + +def test_messages_count_tokens_without_system_or_tools_sends_a_bare_body_to_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, _bare(model)) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_BARE,) + + +def test_utils_token_counter_call_endpoint_reports_the_mantle_tokenizer( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request( + "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"} + ) + payload: Final = _payload(response) + assert (payload["total_tokens"], payload["tokenizer_type"]) == (_MANTLE_COUNT, "bedrock_mantle_api") + assert payload["original_response"] == {"input_tokens": _MANTLE_COUNT}, response.text + assert (payload["request_model"], payload["model_used"]) == (model, _OPUS), response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_FULL,) + + +def test_gemini_count_tokens_route_reports_the_mantle_total(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) + assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == ({"model": _OPUS_BASE, "messages": _GEMINI_MESSAGES},) + + +def test_gemini_count_tokens_route_reports_the_runtime_total_for_a_model_bedrock_counts( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url, _SONNET) + response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) + assert _payload(response) == {"totalTokens": _RUNTIME_COUNT, "promptTokensDetails": []}, response.text + assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",) + assert mantle.drain() == () + + +def test_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request( + "POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"} + ) + assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_BARE,) + + +def test_anthropic_sdk_count_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + counted: Final = _anthropic_client(gateway).messages.count_tokens( + model=model, messages=_SDK_MESSAGES, system=_SYSTEM, tools=_SDK_TOOLS + ) + assert counted.input_tokens == _MANTLE_COUNT, counted + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_FULL,) + + +def test_async_anthropic_sdk_count_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + counted: Final = asyncio.run( + _async_anthropic_client(gateway).messages.count_tokens(model=model, messages=_SDK_MESSAGES, system=_SYSTEM) + ) + assert counted.input_tokens == _MANTLE_COUNT, counted + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_mantle_body(system=_SYSTEM),) + + +def test_model_the_runtime_counts_never_reaches_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url, _SONNET) + assert _payload(_count(gateway, _full(model))) == {"input_tokens": _RUNTIME_COUNT} + detailed: Final = gateway.request( + "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"} + ) + payload: Final = _payload(detailed) + assert (payload["total_tokens"], payload["tokenizer_type"]) == (_RUNTIME_COUNT, "bedrock_api"), detailed.text + assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",) * 2 + assert mantle.drain() == () + + +def test_non_claude_model_the_runtime_cannot_count_falls_back_locally_without_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url, _NOVA) + response: Final = _count(gateway, _bare(model)) + assert _payload(response) == {"input_tokens": _local_count(gateway, _bare(model))}, response.text + assert _runtime_count_targets(runtime) == (f"/model/{_NOVA}/count-tokens",) + assert mantle.drain() == () + + +def test_runtime_403_on_a_claude_model_never_reaches_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + body: Final[dict[str, JsonValue]] = { + "model": model, + "messages": [{"role": "user", "content": "status=403 Count this message"}], + } + response: Final = _count(gateway, body) + assert _payload(response) == {"input_tokens": _local_count(gateway, body)}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert mantle.drain() == () + + +def test_unreachable_runtime_falls_back_locally_without_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario: + model: Final = _deployment(scenario, _closed_port_url()) + response: Final = _count(gateway, _full(model)) + assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text + assert mantle.drain() == () + + +def test_bedrock_passthrough_count_tokens_still_answers_the_runtime_rejection( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request("POST", "/bedrock/v1/messages/count_tokens", _full(model)) + assert response.status_code == 400, response.text + assert _UNSUPPORTED in response.text and "input_tokens" not in response.text, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert mantle.drain() == () + + +def test_responses_input_tokens_with_instructions_still_counts_locally( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": model, "input": "Count this message", "instructions": "Be terse"}, + ) + payload: Final = _payload(response) + assert len(_runtime_count_targets(runtime)) == 1 + (sent,) = _mantle_bodies(mantle) + messages: Final = sent["messages"] + assert isinstance(messages, list) and messages[0] == {"role": "system", "content": "Be terse"}, sent + assert "system" not in sent, sent + local: Final = _local_count(gateway, {"model": model, "messages": messages}) + assert payload == {"object": "response.input_tokens", "input_tokens": local}, response.text + + +@pytest.mark.parametrize("status", [400, 403, 404, 500, 503]) +def test_messages_count_tokens_falls_back_locally_when_mantle_errors( + counting_proxy: OwnedProxy, mantle_port: int, status: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_rejecting(status)), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, _full(model)) + assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_FULL,) + + +@pytest.mark.parametrize( + "body", + [b'{"inputTokens": 7}', b"not json at all", b"{}"], + ids=["runtime_key", "not_json", "empty_object"], +) +def test_messages_count_tokens_falls_back_locally_when_mantle_answers_without_input_tokens( + counting_proxy: OwnedProxy, mantle_port: int, body: bytes +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(lambda _request: Reply(body=body)), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, _bare(model)) + assert _payload(response) == {"input_tokens": _local_count(gateway, _bare(model))}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_BARE,) + + +def test_messages_count_tokens_falls_back_locally_when_mantle_is_unreachable(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_runtime) as runtime, gateway.scenario() as scenario: + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, _full(model)) + assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + + +def test_messages_count_tokens_falls_back_locally_when_mantle_never_answers( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + + def hold(request: Request) -> Reply: + held.put(request.target) + release.wait(timeout=120) + return Reply(drop_connection=True) + + with ExitStack() as stack: + (client,) = _clients(stack, str(gateway.client.base_url), 1, timeout=120) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=1)) + runtime: Final = stack.enter_context(wire_server(_runtime)) + mantle: Final = stack.enter_context(wire_server(_mantle(hold), port=mantle_port)) + stack.callback(release.set) + scenario: Final = stack.enter_context(gateway.scenario()) + model: Final = _deployment(scenario, runtime.url) + body: Final = _full(model) + local: Final = _local_count(gateway, body) + future: Final = pool.submit(_counted_on, client, gateway.key, body) + eventually(held.qsize, lambda size: size == 1, seconds=30) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + assert _local_count(gateway, body) == local + assert not future.done() + assert future.result(timeout=120) == (200, local) + assert len(_runtime_count_targets(runtime)) == 1 + assert len(mantle.drain()) == 1 + + +def test_messages_count_tokens_falls_back_locally_when_mantle_rejects_a_non_text_system( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, {**_bare(model), "system": 5}) + assert _payload(response) == {"input_tokens": _local_count(gateway, _bare(model))}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_mantle_body(system=5),) + + +@pytest.mark.parametrize( + "tools", + [5, "", "x" * 5120, ["get_weather"]], + ids=["int", "empty_string", "5kb_string", "list_of_strings"], +) +def test_messages_count_tokens_rejects_malformed_tools_without_calling_either_peer( + counting_proxy: OwnedProxy, mantle_port: int, tools: JsonValue +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + refused: Final = _count(gateway, {**_bare(model), "tools": tools}) + assert 400 <= refused.status_code < 600, refused.text + assert "input_tokens" not in refused.text, refused.text + assert runtime.drain() == () and mantle.drain() == () + assert _generated_then_counted(gateway.client, gateway.key, model, _bare(model)) == (200, 200, _MANTLE_COUNT) + assert _mantle_bodies(mantle) == (_MANTLE_BARE,) + + +@pytest.mark.parametrize("fields", [{}, {"messages": []}], ids=["missing", "empty"]) +def test_messages_count_tokens_without_messages_is_rejected( + counting_proxy: OwnedProxy, mantle_port: int, fields: dict[str, JsonValue] +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, {"model": model, "system": _SYSTEM, "tools": _TOOLS, **fields}) + assert response.status_code == 400, response.text + assert "messages parameter is required" in response.text, response.text + assert runtime.drain() == () and mantle.drain() == () + + +def test_messages_count_tokens_unauthenticated_request_never_reaches_either_peer( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", _full(model), key="sk-not-a-key-this-proxy-issued" + ) + assert response.status_code == 401, response.text + assert "input_tokens" not in response.text, response.text + assert runtime.drain() == () and mantle.drain() == () + + +def test_messages_count_tokens_forwards_a_5kb_system_verbatim_to_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + system: Final = "Answer in one sentence. " * 214 + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, {**_bare(model), "system": system}) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_mantle_body(system=system),) + + +def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value_each_to_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + fields: Final = f'"system": {json.dumps(_SYSTEM)}, "tools": {json.dumps(_TOOLS)}' + response: Final = gateway.client.post( + "/v1/messages/count_tokens", + content=f'{{"model": "{model}", "messages": {json.dumps(_MESSAGES)}, {fields}, {fields}}}', + headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, + ) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + (sent,) = mantle.drain() + assert _mantle_requests((sent,)) == (_MANTLE_FULL,) + assert (sent.body.count(b'"system"'), sent.body.count(b'"tools"')) == (1, 1), sent.body + + +def test_messages_count_tokens_leaves_empty_tools_and_system_out_of_the_mantle_body( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, {**_bare(model), "tools": [], "system": ""}) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_BARE,) + + +def test_messages_count_tokens_repeated_request_reaches_both_peers_each_time( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2)) + assert answers == ({"input_tokens": _MANTLE_COUNT},) * 2 + assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2 + assert _mantle_bodies(mantle) == (_MANTLE_BARE,) * 2 + + +def test_disabled_token_counter_surfaces_the_mantle_error_instead_of_counting_locally( + gateway: Gateway, mantle_port: int, tmp_path: Path +) -> None: + with ExitStack() as stack: + runtime: Final = stack.enter_context(wire_server(_runtime)) + config: Final = _owned_config( + tmp_path / "disabled-token-counter.yaml", runtime.url, {"disable_token_counter": True} + ) + owned: Final = stack.enter_context( + owned_proxy_process( + gateway, + tmp_path, + _mantle_environment(mantle_port), + config=config, + workers=2, + remove_environment=_INHERITED_BEARER, + ) + ) + with wire_server(_mantle(_rejecting(403)), port=mantle_port) as refusing: + refused: Final = _count(owned.gateway, _full(_OWNED_OPUS)) + assert refused.status_code == 403, refused.text + assert _REJECTION in refused.text and "input_tokens" not in refused.text, refused.text + assert _mantle_bodies(refusing) == (_MANTLE_FULL,) + with wire_server(_mantle(), port=mantle_port) as counting: + assert _payload(_count(owned.gateway, _full(_OWNED_OPUS))) == {"input_tokens": _MANTLE_COUNT} + assert _mantle_bodies(counting) == (_MANTLE_FULL,) + assert _payload(_count(owned.gateway, _full(_OWNED_SONNET))) == {"input_tokens": _RUNTIME_COUNT} + rejected: Final = _count(owned.gateway, _full(_OWNED_NOVA)) + assert rejected.status_code == 400, rejected.text + assert _UNSUPPORTED in rejected.text and "input_tokens" not in rejected.text, rejected.text + assert counting.drain() == () + assert len(_runtime_count_targets(runtime)) == 4 + + +def test_mantle_outage_between_concurrent_waves_falls_back_then_recovers( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ExitStack() as stack: + clients: Final = _clients(stack, str(gateway.client.base_url), 8) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + runtime: Final = stack.enter_context(wire_server(_runtime)) + scenario: Final = stack.enter_context(gateway.scenario()) + model: Final = _deployment(scenario, runtime.url) + body: Final = _full(model) + local: Final = _local_count(gateway, body) + + def generate_then_count(client: httpx.Client) -> tuple[int, int, JsonValue]: + return _generated_then_counted(client, gateway.key, model, body) + + def count_only(client: httpx.Client) -> tuple[int, JsonValue]: + return _counted_on(client, gateway.key, body) + + with wire_server(_mantle(), port=mantle_port) as mantle: + assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _MANTLE_COUNT),) * len(clients) + assert _mantle_bodies(mantle) == (_MANTLE_FULL,) * len(clients) + outage: Final = tuple(pool.map(count_only, clients)) + assert outage == ((200, local),) * len(clients) + with wire_server(_mantle(), port=mantle_port) as revived: + assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _MANTLE_COUNT),) * len(clients) + assert _mantle_bodies(revived) == (_MANTLE_FULL,) * len(clients) + assert len(_runtime_count_targets(runtime)) == 3 * len(clients) + + +def test_slow_mantle_holds_concurrent_counts_without_stalling_the_proxy( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + with ExitStack() as stack: + clients: Final = _clients(stack, str(gateway.client.base_url), 6) + runtime: Final = stack.enter_context(wire_server(_runtime)) + mantle: Final = stack.enter_context(wire_server(_mantle(_holding(held, release, 20)), port=mantle_port)) + scenario: Final = stack.enter_context(gateway.scenario()) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + stack.callback(release.set) + model: Final = _deployment(scenario, runtime.url) + futures: Final = tuple( + pool.submit(_generated_then_counted, client, gateway.key, model, _full(model)) for client in clients + ) + eventually(held.qsize, lambda size: size == len(clients), seconds=30) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + assert _local_count(gateway, _full(model)) > 0 + assert not any(future.done() for future in futures) + release.set() + assert tuple(future.result(timeout=30) for future in futures) == ((200, 200, _MANTLE_COUNT),) * len(clients) + assert _mantle_bodies(mantle) == (_MANTLE_FULL,) * len(clients) + assert len(_runtime_count_targets(runtime)) == len(clients) + + +def test_worker_sigkill_mid_burst_leaves_the_sibling_counting( + gateway: Gateway, mantle_port: int, tmp_path: Path +) -> None: + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + with ExitStack() as stack: + runtime: Final = stack.enter_context(wire_server(_runtime)) + mantle: Final = stack.enter_context(wire_server(_mantle(_holding(held, release, 60)), port=mantle_port)) + config: Final = _owned_config(tmp_path / "worker-kill.yaml", runtime.url, {}) + owned: Final = stack.enter_context( + owned_proxy_process( + gateway, + tmp_path, + _mantle_environment(mantle_port), + config=config, + workers=2, + remove_environment=_INHERITED_BEARER, + ) + ) + body: Final = _full(_OWNED_OPUS) + proxy_url: Final = owned.gateway.client.base_url + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + clients: Final = _clients(stack, str(proxy_url), 12) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + stack.callback(release.set) + probed: Final[SimpleQueue[int]] = SimpleQueue() + futures: Final[tuple[Future[tuple[int, tuple[int, JsonValue] | None]], ...]] = tuple( + pool.submit(_probed_then_counted_or_dropped, client, owned.gateway.key, body, probed) for client in clients + ) + eventually(held.qsize, lambda size: size == len(clients), seconds=30) + ports: Final = frozenset(probed.get_nowait() for _ in clients) + shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & ports for pid in workers} + assert sum(map(len, shares.values())) == len(clients), shares + victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) + psutil.Process(victim).send_signal(signal.SIGKILL) + release.set() + for port, result in (future.result(timeout=60) for future in futures): + assert result == (None if port in shares[victim] else (200, _MANTLE_COUNT)), (port, result, shares) + second_wave: Final = _clients(stack, str(proxy_url), 6) + assert tuple(_counted_on(client, owned.gateway.key, body) for client in second_wave) == ( + (200, _MANTLE_COUNT), + ) * len(second_wave) + assert len(_mantle_bodies(mantle)) == len(clients) + len(second_wave) + eventually(lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda started: started >= 3, 120) + assert owned.process.poll() is None + + +def _streamed_text(text: str) -> str: + events: Final = tuple( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in text.splitlines() + if line.startswith("data: {") + ) + return "".join(_delta_content(event) for event in events) + + +def _delta_content(event: dict[str, JsonValue]) -> str: + choices: Final = event.get("choices") + if not isinstance(choices, list) or not choices: + return "" + delta: Final = object_value(choices[0]).get("delta") + if not isinstance(delta, dict): + return "" + content: Final = delta.get("content") + return content if isinstance(content, str) else "" + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_chat_completions_on_the_same_deployment_still_generate( + counting_proxy: OwnedProxy, mantle_port: int, stream: bool +) -> None: + gateway: Final = counting_proxy.gateway + marker: Final = uuid.uuid4().hex + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"chat control marker-{marker}"}], + "stream": stream, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + generated: Final = _streamed_text(response.text) if stream else response.text + assert answer(marker) in generated, response.text + assert not stream or response.text.rstrip().endswith("data: [DONE]"), response.text + (sent,) = runtime.drain() + assert target_of(sent) == f"/model/{_OPUS}/{'converse-stream' if stream else 'converse'}", sent.target + assert mantle.drain() == () + + +def test_messages_endpoint_on_the_same_deployment_still_generates(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + marker: Final = uuid.uuid4().hex + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request( + "POST", + "/v1/messages", + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": f"marker-{marker}"}]}, + ) + assert response.status_code == 200, response.text + assert answer(marker) in response.text, response.text + (sent,) = runtime.drain() + assert target_of(sent) == f"/model/{_OPUS}/invoke", sent.target + assert mantle.drain() == () + + +def test_responses_endpoint_on_the_same_deployment_still_generates( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + marker: Final = uuid.uuid4().hex + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": f"marker-{marker}"}) + assert response.status_code == 200, response.text + assert answer(marker) in response.text, response.text + (sent,) = runtime.drain() + assert target_of(sent) == f"/model/{_OPUS}/converse", sent.target + assert mantle.drain() == () + + +def test_utils_token_counter_local_mode_never_calls_either_peer(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + assert _local_count(gateway, _full(model)) > 0 + assert runtime.drain() == () and mantle.drain() == () diff --git a/tests/integration/providers/test_gemini_chat_video_metadata_wire.py b/tests/integration/providers/test_gemini_chat_video_metadata_wire.py new file mode 100644 index 00000000000..32914341483 --- /dev/null +++ b/tests/integration/providers/test_gemini_chat_video_metadata_wire.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +import base64 +import json +from collections.abc import Callable +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, Scenario, object_value +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +BACKEND: Final = "gemini-3.7-flash" +PROJECT: Final = "scripted-project" +LOCATION: Final = "us-central1" +MODEL_PATH: Final = f"/v1/projects/{PROJECT}/locations/{LOCATION}/publishers/google/models/{BACKEND}" +CLIP: Final = base64.b64encode(b"\x00\x00\x00\x18ftypmp42" + bytes(24)).decode() +DATA_URI: Final = f"data:video/mp4;base64,{CLIP}" +PROMPT: Final = "describe the clip" +FILE_BLOCK: Final[dict[str, JsonValue]] = { + "type": "file", + "file": {"file_data": DATA_URI, "video_metadata": {"fps": 1.0, "start_offset": "0s", "end_offset": "3s"}}, +} +TEXT_PART: Final[dict[str, JsonValue]] = {"text": PROMPT} +BARE_VIDEO_PART: Final[dict[str, JsonValue]] = {"inline_data": {"mime_type": "video/mp4", "data": CLIP}} +VIDEO_PART: Final[dict[str, JsonValue]] = { + **BARE_VIDEO_PART, + "video_metadata": {"fps": 1.0, "startOffset": "0s", "endOffset": "3s"}, +} +REPLY: Final[dict[str, JsonValue]] = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "a cat"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 2, "totalTokenCount": 7}, +} + + +def chat_peer(path: str, authorization: str | None) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST", request.method + assert authorization is None or request.headers["authorization"] == authorization + target, _, query = request.target.partition("?") + if target == f"{path}:streamGenerateContent": + assert "alt=sse" in query, request.target + return Reply(content_type="text/event-stream", chunks=(f"data: {json.dumps(REPLY)}\n\n".encode(),)) + assert target == f"{path}:generateContent", request.target + return Reply(body=json.dumps(REPLY).encode()) + + return respond + + +def wire_parts(wire: Wire) -> list[JsonValue]: + received: Final = wire.drain() + assert len(received) == 1, [item.target for item in received] + body: Final = object_value(json.loads(received[0].body)) + contents: Final = body["contents"] + assert isinstance(contents, list) and len(contents) == 1, body + parts: Final = object_value(contents[0])["parts"] + assert isinstance(parts, list), body + return parts + + +def gemini_model(scenario: Scenario, url: str) -> str: + return scenario.model(model=f"gemini/{BACKEND}", api_key="scripted-gemini-key", api_base=url) + + +def vertex_model(gateway: Gateway, scenario: Scenario, url: str) -> str: + return scenario.model( + model=f"vertex_ai/{BACKEND}", + api_key=None, + api_base=url, + vertex_project=PROJECT, + vertex_location=LOCATION, + vertex_credentials=service_account_json(PROJECT, gateway.upstream_url), + ) + + +def chat(gateway: Gateway, model: str, stream: bool) -> httpx.Response: + body: Final[dict[str, JsonValue]] = { + "model": model, + "messages": [{"role": "user", "content": [{"type": "text", "text": PROMPT}, FILE_BLOCK]}], + "stream": stream, + } + return gateway.request("POST", "/v1/chat/completions", body) + + +def answered(response: httpx.Response) -> None: + assert response.status_code == 200, response.text + assert "a cat" in response.text, response.text + + +@pytest.mark.parametrize("stream", (False, True), ids=("non-stream", "stream")) +def test_gemini_chat_file_block_video_metadata_reaches_the_wire(gateway: Gateway, stream: bool) -> None: + with wire_server(chat_peer(f"/models/{BACKEND}", None)) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + answered(chat(gateway, model, stream)) + assert wire_parts(wire) == [TEXT_PART, VIDEO_PART] + + +@pytest.mark.parametrize("stream", (False, True), ids=("non-stream", "stream")) +def test_vertex_chat_file_block_video_metadata_reaches_the_wire(gateway: Gateway, stream: bool) -> None: + with wire_server(chat_peer(MODEL_PATH, "Bearer scripted-token")) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url) + answered(chat(gateway, model, stream)) + assert wire_parts(wire) == [TEXT_PART, VIDEO_PART] + + +def test_responses_input_file_video_reaches_the_wire_as_inline_data(gateway: Gateway) -> None: + with wire_server(chat_peer(f"/models/{BACKEND}", None)) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + body: Final[dict[str, JsonValue]] = { + "model": model, + "input": [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": PROMPT}, + {"type": "input_file", "file_data": DATA_URI, "filename": "clip.mp4"}, + ], + } + ], + } + answered(gateway.request("POST", "/v1/responses", body)) + assert wire_parts(wire) == [TEXT_PART, BARE_VIDEO_PART] + + +def test_messages_document_video_reaches_the_wire_as_inline_data(gateway: Gateway) -> None: + with wire_server(chat_peer(f"/models/{BACKEND}", None)) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + body: Final[dict[str, JsonValue]] = { + "model": model, + "max_tokens": 32, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": PROMPT}, + {"type": "document", "source": {"type": "base64", "media_type": "video/mp4", "data": CLIP}}, + ], + } + ], + } + answered(gateway.request("POST", "/v1/messages", body)) + assert wire_parts(wire) == [TEXT_PART, BARE_VIDEO_PART] diff --git a/tests/integration/providers/test_gemini_embedding_file_block_chaos.py b/tests/integration/providers/test_gemini_embedding_file_block_chaos.py new file mode 100644 index 00000000000..cee0758b433 --- /dev/null +++ b/tests/integration/providers/test_gemini_embedding_file_block_chaos.py @@ -0,0 +1,223 @@ +from __future__ import annotations + +import json +import os +import signal +import threading +import uuid +import zlib +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import psutil +import pytest +from integration._support.client import Gateway, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import group_members, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +BACKEND: Final = "gemini-embedding-001" +TARGET: Final = f"/models/{BACKEND}:batchEmbedContents" +END_OFFSETS: Final = {"fast": "3s", "slow": "5s", "malformed": "3s", "dropped": "9s"} +PLAN: Final = tuple(enumerate(("fast",) * 16 + ("slow",) * 8 + ("malformed",) * 8 + ("dropped",) * 8)) +BURST: Final = 24 +SPEND_SQL: Final = 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = %s' + + +@dataclass(frozen=True, slots=True) +class Attempt: + kind: str + index: int + status: int + call_id: str + + +def source(kind: str, index: int) -> str: + return f"gs://scripted-bucket/chaos/{kind}-{index}.mp4" + + +def body(model: str, kind: str, index: int) -> dict[str, JsonValue]: + fps: Final[JsonValue] = "x" if kind == "malformed" else 1.0 + metadata: Final[dict[str, JsonValue]] = {"fps": fps, "start_offset": "0s", "end_offset": END_OFFSETS[kind]} + return { + "model": model, + "input": [{"type": "file", "file": {"file_id": source(kind, index), "video_metadata": metadata}}], + } + + +def first_part(request: Request) -> dict[str, JsonValue]: + requests: Final = object_value(json.loads(request.body))["requests"] + assert isinstance(requests, list) and len(requests) == 1, requests + parts: Final = object_value(object_value(requests[0])["content"])["parts"] + assert isinstance(parts, list) and len(parts) == 1, parts + return object_value(parts[0]) + + +def source_on_the_wire(request: Request) -> str: + return string_value(object_value(first_part(request)["file_data"])["file_uri"]) + + +def embed_reply(request: Request) -> Reply: + assert request.method == "POST", request.method + assert request.target.split("?")[0] == TARGET, request.target + values: Final = [zlib.crc32(source_on_the_wire(request).encode()) / 2**32, 0.5] + return Reply(body=json.dumps({"embeddings": [{"values": values}]}).encode()) + + +def chaos_peer(healed: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + end_offset: Final = object_value(first_part(request)["video_metadata"])["endOffset"] + if end_offset == END_OFFSETS["dropped"] and not healed.is_set(): + return Reply(drop_connection=True) + reply: Final = embed_reply(request) + if end_offset == END_OFFSETS["slow"]: + return Reply(chunks=(reply.body[:8], reply.body[8:]), pause_between_chunks=1.5) + return reply + + return respond + + +def wire_sources(wire: Wire) -> tuple[str, ...]: + return tuple(source_on_the_wire(request) for request in wire.drain()) + + +def spend_row_count(call_id: str) -> int: + return len(read_rows(SPEND_SQL, (call_id,))) + + +def spend_row_counts(call_ids: tuple[str, ...]) -> tuple[int, ...]: + return tuple(spend_row_count(call_id) for call_id in call_ids) + + +def attempt(gateway: Gateway, model: str, index: int, kind: str) -> Attempt: + response: Final = gateway.request("POST", "/v1/embeddings", body(model, kind, index)) + return Attempt(kind, index, response.status_code, response.headers.get("x-litellm-call-id", "")) + + +def of_kind(attempts: tuple[Attempt, ...], *kinds: str) -> tuple[Attempt, ...]: + return tuple(item for item in attempts if item.kind in kinds) + + +@pytest.mark.timeout(240) +def test_mixed_burst_keeps_every_answer_and_spend_row_honest(gateway: Gateway) -> None: + healed: Final = threading.Event() + with wire_server(chaos_peer(healed)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"gemini/{BACKEND}", api_key="scripted-gemini-key", api_base=wire.url) + with ThreadPoolExecutor(max_workers=len(PLAN)) as pool: + attempts: Final = tuple(pool.map(lambda plan: attempt(gateway, model, plan[0], plan[1]), PLAN)) + served: Final = of_kind(attempts, "fast", "slow") + assert all(item.status == 200 for item in served), attempts + assert all(item.status == 400 for item in of_kind(attempts, "malformed")), attempts + assert all(item.status >= 500 for item in of_kind(attempts, "dropped")), attempts + reached: Final = wire_sources(wire) + assert sorted(item for item in reached if "/malformed-" not in item and "/dropped-" not in item) == sorted( + source(item.kind, item.index) for item in served + ), reached + assert not any("/malformed-" in item for item in reached), reached + assert {item for item in reached if "/dropped-" in item} == { + source(item.kind, item.index) for item in of_kind(attempts, "dropped") + }, reached + call_ids: Final = tuple(item.call_id for item in served) + assert all(call_ids), attempts + counts: Final = eventually( + lambda: spend_row_counts(call_ids), lambda found: all(count >= 1 for count in found), seconds=90 + ) + assert counts == (1,) * len(call_ids), counts + healed.set() + with ThreadPoolExecutor(max_workers=len(PLAN)) as pool: + resent: Final = tuple( + pool.map(lambda item: attempt(gateway, model, item.index, item.kind), of_kind(attempts, "dropped")) + ) + assert all(item.status == 200 for item in resent), resent + assert sorted(wire_sources(wire)) == sorted(source(item.kind, item.index) for item in resent) + + +def write_config(directory: Path, url: str, name: str) -> Path: + config: Final = directory / f"gemini_embeddings_{uuid.uuid4().hex}.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": name, + "litellm_params": { + "model": f"gemini/{BACKEND}", + "api_key": "scripted-gemini-key", + "api_base": url, + }, + } + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + }, + } + ) + ) + return config + + +def worker_pids(root_pid: int) -> frozenset[int]: + def is_worker(process: psutil.Process) -> bool: + try: + return "spawn_main" in " ".join(process.cmdline()) + except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess): + return False + + return frozenset(process.pid for process in group_members(root_pid) if is_worker(process)) + + +def outcome(candidate: Gateway, model: str, index: int) -> str: + try: + response: Final = candidate.request("POST", "/v1/embeddings", body(model, "fast", index)) + except httpx.TransportError as error: + return f"transport:{type(error).__name__}" + assert response.status_code == 200, response.text + return f"ok:{response.headers['x-litellm-call-id']}" + + +@pytest.mark.timeout(300) +def test_block_embeddings_survive_a_worker_kill(gateway: Gateway, tmp_path: Path) -> None: + name: Final = f"gemini-embeddings-{uuid.uuid4().hex}" + with wire_server(embed_reply) as wire: + config: Final = write_config(tmp_path, wire.url, name) + with owned_proxy_process(gateway, tmp_path, {"STORE_MODEL_IN_DB": "False"}, config=config, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually(lambda: worker_pids(owned.process.pid), lambda pids: len(pids) == 2, seconds=30) + victim: Final = min(workers) + assert outcome(candidate, name, 0).startswith("ok:") + + def attempt_around_the_kill(index: int) -> str: + if index == 2: + os.kill(victim, signal.SIGKILL) + return outcome(candidate, name, index) + + with ThreadPoolExecutor(max_workers=BURST) as pool: + outcomes: Final = tuple(pool.map(attempt_around_the_kill, range(1, BURST + 1))) + assert outcomes.count("ok") == 0 and any(item.startswith("ok:") for item in outcomes), outcomes + assert all(item.startswith("ok:") or item.startswith("transport:") for item in outcomes), outcomes + respawned: Final = eventually( + lambda: worker_pids(owned.process.pid), + lambda pids: len(pids) == 2 and victim not in pids, + seconds=60, + ) + assert victim not in respawned, respawned + + def settled_burst() -> tuple[str, ...]: + with ThreadPoolExecutor(max_workers=BURST) as pool: + return tuple(pool.map(lambda index: outcome(candidate, name, BURST + 1 + index), range(BURST))) + + final: Final = eventually( + settled_burst, lambda values: all(v.startswith("ok:") for v in values), seconds=40 + ) + call_ids: Final = tuple(item.removeprefix("ok:") for item in (*outcomes, *final) if item.startswith("ok:")) + counts: Final = eventually( + lambda: spend_row_counts(call_ids), lambda found: all(count >= 1 for count in found), seconds=90 + ) + assert counts == (1,) * len(call_ids), counts + assert len(wire.drain()) >= len(call_ids) diff --git a/tests/integration/providers/test_gemini_embedding_file_block_wire.py b/tests/integration/providers/test_gemini_embedding_file_block_wire.py new file mode 100644 index 00000000000..d6439d27cd5 --- /dev/null +++ b/tests/integration/providers/test_gemini_embedding_file_block_wire.py @@ -0,0 +1,294 @@ +from __future__ import annotations + +import asyncio +import base64 +import json +import uuid +import zlib +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, Scenario, object_value, string_value +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from openai import AsyncOpenAI, OpenAI +from openai.types import CreateEmbeddingResponse +from pydantic import JsonValue + +BACKEND: Final = "gemini-embedding-001" +TARGET: Final = f"/models/{BACKEND}:batchEmbedContents" +CLIP: Final = base64.b64encode(b"\x00\x00\x00\x18ftypmp42" + bytes(24)).decode() +DATA_URI: Final = f"data:video/mp4;base64,{CLIP}" +GCS_URI: Final = "gs://scripted-bucket/clips/animals.mp4" +METADATA: Final[dict[str, JsonValue]] = {"fps": 1.0, "start_offset": "0s", "end_offset": "3s"} +WIRE_METADATA: Final[dict[str, JsonValue]] = {"fps": 1.0, "startOffset": "0s", "endOffset": "3s"} +INLINE_DATA: Final[dict[str, JsonValue]] = {"mime_type": "video/mp4", "data": CLIP} +INLINE_PART: Final[dict[str, JsonValue]] = {"inline_data": INLINE_DATA, "video_metadata": WIRE_METADATA} +TEXT_PART: Final[dict[str, JsonValue]] = {"text": "a red bus"} +ACCEPTED_FORMS: Final = "must be a data: URI, a gs:// URL, a files/ reference, or a Gemini Files API URI" +LONG_OFFSET: Final = "9" * 5000 + "s" + + +def block(**file: JsonValue) -> dict[str, JsonValue]: + return {"type": "file", "file": file} + + +def gcs_part(uri: str = GCS_URI, metadata: dict[str, JsonValue] | None = WIRE_METADATA) -> dict[str, JsonValue]: + file_data: Final[dict[str, JsonValue]] = {"file_data": {"mime_type": "video/mp4", "file_uri": uri}} + return file_data if metadata is None else {**file_data, "video_metadata": metadata} + + +GCS_PART: Final = gcs_part() +NOISY_BLOCK: Final[dict[str, JsonValue]] = { + **block(file_id=GCS_URI, mime_type="video/mp4", video_metadata={**METADATA, "frame_rate": 2}), + "caption": "unknown block key", +} + + +def list_value(value: JsonValue) -> list[JsonValue]: + assert isinstance(value, list), value + return value + + +def source_of(part: JsonValue) -> str: + item: Final = object_value(part) + if "text" in item: + return string_value(item["text"]) + if "file_data" in item: + return string_value(object_value(item["file_data"])["file_uri"]) + return string_value(object_value(item["inline_data"])["data"]) + + +def vector(parts: list[JsonValue]) -> list[float]: + return [zlib.crc32("|".join(source_of(part) for part in parts).encode()) / 2**32, 0.5] + + +def request_parts(request: Request) -> list[list[JsonValue]]: + body: Final = object_value(json.loads(request.body)) + return [list_value(object_value(object_value(item)["content"])["parts"]) for item in list_value(body["requests"])] + + +def embed_peer(request: Request) -> Reply: + assert request.method == "POST", request.method + assert request.target.split("?")[0] == TARGET, request.target + embeddings: Final = [{"values": vector(parts)} for parts in request_parts(request)] + return Reply(body=json.dumps({"embeddings": embeddings}).encode()) + + +def wire_parts(wire: Wire) -> list[list[JsonValue]]: + received: Final = wire.drain() + assert len(received) == 1, [item.target for item in received] + return request_parts(received[0]) + + +def embeddings(response: httpx.Response) -> list[JsonValue]: + assert response.status_code == 200, response.text + data: Final = [object_value(item) for item in list_value(object_value(response.json())["data"])] + assert [item["index"] for item in data] == list(range(len(data))), response.text + assert all(item["object"] == "embedding" for item in data), response.text + return [item["embedding"] for item in data] + + +def gemini_model(scenario: Scenario, url: str, **litellm_params: JsonValue) -> str: + return scenario.model(model=f"gemini/{BACKEND}", api_key="scripted-gemini-key", api_base=url, **litellm_params) + + +def embed(gateway: Gateway, model: str, elements: JsonValue, **extra: JsonValue) -> httpx.Response: + return gateway.request("POST", "/v1/embeddings", {"model": model, "input": elements, **extra}) + + +def proxy_root(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def write_config(directory: Path, url: str, name: str) -> Path: + config: Final = directory / f"gemini_embeddings_{uuid.uuid4().hex}.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": name, + "litellm_params": { + "model": f"gemini/{BACKEND}", + "api_key": "scripted-gemini-key", + "api_base": url, + }, + } + ], + "litellm_settings": {"drop_params": True}, + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + }, + } + ) + ) + return config + + +def test_string_input_reaches_the_wire_as_a_text_part(gateway: Gateway) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + assert embeddings(embed(gateway, model, ["a red bus"])) == [vector([TEXT_PART])] + assert wire_parts(wire) == [[TEXT_PART]] + + +@pytest.mark.parametrize( + ("file", "part"), + ( + pytest.param({"file_data": DATA_URI, "video_metadata": METADATA}, INLINE_PART, id="data-uri"), + pytest.param({"file_id": GCS_URI, "video_metadata": METADATA}, GCS_PART, id="gcs"), + pytest.param( + {"file_id": GCS_URI, "video_metadata": {"fps": 1, "start_offset": "", "end_offset": LONG_OFFSET}}, + gcs_part(metadata={"fps": 1, "startOffset": "", "endOffset": LONG_OFFSET}), + id="verbatim-strings-and-int-fps", + ), + pytest.param( + {"file_data": DATA_URI, "format": "video/quicktime"}, + {"inline_data": {"mime_type": "video/quicktime", "data": CLIP}}, + id="format-overrides-the-data-uri-mime", + ), + pytest.param( + {"file_id": "gs://scripted-bucket/clips/clip.bin", "format": "video/mp4"}, + gcs_part("gs://scripted-bucket/clips/clip.bin", None), + id="format-names-an-unlisted-extension", + ), + pytest.param({"file_id": GCS_URI}, gcs_part(metadata=None), id="no-metadata"), + pytest.param({"file_id": GCS_URI, "video_metadata": {}}, gcs_part(metadata=None), id="empty-metadata"), + ), +) +def test_file_block_reaches_the_wire_as_one_part( + gateway: Gateway, file: dict[str, JsonValue], part: dict[str, JsonValue] +) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + assert embeddings(embed(gateway, model, [block(**file)])) == [vector([part])] + assert wire_parts(wire) == [[part]] + + +def test_nested_text_and_block_become_one_request_with_two_parts(gateway: Gateway) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + response: Final = embed(gateway, model, [["a red bus", block(file_id=GCS_URI, video_metadata=METADATA)]]) + assert embeddings(response) == [vector([TEXT_PART, GCS_PART])] + assert wire_parts(wire) == [[TEXT_PART, GCS_PART]] + + +def test_repeated_blocks_become_one_request_each(gateway: Gateway) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + response: Final = embed(gateway, model, [block(file_id=GCS_URI, video_metadata=METADATA)] * 2) + assert embeddings(response) == [vector([GCS_PART])] * 2 + assert wire_parts(wire) == [[GCS_PART], [GCS_PART]] + + +@pytest.mark.parametrize( + ("element", "fragment"), + ( + pytest.param(block(file_id=GCS_URI, video_metadata={"fps": "1"}), "file.video_metadata.fps", id="fps-string"), + pytest.param( + block(file_id=GCS_URI, video_metadata={"start_offset": 0}), + "file.video_metadata.start_offset", + id="offset-int", + ), + pytest.param( + block(file_id=GCS_URI, video_metadata={"end_offset": ["3s"]}), + "file.video_metadata.end_offset", + id="offset-list", + ), + pytest.param(block(file_id=GCS_URI, video_metadata="1fps"), "file.video_metadata", id="metadata-string"), + pytest.param( + block(file_id=GCS_URI, video_metadata={**METADATA, "frame_rate": 2}), + "frame_rate", + id="unknown-metadata-key", + ), + pytest.param(block(file_id=GCS_URI, mime_type="video/mp4"), "mime_type", id="unknown-file-key"), + pytest.param({**block(file_id=GCS_URI), "caption": "x"}, "caption", id="unknown-block-key"), + pytest.param({"type": "video", "file": {"file_id": GCS_URI}}, "type", id="wrong-type"), + pytest.param( + block(file_id=GCS_URI, file_data=DATA_URI), + "takes file.file_id or file.file_data, not both", + id="both-sources", + ), + pytest.param(block(video_metadata=METADATA), "needs file.file_id or file.file_data", id="no-source"), + pytest.param(block(file_data=CLIP), ACCEPTED_FORMS, id="bare-base64"), + pytest.param(block(file_id="https://example.com/clip.mp4"), ACCEPTED_FORMS, id="http-url"), + pytest.param(block(file_id=""), ACCEPTED_FORMS, id="empty-file-id"), + pytest.param(block(file_id=GCS_URI, format=""), "file.format", id="empty-format"), + pytest.param(block(file_data=5), "file.file_data", id="file-data-int"), + pytest.param(block(file_data=[DATA_URI]), "file.file_data", id="file-data-list"), + pytest.param(1, "got int", id="int-element"), + ), +) +def test_malformed_blocks_answer_400_before_any_provider_call( + gateway: Gateway, element: JsonValue, fragment: str +) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + response: Final = embed(gateway, model, [element]) + assert response.status_code == 400, response.text + assert fragment in response.text, response.text + assert wire.drain() == (), "the rejected input reached the provider" + + +def test_bare_block_input_answers_400_before_any_provider_call(gateway: Gateway) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + response: Final = embed(gateway, model, block(file_id=GCS_URI, video_metadata=METADATA)) + assert response.status_code == 400, response.text + assert "input must be a string or a list" in response.text, response.text + assert wire.drain() == (), "the rejected input reached the provider" + + +def test_request_drop_params_strips_unknown_keys_at_every_level(gateway: Gateway) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + assert embeddings(embed(gateway, model, [NOISY_BLOCK], drop_params=True)) == [vector([GCS_PART])] + assert wire_parts(wire) == [[GCS_PART]] + + +def test_deployment_drop_params_strips_unknown_keys_at_every_level(gateway: Gateway) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url, drop_params=True) + assert embeddings(embed(gateway, model, [NOISY_BLOCK])) == [vector([GCS_PART])] + assert wire_parts(wire) == [[GCS_PART]] + + +@pytest.mark.timeout(300) +def test_yaml_drop_params_strips_unknown_keys_at_every_level(gateway: Gateway, tmp_path: Path) -> None: + name: Final = f"gemini-embeddings-{uuid.uuid4().hex}" + with wire_server(embed_peer) as wire: + config: Final = write_config(tmp_path, wire.url, name) + with owned_proxy_process(gateway, tmp_path, {"STORE_MODEL_IN_DB": "False"}, config=config) as owned: + assert embeddings(embed(owned.gateway, name, [NOISY_BLOCK])) == [vector([GCS_PART])] + assert wire_parts(wire) == [[GCS_PART]] + + +def test_openai_sdk_clients_send_blocks_through_the_proxy(gateway: Gateway) -> None: + sync_uri: Final = "gs://scripted-bucket/clips/sync.mp4" + async_uri: Final = "gs://scripted-bucket/clips/async.mp4" + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = gemini_model(scenario, wire.url) + with OpenAI(base_url=f"{proxy_root(gateway)}/v1", api_key=gateway.key, max_retries=0) as client: + served: Final = client.post( + "/embeddings", + body={"model": model, "input": [block(file_id=sync_uri, video_metadata=METADATA)]}, + cast_to=CreateEmbeddingResponse, + ) + assert served.data[0].embedding == vector([gcs_part(sync_uri)]), served.model_dump_json() + assert wire_parts(wire) == [[gcs_part(sync_uri)]] + + async def drive() -> CreateEmbeddingResponse: + async with AsyncOpenAI(base_url=f"{proxy_root(gateway)}/v1", api_key=gateway.key, max_retries=0) as client: + return await client.post( + "/embeddings", + body={"model": model, "input": [block(file_id=async_uri, video_metadata=METADATA)]}, + cast_to=CreateEmbeddingResponse, + ) + + served_async: Final = asyncio.run(drive()) + assert served_async.data[0].embedding == vector([gcs_part(async_uri)]), served_async.model_dump_json() + assert wire_parts(wire) == [[gcs_part(async_uri)]] diff --git a/tests/integration/providers/test_vertex_embedding_batch_file_block_wire.py b/tests/integration/providers/test_vertex_embedding_batch_file_block_wire.py new file mode 100644 index 00000000000..042ca8b0889 --- /dev/null +++ b/tests/integration/providers/test_vertex_embedding_batch_file_block_wire.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import json +from typing import Final +from urllib.parse import unquote + +import httpx +from integration._support.client import Gateway, Scenario +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +BACKEND: Final = "gemini-embedding-2-preview" +PROJECT: Final = "scripted-project" +LOCATION: Final = "us-central1" +BUCKET: Final = "scripted-bucket" +UPLOAD_PREFIX: Final = f"/upload/storage/v1/b/{BUCKET}/o?uploadType=media&name=" +GCS_URI: Final = "gs://scripted-bucket/clips/animals.mp4" +BLOCK: Final[dict[str, JsonValue]] = { + "type": "file", + "file": {"file_id": GCS_URI, "video_metadata": {"fps": 1.0, "start_offset": "0s", "end_offset": "3s"}}, +} +GCS_PART: Final[dict[str, JsonValue]] = { + "file_data": {"mime_type": "video/mp4", "file_uri": GCS_URI}, + "video_metadata": {"fps": 1.0, "startOffset": "0s", "endOffset": "3s"}, +} +TEXT_PART: Final[dict[str, JsonValue]] = {"text": "a red bus"} + + +def row(custom_id: str, model: str, elements: JsonValue) -> dict[str, JsonValue]: + return { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": model, "input": elements}, + } + + +def stored_row(key: str, parts: list[JsonValue]) -> dict[str, JsonValue]: + return {"key": key, "request": {"content": {"parts": parts}}} + + +def gcs_peer(request: Request) -> Reply: + assert request.method == "POST", request.method + assert request.target.startswith(UPLOAD_PREFIX), request.target + assert request.headers["authorization"] == "Bearer scripted-token" + name: Final = unquote(request.target.removeprefix(UPLOAD_PREFIX)) + stored: Final = { + "kind": "storage#object", + "id": f"{BUCKET}/{name}/1759950000000000", + "name": name, + "bucket": BUCKET, + "size": str(len(request.body)), + "timeCreated": "2026-10-08T19:00:00.000Z", + "contentType": "application/octet-stream", + } + return Reply(body=json.dumps(stored).encode()) + + +def vertex_batch_model(gateway: Gateway, scenario: Scenario, url: str) -> str: + return scenario.model( + model=f"vertex_ai/{BACKEND}", + api_key=None, + api_base=url, + vertex_project=PROJECT, + vertex_location=LOCATION, + vertex_credentials=service_account_json(PROJECT, gateway.upstream_url), + gcs_bucket_name=BUCKET, + ) + + +def upload(gateway: Gateway, model: str, rows: tuple[dict[str, JsonValue], ...]) -> httpx.Response: + jsonl: Final = "\n".join(json.dumps(line) for line in rows) + "\n" + return gateway.request_multipart( + "/v1/files", + {"purpose": "batch", "target_model_names": model}, + {"file": ("in.jsonl", jsonl.encode(), "application/jsonl")}, + ) + + +def test_block_rows_upload_as_embed_content_requests(gateway: Gateway) -> None: + with wire_server(gcs_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_batch_model(gateway, scenario, wire.url) + rows: Final = ( + row("mixed", model, [BLOCK, "a red bus"]), + row("bare", model, "a red bus"), + row("nested", model, [[BLOCK, "a red bus"]]), + ) + response: Final = upload(gateway, model, rows) + assert response.status_code == 200, response.text + assert response.json()["object"] == "file" and response.json()["purpose"] == "batch", response.text + uploads: Final = wire.drain() + assert len(uploads) == 1, [upload.target for upload in uploads] + stored: Final = tuple(json.loads(line) for line in uploads[0].body.decode().splitlines() if line.strip()) + assert stored == ( + stored_row("mixed#0/2", [GCS_PART]), + stored_row("mixed#1/2", [TEXT_PART]), + stored_row("bare", [TEXT_PART]), + stored_row("nested", [GCS_PART, TEXT_PART]), + ) + + +def test_bare_block_row_fails_the_upload_before_any_storage_write(gateway: Gateway) -> None: + with wire_server(gcs_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_batch_model(gateway, scenario, wire.url) + response: Final = upload(gateway, model, (row("mixed", model, [BLOCK, "a red bus"]), row("bare", model, BLOCK))) + assert response.status_code >= 400, response.text + assert "got dict" in response.text, response.text + assert wire.drain() == (), "a rejected batch file reached the bucket" diff --git a/tests/integration/providers/test_vertex_embedding_file_block_wire.py b/tests/integration/providers/test_vertex_embedding_file_block_wire.py new file mode 100644 index 00000000000..ad66412a27a --- /dev/null +++ b/tests/integration/providers/test_vertex_embedding_file_block_wire.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +import json +from typing import Final + +import httpx +from integration._support.client import Gateway, Scenario, object_value +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +BACKEND: Final = "gemini-embedding-2-preview" +PROJECT: Final = "scripted-project" +LOCATION: Final = "us-central1" +TARGET: Final = f"/v1/projects/{PROJECT}/locations/{LOCATION}/publishers/google/models/{BACKEND}:embedContent" +GCS_URI: Final = "gs://scripted-bucket/clips/animals.mp4" +METADATA: Final[dict[str, JsonValue]] = {"fps": 1.0, "start_offset": "0s", "end_offset": "3s"} +GCS_PART: Final[dict[str, JsonValue]] = { + "file_data": {"mime_type": "video/mp4", "file_uri": GCS_URI}, + "video_metadata": {"fps": 1.0, "startOffset": "0s", "endOffset": "3s"}, +} +TEXT_PART: Final[dict[str, JsonValue]] = {"text": "a red bus"} +VALUES: Final = [0.75, 0.25] + + +def block(**file: JsonValue) -> dict[str, JsonValue]: + return {"type": "file", "file": file} + + +def embed_peer(request: Request) -> Reply: + assert request.method == "POST", request.method + assert request.target == TARGET, request.target + assert request.headers["authorization"] == "Bearer scripted-token" + return Reply(body=json.dumps({"embedding": {"values": VALUES}}).encode()) + + +def wire_parts(wire: Wire) -> list[JsonValue]: + received: Final = wire.drain() + assert len(received) == 1, [item.target for item in received] + body: Final = object_value(json.loads(received[0].body)) + parts: Final = object_value(body["content"])["parts"] + assert isinstance(parts, list), body + return parts + + +def vertex_model(gateway: Gateway, scenario: Scenario, url: str, **litellm_params: JsonValue) -> str: + return scenario.model( + model=f"vertex_ai/{BACKEND}", + api_key=None, + api_base=url, + vertex_project=PROJECT, + vertex_location=LOCATION, + vertex_credentials=service_account_json(PROJECT, gateway.upstream_url), + **litellm_params, + ) + + +def embed(gateway: Gateway, model: str, elements: JsonValue, **extra: JsonValue) -> httpx.Response: + return gateway.request("POST", "/v1/embeddings", {"model": model, "input": elements, **extra}) + + +def single_embedding(response: httpx.Response) -> JsonValue: + assert response.status_code == 200, response.text + data: Final = object_value(response.json())["data"] + assert isinstance(data, list) and len(data) == 1, response.text + item: Final = object_value(data[0]) + assert item["index"] == 0 and item["object"] == "embedding", response.text + return item["embedding"] + + +def rejected(gateway: Gateway, wire: Wire, model: str, elements: JsonValue, fragment: str) -> None: + response: Final = embed(gateway, model, elements) + assert response.status_code == 400, response.text + assert fragment in response.text, response.text + assert wire.drain() == (), "the rejected input reached the provider" + + +def test_string_input_reaches_embed_content_as_a_text_part(gateway: Gateway) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url) + assert single_embedding(embed(gateway, model, ["a red bus"])) == VALUES + assert wire_parts(wire) == [TEXT_PART] + + +def test_file_block_with_video_metadata_reaches_embed_content(gateway: Gateway) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url) + assert single_embedding(embed(gateway, model, [block(file_id=GCS_URI, video_metadata=METADATA)])) == VALUES + assert wire_parts(wire) == [GCS_PART] + + +def test_request_drop_params_strips_unknown_keys_at_every_level(gateway: Gateway) -> None: + noisy: Final[dict[str, JsonValue]] = { + **block(file_id=GCS_URI, mime_type="video/mp4", video_metadata={**METADATA, "frame_rate": 2}), + "caption": "unknown block key", + } + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url) + assert single_embedding(embed(gateway, model, [noisy], drop_params=True)) == VALUES + assert wire_parts(wire) == [GCS_PART] + + +def test_bad_fps_answers_400_before_any_provider_call(gateway: Gateway) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url) + rejected(gateway, wire, model, [block(file_id=GCS_URI, video_metadata={"fps": "1"})], "file.video_metadata.fps") + + +def test_files_api_reference_answers_400_naming_the_gemini_provider(gateway: Gateway) -> None: + with wire_server(embed_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url) + rejected( + gateway, + wire, + model, + [block(file_id="files/abc", video_metadata=METADATA)], + "Gemini Files API references are only supported through the gemini/ provider", + ) diff --git a/tests/integration/spend/test_batch_completion_accounting.py b/tests/integration/spend/test_batch_completion_accounting.py index 4ecab10f942..bbc0300b3f0 100644 --- a/tests/integration/spend/test_batch_completion_accounting.py +++ b/tests/integration/spend/test_batch_completion_accounting.py @@ -63,7 +63,7 @@ def _failed_line(index: int) -> str: ) -def _batch_routes(model: str) -> RoutedResponse: +def batch_routes(model: str) -> RoutedResponse: output_lines: Final = ( _succeeded_line(1, model, **FIRST_LINE), _succeeded_line(2, model, **SECOND_LINE), @@ -139,7 +139,7 @@ def _batches_total_tokens(endpoints: dict[str, object] | None) -> int | None: return int(total_tokens) if isinstance(total_tokens, (int, float, str)) else None -def _input_file(model: str) -> bytes: +def batch_input_file(model: str) -> bytes: return ( "\n".join( json.dumps( @@ -166,13 +166,13 @@ def test_completed_batch_spend_row_records_reasoning_tokens_and_error_file_failu with gateway.scenario() as scenario: key: Final = scenario.key() scenario_id: Final = f"batch-accounting-{uuid.uuid4().hex[:12]}" - handle: Final = register_scenario(scenario_id, _batch_routes("gpt-4o-mini")) + handle: Final = register_scenario(scenario_id, batch_routes("gpt-4o-mini")) scenario.cleanups.callback(delete_scenario, handle) model: Final = scenario.model(api_base=handle.api_base()) file_response: Final = gateway.request_multipart( "/v1/files", {"purpose": "batch", "model": model}, - {"file": ("in.jsonl", _input_file(model), "application/jsonl")}, + {"file": ("in.jsonl", batch_input_file(model), "application/jsonl")}, key=key, ) assert file_response.status_code == 200, file_response.text @@ -235,7 +235,7 @@ BATCH_SPEND: Final = (BATCH_PROMPT_TOKENS * INPUT_COST_PER_TOKEN + BATCH_COMPLET def test_completed_batch_spend_lands_under_batches_in_team_endpoint_activity(gateway: Gateway) -> None: with gateway.scenario() as scenario: scenario_id: Final = f"batch-endpoint-{uuid.uuid4().hex[:12]}" - handle: Final = register_scenario(scenario_id, _batch_routes("gpt-4o-mini")) + handle: Final = register_scenario(scenario_id, batch_routes("gpt-4o-mini")) scenario.cleanups.callback(delete_scenario, handle) model: Final = scenario.model( api_base=handle.api_base(), @@ -247,7 +247,7 @@ def test_completed_batch_spend_lands_under_batches_in_team_endpoint_activity(gat file_response: Final = gateway.request_multipart( "/v1/files", {"purpose": "batch", "model": model}, - {"file": ("in.jsonl", _input_file(model), "application/jsonl")}, + {"file": ("in.jsonl", batch_input_file(model), "application/jsonl")}, key=key, ) assert file_response.status_code == 200, file_response.text diff --git a/tests/integration/spend/test_budget_limit_envelopes.py b/tests/integration/spend/test_budget_limit_envelopes.py new file mode 100644 index 00000000000..6ac2a11e657 --- /dev/null +++ b/tests/integration/spend/test_budget_limit_envelopes.py @@ -0,0 +1,172 @@ +from __future__ import annotations + +import json +import re +import uuid +from collections.abc import Sequence +from hashlib import sha256 +from typing import Final, Literal + +import httpx +import pytest +from pydantic import JsonValue + +from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value +from tests.integration._support.database import read_rows +from tests.integration._support.provider import PROVIDER_URL, SharedProvider +from tests.integration._support.wire import Reply + +_CALL_COST: Final = 0.02 +_TINY_BUDGET: Final = 0.0000000005 +_LIMIT_FIELDS: Final = ("max_budget", "rpm_limit", "tpm_limit") + + +def _completion() -> Reply: + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "scripted"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + +def _priced_model(scenario: Scenario) -> str: + return scenario.model(api_base=f"{PROVIDER_URL}/v1", input_cost_per_token=0.001, output_cost_per_token=0.002) + + +def _ask(gateway: Gateway, model: str, key: str) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"budget {uuid.uuid4().hex}"}]}, + key=key, + ) + + +def _served(gateway: Gateway, provider: SharedProvider, model: str, key: str) -> None: + provider.expect(_completion()) + response: Final = _ask(gateway, model, key) + assert response.status_code == 200, response.text + assert len(provider.received()) == 1 + + +def _spend_reaches(table: Literal["key", "team"], identity: str, amount: float) -> None: + query: Final = ( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s' + if table == "key" + else 'SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id = %s' + ) + eventually( + lambda: read_rows(query, (identity,)), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= amount, + seconds=70, + ) + + +_COST_AND_LIMIT: Final = re.compile(r"Current cost: ([^\s,]+), Max budget: ([^\s,]+)") + + +def _budget_refusal( + gateway: Gateway, provider: SharedProvider, model: str, key: str, *, spent: float, limit: float +) -> str: + refused: Final = _ask(gateway, model, key) + assert refused.status_code == 422, f"{refused.status_code} {refused.text}" + error: Final = object_value(JSON_OBJECT.validate_json(refused.content)["error"]) + assert error["type"] == "budget_exceeded", refused.text + assert error["code"] == "422", refused.text + message: Final = string_value(error["message"]) + assert "Budget has been exceeded!" in message, refused.text + figures: Final = _COST_AND_LIMIT.search(message) + assert figures is not None, message + assert float(figures[1]) == pytest.approx(spent), message + assert float(figures[2]) == pytest.approx(limit), message + assert provider.received() == () + return message + + +def _hashed(key: str) -> str: + return sha256(key.encode()).hexdigest() + + +def test_a_key_with_a_tiny_budget_serves_once_and_then_answers_budget_exceeded( + gateway: Gateway, provider: SharedProvider +) -> None: + with gateway.scenario() as scenario: + model: Final = _priced_model(scenario) + key: Final = scenario.key(models=[model], max_budget=_TINY_BUDGET) + _served(gateway, provider, model, key) + _spend_reaches("key", _hashed(key), _CALL_COST) + _budget_refusal(gateway, provider, model, key, spent=_CALL_COST, limit=_TINY_BUDGET) + + +def test_a_key_with_a_zero_budget_is_refused_before_the_provider_is_called( + gateway: Gateway, provider: SharedProvider +) -> None: + with gateway.scenario() as scenario: + model: Final = _priced_model(scenario) + key: Final = scenario.key(models=[model], max_budget=0) + _budget_refusal(gateway, provider, model, key, spent=0.0, limit=0.0) + + +def test_a_key_with_room_for_two_calls_serves_both_before_answering_budget_exceeded( + gateway: Gateway, provider: SharedProvider +) -> None: + with gateway.scenario() as scenario: + model: Final = _priced_model(scenario) + limit: Final = _CALL_COST * 1.5 + key: Final = scenario.key(models=[model], max_budget=limit) + _served(gateway, provider, model, key) + _spend_reaches("key", _hashed(key), _CALL_COST) + _served(gateway, provider, model, key) + _spend_reaches("key", _hashed(key), _CALL_COST * 2) + _budget_refusal(gateway, provider, model, key, spent=_CALL_COST * 2, limit=limit) + + +def test_a_team_key_serves_once_and_then_answers_the_team_budget_envelope( + gateway: Gateway, provider: SharedProvider +) -> None: + with gateway.scenario() as scenario: + model: Final = _priced_model(scenario) + team: Final = scenario.team(models=[model], max_budget=_TINY_BUDGET) + key: Final = scenario.key(team_id=team, models=[model]) + _served(gateway, provider, model, key) + _spend_reaches("team", team, _CALL_COST) + message: Final = _budget_refusal(gateway, provider, model, key, spent=_CALL_COST, limit=_TINY_BUDGET) + assert f"Team={team}" in message, message + + +def _limits(record: dict[str, JsonValue]) -> Sequence[JsonValue]: + return [record[field] for field in _LIMIT_FIELDS] + + +@pytest.mark.parametrize("field", _LIMIT_FIELDS) +def test_a_key_limit_is_set_by_update_and_reset_to_null(gateway: Gateway, field: str) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key(max_budget=None, rpm_limit=None, tpm_limit=None) + raised: Final = gateway.post("/key/update", {"key": key, field: 10}) + assert raised[field] == 10, raised + assert [value for name, value in zip(_LIMIT_FIELDS, _limits(raised)) if name != field] == [None, None] + cleared: Final = gateway.post("/key/update", {"key": key, field: None}) + assert _limits(cleared) == [None, None, None], cleared + saved: Final = object_value(gateway.get("/key/info", {"key": key})["info"]) + assert _limits(saved) == [None, None, None], saved + + +@pytest.mark.parametrize("field", _LIMIT_FIELDS) +def test_a_team_limit_is_set_by_update_and_reset_to_null(gateway: Gateway, field: str) -> None: + with gateway.scenario() as scenario: + team: Final = scenario.team(max_budget=None, rpm_limit=None, tpm_limit=None) + raised: Final = object_value(gateway.post("/team/update", {"team_id": team, field: 10})["data"]) + assert raised[field] == 10, raised + cleared: Final = object_value(gateway.post("/team/update", {"team_id": team, field: None})["data"]) + assert _limits(cleared) == [None, None, None], cleared + saved: Final = object_value(gateway.get("/team/info", {"team_id": team})["team_info"]) + assert _limits(saved) == [None, None, None], saved diff --git a/tests/integration/spend/test_spend_logs_metadata_fields.py b/tests/integration/spend/test_spend_logs_metadata_fields.py new file mode 100644 index 00000000000..d7f1dbfd998 --- /dev/null +++ b/tests/integration/spend/test_spend_logs_metadata_fields.py @@ -0,0 +1,509 @@ +import json +from collections.abc import Callable, Generator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import Final +from uuid import uuid4 + +import httpx +import pytest +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows, scratch_database +from integration._support.process import owned_proxy +from integration._support.upstream import delete_scenario, register_scenario +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.spend.test_batch_completion_accounting import batch_input_file, batch_routes +from integration.streaming.test_stream_contracts import text_stream +from pydantic import JsonValue + +CACHED_PROMPT_TOKENS: Final = 4 + + +def _config( + tmp_path: Path, + spend_logs_metadata_fields: Mapping[str, JsonValue] | None, + *, + model_list: tuple[Mapping[str, JsonValue], ...] = (), + litellm_settings: Mapping[str, JsonValue] | None = None, + guardrails: tuple[Mapping[str, JsonValue], ...] = (), +) -> Path: + retention: Final = ( + {} if spend_logs_metadata_fields is None else {"spend_logs_metadata_fields": dict(spend_logs_metadata_fields)} + ) + config: Final = tmp_path / f"spend-logs-metadata-fields-{uuid4()}.json" + config.write_text( + json.dumps( + { + "model_list": [dict(model) for model in model_list], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "store_model_in_db": True, + "proxy_batch_write_at": 1, + "proxy_batch_polling_interval": 1, + **retention, + }, + "litellm_settings": dict(litellm_settings or {}), + "guardrails": [dict(guardrail) for guardrail in guardrails], + } + ) + ) + return config + + +def _respond(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=b'{"object":"list","data":[]}') + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid4()}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 2, + "total_tokens": 12, + "prompt_tokens_details": {"cached_tokens": CACHED_PROMPT_TOKENS}, + }, + } + ).encode() + ) + + +def _stored_row_after_one_chat( + gateway: Gateway, tmp_path: Path, spend_logs_metadata_fields: Mapping[str, JsonValue] +) -> tuple[dict[str, JsonValue], str]: + with ( + wire_server(_respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, spend_logs_metadata_fields)) as isolated, + isolated.scenario() as scenario, + ): + model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") + api_key: Final = scenario.key(key_alias=f"metadata-fields-{uuid4()}", models=[model]) + response_id: Final = string_value(isolated.chat(model, key=api_key)["id"]) + rows: Final = eventually( + lambda: read_rows( + 'SELECT metadata, proxy_server_request, response FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (response_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + return rows[0], sha256(api_key.encode()).hexdigest() + + +def test_excluded_metadata_fields_are_not_stored_but_still_reach_daily_spend(gateway: Gateway, tmp_path: Path) -> None: + row, hashed_key = _stored_row_after_one_chat( + gateway, tmp_path, {"exclude": ["model_map_information", "usage_object"]} + ) + + metadata: Final = object_value(row["metadata"]) + assert "model_map_information" not in metadata, metadata + assert "usage_object" not in metadata, metadata + assert {"status", "cold_storage_object_key"} <= set(metadata), metadata + assert string_value(metadata["user_api_key_alias"]).startswith("metadata-fields-") + assert row["proxy_server_request"] == {} + assert row["response"] == {} + daily: Final = eventually( + lambda: read_rows( + 'SELECT prompt_tokens, cache_read_input_tokens FROM "LiteLLM_DailyUserSpend" WHERE api_key=%s', + (hashed_key,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert daily[0]["cache_read_input_tokens"] == CACHED_PROMPT_TOKENS, daily + + +def test_included_metadata_fields_are_the_only_ones_stored_besides_always_kept( + gateway: Gateway, tmp_path: Path +) -> None: + row, _ = _stored_row_after_one_chat(gateway, tmp_path, {"include": ["user_api_key_alias"]}) + + assert set(object_value(row["metadata"])) == {"status", "cold_storage_object_key", "user_api_key_alias"} + + +def test_guardrail_usage_is_tracked_when_guardrail_information_is_not_stored(gateway: Gateway, tmp_path: Path) -> None: + guardrail_name: Final = f"metadata-fields-guardrail-{uuid4()}" + with ( + wire_server(lambda _: Reply(body=b'{"action":"NONE"}')) as policy, + wire_server(_respond) as wire, + owned_proxy( + gateway, + tmp_path, + {}, + config=_config( + tmp_path, + {"exclude": ["guardrail_information"]}, + guardrails=( + { + "guardrail_name": guardrail_name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": True, + "api_base": policy.url, + "api_key": "synthetic-guardrail-key", + }, + }, + ), + ), + ) as isolated, + isolated.scenario() as scenario, + ): + model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") + response_id: Final = string_value(isolated.chat(model, key=scenario.key(models=[model]))["id"]) + assert len(policy.drain()) == 1 + rows: Final = eventually( + lambda: read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + assert "guardrail_information" not in object_value(rows[0]["metadata"]), rows[0] + indexed: Final = eventually( + lambda: read_rows( + 'SELECT guardrail_id FROM "LiteLLM_SpendLogGuardrailIndex" WHERE request_id=%s', (response_id,) + ), + lambda values: len(values) == 1, + seconds=70, + ) + metrics: Final = eventually( + lambda: read_rows( + 'SELECT requests_evaluated, passed_count FROM "LiteLLM_DailyGuardrailMetrics" WHERE guardrail_id=%s', + (string_value(indexed[0]["guardrail_id"]),), + ), + lambda values: values == [{"requests_evaluated": 1, "passed_count": 1}], + seconds=70, + ) + assert metrics == [{"requests_evaluated": 1, "passed_count": 1}], metrics + + +EXCLUDED: Final = ("model_map_information", "user_api_key_alias") + + +@dataclass(frozen=True, slots=True) +class _Isolated: + proxy: Gateway + scenario: Scenario + key: str + database_url: str | None = None + + @property + def hashed_key(self) -> str: + return sha256(self.key.encode()).hexdigest() + + def rows(self, count: int, where: str = "TRUE", seconds: int = 70) -> tuple[dict[str, JsonValue], ...]: + return tuple( + eventually( + lambda: read_rows( + "SELECT request_id, call_type, status, spend, cache_hit, metadata " + f'FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND {where} ORDER BY "startTime"', + (self.hashed_key,), + database_url=self.database_url, + ), + lambda values: len(values) == count, + seconds=seconds, + ) + ) + + +@contextmanager +def _isolated(gateway: Gateway, config: Path, tmp_path: Path, database_url: str | None = None) -> Generator[_Isolated]: + database: Final = {} if database_url is None else {"DATABASE_URL": database_url} + replica: Final = () if database_url is None else ("DATABASE_URL_READ_REPLICA",) + with ( + owned_proxy(gateway, tmp_path, database, config=config, remove_environment=replica) as proxy, + proxy.scenario() as scenario, + ): + key: Final = scenario.key(key_alias=f"metadata-fields-{uuid4()}") + yield _Isolated(proxy, scenario, key, database_url) + + +def _assert_filtered(row: Mapping[str, JsonValue], *kept: str) -> dict[str, JsonValue]: + metadata: Final = object_value(row["metadata"]) + assert not set(EXCLUDED) & set(metadata), metadata + assert {"status", "cold_storage_object_key", *kept} <= set(metadata), metadata + return metadata + + +def _anthropic_event(event: Mapping[str, JsonValue]) -> bytes: + return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def _anthropic_stream(model: JsonValue) -> tuple[bytes, ...]: + message: Final = {**_anthropic_message(model), "content": [], "stop_reason": None} + events: Final[tuple[Mapping[str, JsonValue], ...]] = ( + {"type": "message_start", "message": message}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 2}}, + {"type": "message_stop"}, + ) + return tuple(_anthropic_event(event) for event in events) + + +def _respond_streaming_or_not(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=b'{"object":"list","data":[]}') + if request.target.startswith("/v1/messages"): + return _anthropic(request) + if JSON_OBJECT.validate_json(request.body).get("stream") is True: + return Reply(content_type="text/event-stream", chunks=text_stream(f"chatcmpl-{uuid4()}")) + return _respond(request) + + +def _stream(proxy: Gateway, key: str, path: str, body: Mapping[str, JsonValue]) -> None: + with proxy.client.stream("POST", path, json=dict(body), headers={"Authorization": f"Bearer {key}"}) as response: + assert response.status_code == 200, response.read().decode() + lines: Final = tuple(response.iter_lines()) + assert any(line.startswith("data:") for line in lines), lines + + +def _call(proxy: Gateway, key: str, path: str, body: Mapping[str, JsonValue]) -> None: + response: Final = proxy.request("POST", path, body, key=key) + assert response.status_code == 200, response.text + + +def test_every_inference_surface_and_stream_stores_filtered_metadata(gateway: Gateway, tmp_path: Path) -> None: + with ( + wire_server(_respond_streaming_or_not) as wire, + _isolated(gateway, _config(tmp_path, {"exclude": list(EXCLUDED)}), tmp_path) as isolated, + ): + model: Final = isolated.scenario.model( + model="deepseek/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-key" + ) + anthropic: Final = isolated.scenario.model( + model="anthropic/claude-sonnet-4-6", api_base=wire.url, api_key="synthetic-key" + ) + prompt: Final[list[JsonValue]] = [{"role": "user", "content": f"surfaces {uuid4()}"}] + calls: Final[tuple[Callable[[Gateway, str, str, Mapping[str, JsonValue]], None], ...]] = ( + _stream, + _call, + _stream, + _call, + _stream, + ) + requests: Final[tuple[tuple[str, Mapping[str, JsonValue]], ...]] = ( + ("/v1/chat/completions", {"model": model, "messages": prompt, "stream": True}), + ("/v1/messages", {"model": anthropic, "max_tokens": 16, "messages": prompt}), + ("/v1/messages", {"model": anthropic, "max_tokens": 16, "messages": prompt, "stream": True}), + ("/v1/responses", {"model": model, "input": f"surfaces {uuid4()}"}), + ("/v1/responses", {"model": model, "input": f"surfaces {uuid4()}", "stream": True}), + ) + for send, (path, body) in zip(calls, requests, strict=True): + send(isolated.proxy, isolated.key, path, body) + + rows: Final = isolated.rows(len(requests)) + for row in rows: + assert row["status"] == "success", row + _assert_filtered(row, "usage_object") + + +def test_failed_request_keeps_its_error_and_status_but_drops_excluded_metadata( + gateway: Gateway, tmp_path: Path +) -> None: + def fail(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=b'{"object":"list","data":[]}') + return Reply(status=500, body=b'{"error":{"message":"scripted upstream failure","type":"server_error"}}') + + with ( + wire_server(fail) as wire, + _isolated(gateway, _config(tmp_path, {"exclude": list(EXCLUDED)}), tmp_path) as isolated, + ): + model: Final = isolated.scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic") + response: Final = isolated.proxy.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "fail"}]}, + key=isolated.key, + ) + assert response.status_code >= 500, response.text + + row: Final = isolated.rows(1)[0] + assert row["status"] == "failure", row + metadata: Final = _assert_filtered(row, "error_information") + assert metadata["status"] == "failure", metadata + assert "scripted upstream failure" in json.dumps(metadata["error_information"]), metadata + + +def test_response_cache_hit_row_is_filtered_and_charged_nothing(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _config( + tmp_path, {"exclude": list(EXCLUDED)}, litellm_settings={"cache": True, "cache_params": {"type": "local"}} + ) + with wire_server(_respond) as wire, _isolated(gateway, config, tmp_path) as isolated: + model: Final = isolated.scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic") + text: Final = f"cache {uuid4()}" + first: Final = string_value(isolated.proxy.chat(model, key=isolated.key, text=text)["id"]) + isolated.proxy.chat(model, key=isolated.key, text=text) + + rows: Final = isolated.rows(2) + assert len([request for request in wire.drain() if request.method == "POST"]) == 1 + paid, hit = sorted(rows, key=lambda row: row["cache_hit"] == "True") + assert paid["request_id"] == first and float(str(paid["spend"])) > 0, rows + assert hit["cache_hit"] == "True" and string_value(hit["request_id"]).startswith(first + "_cache_hit"), rows + assert float(str(hit["spend"])) == 0, rows + for row in rows: + _assert_filtered(row, "usage_object") + + +def test_batch_cost_row_is_filtered_and_charged_once(gateway: Gateway, tmp_path: Path) -> None: + with _isolated(gateway, _config(tmp_path, {"exclude": list(EXCLUDED)}), tmp_path) as isolated: + handle: Final = register_scenario(f"metadata-fields-batch-{uuid4().hex[:12]}", batch_routes("gpt-4o-mini")) + isolated.scenario.cleanups.callback(delete_scenario, handle) + model: Final = isolated.scenario.model(api_base=handle.api_base()) + uploaded: Final = isolated.proxy.request_multipart( + "/v1/files", + {"purpose": "batch", "model": model}, + {"file": ("in.jsonl", batch_input_file(model), "application/jsonl")}, + key=isolated.key, + ) + assert uploaded.status_code == 200, uploaded.text + batch: Final = isolated.proxy.post( + "/v1/batches", + { + "input_file_id": string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": model, + }, + key=isolated.key, + ) + batch_path: Final = f"/v1/batches/{string_value(batch['id'])}" + retrievals: Final = tuple(isolated.proxy.request("GET", batch_path, key=isolated.key) for _ in range(2)) + assert all(r.status_code == 200 and r.json()["status"] == "completed" for r in retrievals), retrievals + + row: Final = isolated.rows(1, "call_type='aretrieve_batch'")[0] + assert float(str(row["spend"])) > 0, row + metadata: Final = _assert_filtered(row, "usage_object") + assert (metadata["batch_successful_requests"], metadata["batch_failed_requests"]) == (2, 3), metadata + + +def _update_retention(proxy: Gateway, value: JsonValue) -> httpx.Response: + return proxy.request( + "POST", + "/config/field/update", + {"field_name": "spend_logs_metadata_fields", "field_value": value, "config_type": "general_settings"}, + ) + + +def test_runtime_retention_update_rejects_invalid_values_and_applies_valid_ones_without_restart( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + scratch_database() as database_url, + wire_server(_respond) as wire, + _isolated(gateway, _config(tmp_path, None), tmp_path, database_url) as isolated, + ): + model: Final = isolated.scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic") + invalid: Final[tuple[JsonValue, ...]] = ( + {"include": ["usage_object"], "exclude": ["model_map_information"]}, + {}, + {"exclude": ["not_a_metadata_field"]}, + {"exclude": ["status"]}, + ) + + isolated.proxy.chat(model, key=isolated.key) + assert set(EXCLUDED) <= set(object_value(isolated.rows(1)[0]["metadata"])) + for value in invalid: + assert _update_retention(isolated.proxy, value).status_code == 400, value + isolated.proxy.chat(model, key=isolated.key) + assert set(EXCLUDED) <= set(object_value(isolated.rows(2)[-1]["metadata"])) + + assert _update_retention(isolated.proxy, {"exclude": list(EXCLUDED)}).status_code == 200 + isolated.proxy.chat(model, key=isolated.key) + _assert_filtered(isolated.rows(3)[-1]) + + for value in invalid: + assert _update_retention(isolated.proxy, value).status_code == 400, value + isolated.proxy.chat(model, key=isolated.key) + _assert_filtered(isolated.rows(4)[-1]) + + +SAVINGS_FIELDS: Final = ("autorouter_savings_estimate", "autorouter_savings") + + +def _anthropic_message(model: JsonValue) -> dict[str, JsonValue]: + return { + "id": f"msg_{uuid4().hex}", + "type": "message", + "role": "assistant", + "model": model, + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 12, + "output_tokens": 2, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + } + + +def _anthropic(request: Request) -> Reply: + if request.target.startswith("/v1/messages/count_tokens"): + return Reply(body=b'{"input_tokens":12}') + body: Final = JSON_OBJECT.validate_json(request.body) + if body.get("stream") is True: + return Reply(content_type="text/event-stream", chunks=_anthropic_stream(body["model"])) + return Reply(body=json.dumps(_anthropic_message(body["model"])).encode()) + + +def _router_models(wire: Wire, router: str) -> tuple[Mapping[str, JsonValue], ...]: + tiers: Final = {"cheap": "anthropic/claude-sonnet-4-6", "frontier": "anthropic/claude-opus-4-8"} + return ( + *( + {"model_name": name, "litellm_params": {"model": model, "api_base": wire.url, "api_key": "synthetic"}} + for name, model in tiers.items() + ), + { + "model_name": router, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "cheap", "MEDIUM": "cheap", "COMPLEX": "frontier", "REASONING": "frontier"}, + }, + }, + }, + ) + + +@pytest.mark.parametrize( + ("retention", "stored"), + [(None, SAVINGS_FIELDS), ({"exclude": list(SAVINGS_FIELDS)}, ())], + ids=["unset-keeps-savings", "excluded-savings-stay-out"], +) +@pytest.mark.timeout(240) +def test_delayed_autorouter_savings_publication_respects_retention( + gateway: Gateway, tmp_path: Path, retention: Mapping[str, JsonValue] | None, stored: tuple[str, ...] +) -> None: + router: Final = f"router-{uuid4().hex[:8]}" + with wire_server(_anthropic) as wire: + config: Final = _config(tmp_path, retention, model_list=_router_models(wire, router)) + with _isolated(gateway, config, tmp_path) as isolated: + response: Final = isolated.proxy.request( + "POST", + "/v1/messages", + {"model": router, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]}, + key=isolated.key, + headers={"x-litellm-session-id": f"session-{uuid4()}"}, + ) + assert response.status_code == 200, response.text + + published: Final = isolated.rows( + 1, + "metadata::jsonb ? 'routing_decision' AND NOT metadata::jsonb ? 'autorouter_baseline_observation' " + 'AND request_id IN (SELECT request_id FROM "LiteLLM_AutoRouterBaselineObservation" ' + "WHERE publication IS NOT NULL)", + seconds=150, + )[0] + metadata: Final = object_value(published["metadata"]) + assert {name for name in SAVINGS_FIELDS if name in metadata} == set(stored), metadata diff --git a/tests/integration/translation/decisions/bases/cloudflare.py b/tests/integration/translation/decisions/bases/cloudflare.py index 1668dc215a8..ae86a2ccd87 100644 --- a/tests/integration/translation/decisions/bases/cloudflare.py +++ b/tests/integration/translation/decisions/bases/cloudflare.py @@ -6,7 +6,7 @@ from integration.translation.case import TranslationTestCase """ CLEF_TEST_CASE: Final = TranslationTestCase( scenario="basic", - litellm_endpoint="/v1/decisions", + litellm_endpoint="/v1/systemone", litellm_request={ "model": "cloudflare/@cf/cloudflare/clef", "state": "Ticket (billing): The export job hangs at 99% and never finishes", diff --git a/tests/integration/translation/decisions/bases/openrouter.py b/tests/integration/translation/decisions/bases/openrouter.py index f6217bcf818..b65acff6b06 100644 --- a/tests/integration/translation/decisions/bases/openrouter.py +++ b/tests/integration/translation/decisions/bases/openrouter.py @@ -6,7 +6,7 @@ from integration.translation.case import TranslationTestCase """ TYPESAFE_JEV_1_13_TEST_CASE: Final = TranslationTestCase( scenario="basic", - litellm_endpoint="/v1/decisions", + litellm_endpoint="/v1/systemone", litellm_request={ "model": "openrouter/typesafe/jev-1.13", "state": "Ticket (billing): The export job hangs at 99% and never finishes", diff --git a/tests/integration/translation/decisions/bases/perplexity.py b/tests/integration/translation/decisions/bases/perplexity.py index 71ef7a85220..02b1de51f24 100644 --- a/tests/integration/translation/decisions/bases/perplexity.py +++ b/tests/integration/translation/decisions/bases/perplexity.py @@ -6,7 +6,7 @@ from integration.translation.case import TranslationTestCase """ PPLX_DECIDER_V1_27B_TEST_CASE: Final = TranslationTestCase( scenario="basic", - litellm_endpoint="/v1/decisions", + litellm_endpoint="/v1/systemone", litellm_request={ "model": "perplexity/pplx-decider-v1-27b", "state": "Ticket (billing): The export job hangs at 99% and never finishes", diff --git a/tests/integration/translation/decisions/bases/strands_decider.py b/tests/integration/translation/decisions/bases/strands_decider.py index 83d7ae5e33a..8a1a51f8c3e 100644 --- a/tests/integration/translation/decisions/bases/strands_decider.py +++ b/tests/integration/translation/decisions/bases/strands_decider.py @@ -6,7 +6,7 @@ from integration.translation.case import TranslationTestCase """ STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE: Final = TranslationTestCase( scenario="basic", - litellm_endpoint="/v1/decisions", + litellm_endpoint="/v1/systemone", litellm_request={ "model": "strands_decider/strands-decider-2B-hobson-v19", "state": "Help! My payouts have been failing for 3 days!", diff --git a/tests/integration/translation/decisions/bases/typesafe.py b/tests/integration/translation/decisions/bases/typesafe.py index 2fb384a94ae..5da5b05cb09 100644 --- a/tests/integration/translation/decisions/bases/typesafe.py +++ b/tests/integration/translation/decisions/bases/typesafe.py @@ -6,7 +6,7 @@ from integration.translation.case import TranslationTestCase """ JEV_1_13_0_TEST_CASE: Final = TranslationTestCase( scenario="basic", - litellm_endpoint="/v1/decisions", + litellm_endpoint="/v1/systemone", litellm_request={ "model": "typesafe/jev-1.13.0", "state": "Ticket (billing): The export job hangs at 99% and never finishes", diff --git a/tests/llm_translation/base_audio_transcription_unit_tests.py b/tests/llm_translation/base_audio_transcription_unit_tests.py index 0234c05f853..bdf9f998d57 100644 --- a/tests/llm_translation/base_audio_transcription_unit_tests.py +++ b/tests/llm_translation/base_audio_transcription_unit_tests.py @@ -1,20 +1,11 @@ import httpx import json -import pytest from typing import Any, Dict, List from unittest.mock import MagicMock, Mock, patch import os from litellm._uuid import uuid import litellm -from litellm import transcription -from litellm.litellm_core_utils.get_supported_openai_params import ( - get_supported_openai_params, -) -from litellm.llms.base_llm.audio_transcription.transformation import ( - BaseAudioTranscriptionConfig, -) -from litellm.utils import ProviderConfigManager from abc import ABC, abstractmethod pwd = os.path.dirname(os.path.realpath(__file__)) @@ -24,7 +15,6 @@ file_path = os.path.join(pwd, "gettysburg.wav") audio_file = open(file_path, "rb") - class BaseLLMAudioTranscriptionTest(ABC): @abstractmethod def get_base_audio_transcription_call_args(self) -> dict: @@ -36,65 +26,3 @@ class BaseLLMAudioTranscriptionTest(ABC): """Must return the custom llm provider""" pass - def test_audio_transcription(self): - """ - Test that the audio transcription is translated correctly. - """ - litellm.set_verbose = True - transcription_call_args = self.get_base_audio_transcription_call_args() - transcript = transcription(**transcription_call_args, file=audio_file) - print(f"transcript: {transcript.model_dump()}") - print(f"transcript hidden params: {transcript._hidden_params}") - - assert transcript.text is not None - - @pytest.mark.asyncio - async def test_audio_transcription_async(self): - """ - Test that the audio transcription is translated correctly. - """ - - litellm.set_verbose = True - litellm.turn_on_debug() - AUDIO_FILE = open(file_path, "rb") - transcription_call_args = self.get_base_audio_transcription_call_args() - transcript = await litellm.atranscription( - **transcription_call_args, file=AUDIO_FILE - ) - print(f"transcript: {transcript.model_dump()}") - print(f"transcript hidden params: {transcript._hidden_params}") - - assert transcript.text is not None - - def test_audio_transcription_optional_params(self): - """ - Test that the audio transcription is translated correctly. - """ - transcription_args = self.get_base_audio_transcription_call_args() - model = transcription_args["model"] - custom_llm_provider = self.get_custom_llm_provider() - optional_params = get_supported_openai_params( - model=model, - custom_llm_provider=custom_llm_provider.value, - request_type="transcription", - ) - print(f"optional_params: {optional_params}") - assert optional_params is not None - assert ( - "max_completion_tokens" not in optional_params - ) # assert default chat completion response not returned - - def test_audio_transcription_config(self): - """ - Test that the audio transcription config is implemented and correctly instrumented. - """ - transcription_args = self.get_base_audio_transcription_call_args() - model = transcription_args["model"] - custom_llm_provider = self.get_custom_llm_provider() - config = ProviderConfigManager.get_provider_audio_transcription_config( - model=model, - provider=custom_llm_provider, - ) - print(f"config: {config}") - assert config is not None - assert isinstance(config, BaseAudioTranscriptionConfig) diff --git a/tests/llm_translation/base_embedding_unit_tests.py b/tests/llm_translation/base_embedding_unit_tests.py index 469416fc0cf..7efd08896c8 100644 --- a/tests/llm_translation/base_embedding_unit_tests.py +++ b/tests/llm_translation/base_embedding_unit_tests.py @@ -4,17 +4,14 @@ import json import pytest from typing import Any, Dict, List from unittest.mock import MagicMock, Mock, patch -import os import litellm -from litellm import embedding from litellm.exceptions import BadRequestError from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.utils import ( CustomStreamWrapper, get_supported_openai_params, get_optional_params, - get_optional_params_embeddings, ) import base64 from pathlib import Path @@ -26,7 +23,6 @@ file_data = (Path(__file__).parent.parent / "white_100x100.png").read_bytes() encoded_file = base64.b64encode(file_data).decode("utf-8") base64_image = f"data:image/png;base64,{encoded_file}" - class BaseLLMEmbeddingTest(ABC): """ Abstract base test class that enforces a common test across all test classes. @@ -66,23 +62,3 @@ class BaseLLMEmbeddingTest(ABC): CreateEmbeddingResponse.model_validate(response.model_dump()) - def test_embedding_optional_params_max_retries(self): - embedding_call_args = self.get_base_embedding_call_args() - optional_params = get_optional_params_embeddings( - **embedding_call_args, max_retries=20 - ) - assert optional_params["max_retries"] == 20 - - def test_image_embedding(self): - litellm.set_verbose = True - from litellm.utils import supports_embedding_image_input - - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - base_embedding_call_args = self.get_base_embedding_call_args() - if not supports_embedding_image_input(base_embedding_call_args["model"], None): - print("Model does not support embedding image input") - pytest.skip("Model does not support embedding image input") - - embedding(**base_embedding_call_args, input=[base64_image]) diff --git a/tests/llm_translation/base_llm_unit_tests.py b/tests/llm_translation/base_llm_unit_tests.py index e4ccb589ff4..d8d53dcc45c 100644 --- a/tests/llm_translation/base_llm_unit_tests.py +++ b/tests/llm_translation/base_llm_unit_tests.py @@ -1,11 +1,9 @@ -import httpx import json import pytest import sys from typing import Any, Dict, List -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import MagicMock, Mock import os -import base64 import inspect import litellm @@ -14,10 +12,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.utils import ( CustomStreamWrapper, get_supported_openai_params, - get_optional_params, - ProviderConfigManager, ) -from litellm.main import stream_chunk_builder from typing import Union from litellm.types.utils import Usage, ModelResponse @@ -27,8 +22,6 @@ from openai import OpenAI sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) - - class BaseLLMChatTest(ABC): """ Abstract base test class that enforces a common test across all test classes. @@ -61,155 +54,6 @@ class BaseLLMChatTest(ABC): except litellm.InternalServerError: pytest.skip("Model is overloaded") - def test_developer_role_translation(self): - """ - Test that the developer role is translated correctly for non-OpenAI providers. - - Translate `developer` role to `system` role for non-OpenAI providers. - """ - base_completion_call_args = self.get_base_completion_call_args() - messages = [ - { - "role": "developer", - "content": "Be a good bot!", - }, - { - "role": "user", - "content": [{"type": "text", "text": "Hello, how are you?"}], - }, - ] - try: - response = self.completion_function( - **base_completion_call_args, - messages=messages, - ) - assert response is not None - except litellm.InternalServerError: - pytest.skip("Model is overloaded") - - assert response.choices[0].message.content is not None - - def test_content_list_handling(self): - """Check if content list is supported by LLM API""" - base_completion_call_args = self.get_base_completion_call_args() - messages = [ - { - "role": "user", - "content": [{"type": "text", "text": "Hello, how are you?"}], - } - ] - try: - response = self.completion_function( - **base_completion_call_args, - messages=messages, - ) - assert response is not None - except litellm.InternalServerError: - pytest.skip("Model is overloaded") - - # for OpenAI the content contains the JSON schema, so we need to assert that the content is not None - assert response.choices[0].message.content is not None - - def test_tool_call_with_property_type_array(self): - litellm.turn_on_debug() - from litellm.utils import supports_function_calling - - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - base_completion_call_args = self.get_base_completion_call_args() - if not supports_function_calling(base_completion_call_args["model"], None): - print("Model does not support function calling") - pytest.skip("Model does not support function calling") - base_completion_call_args = self.get_base_completion_call_args() - response = self.completion_function( - **base_completion_call_args, - messages=[ - { - "role": "user", - "content": "Tell me if the shoe brand Air Jordan has more models than the shoe brand Nike.", - } - ], - tools=[ - { - "type": "function", - "function": { - "name": "shoe_get_id", - "description": "Get information about a show by its ID or name", - "parameters": { - "type": "object", - "properties": { - "shoe_id": { - "type": ["string", "number"], - "description": "The shoe ID or name", - } - }, - "required": ["shoe_id"], - "additionalProperties": False, - "$schema": "http://json-schema.org/draft-07/schema#", - }, - }, - }, - ], - ) - print(response) - print(json.dumps(response, indent=4, default=str)) - - @pytest.mark.flaky(retries=3, delay=1) - def test_tool_call_with_empty_enum_property(self): - litellm.turn_on_debug() - from litellm.utils import supports_function_calling - - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - base_completion_call_args = self.get_base_completion_call_args() - if not supports_function_calling(base_completion_call_args["model"], None): - print("Model does not support function calling") - pytest.skip("Model does not support function calling") - base_completion_call_args = self.get_base_completion_call_args() - response = self.completion_function( - **base_completion_call_args, - messages=[ - { - "role": "user", - "content": "Search for the latest iPhone models and tell me which storage options are available.", - } - ], - tools=[ - { - "type": "function", - "function": { - "name": "litellm_product_search", - "description": "Search for product information and specifications.\n\nSupports filtering by category, brand, price range, and availability.\nCan retrieve detailed product specifications, pricing, and stock information.\nSupports different search modes and result formatting options.\n", - "parameters": { - "properties": { - "search_mode": { - "default": "", - "description": "The search strategy to use for finding products.", - "enum": [ - "", - "product_search", - "product_search_with_filters", - "product_search_with_sorting", - "product_search_with_pagination", - "product_search_with_aggregation", - ], - "title": "Search Mode", - "type": "string", - }, - }, - "required": ["search_mode"], - "title": "product_search_arguments", - "type": "object", - }, - }, - } - ], - ) - print(response) - print(json.dumps(response, indent=4, default=str)) - def test_streaming(self): """Check if litellm handles streaming correctly""" from litellm.types.utils import ModelResponseStream @@ -252,16 +96,6 @@ class BaseLLMChatTest(ABC): # assert resp.usage.completion_tokens > 0 # assert resp.usage.total_tokens > 0 - def test_pydantic_model_input(self): - litellm.set_verbose = True - - from litellm import completion, Message - - base_completion_call_args = self.get_base_completion_call_args() - messages = [Message(content="Hello, how are you?", role="user")] - - self.completion_function(**base_completion_call_args, messages=messages) - def test_web_search(self): from litellm.utils import supports_web_search @@ -388,407 +222,6 @@ class BaseLLMChatTest(ABC): assert response is not None - def test_file_data_unit_test(self, pdf_messages): - from litellm.utils import supports_pdf_input, return_raw_request - from litellm.types.utils import CallTypes - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_anthropic_image_obj, - ) - - media_chunk = convert_to_anthropic_image_obj( - openai_image_url=pdf_messages, - format=None, - ) - - file_content = [ - {"type": "text", "text": "What's this file about?"}, - { - "type": "file", - "file": { - "file_data": pdf_messages, - }, - }, - ] - - image_messages = [{"role": "user", "content": file_content}] - - base_completion_call_args = self.get_base_completion_call_args() - - if not supports_pdf_input(base_completion_call_args["model"], None): - pytest.skip("Model does not support image input") - - raw_request = return_raw_request( - endpoint=CallTypes.completion, - kwargs={**base_completion_call_args, "messages": image_messages}, - ) - - print("RAW REQUEST", raw_request) - - assert media_chunk["data"] in json.dumps(raw_request) - - def test_message_with_name(self): - try: - litellm.set_verbose = True - base_completion_call_args = self.get_base_completion_call_args() - messages = [ - {"role": "user", "content": "Hello", "name": "test_name"}, - ] - response = self.completion_function( - **base_completion_call_args, messages=messages - ) - assert response is not None - except litellm.RateLimitError: - pass - - @pytest.mark.parametrize( - "response_format", - [ - {"type": "json_object"}, - {"type": "text"}, - ], - ) - @pytest.mark.flaky(retries=6, delay=1) - def test_json_response_format(self, response_format): - """ - Test that the JSON response format is supported by the LLM API - """ - from litellm.utils import supports_response_schema - - base_completion_call_args = self.get_base_completion_call_args() - litellm.set_verbose = True - - if not supports_response_schema(base_completion_call_args["model"], None): - pytest.skip("Model does not support response schema") - - messages = [ - { - "role": "system", - "content": "Your output should be a JSON object with no additional properties. ", - }, - { - "role": "user", - "content": "Respond with this in json. city=San Francisco, state=CA, weather=sunny, temp=60", - }, - ] - - response = self.completion_function( - **base_completion_call_args, - messages=messages, - response_format=response_format, - ) - - print(f"response={response}") - - # OpenAI guarantees that the JSON schema is returned in the content - # relevant issue: https://github.com/BerriAI/litellm/issues/6741 - assert response.choices[0].message.content is not None - - @pytest.mark.parametrize( - "response_format", - [ - {"type": "text"}, - ], - ) - @pytest.mark.flaky(retries=6, delay=1) - def test_response_format_type_text_with_tool_calls_no_tool_choice( - self, response_format - ): - base_completion_call_args = self.get_base_completion_call_args() - messages = [ - {"role": "user", "content": "What's the weather like in Boston today?"}, - ] - tools = [ - { - "type": "function", - "function": { - "name": "get_current_weather", - "description": "Get the current weather in a given location", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. San Francisco, CA", - }, - "unit": { - "type": "string", - "enum": ["celsius", "fahrenheit"], - }, - }, - "required": ["location"], - }, - }, - } - ] - try: - print(f"MAKING LLM CALL") - response = self.completion_function( - **base_completion_call_args, - messages=messages, - response_format=response_format, - tools=tools, - drop_params=True, - ) - print(f"RESPONSE={response}") - except litellm.ContextWindowExceededError: - pytest.skip("Model exceeded context window") - assert response is not None - - def test_response_format_type_text(self): - """ - Test that the response format type text does not lead to tool calls - """ - from litellm import LlmProviders - - base_completion_call_args = self.get_base_completion_call_args() - litellm.set_verbose = True - - _, provider, _, _ = litellm.get_llm_provider( - model=base_completion_call_args["model"] - ) - - provider_config = ProviderConfigManager.get_provider_chat_config( - base_completion_call_args["model"], LlmProviders(provider) - ) - - print(f"provider_config={provider_config}") - - translated_params = provider_config.map_openai_params( - non_default_params={"response_format": {"type": "text"}}, - optional_params={}, - model=base_completion_call_args["model"], - drop_params=False, - ) - - assert "tool_choice" not in translated_params - assert ( - "tools" not in translated_params - ), f"Got tools={translated_params['tools']}, expected no tools" - - print(f"translated_params={translated_params}") - - @pytest.mark.flaky(retries=6, delay=1) - def test_json_response_pydantic_obj(self): - litellm.turn_on_debug() - from pydantic import BaseModel - from litellm.utils import supports_response_schema - - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - class TestModel(BaseModel): - first_response: str - - base_completion_call_args = self.get_base_completion_call_args() - if not supports_response_schema(base_completion_call_args["model"], None): - pytest.skip("Model does not support response schema") - - try: - res = self.completion_function( - **base_completion_call_args, - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - { - "role": "user", - "content": "What is the capital of France?", - }, - ], - response_format=TestModel, - timeout=5, - ) - assert res is not None - - print(res.choices[0].message) - - assert res.choices[0].message.content is not None - assert res.choices[0].message.tool_calls is None - except litellm.Timeout: - pytest.skip("Model took too long to respond") - except litellm.InternalServerError: - pytest.skip("Model is overloaded") - - @pytest.mark.flaky(retries=6, delay=1) - def test_json_response_nested_pydantic_obj(self): - from pydantic import BaseModel - from litellm.utils import supports_response_schema - - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - class CalendarEvent(BaseModel): - name: str - date: str - participants: list[str] - - class EventsList(BaseModel): - events: list[CalendarEvent] - - messages = [ - {"role": "user", "content": "List 5 important events in the XIX century"} - ] - - base_completion_call_args = self.get_base_completion_call_args() - if not supports_response_schema(base_completion_call_args["model"], None): - pytest.skip( - f"Model={base_completion_call_args['model']} does not support response schema" - ) - - try: - res = self.completion_function( - **base_completion_call_args, - messages=messages, - response_format=EventsList, - timeout=60, - ) - assert res is not None - - print(res.choices[0].message) - - assert res.choices[0].message.content is not None - assert res.choices[0].message.tool_calls is None - except litellm.Timeout: - pytest.skip("Model took too long to respond") - except litellm.InternalServerError: - pytest.skip("Model is overloaded") - - @pytest.mark.flaky(retries=6, delay=1) - def test_json_response_nested_json_schema(self): - """ - PROD Test: ensure nested json schema sent to proxy works as expected. - """ - litellm.turn_on_debug() - from pydantic import BaseModel - from litellm.utils import supports_response_schema - from litellm.llms.base_llm.base_utils import type_to_response_format_param - - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - class CalendarEvent(BaseModel): - name: str - date: str - participants: list[str] - - class EventsList(BaseModel): - events: list[CalendarEvent] - - response_format = type_to_response_format_param(EventsList) - - messages = [ - {"role": "user", "content": "List 5 important events in the XIX century"} - ] - - base_completion_call_args = self.get_base_completion_call_args() - if not supports_response_schema(base_completion_call_args["model"], None): - pytest.skip( - f"Model={base_completion_call_args['model']} does not support response schema" - ) - - try: - res = self.completion_function( - **base_completion_call_args, - messages=messages, - response_format=response_format, - timeout=60, - ) - assert res is not None - - print(res.choices[0].message) - - assert res.choices[0].message.content is not None - assert res.choices[0].message.tool_calls is None - except litellm.Timeout: - pytest.skip("Model took too long to respond") - except litellm.InternalServerError: - pytest.skip("Model is overloaded") - - @pytest.mark.flaky(retries=6, delay=1) - def test_audio_input(self): - """ - Test that audio input is supported by the LLM API - """ - from litellm.utils import supports_audio_input - - litellm.turn_on_debug() - base_completion_call_args = self.get_base_completion_call_args() - if not supports_audio_input(base_completion_call_args["model"], None): - pytest.skip( - f"Model={base_completion_call_args['model']} does not support audio input" - ) - - url = "https://openaiassets.blob.core.windows.net/$web/API/docs/audio/alloy.wav" - response = httpx.get(url) - response.raise_for_status() - wav_data = response.content - encoded_string = base64.b64encode(wav_data).decode("utf-8") - - completion = self.completion_function( - **base_completion_call_args, - messages=[ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is in this recording?"}, - { - "type": "input_audio", - "input_audio": {"data": encoded_string, "format": "wav"}, - }, - ], - }, - ], - ) - - print(completion.choices[0].message) - - @pytest.mark.flaky(retries=6, delay=1) - def test_json_response_format_stream(self): - """ - Test that the JSON response format with streaming is supported by the LLM API - """ - from litellm.utils import supports_response_schema - - base_completion_call_args = self.get_base_completion_call_args() - litellm.set_verbose = True - - base_completion_call_args = self.get_base_completion_call_args() - if not supports_response_schema(base_completion_call_args["model"], None): - pytest.skip("Model does not support response schema") - - messages = [ - { - "role": "system", - "content": "Your output should be a JSON object with no additional properties. ", - }, - { - "role": "user", - "content": "Respond with this in json. city=San Francisco, state=CA, weather=sunny, temp=60", - }, - ] - - try: - response = self.completion_function( - **base_completion_call_args, - messages=messages, - response_format={"type": "json_object"}, - stream=True, - ) - except litellm.InternalServerError: - pytest.skip("Model is overloaded") - - print(response) - - content = "" - for chunk in response: - content += chunk.choices[0].delta.content or "" - - print(f"content={content}") - - # OpenAI guarantees that the JSON schema is returned in the content - # relevant issue: https://github.com/BerriAI/litellm/issues/6741 - # we need to assert that the JSON schema was returned in the content, (for Anthropic we were returning it as part of the tool call) - assert content is not None - assert len(content) > 0 - @pytest.fixture def tool_call_no_arguments(self): return { @@ -803,118 +236,6 @@ class BaseLLMChatTest(ABC): ], } - @pytest.mark.parametrize("detail", [None, "low", "high"]) - @pytest.mark.parametrize( - "image_url", - [ - # In-repo logo served via jsdelivr (sha-pinned, immutable). - # Bedrock fetches the URL and base64-embeds it in the - # Converse request body; using a multi-MB hosted product - # photo here previously bloated cassettes to ~60 MB each. - "https://cdn.jsdelivr.net/gh/BerriAI/litellm@d769e81c90d453240c61fc572cdb27fae06a89d0/ui/litellm-dashboard/public/assets/logos/litellm_logo.jpg", - "https://awsmp-logos.s3.amazonaws.com/seller-xw5kijmvmzasy/c233c9ade2ccb5491072ae232c814942.png", - ], - ) - @pytest.mark.flaky(retries=4, delay=2) - def test_image_url(self, detail, image_url): - litellm.set_verbose = True - from litellm.utils import supports_vision - - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - base_completion_call_args = self.get_base_completion_call_args() - if not supports_vision(base_completion_call_args["model"], None): - pytest.skip("Model does not support image input") - elif "http://" in image_url and ( - "fireworks_ai" in base_completion_call_args.get("model", "") - or "mistral" in base_completion_call_args.get("model", "") - ): - pytest.skip("Model does not support http:// input") - - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What's in this image?"}, - { - "type": "image_url", - "image_url": { - "url": image_url, - }, - }, - ], - } - ] - - if detail is not None: - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What's in this image?"}, - { - "type": "image_url", - "image_url": { - # sha-pinned in-repo logo via jsdelivr; gstatic's - # robots.txt blocks server-side fetchers (e.g. - # Anthropic), which 400s the request. - "url": "https://cdn.jsdelivr.net/gh/BerriAI/litellm@d769e81c90d453240c61fc572cdb27fae06a89d0/ui/litellm-dashboard/public/assets/logos/litellm_logo.jpg", - "detail": detail, - }, - }, - ], - } - ] - try: - response = self.completion_function( - **base_completion_call_args, messages=messages - ) - except litellm.InternalServerError: - pytest.skip("Model is overloaded") - - assert response is not None - - def test_image_url_string(self): - litellm.set_verbose = True - from litellm.utils import supports_vision - - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - image_url = "https://awsmp-logos.s3.amazonaws.com/seller-xw5kijmvmzasy/c233c9ade2ccb5491072ae232c814942.png" - - base_completion_call_args = self.get_base_completion_call_args() - if not supports_vision(base_completion_call_args["model"], None): - pytest.skip("Model does not support image input") - elif "http://" in image_url and "fireworks_ai" in base_completion_call_args.get( - "model" - ): - pytest.skip("Model does not support http:// input") - - image_url_param = image_url - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What's in this image?"}, - { - "type": "image_url", - "image_url": image_url_param, - }, - ], - } - ] - - try: - response = self.completion_function( - **base_completion_call_args, messages=messages - ) - except litellm.InternalServerError: - pytest.skip("Model is overloaded") - - assert response is not None - @pytest.fixture def pdf_messages(self): import base64 @@ -932,42 +253,6 @@ class BaseLLMChatTest(ABC): return url - @pytest.mark.flaky(retries=3, delay=1) - def test_empty_tools(self): - """ - Related Issue: https://github.com/BerriAI/litellm/issues/9080 - """ - try: - from litellm import completion, ModelResponse - - litellm.set_verbose = True - litellm.turn_on_debug() - from litellm.utils import supports_function_calling - - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - base_completion_call_args = self.get_base_completion_call_args() - if not supports_function_calling(base_completion_call_args["model"], None): - print("Model does not support function calling") - pytest.skip("Model does not support function calling") - - response = completion( - **base_completion_call_args, - messages=[{"role": "user", "content": "Hello, how are you?"}], - tools=[], - ) # just make sure call doesn't fail - print("response: ", response) - assert response is not None - except litellm.ContentPolicyViolationError: - pass - except litellm.InternalServerError: - pytest.skip("Model is overloaded") - except litellm.RateLimitError: - pass - except Exception as e: - pytest.fail(f"Error occurred: {e}") - @pytest.mark.flaky(retries=3, delay=1) def test_basic_tool_calling(self): try: @@ -1091,94 +376,6 @@ class BaseLLMChatTest(ABC): except Exception as e: pytest.fail(f"Error occurred: {e}") - @pytest.mark.flaky(retries=3, delay=1) - @pytest.mark.asyncio - async def test_completion_cost(self): - from litellm import completion_cost - - litellm.turn_on_debug() - - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - litellm.set_verbose = True - response = await self.async_completion_function( - **self.get_base_completion_call_args(), - messages=[{"role": "user", "content": "Hello, how are you?"}], - ) - print(response._hidden_params["response_cost"]) - - assert response._hidden_params["response_cost"] > 0 - - @pytest.mark.parametrize("input_type", ["input_audio", "audio_url"]) - @pytest.mark.parametrize("format_specified", [True]) - def test_supports_audio_input(self, input_type, format_specified): - from litellm.utils import return_raw_request, supports_audio_input - from litellm.types.utils import CallTypes - - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - litellm.drop_params = True - base_completion_call_args = self.get_base_completion_call_args() - if not supports_audio_input(base_completion_call_args["model"], None): - print("Model does not support audio input") - pytest.skip("Model does not support audio input") - - url = "https://openaiassets.blob.core.windows.net/$web/API/docs/audio/alloy.wav" - response = httpx.get(url) - response.raise_for_status() - wav_data = response.content - audio_format = "wav" - encoded_string = base64.b64encode(wav_data).decode("utf-8") - - audio_content = [{"type": "text", "text": "What is in this recording?"}] - - test_file_id = "gs://bucket/file.wav" - - if input_type == "input_audio": - audio_content.append( - { - "type": "input_audio", - "input_audio": {"data": encoded_string, "format": audio_format}, - } - ) - elif input_type == "audio_url": - audio_content.append( - { - "type": "file", - "file": { - "file_id": test_file_id, - "filename": "my-sample-audio-file", - }, - } - ) - - raw_request = return_raw_request( - endpoint=CallTypes.completion, - kwargs={ - **base_completion_call_args, - "modalities": ["text", "audio"], - "audio": {"voice": "alloy", "format": audio_format}, - "messages": [ - { - "role": "user", - "content": audio_content, - }, - ], - }, - ) - print("raw_request: ", raw_request) - - if input_type == "input_audio": - assert encoded_string in json.dumps( - raw_request - ), "Audio data not sent to gemini" - elif input_type == "audio_url": - assert test_file_id in json.dumps( - raw_request - ), "Audio URL not sent to gemini" - def test_function_calling_with_tool_response(self): from litellm.utils import supports_function_calling from litellm import completion @@ -1277,53 +474,6 @@ class BaseLLMChatTest(ABC): except litellm.ServiceUnavailableError: pass - def test_reasoning_effort(self): - """Test that reasoning_effort is passed correctly to the model""" - from litellm.utils import supports_reasoning - from litellm import completion - - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - base_completion_call_args = ( - self.get_base_completion_call_args_with_reasoning_model() - ) - if len(base_completion_call_args) == 0: - print("base_completion_call_args is empty") - pytest.skip("Model does not support reasoning") - if not supports_reasoning(base_completion_call_args["model"], None): - print("Model does not support reasoning") - pytest.skip("Model does not support reasoning") - - _, provider, _, _ = litellm.get_llm_provider( - model=base_completion_call_args["model"] - ) - - ## CHECK PARAM MAPPING - optional_params = get_optional_params( - model=base_completion_call_args["model"], - custom_llm_provider=provider, - reasoning_effort="high", - ) - # either accepts reasoning effort or thinking budget - from litellm.constants import DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET - - assert "reasoning_effort" in optional_params or str( - DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET - ) in json.dumps(optional_params) - - try: - litellm.turn_on_debug() - response = completion( - **base_completion_call_args, - reasoning_effort="low", - messages=[{"role": "user", "content": "Hello!"}], - ) - print(f"response: {response}") - except Exception as e: - pytest.fail(f"Error: {e}") - - class BaseOSeriesModelsTest(ABC): # test across azure/openai @abstractmethod def get_base_completion_call_args(self): @@ -1333,105 +483,6 @@ class BaseOSeriesModelsTest(ABC): # test across azure/openai def get_client(self) -> OpenAI: pass - def test_reasoning_effort(self): - """Test that reasoning_effort is passed correctly to the model""" - - from litellm import completion - - client = self.get_client() - - completion_args = self.get_base_completion_call_args() - - with patch.object( - client.chat.completions.with_raw_response, "create" - ) as mock_client: - try: - completion( - **completion_args, - reasoning_effort="low", - messages=[{"role": "user", "content": "Hello!"}], - client=client, - ) - except Exception as e: - print(f"Error: {e}") - - mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs - print("request_body: ", request_body) - assert request_body["reasoning_effort"] == "low" - - def test_developer_role_translation(self): - """Test that developer role is translated correctly to system role for non-OpenAI providers""" - from litellm import completion - - client = self.get_client() - - completion_args = self.get_base_completion_call_args() - - with patch.object( - client.chat.completions.with_raw_response, "create" - ) as mock_client: - try: - completion( - **completion_args, - reasoning_effort="low", - messages=[ - {"role": "developer", "content": "Be a good bot!"}, - {"role": "user", "content": "Hello!"}, - ], - client=client, - ) - except Exception as e: - print(f"Error: {e}") - - mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs - print("request_body: ", request_body) - assert ( - request_body["messages"][0]["role"] == "developer" - ), "Got={} instead of system".format(request_body["messages"][0]["role"]) - assert request_body["messages"][0]["content"] == "Be a good bot!" - - def test_completion_o_series_models_temperature(self): - """ - Test that temperature is not passed to O-series models - """ - try: - from litellm import completion - - client = self.get_client() - - completion_args = self.get_base_completion_call_args() - - with patch.object( - client.chat.completions.with_raw_response, "create" - ) as mock_client: - try: - completion( - **completion_args, - temperature=0.0, - messages=[ - { - "role": "user", - "content": "Hello, world!", - } - ], - drop_params=True, - client=client, - ) - except Exception as e: - print(f"Error: {e}") - - mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs - print("request_body: ", request_body) - assert ( - "temperature" not in request_body - ), "temperature should not be in the request body" - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - class BaseAnthropicChatTest(ABC): """ Ensures consistent result across anthropic model usage @@ -1451,178 +502,6 @@ class BaseAnthropicChatTest(ABC): def completion_function(self): return litellm.completion - def test_anthropic_response_format_streaming_vs_non_streaming(self): - args = { - "messages": [ - { - "content": "Your goal is to summarize the previous agent's thinking process into short descriptions to let user better understand the research progress. If no information is available, just say generic phrase like 'Doing some research...' with the given output format. Make sure to adhere to the output format no matter what, even if you don't have any information or you are not allowed to respond to the given input information (then just say generic phrase like 'Doing some research...').", - "role": "system", - }, - { - "role": "user", - "content": "Here is the input data (previous agent's output): \n\n Let's try to refine our search further, focusing more on the technical aspects of home automation and home energy system management:", - }, - ], - "response_format": { - "type": "json_schema", - "json_schema": { - "name": "final_output", - "strict": True, - "schema": { - "description": 'Progress report for the thinking process\n\nThis model represents a snapshot of the agent\'s current progress during\nthe thinking process, providing a brief description of the current activity.\n\nAttributes:\n agent_doing: Brief description of what the agent is currently doing.\n Should be kept under 10 words. Example: "Learning about home automation"', - "properties": { - "agent_doing": {"title": "Agent Doing", "type": "string"} - }, - "required": ["agent_doing"], - "title": "ThinkingStep", - "type": "object", - "additionalProperties": False, - }, - }, - }, - } - - base_completion_call_args = self.get_base_completion_call_args() - - response = self.completion_function( - **base_completion_call_args, **args, stream=True - ) - - chunks = [] - for chunk in response: - print(f"chunk: {chunk}") - chunks.append(chunk) - - print(f"chunks: {chunks}") - built_response = stream_chunk_builder(chunks=chunks) - - non_stream_response = self.completion_function( - **base_completion_call_args, **args, stream=False - ) - - print( - "built_response.choices[0].message.content", - built_response.choices[0].message.content, - ) - print( - "non_stream_response.choices[0].message.content", - non_stream_response.choices[0].message.content, - ) - assert ( - json.loads(built_response.choices[0].message.content).keys() - == json.loads(non_stream_response.choices[0].message.content).keys() - ), f"Got={json.loads(built_response.choices[0].message.content)}, Expected={json.loads(non_stream_response.choices[0].message.content)}" - - def test_completion_thinking_with_response_format(self): - from pydantic import BaseModel - - litellm.turn_on_debug() - - class RFormat(BaseModel): - question: str - answer: str - - base_completion_call_args = self.get_base_completion_call_args_with_thinking() - - messages = [{"role": "user", "content": "Generate 5 question + answer pairs"}] - response = self.completion_function( - **base_completion_call_args, - messages=messages, - response_format=RFormat, - ) - - print(response) - - def test_completion_thinking_with_max_tokens(self): - from pydantic import BaseModel - - litellm.turn_on_debug() - - base_completion_call_args = self.get_base_completion_call_args_with_thinking() - - messages = [{"role": "user", "content": "Generate 5 question + answer pairs"}] - response = self.completion_function( - **base_completion_call_args, - messages=messages, - max_completion_tokens=20000, - ) - - print(response) - - def test_completion_thinking_without_max_tokens(self): - from pydantic import BaseModel - - litellm.turn_on_debug() - - base_completion_call_args = self.get_base_completion_call_args_with_thinking() - - messages = [{"role": "user", "content": "Generate 5 question + answer pairs"}] - response = self.completion_function( - **base_completion_call_args, - messages=messages, - ) - - print(response) - - def test_anthropic_thinking_output_stream(self): - # litellm.set_verbose = True - try: - base_completion_call_args = ( - self.get_base_completion_call_args_with_thinking() - ) - resp = litellm.completion( - **base_completion_call_args, - messages=[{"role": "user", "content": "Tell me a joke."}], - stream=True, - timeout=10, - ) - - reasoning_content_exists = False - signature_block_exists = False - tool_call_exists = False - for chunk in resp: - print(f"chunk 2: {chunk}") - if chunk.choices[0].delta.tool_calls: - tool_call_exists = True - if ( - hasattr(chunk.choices[0].delta, "thinking_blocks") - and chunk.choices[0].delta.thinking_blocks is not None - and chunk.choices[0].delta.reasoning_content is not None - and isinstance(chunk.choices[0].delta.thinking_blocks, list) - and len(chunk.choices[0].delta.thinking_blocks) > 0 - and isinstance(chunk.choices[0].delta.reasoning_content, str) - ): - reasoning_content_exists = True - print(chunk.choices[0].delta.thinking_blocks[0]) - if chunk.choices[0].delta.thinking_blocks[0].get("signature"): - signature_block_exists = True - assert not tool_call_exists - assert reasoning_content_exists - assert signature_block_exists - except litellm.Timeout: - pytest.skip("Model is timing out") - - def test_anthropic_reasoning_effort_thinking_translation(self): - base_completion_call_args = self.get_base_completion_call_args_with_thinking() - _, provider, _, _ = litellm.get_llm_provider( - model=base_completion_call_args["model"] - ) - - from litellm.constants import DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET - - optional_params = get_optional_params( - model=base_completion_call_args.get("model"), - custom_llm_provider=provider, - reasoning_effort="high", - ) - assert optional_params["thinking"] == { - "type": "enabled", - "budget_tokens": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, - } - - assert "reasoning_effort" not in optional_params - - class BaseReasoningLLMTests(ABC): """ Base class for testing reasoning llms diff --git a/tests/llm_translation/base_rerank_unit_tests.py b/tests/llm_translation/base_rerank_unit_tests.py index e3c66dcc0be..dbf7d0fd1ff 100644 --- a/tests/llm_translation/base_rerank_unit_tests.py +++ b/tests/llm_translation/base_rerank_unit_tests.py @@ -1,10 +1,8 @@ import asyncio import httpx import json -import pytest from typing import Any, Dict, List from unittest.mock import MagicMock, Mock, patch -import os import litellm from litellm.exceptions import BadRequestError @@ -18,7 +16,6 @@ from litellm.utils import ( # test_example.py from abc import ABC, abstractmethod - def assert_response_shape(response, custom_llm_provider): expected_response_shape = {"id": str, "results": list, "meta": dict} @@ -63,7 +60,6 @@ def assert_response_shape(response, custom_llm_provider): expected_billed_units_shape["search_units"], ) - class BaseLLMRerankTest(ABC): """ Abstract base test class that enforces a common test across all test classes. @@ -87,57 +83,3 @@ class BaseLLMRerankTest(ABC): """ return None - @pytest.mark.asyncio() - @pytest.mark.parametrize("sync_mode", [True, False]) - async def test_basic_rerank(self, sync_mode): - litellm.turn_on_debug() - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - rerank_call_args = self.get_base_rerank_call_args() - custom_llm_provider = self.get_custom_llm_provider() - if sync_mode is True: - response = litellm.rerank( - **rerank_call_args, - query="hello", - documents=["hello", "world"], - top_n=2, - ) - - print("re rank response: ", response) - - assert response.id is not None - assert response.results is not None - - assert response._hidden_params["response_cost"] is not None - - # Check expected cost - expected_cost = self.get_expected_cost() - if expected_cost is not None: - # If expected cost is specified, check exact match or >= for 0 - if expected_cost == 0.0: - assert response._hidden_params["response_cost"] >= 0 - else: - assert response._hidden_params["response_cost"] == expected_cost - else: - # Default behavior: cost should be greater than 0 - assert response._hidden_params["response_cost"] > 0 - - assert_response_shape( - response=response, custom_llm_provider=custom_llm_provider.value - ) - else: - response = await litellm.arerank( - **rerank_call_args, - query="hello", - documents=["hello", "world"], - top_n=2, - ) - - print("async re rank response: ", response) - - assert response.id is not None - assert response.results is not None - - assert_response_shape( - response=response, custom_llm_provider=custom_llm_provider.value - ) diff --git a/tests/llm_translation/interactions/base_interactions_test.py b/tests/llm_translation/interactions/base_interactions_test.py index 22ecce3a57b..ed09ecb2b7b 100644 --- a/tests/llm_translation/interactions/base_interactions_test.py +++ b/tests/llm_translation/interactions/base_interactions_test.py @@ -12,7 +12,6 @@ import pytest import litellm.interactions as interactions - class BaseInteractionsTest(ABC): """Abstract base class for interactions API tests. @@ -102,17 +101,3 @@ class BaseInteractionsTest(ABC): assert len(chunks) > 0 - @pytest.mark.asyncio - async def test_acreate_simple(self): - """Test async interaction creation.""" - api_key = self.get_api_key() - if not api_key: - pytest.skip(f"API key not set for {self.__class__.__name__}") - - response = await interactions.acreate( - model=self.get_model(), - input="What is the speed of light?", - api_key=api_key, - ) - assert response is not None - assert response.id is not None or response.status is not None diff --git a/tests/llm_translation/test_anthropic_completion.py b/tests/llm_translation/test_anthropic_completion.py index af7ba989b07..73872728a59 100644 --- a/tests/llm_translation/test_anthropic_completion.py +++ b/tests/llm_translation/test_anthropic_completion.py @@ -32,7 +32,6 @@ from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion from httpx import Headers from base_llm_unit_tests import BaseLLMChatTest, BaseAnthropicChatTest - def streaming_format_tests(chunk: dict, idx: int): """ 1st chunk - chunk.get("type") == "message_start" @@ -46,7 +45,6 @@ def streaming_format_tests(chunk: dict, idx: int): elif idx == 2: assert chunk.get("type") == "content_block_delta" - anthropic_chunk_list = [ { "type": "content_block_start", @@ -234,17 +232,6 @@ anthropic_chunk_list = [ {"type": "message_stop"}, ] - - - - - - - - - - - @pytest.mark.parametrize( "tool_type, tool_config, message_content", [ @@ -295,16 +282,8 @@ def test_anthropic_tool_use(tool_type, tool_config, message_content): except litellm.InternalServerError: pass - - - - - - - from litellm import completion - class TestAnthropicCompletion(BaseLLMChatTest, BaseAnthropicChatTest): def get_base_completion_call_args(self) -> dict: return {"model": "anthropic/claude-sonnet-4-5-20250929"} @@ -380,33 +359,10 @@ class TestAnthropicCompletion(BaseLLMChatTest, BaseAnthropicChatTest): @pytest.mark.asyncio async def test_pdf_handling(self, pdf_messages, sync_mode): await super().test_pdf_handling(pdf_messages, sync_mode) - test_content_list_handling = None - test_image_url = None - test_image_url_string = None test_web_search = None - - - - - - - - - - - - - - - from litellm.constants import RESPONSE_FORMAT_TOOL_NAME - - - - - def test_anthropic_citations_api(): """ Test the citations API @@ -454,7 +410,6 @@ def test_anthropic_citations_api(): assert "start_char_index" in citation assert "end_char_index" in citation - def test_anthropic_citations_api_streaming(): resp = completion( @@ -493,7 +448,6 @@ def test_anthropic_citations_api_streaming(): assert has_citations - @pytest.mark.parametrize( "model", [ @@ -538,7 +492,6 @@ def test_anthropic_thinking_output_stream(model): except litellm.Timeout: pytest.skip("Model is timing out") - def test_anthropic_custom_headers(): from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -576,9 +529,6 @@ def test_anthropic_custom_headers(): headers = mock_post.call_args[1]["headers"] assert "computer-use-2025-01-24" in headers["anthropic-beta"] - - - @pytest.mark.parametrize( "optional_params", [ @@ -617,7 +567,6 @@ def test_anthropic_websearch(optional_params: dict): assert response.usage.server_tool_use is not None assert response.usage.server_tool_use.web_search_requests >= 1 - def test_anthropic_text_editor(): litellm.turn_on_debug() params = { @@ -640,7 +589,6 @@ def test_anthropic_text_editor(): assert response is not None - @pytest.mark.parametrize("spec", ["anthropic", "openai"]) @pytest.mark.skipif( os.getenv("ZAPIER_CI_CD_MCP_TOKEN") is None, reason="ZAPIER_CI_CD_MCP_TOKEN not set" @@ -682,7 +630,6 @@ def test_anthropic_mcp_server_tool_use(spec: str): except litellm.InternalServerError as e: pytest.skip(f"Skipping test due to internal server error: {e}") - @pytest.mark.parametrize( "model", ["openai/gpt-4.1", "anthropic/claude-sonnet-4-5-20250929"] ) @@ -714,7 +661,6 @@ def test_anthropic_mcp_server_responses_api(model: str): assert response is not None - def test_anthropic_prefix_prompt(): params = { "model": "anthropic/claude-sonnet-4-5-20250929", @@ -729,7 +675,6 @@ def test_anthropic_prefix_prompt(): assert response is not None assert response.choices[0].message.content.startswith("Argentina") - @pytest.mark.asyncio async def test_claude_tool_use_with_anthropic_acreate(): response = await litellm.anthropic.messages.acreate( @@ -754,9 +699,6 @@ async def test_claude_tool_use_with_anthropic_acreate(): async for chunk in response: print(chunk) - - - def test_anthropic_streaming(): request_data = { @@ -811,7 +753,6 @@ def test_anthropic_streaming(): assert role_set_count == 1 - def test_anthropic_via_responses_api(): from litellm.types.llms.openai import ResponsesAPIStreamEvents @@ -934,11 +875,6 @@ def test_anthropic_via_responses_api(): print(f"✓ All {len(events_seen)} events matched expected structure") print(f"✓ Received {text_delta_count} text delta chunks") - - - - - def _make_transform_request(optional_params: dict, litellm_params: dict) -> dict: from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -950,19 +886,6 @@ def _make_transform_request(optional_params: dict, litellm_params: dict) -> dict headers={}, ) - - - - - - - - - - - - - def test_anthropic_basic_completion_replay(): response = litellm.completion( model="anthropic/claude-sonnet-4-5-20250929", @@ -976,7 +899,6 @@ def test_anthropic_basic_completion_replay(): assert response.usage.completion_tokens > 0 assert response.choices[0].finish_reason in {"stop", "length"} - def test_anthropic_streaming_completion_replay(): stream = litellm.completion( model="anthropic/claude-sonnet-4-5-20250929", diff --git a/tests/llm_translation/test_azure_o_series.py b/tests/llm_translation/test_azure_o_series.py index 4d5ac43ab6b..ef774d5f7a3 100644 --- a/tests/llm_translation/test_azure_o_series.py +++ b/tests/llm_translation/test_azure_o_series.py @@ -1,15 +1,11 @@ import os - import pytest import litellm from base_llm_unit_tests import BaseLLMChatTest, BaseOSeriesModelsTest - class TestAzureOpenAIO3Mini(BaseOSeriesModelsTest, BaseLLMChatTest): - test_content_list_handling = None - test_empty_tools = None test_function_calling_with_tool_response = None def get_base_completion_call_args(self): @@ -34,8 +30,6 @@ class TestAzureOpenAIO3Mini(BaseOSeriesModelsTest, BaseLLMChatTest): def test_basic_tool_calling(self): pass - - class TestAzureOpenAIO3(BaseOSeriesModelsTest): def get_base_completion_call_args(self): return { diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index 768be9b1080..3c8e7ff416d 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -35,7 +35,6 @@ litellm.success_callback = [] user_message = "Write a short poem about the sky" messages = [{"content": user_message, "role": "user"}] - @pytest.fixture(autouse=True) def reset_callbacks(): print("\npytest fixture - resetting callbacks") @@ -44,7 +43,6 @@ def reset_callbacks(): litellm.failure_callback = [] litellm.callbacks = [] - def test_completion_bedrock_claude_completion_auth(monkeypatch): print("calling bedrock claude completion params auth") @@ -72,16 +70,13 @@ def test_completion_bedrock_claude_completion_auth(monkeypatch): except Exception as e: pytest.fail(f"Error occurred: {e}") - # test_completion_bedrock_claude_completion_auth() - @pytest.mark.parametrize("streaming", [True, False]) def test_completion_bedrock_guardrails(streaming): litellm.set_verbose = True - # verbose_logger.setLevel(logging.DEBUG) try: if streaming is False: @@ -146,10 +141,8 @@ def test_completion_bedrock_guardrails(streaming): except Exception as e: pytest.fail(f"Error occurred: {e}") - # test_completion_bedrock_claude_2_1_completion_auth() - def test_completion_bedrock_claude_external_client_auth(monkeypatch): print("\ncalling bedrock claude external client auth") @@ -187,21 +180,10 @@ def test_completion_bedrock_claude_external_client_auth(monkeypatch): except Exception as e: pytest.fail(f"Error occurred: {e}") - # test_completion_bedrock_claude_external_client_auth() - - - - - - - - - # test_completion_bedrock_claude_sts_client_auth() - @pytest.mark.parametrize( "stop", [""], @@ -240,7 +222,6 @@ def test_bedrock_stop_value(stop, model): except Exception as e: pytest.fail(f"Error occurred: {e}") - @pytest.mark.parametrize( "system", ["You are an AI", [{"type": "text", "text": "You are an AI"}], ""], @@ -282,7 +263,6 @@ def test_bedrock_system_prompt(system, model): def test_completion_bedrock_mistral_completion_auth(): print("calling bedrock mistral completion params auth") - litellm.turn_on_debug() # aws_access_key_id = os.environ["AWS_ACCESS_KEY_ID"] @@ -311,10 +291,8 @@ def test_completion_bedrock_mistral_completion_auth(): except Exception as e: pytest.fail(f"Error occurred: {e}") - # test_completion_bedrock_mistral_completion_auth() - def test_bedrock_ptu(): """ Check if a url with 'modelId' passed in, is created correctly @@ -346,7 +324,6 @@ def test_bedrock_ptu(): ) mock_client_post.assert_called_once() - @pytest.mark.asyncio async def test_bedrock_custom_api_base(): """ @@ -384,7 +361,6 @@ async def test_bedrock_custom_api_base(): ) mock_client_post.assert_called_once() - @pytest.mark.parametrize( "model", [ @@ -421,7 +397,6 @@ async def test_bedrock_extra_headers(model): ) mock_client_post.assert_called_once() - @pytest.mark.asyncio async def test_bedrock_custom_prompt_template(): """ @@ -466,7 +441,6 @@ async def test_bedrock_custom_prompt_template(): assert prompt == "<|im_start|>user\nWhat's AWS?<|im_end|>" mock_client_post.assert_called_once() - def test_completion_bedrock_external_client_region(monkeypatch): print("\ncalling bedrock claude external client auth") @@ -515,32 +489,15 @@ def test_completion_bedrock_external_client_region(monkeypatch): except Exception as e: pytest.fail(f"Error occurred: {e}") - - - - - - - - - - - - - from litellm.litellm_core_utils.prompt_templates.factory import ( _bedrock_converse_messages_pt, ) - - - def test_base_aws_llm_get_credentials(): import time import boto3 - start_time = time.time() session = boto3.Session( aws_access_key_id="test", @@ -571,15 +528,6 @@ def test_base_aws_llm_get_credentials(): ) ) - - - - - - - - - def test_bedrock_converse_route(): litellm.set_verbose = True try: @@ -593,7 +541,6 @@ def test_bedrock_converse_route(): else: raise - def test_bedrock_mapped_converse_models(): litellm.set_verbose = True os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" @@ -604,23 +551,8 @@ def test_bedrock_mapped_converse_models(): messages=[{"role": "user", "content": "Hello, world!"}], ) - - - - - - - - - class TestBedrockConverseChatCrossRegion(BaseLLMChatTest): - test_content_list_handling = None - test_developer_role_translation = None test_function_calling_with_tool_response = None - test_image_url = None - test_json_response_format_stream = None - test_tool_call_with_empty_enum_property = None - test_tool_call_with_property_type_array = None def get_base_completion_call_args(self) -> dict: os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" @@ -649,10 +581,7 @@ class TestBedrockConverseChatCrossRegion(BaseLLMChatTest): assert cost > 0 - class TestBedrockConverseAnthropicUnitTests(BaseAnthropicChatTest): - test_completion_thinking_with_max_tokens = None - test_completion_thinking_without_max_tokens = None def get_base_completion_call_args(self) -> dict: return { @@ -665,12 +594,8 @@ class TestBedrockConverseAnthropicUnitTests(BaseAnthropicChatTest): "thinking": {"type": "enabled", "budget_tokens": 16000}, } - class TestBedrockConverseChatNormal(BaseLLMChatTest): - test_content_list_handling = None - test_empty_tools = None test_function_calling_with_tool_response = None - test_image_url = None def get_base_completion_call_args(self) -> dict: os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" @@ -681,12 +606,8 @@ class TestBedrockConverseChatNormal(BaseLLMChatTest): "aws_region_name": "us-east-1", } - - class TestBedrockConverseNovaTestSuite(BaseLLMChatTest): - test_content_list_handling = None test_function_calling_with_tool_response = None - test_image_url = None def get_base_completion_call_args(self) -> dict: os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" @@ -697,9 +618,6 @@ class TestBedrockConverseNovaTestSuite(BaseLLMChatTest): "aws_region_name": "us-east-1", } - - - class TestBedrockRerank(BaseLLMRerankTest): def get_custom_llm_provider(self) -> litellm.LlmProviders: return litellm.LlmProviders.BEDROCK @@ -709,7 +627,6 @@ class TestBedrockRerank(BaseLLMRerankTest): "model": "bedrock/arn:aws:bedrock:us-west-2::foundation-model/amazon.rerank-v1:0", } - class TestBedrockCohereRerank(BaseLLMRerankTest): def get_custom_llm_provider(self) -> litellm.LlmProviders: return litellm.LlmProviders.BEDROCK @@ -719,13 +636,6 @@ class TestBedrockCohereRerank(BaseLLMRerankTest): "model": "bedrock/arn:aws:bedrock:us-west-2::foundation-model/cohere.rerank-v3-5:0", } - - - - - - - @pytest.mark.parametrize("top_k_param", ["top_k", "topK"]) def test_bedrock_nova_topk(top_k_param): litellm.set_verbose = True @@ -751,7 +661,6 @@ def test_bedrock_nova_topk(top_k_param): assert "inferenceConfig" in captured_data["additionalModelRequestFields"] assert captured_data["additionalModelRequestFields"]["inferenceConfig"]["topK"] == 10 - def test_bedrock_cross_region_inference(monkeypatch): from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -777,7 +686,6 @@ def test_bedrock_cross_region_inference(monkeypatch): == "https://bedrock-runtime.us-west-2.amazonaws.com/model/us.meta.llama3-3-70b-instruct-v1%3A0/converse" ) - def test_bedrock_empty_content_real_call(): completion( model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", @@ -795,11 +703,6 @@ def test_bedrock_empty_content_real_call(): ], ) - - - - - class TestBedrockEmbedding(BaseLLMEmbeddingTest): def get_base_embedding_call_args(self) -> dict: return { @@ -809,8 +712,6 @@ class TestBedrockEmbedding(BaseLLMEmbeddingTest): def get_custom_llm_provider(self) -> litellm.LlmProviders: return litellm.LlmProviders.BEDROCK - - @pytest.mark.asyncio async def test_bedrock_image_url_sync_client(): import logging @@ -849,9 +750,6 @@ async def test_bedrock_image_url_sync_client(): print(e) mock_post.assert_called_once() - - - def test_bedrock_custom_proxy(): from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -874,7 +772,6 @@ def test_bedrock_custom_proxy(): assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer Token" - def test_bedrock_custom_deepseek(): import json @@ -924,13 +821,6 @@ def test_bedrock_custom_deepseek(): print(f"Error: {str(e)}") raise e - - - - - - - def test_bedrock_description_param(): from litellm import completion from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -969,7 +859,6 @@ def test_bedrock_description_param(): "Find the meaning inside a poem" in request_body_str ) # assert description is passed - @pytest.mark.parametrize( "sync_mode", [ @@ -1029,7 +918,6 @@ async def test_bedrock_thinking_in_assistant_message(sync_mode): in json_data ) - @pytest.mark.asyncio async def test_bedrock_stream_thinking_content_openwebui(): """ @@ -1098,7 +986,6 @@ async def test_bedrock_stream_thinking_content_openwebui(): len(response_content) > 0 ), "There should be non-empty content after thinking tags" - def test_bedrock_application_inference_profile(): from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -1163,7 +1050,6 @@ def test_bedrock_application_inference_profile(): ) assert mock_post2.call_args.kwargs["url"] == mock_post.call_args.kwargs["url"] - def return_mocked_response(model: str): if model == "bedrock/mistral.mistral-large-2407-v1:0": return { @@ -1178,7 +1064,6 @@ def return_mocked_response(model: str): "usage": {"inputTokens": 5, "outputTokens": 10, "totalTokens": 15}, } - @pytest.mark.parametrize( "model", [ @@ -1222,9 +1107,6 @@ async def test_bedrock_max_completion_tokens(model: str): "inferenceConfig": {"maxTokens": 10}, } - - - @pytest.mark.asyncio async def test_bedrock_passthrough_router(): """ @@ -1278,7 +1160,6 @@ async def test_bedrock_passthrough_router(): assert response.status_code == 200 - @pytest.mark.asyncio async def test_bedrock_converse__streaming_passthrough(monkeypatch): import asyncio @@ -1332,7 +1213,6 @@ async def test_bedrock_converse__streaming_passthrough(monkeypatch): assert response_cost is not None and response_cost > 0 assert "standard_logging_object" in mock_callback.call_args.kwargs["kwargs"] - @pytest.mark.asyncio async def test_bedrock_streaming_passthrough_test2(monkeypatch): import asyncio @@ -1383,7 +1263,6 @@ async def test_bedrock_streaming_passthrough_test2(monkeypatch): assert "standard_logging_object" in mock_callback.call_args.kwargs["kwargs"] assert "response_cost" in mock_callback.call_args.kwargs["kwargs"] - def test_bedrock_openai_imported_model(): """ Test that Bedrock imported models using OpenAI format work correctly. @@ -1485,25 +1364,6 @@ def test_bedrock_openai_imported_model(): assert request_body["max_tokens"] == 300 assert request_body["temperature"] == 0.5 - - - - - - - - - - - - - - - - - - - def test_bedrock_openai_multiple_message_types(): """ Test that various message content types are handled correctly. @@ -1552,14 +1412,10 @@ def test_bedrock_openai_multiple_message_types(): print("✓ Multiple message types handled correctly") - - - # ============================================================================ # Nova Grounding (web_search_options) Unit Tests (Mocked) # ============================================================================ - def test_bedrock_nova_grounding_web_search_options_non_streaming(): """ Unit test for Nova grounding using web_search_options parameter (non-streaming). @@ -1618,7 +1474,6 @@ def test_bedrock_nova_grounding_web_search_options_non_streaming(): f"✓ web_search_options correctly transformed to systemTool (non-streaming)" ) - def test_bedrock_nova_grounding_with_function_tools(): """ Unit test for Nova grounding combined with regular function tools. @@ -1703,7 +1558,6 @@ def test_bedrock_nova_grounding_with_function_tools(): assert system_tool_found, "systemTool (nova_grounding) should be present" print(f"✓ Both function tools and web_search_options correctly combined") - @pytest.mark.asyncio async def test_bedrock_nova_grounding_async(): """ @@ -1757,9 +1611,6 @@ async def test_bedrock_nova_grounding_async(): assert system_tool_found, "systemTool with nova_grounding should be present" print(f"✓ Async web_search_options correctly transformed to systemTool") - - - def test_bedrock_nova_grounding_request_transformation(): """ Unit test to verify that web_search_options transforms to systemTool in the request. diff --git a/tests/llm_translation/test_bedrock_gpt_oss.py b/tests/llm_translation/test_bedrock_gpt_oss.py index e0d7b1b904f..b5252847805 100644 --- a/tests/llm_translation/test_bedrock_gpt_oss.py +++ b/tests/llm_translation/test_bedrock_gpt_oss.py @@ -1,8 +1,6 @@ from base_llm_unit_tests import BaseLLMChatTest - class TestBedrockGPTOSS(BaseLLMChatTest): - test_json_response_format = None def get_base_completion_call_args(self) -> dict: return { @@ -19,8 +17,3 @@ class TestBedrockGPTOSS(BaseLLMChatTest): """ pass - async def test_completion_cost(self): - """ - Bedrock GPT-OSS models are flaky and occasionally report 0 token counts in api response - """ - pass diff --git a/tests/llm_translation/test_bedrock_invoke_tests.py b/tests/llm_translation/test_bedrock_invoke_tests.py index 584b0ef341f..8f61a359f6a 100644 --- a/tests/llm_translation/test_bedrock_invoke_tests.py +++ b/tests/llm_translation/test_bedrock_invoke_tests.py @@ -5,7 +5,6 @@ import os import litellm from litellm.types.llms.bedrock import BedrockInvokeNovaRequest - _LITELLM_LOGO_IMAGE_URL = ( "https://cdn.jsdelivr.net/gh/BerriAI/litellm@d769e81c90d453240c61fc572cdb27fae06a89d0/" "ui/litellm-dashboard/public/assets/logos/litellm_logo.jpg" @@ -15,7 +14,6 @@ _AWSMP_LOGO_IMAGE_URL = ( "c233c9ade2ccb5491072ae232c814942.png" ) - @pytest.mark.flaky(retries=3, delay=5) class TestBedrockInvokeClaudeJson(BaseLLMChatTest): def get_base_completion_call_args(self) -> dict: @@ -24,26 +22,9 @@ class TestBedrockInvokeClaudeJson(BaseLLMChatTest): "model": "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", } - @pytest.mark.parametrize( - "image_url, detail", - [ - (_LITELLM_LOGO_IMAGE_URL, None), - (_LITELLM_LOGO_IMAGE_URL, "low"), - (_LITELLM_LOGO_IMAGE_URL, "high"), - (_AWSMP_LOGO_IMAGE_URL, "low"), - (_AWSMP_LOGO_IMAGE_URL, "high"), - ], - ) - @pytest.mark.flaky(retries=4, delay=2) - def test_image_url(self, image_url, detail): - super().test_image_url(detail=detail, image_url=image_url) - test_content_list_handling = None - test_image_url_string = None test_pdf_handling = None - class TestBedrockInvokeNovaJson(BaseLLMChatTest): - test_json_response_format = None def get_base_completion_call_args(self) -> dict: return { @@ -57,11 +38,3 @@ class TestBedrockInvokeNovaJson(BaseLLMChatTest): f"Skipping non-JSON test: {request.function.__name__} does not contain 'json'" ) - def test_json_response_pydantic_obj(self): - if os.environ.get("LITELLM_RUN_LIVE_BEDROCK_NOVA_JSON_TESTS") != "1": - pytest.skip("Live Bedrock Nova response-schema E2E tests are opt-in") - if os.environ.get("CASSETTE_REDIS_URL"): - pytest.skip( - "Live Bedrock Nova response-schema E2E tests cannot run under VCR replay" - ) - super().test_json_response_pydantic_obj() diff --git a/tests/llm_translation/test_bedrock_llama.py b/tests/llm_translation/test_bedrock_llama.py index 9c83c28bbe0..3a3f0064c52 100644 --- a/tests/llm_translation/test_bedrock_llama.py +++ b/tests/llm_translation/test_bedrock_llama.py @@ -3,10 +3,7 @@ import pytest import litellm - class TestBedrockTestSuite(BaseLLMChatTest): - test_content_list_handling = None - test_empty_tools = None test_function_calling_with_tool_response = None def get_base_completion_call_args(self) -> dict: diff --git a/tests/llm_translation/test_bedrock_moonshot.py b/tests/llm_translation/test_bedrock_moonshot.py index cf5e554e91b..b2026b9ed1f 100644 --- a/tests/llm_translation/test_bedrock_moonshot.py +++ b/tests/llm_translation/test_bedrock_moonshot.py @@ -15,20 +15,12 @@ from base_llm_unit_tests import BaseLLMChatTest import litellm - class TestBedrockMoonshotInvoke(BaseLLMChatTest): """ Test suite for Bedrock Moonshot via invoke route. Inherits all standard LLM tests from BaseLLMChatTest. """ - test_json_response_format_stream = None - test_completion_cost = None - test_content_list_handling = None - test_developer_role_translation = None - test_message_with_name = None - test_pydantic_model_input = None - test_response_format_type_text_with_tool_calls_no_tool_choice = None test_streaming = None def get_base_completion_call_args(self) -> dict: diff --git a/tests/llm_translation/test_bedrock_nova_json.py b/tests/llm_translation/test_bedrock_nova_json.py index 8adfef50618..7c1023e0671 100644 --- a/tests/llm_translation/test_bedrock_nova_json.py +++ b/tests/llm_translation/test_bedrock_nova_json.py @@ -3,15 +3,8 @@ import pytest import litellm - class TestBedrockNovaJson(BaseLLMChatTest): - test_content_list_handling = None - test_developer_role_translation = None - test_empty_tools = None test_function_calling_with_tool_response = None - test_json_response_format_stream = None - test_tool_call_with_empty_enum_property = None - test_tool_call_with_property_type_array = None def get_base_completion_call_args(self) -> dict: litellm.turn_on_debug() @@ -19,14 +12,6 @@ class TestBedrockNovaJson(BaseLLMChatTest): "model": "bedrock/converse/us.amazon.nova-micro-v1:0", } - def test_json_response_nested_pydantic_obj(self): - pass - - def test_json_response_nested_json_schema(self): - pass - - - # @pytest.fixture(autouse=True) # def skip_non_json_tests(self, request): # if not "json" in request.function.__name__.lower(): diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index d05afe92626..acbad4fb700 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -2,7 +2,6 @@ import os import pytest - from base_llm_unit_tests import BaseLLMChatTest from litellm.llms.vertex_ai.context_caching.transformation import ( separate_cached_messages, @@ -12,18 +11,9 @@ import litellm from litellm import completion import json - - - class TestGoogleAIStudioGemini(BaseLLMChatTest): test_async_pdf_handling_with_file_id = None - test_content_list_handling = None - test_developer_role_translation = None test_function_calling_with_tool_response = None - test_image_url = None - test_json_response_nested_json_schema = None - test_json_response_nested_pydantic_obj = None - test_json_response_pydantic_obj = None test_web_search = None def get_base_completion_call_args(self) -> dict: @@ -32,7 +22,6 @@ class TestGoogleAIStudioGemini(BaseLLMChatTest): def get_base_completion_call_args_with_reasoning_model(self) -> dict: return {"model": "gemini/gemini-2.5-flash"} - @pytest.mark.flaky(retries=3, delay=2) def test_url_context(self): from litellm.utils import supports_url_context @@ -64,11 +53,6 @@ class TestGoogleAIStudioGemini(BaseLLMChatTest): ), "URL context metadata should be present" print(f"response={response}") - - - - - def test_gemini_image_generation(): # litellm.turn_on_debug() response = completion( @@ -90,23 +74,6 @@ def test_gemini_image_generation(): .startswith("data:image/png;base64,") ) - - - - - - - - - - - - - - - - - def test_gemini_thinking(): litellm.turn_on_debug() from litellm.types.utils import Message, CallTypes @@ -146,9 +113,6 @@ def test_gemini_thinking(): print(response.choices[0].message) assert response.choices[0].message.content is not None - - - def test_gemini_finish_reason(): import os from litellm import completion @@ -163,7 +127,6 @@ def test_gemini_finish_reason(): assert response.choices[0].finish_reason is not None assert response.choices[0].finish_reason == "length" - @pytest.mark.flaky(retries=3, delay=2) def test_gemini_url_context(): from litellm import completion @@ -189,7 +152,6 @@ def test_gemini_url_context(): assert urlMetadata["retrievedUrl"] == URL1 assert urlMetadata["urlRetrievalStatus"] == "URL_RETRIEVAL_STATUS_SUCCESS" - @pytest.mark.flaky(retries=3, delay=2) def test_gemini_with_grounding(): from litellm import completion, Usage, stream_chunk_builder @@ -226,7 +188,6 @@ def test_gemini_with_grounding(): assert usage.prompt_tokens_details.web_search_requests is not None assert usage.prompt_tokens_details.web_search_requests > 0 - def test_gemini_with_empty_function_call_arguments(): from litellm import completion @@ -248,9 +209,6 @@ def test_gemini_with_empty_function_call_arguments(): print(response) assert response.choices[0].message.content is not None - - - def test_gemini_tool_use(): data = { "max_tokens": 8192, @@ -299,7 +257,6 @@ def test_gemini_tool_use(): assert stop_reason is not None assert stop_reason == "tool_calls" - @pytest.mark.asyncio async def test_gemini_image_generation_async(): litellm.turn_on_debug() @@ -332,7 +289,6 @@ async def test_gemini_image_generation_async(): assert IMAGE_URL["url"] is not None, "IMAGE_URL['url'] is not None" assert IMAGE_URL["url"].startswith("data:image/png;base64,") - @pytest.mark.asyncio async def test_gemini_image_generation_async_stream(): # litellm.turn_on_debug() @@ -367,7 +323,6 @@ async def test_gemini_image_generation_async_stream(): assert model_response_image is not None assert model_response_image["url"].startswith("data:image/png;base64,") - def test_system_message_with_no_user_message(): """ Test that the system message is translated correctly for non-OpenAI providers. @@ -387,7 +342,6 @@ def test_system_message_with_no_user_message(): assert response.choices[0].message.content is not None - def get_current_weather(location, unit="fahrenheit"): """Get the current weather in a given location""" if "tokyo" in location.lower(): @@ -401,7 +355,6 @@ def get_current_weather(location, unit="fahrenheit"): else: return json.dumps({"location": location, "temperature": "unknown"}) - def test_gemini_with_thinking(): from litellm import completion @@ -493,11 +446,6 @@ def test_gemini_with_thinking(): ) # get a new response from the model where it can see the function response print("second response\n", second_response) - - - - - @pytest.mark.parametrize( "status_code,expected_exception", [ @@ -582,7 +530,6 @@ def l(status_code, expected_exception): "VertexAIException" not in error_message ), f"Should not contain 'VertexAIException' for status {status_code}, got: {error_message}" - def test_gemini_embedding(): litellm.turn_on_debug() response = litellm.embedding( @@ -592,23 +539,6 @@ def test_gemini_embedding(): print("response: ", response) assert response is not None - - - - - - - - - - - - - - - - - @pytest.mark.asyncio async def test_gemini_openai_web_search_tool_to_google_search(): """ diff --git a/tests/llm_translation/test_groq.py b/tests/llm_translation/test_groq.py index 55b1b9bc7eb..ea633b5cc84 100644 --- a/tests/llm_translation/test_groq.py +++ b/tests/llm_translation/test_groq.py @@ -1,6 +1,5 @@ - # sys.path.insert( # 0, os.path.abspath("../..") # ) # noqa @@ -8,10 +7,7 @@ from base_llm_unit_tests import BaseLLMChatTest - class TestGroq(BaseLLMChatTest): - test_content_list_handling = None - test_empty_tools = None test_web_search = None def get_base_completion_call_args(self) -> dict: @@ -19,5 +15,3 @@ class TestGroq(BaseLLMChatTest): "model": "groq/openai/gpt-oss-120b", } - def test_tool_call_with_empty_enum_property(self): - pass diff --git a/tests/llm_translation/test_huggingface_chat_completion.py b/tests/llm_translation/test_huggingface_chat_completion.py index 57c2d4df71c..6e076e6de63 100644 --- a/tests/llm_translation/test_huggingface_chat_completion.py +++ b/tests/llm_translation/test_huggingface_chat_completion.py @@ -7,7 +7,6 @@ from unittest.mock import AsyncMock, MagicMock, patch from base_llm_unit_tests import BaseLLMChatTest - import pytest import litellm @@ -126,7 +125,6 @@ MOCK_STREAMING_CHUNKS = [ }, ] - PROVIDER_MAPPING_RESPONSE = { "fireworks-ai": { "status": "live", @@ -145,14 +143,12 @@ PROVIDER_MAPPING_RESPONSE = { }, } - @pytest.fixture def mock_provider_mapping(): with patch("litellm.llms.huggingface.chat.transformation.fetch_inference_provider_mapping") as mock: mock.return_value = PROVIDER_MAPPING_RESPONSE yield mock - @pytest.fixture(autouse=True) def clear_lru_cache(): from litellm.llms.huggingface.common_utils import fetch_inference_provider_mapping @@ -161,7 +157,6 @@ def clear_lru_cache(): yield fetch_inference_provider_mapping.cache_clear() - @pytest.fixture def mock_http_handler(): """Fixture to mock the HTTP handler""" @@ -188,7 +183,6 @@ def mock_http_handler(): mock.side_effect = mock_side_effect yield mock - @pytest.fixture def mock_http_async_handler(): """Fixture to mock the async HTTP handler""" @@ -220,7 +214,6 @@ def mock_http_async_handler(): mock.side_effect = mock_side_effect yield mock - class TestHuggingFace(BaseLLMChatTest): @pytest.fixture(autouse=True) def setup(self, mock_provider_mapping, mock_http_handler, mock_http_async_handler): @@ -355,8 +348,6 @@ class TestHuggingFace(BaseLLMChatTest): == tool_call_no_arguments["tool_calls"][0]["function"]["arguments"] ) - - def test_completion_with_api_base(self): messages = [{"role": "user", "content": "This is a test message"}] api_base = "https://abcd123.us-east-1.aws.endpoints.huggingface.cloud" @@ -421,9 +412,3 @@ class TestHuggingFace(BaseLLMChatTest): called_url = call_args[1]["url"] assert called_url == f"{api_base}/v1/chat/completions" - - - - @pytest.mark.asyncio - async def test_completion_cost(self): - pass diff --git a/tests/llm_translation/test_mistral_audio_transcription_transformation.py b/tests/llm_translation/test_mistral_audio_transcription_transformation.py index db77eabba23..6328e52cbed 100644 --- a/tests/llm_translation/test_mistral_audio_transcription_transformation.py +++ b/tests/llm_translation/test_mistral_audio_transcription_transformation.py @@ -8,7 +8,6 @@ from tests.llm_translation.base_audio_transcription_unit_tests import ( BaseLLMAudioTranscriptionTest, ) - @pytest.mark.skipif( not os.getenv("MISTRAL_API_KEY"), reason="MISTRAL_API_KEY not set, skipping Mistral audio transcription tests", @@ -22,8 +21,3 @@ class TestMistralAudioTranscription(BaseLLMAudioTranscriptionTest): def get_custom_llm_provider(self) -> litellm.LlmProviders: return litellm.LlmProviders.MISTRAL - def test_audio_transcription_async(self): # type: ignore[override] - pytest.skip( - "Async audio transcription test for Mistral is skipped in this suite; " - "async test plugins (e.g. pytest-asyncio/anyio) are not configured here." - ) diff --git a/tests/llm_translation/test_nvidia_nim.py b/tests/llm_translation/test_nvidia_nim.py index 2de1d0f3b84..402035444b1 100644 --- a/tests/llm_translation/test_nvidia_nim.py +++ b/tests/llm_translation/test_nvidia_nim.py @@ -3,7 +3,6 @@ from datetime import datetime from typing import Final from unittest.mock import AsyncMock - import httpx import pytest from openai.types import CreateEmbeddingResponse, Embedding @@ -16,7 +15,6 @@ from litellm import completion from base_rerank_unit_tests import BaseLLMRerankTest from tests.capturing_transport import CapturingTransport - def test_completion_nvidia_nim(): from openai import OpenAI @@ -59,7 +57,6 @@ def test_completion_nvidia_nim(): assert request_body["frequency_penalty"] == 0.1 assert request_body["presence_penalty"] == 0.5 - class TestNvidiaNim(BaseLLMRerankTest): def get_custom_llm_provider(self) -> litellm.LlmProviders: return litellm.LlmProviders.NVIDIA_NIM @@ -73,43 +70,3 @@ class TestNvidiaNim(BaseLLMRerankTest): """Nvidia NIM rerank models are free (cost = 0.0)""" return 0.0 - @pytest.mark.asyncio() - @pytest.mark.parametrize("sync_mode", [True, False]) - async def test_basic_rerank(self, sync_mode, monkeypatch): - """ - Override the base live rerank test with a mocked HTTP layer. - - NVIDIA reached end-of-life for the hosted - nvidia/llama-3.2-nv-rerankqa-1b-v2 rerank API on 2026-05-18 and - published no replacement model, so a live call now returns HTTP 410 - ("Gone"). NVIDIA's hosted catalog rotates on a schedule, so pointing - at another live model would only defer the same failure. Mock the - transport instead (same pattern as - test_nvidia_nim_rerank_ranking_endpoint above) so the request/response - transformation and cost calculation stay covered offline. - """ - monkeypatch.setenv("NVIDIA_NIM_API_KEY", "fake-api-key") - - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {} - mock_response.text = "" - mock_response.json.return_value = { - "rankings": [ - {"index": 0, "logit": 0.95}, - {"index": 1, "logit": 0.75}, - ], - "usage": {"total_tokens": 7}, - } - - with ( - patch( - "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", - return_value=mock_response, - ), - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - return_value=mock_response, - ), - ): - await super().test_basic_rerank(sync_mode=sync_mode) diff --git a/tests/llm_translation/test_openai_o1.py b/tests/llm_translation/test_openai_o1.py index 61b17dfecf1..6b597b016fc 100644 --- a/tests/llm_translation/test_openai_o1.py +++ b/tests/llm_translation/test_openai_o1.py @@ -1,18 +1,13 @@ import os from unittest.mock import patch - import pytest import litellm from litellm import ModelResponse from base_llm_unit_tests import BaseLLMChatTest, BaseOSeriesModelsTest - class TestOpenAIO1(BaseOSeriesModelsTest, BaseLLMChatTest): - test_empty_tools = None - test_tool_call_with_empty_enum_property = None - test_tool_call_with_property_type_array = None def get_base_completion_call_args(self): return { @@ -24,9 +19,6 @@ class TestOpenAIO1(BaseOSeriesModelsTest, BaseLLMChatTest): return OpenAI(api_key="fake-api-key") - - - class TestOpenAIO3(BaseOSeriesModelsTest, BaseLLMChatTest): test_basic_tool_calling = None test_function_calling_with_tool_response = None @@ -41,9 +33,6 @@ class TestOpenAIO3(BaseOSeriesModelsTest, BaseLLMChatTest): return OpenAI(api_key="fake-api-key") - - - def test_o3_reasoning_effort(): resp = litellm.completion( model="o3-mini", diff --git a/tests/llm_translation/test_ovhcloud_audio_transcription_transformation.py b/tests/llm_translation/test_ovhcloud_audio_transcription_transformation.py index 8cc46dc98d0..a1e03ae5e6f 100644 --- a/tests/llm_translation/test_ovhcloud_audio_transcription_transformation.py +++ b/tests/llm_translation/test_ovhcloud_audio_transcription_transformation.py @@ -12,7 +12,6 @@ from tests.llm_translation.base_audio_transcription_unit_tests import ( BaseLLMAudioTranscriptionTest, ) - @pytest.mark.skipif( not os.getenv("OVHCLOUD_API_KEY"), reason="OVHCLOUD_API_KEY not set, skipping OVHCloud audio transcription tests", @@ -29,12 +28,6 @@ class TestOVHCloudAudioTranscription(BaseLLMAudioTranscriptionTest): # Override the async base test with a sync no-op to avoid # 'async def functions are not natively supported' failures when # running this file in isolation without pytest-asyncio. - def test_audio_transcription_async(self): # type: ignore[override] - pytest.skip( - "Async audio transcription test for OVHCloud is skipped in this suite; " - "async test plugins (e.g. pytest-asyncio/anyio) are not configured here." - ) - @pytest.mark.skipif( not os.getenv("OVHCLOUD_API_KEY"), diff --git a/tests/llm_translation/test_together_ai.py b/tests/llm_translation/test_together_ai.py index a203c9edcfe..b78fbeca9e8 100644 --- a/tests/llm_translation/test_together_ai.py +++ b/tests/llm_translation/test_together_ai.py @@ -8,21 +8,12 @@ import json from datetime import datetime from unittest.mock import AsyncMock - import litellm import pytest - class TestTogetherAI(BaseLLMChatTest): test_basic_tool_calling = None - test_empty_tools = None test_function_calling_with_tool_response = None - test_json_response_format = None - test_json_response_nested_json_schema = None - test_json_response_nested_pydantic_obj = None - test_json_response_pydantic_obj = None - test_tool_call_with_empty_enum_property = None - test_tool_call_with_property_type_array = None def get_base_completion_call_args(self) -> dict: litellm.set_verbose = True diff --git a/tests/logging_callback_tests/test_alerting.py b/tests/logging_callback_tests/test_alerting.py deleted file mode 100644 index 691bb58b998..00000000000 --- a/tests/logging_callback_tests/test_alerting.py +++ /dev/null @@ -1,338 +0,0 @@ -# What is this? -## Tests slack alerting on proxy logging object - -import asyncio - -# import logging -# logging.basicConfig(level=logging.DEBUG) -from datetime import datetime -from unittest.mock import AsyncMock, patch - -import pytest - -import litellm -from litellm.caching.caching import DualCache -from litellm.integrations.SlackAlerting.slack_alerting import ( - SlackAlerting, -) -from litellm.proxy._types import CallInfo, Litellm_EntityType -from litellm.proxy.utils import ProxyLogging -from litellm.types.integrations.slack_alerting import AlertType - - -@pytest.mark.asyncio -async def test_get_api_base(): - _pl = ProxyLogging(user_api_key_cache=DualCache()) - _pl.update_values(alerting=["slack"], alerting_threshold=100, redis_cache=None) - model = "chatgpt-v-3" - messages = [{"role": "user", "content": "Hey how's it going?"}] - litellm_params = { - "acompletion": True, - "api_key": None, - "api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/", - "force_timeout": 600, - "logger_fn": None, - "verbose": False, - "custom_llm_provider": "azure", - "litellm_call_id": "68f46d2d-714d-4ad8-8137-69600ec8755c", - "model_alias_map": {}, - "completion_call_id": None, - "metadata": None, - "model_info": None, - "proxy_server_request": None, - "preset_cache_key": None, - "no-log": False, - "stream_response": {}, - } - start_time = datetime.now() - end_time = datetime.now() - - time_difference_float, model, api_base, messages = ( - _pl.slack_alerting_instance._response_taking_too_long_callback_helper( - kwargs={ - "model": model, - "messages": messages, - "litellm_params": litellm_params, - }, - start_time=start_time, - end_time=end_time, - ) - ) - - assert api_base is not None - assert isinstance(api_base, str) - assert len(api_base) > 0 - request_info = ( - f"\nRequest Model: `{model}`\nAPI Base: `{api_base}`\nMessages: `{messages}`" - ) - slow_message = f"`Responses are slow - {round(time_difference_float,2)}s response time > Alerting threshold: {100}s`" - await _pl.alerting_handler( - message=slow_message + request_info, - level="Low", - alert_type=AlertType.llm_too_slow, - ) - print("passed test_get_api_base") - - -# Create a mock environment for testing -@pytest.fixture -def mock_env(monkeypatch): - monkeypatch.setenv("SLACK_WEBHOOK_URL", "https://example.com/webhook") - monkeypatch.setenv("LANGFUSE_HOST", "https://cloud.langfuse.com") - monkeypatch.setenv("LANGFUSE_PROJECT_ID", "test-project-id") - - -# Test the __init__ method - - -@pytest.fixture -def slack_alerting(): - return SlackAlerting( - alerting_threshold=1, internal_usage_cache=DualCache(), alerting=["slack"] - ) - - -# Test for slow LLM responses - - - - -# Test for budget crossed - - -# Test for budget crossed again (should not fire alert 2nd time) - - -# Test for send_alert - should be called once -@pytest.mark.asyncio -async def test_send_alert(slack_alerting): - import logging - - from litellm._logging import verbose_logger - - asyncio.create_task(slack_alerting.periodic_flush()) - verbose_logger.setLevel(level=logging.DEBUG) - with patch.object( - slack_alerting.async_http_handler, "post", new=AsyncMock() - ) as mock_post: - mock_post.return_value.status_code = 200 - await slack_alerting.send_alert( - "Test message", "Low", "budget_alerts", alerting_metadata={} - ) - - await asyncio.sleep(6) - mock_post.assert_awaited_once() - - - - -@pytest.mark.asyncio -async def test_daily_reports_completion(slack_alerting): - with patch.object(slack_alerting, "send_alert", new=AsyncMock()) as mock_send_alert: - litellm.callbacks = [slack_alerting] - - # on async success - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-5.5", - "litellm_params": { - "model": "gpt-5-mini", - }, - } - ] - ) - - await router.acompletion( - model="gpt-5-mini", - messages=[{"role": "user", "content": "Hey, how's it going?"}], - ) - - await asyncio.sleep(3) - response_val = await slack_alerting.send_daily_reports(router=router) - - assert response_val is True - - mock_send_alert.assert_awaited_once() - - # on async failure - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-5.5", - "litellm_params": {"model": "gpt-5-mini", "api_key": "bad_key"}, - } - ] - ) - - try: - await router.acompletion( - model="gpt-5-mini", - messages=[{"role": "user", "content": "Hey, how's it going?"}], - ) - except Exception as e: - pass - - await asyncio.sleep(3) - response_val = await slack_alerting.send_daily_reports(router=router) - - assert response_val is True - - mock_send_alert.assert_awaited() - - - - -# test models with 0 metrics are ignored - - -# test no alert is sent if all None or 0 metrics - - -# test user budget crossed alert sent only once, even if user makes multiple calls - - - - -# @pytest.mark.asyncio -# async def test_webhook_customer_spend_event(): -# """ -# Test if customer spend is working as expected -# """ -# slack_alerting = SlackAlerting(alerting=["webhook"]) - -# with patch.object( -# slack_alerting, "send_webhook_alert", new=AsyncMock() -# ) as mock_send_alert: -# user_info = { -# "token": "sk-test-mock-token-606", -# "spend": 1, -# "max_budget": 0, -# "user_id": "ishaan@berri.ai", -# "user_email": "ishaan@berri.ai", -# "key_alias": "my-test-key", -# "projected_exceeded_date": "10/20/2024", -# "projected_spend": 200, -# } - -# user_info = CallInfo(**user_info) -# for _ in range(50): -# await slack_alerting.budget_alerts( -# type=alerting_type, -# user_info=user_info, -# ) -# mock_send_alert.assert_awaited_once() - - - - - - -@pytest.mark.asyncio -async def test_langfuse_trace_id(): - """ - - Unit test for `_add_langfuse_trace_id_to_alert` function in slack_alerting.py - """ - from litellm.integrations.SlackAlerting.utils import add_langfuse_trace_id_to_alert - from litellm.litellm_core_utils.litellm_logging import Logging - - litellm.success_callback = ["langfuse"] - - litellm_logging_obj = Logging( - model="gpt-5-mini", - messages=[{"role": "user", "content": "hi"}], - stream=False, - call_type="acompletion", - litellm_call_id="1234", - start_time=datetime.now(), - function_id="1234", - ) - - litellm.completion( - model="gpt-5-mini", - messages=[{"role": "user", "content": "Hey how's it going?"}], - mock_response="Hey!", - litellm_logging_obj=litellm_logging_obj, - ) - - await asyncio.sleep(3) - - assert litellm_logging_obj.get_trace_id(service_name="langfuse") is not None - - slack_alerting = SlackAlerting( - alerting_threshold=32, - alerting=["slack"], - alert_types=[AlertType.llm_exceptions], - internal_usage_cache=DualCache(), - ) - - trace_url = await add_langfuse_trace_id_to_alert( - request_data={"litellm_logging_obj": litellm_logging_obj} - ) - - assert trace_url is not None - - returned_trace_id = trace_url.split("/")[-1] - - assert returned_trace_id == litellm_logging_obj.get_trace_id( - service_name="langfuse" - ) - - - - -@pytest.mark.parametrize("report_type", ["weekly", "monthly"]) -@pytest.mark.asyncio -async def test_spend_report_cache(report_type): - """ - Test that spend reports are only sent once within their period - """ - # Mock prisma client response - mock_spend_data = [ - {"team_alias": "team1", "total_spend": 100.0}, - {"team_alias": "team2", "total_spend": 200.0}, - ] - - mock_tag_data = [ - {"individual_request_tag": "tag1", "total_spend": 150.0}, - {"individual_request_tag": "tag2", "total_spend": 150.0}, - ] - - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: - # Setup mock for database query - mock_prisma.db.query_raw = AsyncMock( - side_effect=[mock_spend_data, mock_tag_data] - ) - - slack_alerting = SlackAlerting( - alerting=["webhook"], internal_usage_cache=DualCache() - ) - - user_info = CallInfo( - token="test_token", - spend=100, - max_budget=1000, - user_id="test@test.com", - user_email="test@test.com", - key_alias="test-key", - event_group=Litellm_EntityType.KEY, - ) - - with patch.object( - slack_alerting, "send_alert", new=AsyncMock() - ) as mock_send_alert: - # First call should send alert - if report_type == "weekly": - await slack_alerting.send_weekly_spend_report() - else: - await slack_alerting.send_monthly_spend_report() - - mock_send_alert.assert_called_once() - mock_send_alert.reset_mock() - - # Second call should not send alert (cached) - if report_type == "weekly": - await slack_alerting.send_weekly_spend_report() - else: - await slack_alerting.send_monthly_spend_report() - mock_send_alert.assert_not_called() diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index bf6294baa00..8c17ee132d8 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -6,24 +6,15 @@ import litellm import litellm.vector_stores.main import json from typing import Optional -from unittest.mock import AsyncMock, patch, Mock import pytest import litellm -from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( - VectorStorePreCallHook, -) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import ( StandardLoggingPayload, ) -from litellm.types.vector_stores import ( - VectorStoreSearchResponse, - VectorStoreResultContent, - VectorStoreSearchResult, -) class MockCustomLogger(CustomLogger): @@ -63,113 +54,6 @@ def setup_vector_store_registry(): ) -@pytest.mark.asyncio -async def test_vector_store_hook_routes_search_through_proxy_router( - setup_vector_store_registry, -): - proxy_router = Mock() - proxy_router.avector_store_search = AsyncMock( - return_value=VectorStoreSearchResponse( - object="vector_store.search_results.page", - search_query="what is litellm?", - data=[ - VectorStoreSearchResult( - score=1.0, - content=[VectorStoreResultContent(text="routed context", type="text")], - ) - ], - ) - ) - logging_obj = Mock() - logging_obj.model_call_details = { - "litellm_params": {"metadata": {"user_api_key_team_id": "team-a"}} - } - - with patch("litellm.proxy.proxy_server.llm_router", proxy_router): - _, messages, _ = await VectorStorePreCallHook().async_get_chat_completion_prompt( - model="chat-model", - messages=[{"role": "user", "content": "what is litellm?"}], - non_default_params={"vector_store_ids": ["T37J8R4WTM"]}, - prompt_id=None, - prompt_variables=None, - dynamic_callback_params={}, - litellm_logging_obj=logging_obj, - ) - - proxy_router.avector_store_search.assert_awaited_once_with( - vector_store_id="T37J8R4WTM", - query="what is litellm?", - custom_llm_provider="bedrock", - metadata={"user_api_key_team_id": "team-a"}, - ) - assert messages[0]["content"] == "Context:\n\nrouted context\n\n" - - -@pytest.mark.asyncio -async def test_e2e_bedrock_knowledgebase_retrieval_with_completion( - setup_vector_store_registry, -): - litellm.turn_on_debug() - client = AsyncHTTPHandler() - print("value of litellm.vector_store_registry:", litellm.vector_store_registry) - - with patch.object(client, "post") as mock_post: - # Mock the response for the LLM call - mock_response = Mock() - mock_response.status_code = 200 - mock_response.headers = {"Content-Type": "application/json"} - # Provide proper JSON response content - mock_response.text = json.dumps( - { - "id": "msg_01ABC123", - "type": "message", - "role": "assistant", - "content": [ - { - "type": "text", - "text": "LiteLLM is a library that simplifies LLM API access.", - } - ], - "model": "claude-3.5-sonnet", - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 100, "output_tokens": 50}, - } - ) - mock_response.json = lambda: json.loads(mock_response.text) - mock_post.return_value = mock_response - - try: - response = await litellm.acompletion( - model="anthropic/claude-3.5-sonnet", - messages=[{"role": "user", "content": "what is litellm?"}], - vector_store_ids=["T37J8R4WTM"], - client=client, - ) - except Exception as e: - print(f"Error: {e}") - - # Verify the LLM request was made - mock_post.assert_called_once() - - # Verify the request body - print("call args:", mock_post.call_args) - request_body = mock_post.call_args.kwargs["json"] - print("Request body:", json.dumps(request_body, indent=4, default=str)) - - # Assert content from the knowedge base was applied to the request - - # 1. we should have 2 content blocks, the first is the context from the knowledge base, the second is the user message - content = request_body["messages"][0]["content"] - assert len(content) == 2 - assert content[0]["type"] == "text" - assert content[1]["type"] == "text" - - # 2. the first content block should have the bedrock knowledge base prefix string - # this helps confirm that the context from the knowledge base was applied to the request - assert VectorStorePreCallHook.CONTENT_PREFIX_STRING in content[0]["text"] - - @pytest.mark.asyncio async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call( setup_vector_store_registry, @@ -214,65 +98,6 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call( print(f"First search result has {len(first_search_result['data'])} items") -@pytest.mark.asyncio -async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_streaming( - setup_vector_store_registry, -): - """ - Test that the Bedrock Knowledge Base Hook works with streaming and returns search_results in chunks. - """ - - # Init client - # litellm.turn_on_debug() - async_client = AsyncHTTPHandler() - response = await litellm.acompletion( - model=f"anthropic/{os.environ.get('CI_CD_DEFAULT_ANTHROPIC_MODEL', 'claude-haiku-4-5-20251001')}", - messages=[{"role": "user", "content": "what is litellm?"}], - vector_store_ids=["T37J8R4WTM"], - stream=True, - client=async_client, - ) - - # Collect chunks - chunks = [] - search_results_found = False - async for chunk in response: - chunks.append(chunk) - print(f"Chunk: {chunk}") - - # Check if this chunk has search_results in provider_specific_fields - if hasattr(chunk, "choices") and chunk.choices: - for choice in chunk.choices: - if hasattr(choice, "delta") and choice.delta: - provider_fields = getattr( - choice.delta, "provider_specific_fields", None - ) - if provider_fields and "search_results" in provider_fields: - search_results = provider_fields["search_results"] - print( - f"Found search_results in streaming chunk: {len(search_results)} results" - ) - - # Verify structure - assert search_results is not None - assert len(search_results) > 0 - - first_search_result = search_results[0] - assert "object" in first_search_result - assert ( - first_search_result["object"] - == "vector_store.search_results.page" - ) - assert "data" in first_search_result - assert len(first_search_result["data"]) > 0 - - search_results_found = True - - print(f"Total chunks received: {len(chunks)}") - assert len(chunks) > 0 - assert search_results_found, "search_results should be present in streaming chunks" - - @pytest.mark.asyncio async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools( setup_vector_store_registry, @@ -345,328 +170,6 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools_ print(f" Search was performed and {len(search_results)} result(s) returned") -@pytest.mark.asyncio -async def test_bedrock_kb_request_body_has_transformed_filters( - setup_vector_store_registry, -): - """ - Validate that the Bedrock Knowledge Base request body contains the transformed filters. - """ - captured_request_body: dict = {} - - async def fake_async_vector_store_search_handler( - vector_store_id, - query, - vector_store_search_optional_params, - vector_store_provider_config, - custom_llm_provider, - litellm_params, - logging_obj, - embedding_executor=None, - extra_headers=None, - extra_body=None, - timeout=None, - client=None, - _is_async=False, - ): - litellm_params_dict = ( - litellm_params.model_dump(exclude_none=False) - if hasattr(litellm_params, "model_dump") - else dict(litellm_params) - ) - api_base = vector_store_provider_config.get_complete_url( - api_base=litellm_params_dict.get("api_base"), - litellm_params=litellm_params_dict, - ) - - url, request_body = ( - vector_store_provider_config.transform_search_vector_store_request( - vector_store_id=vector_store_id, - query=query, - vector_store_search_optional_params=vector_store_search_optional_params, - api_base=api_base, - litellm_logging_obj=logging_obj, - litellm_params=litellm_params_dict, - extra_body=None, - ) - ) - captured_request_body["url"] = url - captured_request_body["body"] = request_body - - return VectorStoreSearchResponse( - object="vector_store.search_results.page", - search_query=query if isinstance(query, str) else " ".join(query), - data=[ - VectorStoreSearchResult( - score=0.9, - content=[ - VectorStoreResultContent( - text="LiteLLM is a library", type="text" - ) - ], - ) - ], - ) - - with patch.object( - litellm.vector_stores.main.base_llm_http_handler, - "async_vector_store_search_handler", - new=AsyncMock(side_effect=fake_async_vector_store_search_handler), - ): - response = await litellm.acompletion( - model=f"anthropic/{os.environ.get('CI_CD_DEFAULT_ANTHROPIC_MODEL', 'claude-haiku-4-5-20251001')}", - messages=[{"role": "user", "content": "what is litellm?"}], - max_tokens=10, - tools=[ - { - "type": "file_search", - "vector_store_ids": ["T37J8R4WTM"], - "filters": { - "key": "user_id", - "value": "fake-user-id", - "operator": "eq", - }, - } - ], - ) - - assert response is not None - print( - "captured_request_body:", - json.dumps(captured_request_body, indent=4, default=str), - ) - assert "body" in captured_request_body, "Bedrock KB request body was not captured" - - vector_search = captured_request_body["body"]["retrievalConfiguration"][ - "vectorSearchConfiguration" - ] - aws_filter = vector_search["filter"] - assert "equals" in aws_filter, f"Expected 'equals' in AWS format, got: {aws_filter}" - assert aws_filter["equals"]["key"] == "user_id" - assert aws_filter["equals"]["value"] == "fake-user-id" - - print("✅ Filters transformed correctly: OpenAI format -> AWS Bedrock format") - - -@pytest.mark.asyncio -async def test_openai_with_knowledge_base_mock_openai(setup_vector_store_registry): - """ - Tests that knowledge base content is correctly passed to the OpenAI API call - """ - litellm.set_verbose = True - from openai import AsyncOpenAI - - client = AsyncOpenAI(api_key="fake-api-key") - - # Variable to capture the request - captured_request = {} - - with patch.object( - client.chat.completions.with_raw_response, "create" - ) as mock_client: - # Create async mock that returns proper structure - async def mock_create(**kwargs): - mock_response = Mock() - mock_response.choices = [ - Mock( - message=Mock(content="Mock response from OpenAI", role="assistant") - ) - ] - mock_response.usage = Mock( - prompt_tokens=100, completion_tokens=50, total_tokens=150 - ) - mock_response.id = "chatcmpl-123" - mock_response.object = "chat.completion" - mock_response.created = 1234567890 - mock_response.model = "gpt-5.5" - - # Store the request for verification - captured_request.update(kwargs) - - # Return wrapper with parse method - wrapper = Mock() - wrapper.parse.return_value = mock_response - return wrapper - - mock_client.side_effect = mock_create - - try: - await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": "what is litellm?"}], - vector_store_ids=["T37J8R4WTM"], - client=client, - ) - except Exception as e: - print(f"Error: {e}") - - # Verify the API was called - mock_client.assert_called_once() - request_body = captured_request - - # Verify the request contains messages with knowledge base context - assert "messages" in request_body - messages = request_body["messages"] - - # We expect at least 2 messages: - # 1. User message with the knowledge base context - # 2. User message with the question - assert len(messages) >= 2 - - print("request messages:", json.dumps(messages, indent=4, default=str)) - - # assert message[0] is the user message with the knowledge base context - assert messages[0]["role"] == "user" - assert VectorStorePreCallHook.CONTENT_PREFIX_STRING in messages[0]["content"] - - -@pytest.mark.asyncio -async def test_openai_with_vector_store_ids_in_tool_call_mock_openai( - setup_vector_store_registry, -): - """ - Tests that vector store ids can be passed as tools - - This is the OpenAI format - """ - litellm.set_verbose = True - from openai import AsyncOpenAI - - client = AsyncOpenAI(api_key="fake-api-key") - - # Variable to capture the request - captured_request = {} - - with patch.object( - client.chat.completions.with_raw_response, "create" - ) as mock_client: - # Create async mock that returns proper structure - async def mock_create(**kwargs): - mock_response = Mock() - mock_response.choices = [ - Mock( - message=Mock(content="Mock response from OpenAI", role="assistant") - ) - ] - mock_response.usage = Mock( - prompt_tokens=100, completion_tokens=50, total_tokens=150 - ) - mock_response.id = "chatcmpl-123" - mock_response.object = "chat.completion" - mock_response.created = 1234567890 - mock_response.model = "gpt-5.5" - - # Store the request for verification - captured_request.update(kwargs) - - # Return wrapper with parse method - wrapper = Mock() - wrapper.parse.return_value = mock_response - return wrapper - - mock_client.side_effect = mock_create - - try: - await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": "what is litellm?"}], - tools=[{"type": "file_search", "vector_store_ids": ["T37J8R4WTM"]}], - client=client, - ) - except Exception as e: - print(f"Error: {e}") - - # Verify the API was called - mock_client.assert_called_once() - request_body = captured_request - print("request body:", json.dumps(request_body, indent=4, default=str)) - - # Verify the request contains messages with knowledge base context - assert "messages" in request_body - messages = request_body["messages"] - - # We expect at least 2 messages: - # 1. User message with the knowledge base context - # 2. User message with the question - assert len(messages) >= 2 - - print("request messages:", json.dumps(messages, indent=4, default=str)) - - # assert message[0] is the user message with the knowledge base context - assert messages[0]["role"] == "user" - assert VectorStorePreCallHook.CONTENT_PREFIX_STRING in messages[0]["content"] - - # assert that the tool call was not sent to the upstream llm API if it's a litellm vector store - assert "tools" not in request_body - - -@pytest.mark.asyncio -async def test_openai_with_mixed_tool_call_mock_openai(setup_vector_store_registry): - """Ensure unrecognized vector store tools are forwarded to the provider""" - from openai import AsyncOpenAI - - client = AsyncOpenAI(api_key="fake-api-key") - - # Variable to capture the request - captured_request = {} - - with patch.object( - client.chat.completions.with_raw_response, "create" - ) as mock_client: - # Create async mock that returns proper structure - async def mock_create(**kwargs): - mock_response = Mock() - mock_response.choices = [ - Mock( - message=Mock(content="Mock response from OpenAI", role="assistant") - ) - ] - mock_response.usage = Mock( - prompt_tokens=100, completion_tokens=50, total_tokens=150 - ) - mock_response.id = "chatcmpl-123" - mock_response.object = "chat.completion" - mock_response.created = 1234567890 - mock_response.model = "gpt-5.5" - - # Store the request for verification - captured_request.update(kwargs) - - # Return wrapper with parse method - wrapper = Mock() - wrapper.parse.return_value = mock_response - return wrapper - - mock_client.side_effect = mock_create - - try: - await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": "what is litellm?"}], - tools=[ - {"type": "file_search", "vector_store_ids": ["T37J8R4WTM"]}, - {"type": "file_search", "vector_store_ids": ["unknownVS"]}, - ], - client=client, - ) - except Exception as e: - print(f"Error: {e}") - - mock_client.assert_called_once() - request_body = captured_request - - assert "messages" in request_body - messages = request_body["messages"] - assert len(messages) >= 2 - assert messages[0]["role"] == "user" - assert VectorStorePreCallHook.CONTENT_PREFIX_STRING in messages[0]["content"] - - assert "tools" in request_body - tools = request_body["tools"] - assert len(tools) == 1 - assert tools[0]["vector_store_ids"] == ["unknownVS"] - - # @pytest.mark.asyncio # async def test_logging_with_knowledge_base_hook(setup_vector_store_registry): # """ @@ -723,118 +226,3 @@ async def test_openai_with_mixed_tool_call_mock_openai(setup_vector_store_regist -@pytest.mark.asyncio -async def test_provider_specific_fields_in_proxy_http_response( - setup_vector_store_registry, -): - """ - Test that provider_specific_fields (like search_results) are included - in the proxy HTTP JSON response, not just in Python SDK objects. - - This test catches serialization bugs where exclude=True would strip - provider_specific_fields from the HTTP response. - """ - from fastapi.testclient import TestClient - from litellm.proxy.proxy_server import app, initialize - from unittest.mock import patch as mock_patch - - # Initialize proxy - await initialize( - model="gpt-5-mini", - alias=None, - api_base=None, - debug=False, - temperature=None, - max_tokens=None, - request_timeout=600, - max_budget=None, - drop_params=True, - add_function_to_prompt=False, - headers=None, - save=False, - use_queue=False, - config=None, - ) - - # Create test client - client = TestClient(app) - - # Create mock response with provider_specific_fields - mock_response = litellm.ModelResponse( - id="test-123", - model="gpt-5-mini", - created=1234567890, - object="chat.completion", - ) - - # Create message with provider_specific_fields - mock_message = litellm.Message( - content="LiteLLM is a tool that simplifies working with multiple LLMs.", - role="assistant", - provider_specific_fields={ - "search_results": [ - { - "object": "vector_store.search_results.page", - "search_query": "what is litellm?", - "data": [ - { - "score": 0.95, - "content": [{"text": "Test content", "type": "text"}], - "file_id": "test-file", - "filename": "test.txt", - } - ], - } - ] - }, - ) - - mock_choice = litellm.Choices(finish_reason="stop", index=0, message=mock_message) - - mock_response.choices = [mock_choice] - mock_response.usage = litellm.Usage( - prompt_tokens=10, completion_tokens=20, total_tokens=30 - ) - - # Patch the completion call at the proxy level - with mock_patch("litellm.acompletion", new=AsyncMock(return_value=mock_response)): - # Make HTTP request to proxy - response = client.post( - "/v1/chat/completions", - json={ - "model": "gpt-5-mini", - "messages": [{"role": "user", "content": "What is litellm?"}], - }, - ) - - # Check HTTP response - assert response.status_code == 200 - result = response.json() - - print("HTTP Response JSON:", json.dumps(result, indent=2)) - - # THE KEY ASSERTIONS - These would FAIL with exclude=True! - assert "choices" in result - assert len(result["choices"]) > 0 - - choice = result["choices"][0] - assert "message" in choice - - message = choice["message"] - - # Verify provider_specific_fields is in the JSON response - assert ( - "provider_specific_fields" in message - ), "provider_specific_fields missing from HTTP JSON response! This means exclude=True is preventing serialization." - - assert "search_results" in message["provider_specific_fields"] - search_results = message["provider_specific_fields"]["search_results"] - assert len(search_results) > 0 - - # Verify search result structure - first_result = search_results[0] - assert first_result["object"] == "vector_store.search_results.page" - assert "data" in first_result - assert len(first_result["data"]) > 0 - - print("✅ provider_specific_fields successfully serialized in HTTP response") diff --git a/tests/logging_callback_tests/test_built_in_tools_cost_tracking.py b/tests/logging_callback_tests/test_built_in_tools_cost_tracking.py deleted file mode 100644 index 2d3933c3dad..00000000000 --- a/tests/logging_callback_tests/test_built_in_tools_cost_tracking.py +++ /dev/null @@ -1,160 +0,0 @@ -import traceback -from litellm._uuid import uuid -import pytest -from dotenv import load_dotenv -from fastapi import Request -from fastapi.routing import APIRoute - -load_dotenv() -import io -import time -import json - -# this file is to test litellm/proxy - -import litellm -import asyncio -from typing import Optional -from litellm.types.utils import StandardLoggingPayload, Usage, ModelInfoBase -from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( - StandardBuiltInToolCostTracking, -) - - -class TestCustomLogger(CustomLogger): - def __init__(self): - self.recorded_usage: Optional[Usage] = None - self.standard_logging_payload: Optional[StandardLoggingPayload] = None - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - standard_logging_payload = kwargs.get("standard_logging_object") - self.standard_logging_payload = standard_logging_payload - print( - "standard_logging_payload", - json.dumps(standard_logging_payload, indent=4, default=str), - ) - - self.recorded_usage = Usage( - prompt_tokens=standard_logging_payload.get("prompt_tokens"), - completion_tokens=standard_logging_payload.get("completion_tokens"), - total_tokens=standard_logging_payload.get("total_tokens"), - ) - pass - - -async def _setup_web_search_test(): - """Helper function to setup common test requirements""" - litellm.turn_on_debug() - test_custom_logger = TestCustomLogger() - litellm.callbacks = [test_custom_logger] - return test_custom_logger - - -async def _verify_web_search_cost(test_custom_logger, expected_context_size): - """Helper function to verify web search costs""" - await asyncio.sleep(1) - - standard_logging_payload = test_custom_logger.standard_logging_payload - response = standard_logging_payload.get("response") - response_cost = standard_logging_payload.get("response_cost") - assert response_cost is not None - - # Calculate token cost - model_map_information = standard_logging_payload["model_map_information"] - model_map_value: ModelInfoBase = model_map_information["model_map_value"] - total_token_cost = ( - standard_logging_payload["prompt_tokens"] - * model_map_value["input_cost_per_token"] - ) + ( - standard_logging_payload["completion_tokens"] - * model_map_value["output_cost_per_token"] - ) - - # Verify total cost - if StandardBuiltInToolCostTracking.response_object_includes_web_search_call( - response - ): - assert ( - response_cost - == total_token_cost - + model_map_value["search_context_cost_per_query"][expected_context_size] - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "web_search_options,expected_context_size", - [ - (None, "search_context_size_medium"), - ({"search_context_size": "low"}, "search_context_size_low"), - ({"search_context_size": "high"}, "search_context_size_high"), - ], -) -async def test_openai_web_search_logging_cost_tracking( - web_search_options, expected_context_size -): - """Test web search cost tracking with different search context sizes""" - test_custom_logger = await _setup_web_search_test() - - request_kwargs = { - "model": "openai/gpt-5-search-api", - "messages": [ - { - "role": "user", - "content": f"What was a positive news story from today? {uuid.uuid4()}", - } - ], - } - if web_search_options is not None: - request_kwargs["web_search_options"] = web_search_options - - response = await litellm.acompletion(**request_kwargs) - - await _verify_web_search_cost(test_custom_logger, expected_context_size) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "tools_config,expected_context_size,stream", - [ - ( - [{"type": "web_search_preview", "search_context_size": "low"}], - "search_context_size_low", - True, - ), - ( - [{"type": "web_search_preview", "search_context_size": "low"}], - "search_context_size_low", - False, - ), - ([{"type": "web_search_preview"}], "search_context_size_medium", True), - ([{"type": "web_search_preview"}], "search_context_size_medium", False), - ], -) -async def test_openai_responses_api_web_search_cost_tracking( - tools_config, expected_context_size, stream -): - """Test web search cost tracking with different search context sizes and streaming options""" - test_custom_logger = await _setup_web_search_test() - - response = await litellm.aresponses( - model="openai/gpt-4o", - input=[ - {"role": "user", "content": "What was a positive news story from today?"} - ], - tools=tools_config, - stream=stream, - ) - if stream is True: - async for chunk in response: - print("chunk", chunk) - else: - print("response", response) - - await asyncio.sleep(1) - - if StandardBuiltInToolCostTracking.response_object_includes_web_search_call( - test_custom_logger.standard_logging_payload.get("response") - ): - await _verify_web_search_cost(test_custom_logger, expected_context_size) diff --git a/tests/logging_callback_tests/test_custom_callback_router.py b/tests/logging_callback_tests/test_custom_callback_router.py deleted file mode 100644 index 4a12f8d536d..00000000000 --- a/tests/logging_callback_tests/test_custom_callback_router.py +++ /dev/null @@ -1,754 +0,0 @@ -### What this tests #### -## This test asserts the type of data passed into each method of the custom callback handler -import asyncio -import inspect -import os -import time -import traceback -from datetime import datetime - -import pytest - -from typing import List, Literal, Optional - -import litellm -from litellm import Cache, Router -from litellm.integrations.custom_logger import CustomLogger - -# Test Scenarios (test across completion, streaming, embedding) -## 1: Pre-API-Call -## 2: Post-API-Call -## 3: On LiteLLM Call success -## 4: On LiteLLM Call failure -## fallbacks -## retries - -# Test cases -## 1. Simple Azure OpenAI acompletion + streaming call -## 2. Simple Azure OpenAI aembedding call -## 3. Azure OpenAI acompletion + streaming call with retries -## 4. Azure OpenAI aembedding call with retries -## 5. Azure OpenAI acompletion + streaming call with fallbacks -## 6. Azure OpenAI aembedding call with fallbacks - -## Test interfaces -## 1. router.completion() + router.embeddings() -## 2. proxy.completions + proxy.embeddings - -litellm.num_retries = 0 - - -class CompletionCustomHandler( - CustomLogger -): # https://docs.litellm.ai/docs/observability/custom_callback#callback-class - """ - The set of expected inputs to a custom handler for a - """ - - # Class variables or attributes - def __init__(self): - self.errors = [] - self.states: Optional[ - List[ - Literal[ - "sync_pre_api_call", - "async_pre_api_call", - "post_api_call", - "sync_stream", - "async_stream", - "sync_success", - "async_success", - "sync_failure", - "async_failure", - ] - ] - ] = [] - - def log_pre_api_call(self, model, messages, kwargs): - try: - print(f"received kwargs in pre-input: {kwargs}") - self.states.append("sync_pre_api_call") - ## MODEL - assert isinstance(model, str) - ## MESSAGES - assert isinstance(messages, list) - ## KWARGS - assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) - assert isinstance(kwargs["optional_params"], dict) - assert isinstance(kwargs["litellm_params"], dict) - assert isinstance(kwargs["start_time"], (datetime, type(None))) - assert isinstance(kwargs["stream"], bool) - assert isinstance(kwargs["user"], (str, type(None))) - ### ROUTER-SPECIFIC KWARGS - assert isinstance(kwargs["litellm_params"]["metadata"], dict) - assert isinstance(kwargs["litellm_params"]["metadata"]["model_group"], str) - assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) - assert isinstance(kwargs["litellm_params"]["model_info"], dict) - assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) - assert isinstance(kwargs["litellm_params"]["stream_response"], dict) - except Exception as e: - print(f"Assertion Error: {traceback.format_exc()}") - self.errors.append(traceback.format_exc()) - - def log_post_api_call(self, kwargs, response_obj, start_time, end_time): - try: - self.states.append("post_api_call") - ## START TIME - assert isinstance(start_time, datetime) - ## END TIME - assert end_time == None - ## RESPONSE OBJECT - assert response_obj == None - ## KWARGS - assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) - assert isinstance(kwargs["optional_params"], dict) - assert isinstance(kwargs["litellm_params"], dict) - assert isinstance(kwargs["start_time"], (datetime, type(None))) - assert isinstance(kwargs["stream"], bool) - assert isinstance(kwargs["user"], (str, type(None))) - assert isinstance(kwargs["input"], (list, dict, str)) - assert isinstance(kwargs["api_key"], (str, type(None))) - assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) - or inspect.iscoroutine(kwargs["original_response"]) - or inspect.isasyncgen(kwargs["original_response"]) - ) - assert isinstance(kwargs["additional_args"], (dict, type(None))) - assert isinstance(kwargs["log_event_type"], str) - ### ROUTER-SPECIFIC KWARGS - assert isinstance(kwargs["litellm_params"]["metadata"], dict) - assert isinstance(kwargs["litellm_params"]["metadata"]["model_group"], str) - assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) - assert isinstance(kwargs["litellm_params"]["model_info"], dict) - assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) - assert isinstance(kwargs["litellm_params"]["stream_response"], dict) - except Exception: - print(f"Assertion Error: {traceback.format_exc()}") - self.errors.append(traceback.format_exc()) - - async def async_log_stream_event(self, kwargs, response_obj, start_time, end_time): - try: - self.states.append("async_stream") - ## START TIME - assert isinstance(start_time, datetime) - ## END TIME - assert isinstance(end_time, datetime) - ## RESPONSE OBJECT - assert isinstance(response_obj, litellm.ModelResponseStream) - ## KWARGS - assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) - assert isinstance(kwargs["optional_params"], dict) - assert isinstance(kwargs["litellm_params"], dict) - assert isinstance(kwargs["start_time"], (datetime, type(None))) - assert isinstance(kwargs["stream"], bool) - assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) - assert isinstance(kwargs["api_key"], (str, type(None))) - assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) - or inspect.isasyncgen(kwargs["original_response"]) - or inspect.iscoroutine(kwargs["original_response"]) - ) - assert isinstance(kwargs["additional_args"], (dict, type(None))) - assert isinstance(kwargs["log_event_type"], str) - except Exception: - print(f"Assertion Error: {traceback.format_exc()}") - self.errors.append(traceback.format_exc()) - - def log_success_event(self, kwargs, response_obj, start_time, end_time): - try: - self.states.append("sync_success") - ## START TIME - assert isinstance(start_time, datetime) - ## END TIME - assert isinstance(end_time, datetime) - ## RESPONSE OBJECT - assert isinstance(response_obj, litellm.ModelResponse) - ## KWARGS - assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) - assert isinstance(kwargs["optional_params"], dict) - assert isinstance(kwargs["litellm_params"], dict) - assert isinstance(kwargs["start_time"], (datetime, type(None))) - assert isinstance(kwargs["stream"], bool) - assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) - assert isinstance(kwargs["api_key"], (str, type(None))) - assert isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) - assert isinstance(kwargs["additional_args"], (dict, type(None))) - assert isinstance(kwargs["log_event_type"], str) - assert kwargs["cache_hit"] is None or isinstance(kwargs["cache_hit"], bool) - except Exception: - print(f"Assertion Error: {traceback.format_exc()}") - self.errors.append(traceback.format_exc()) - - def log_failure_event(self, kwargs, response_obj, start_time, end_time): - try: - self.states.append("sync_failure") - ## START TIME - assert isinstance(start_time, datetime) - ## END TIME - assert isinstance(end_time, datetime) - ## RESPONSE OBJECT - assert response_obj == None - ## KWARGS - assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) - assert isinstance(kwargs["optional_params"], dict) - assert isinstance(kwargs["litellm_params"], dict) - assert isinstance(kwargs["start_time"], (datetime, type(None))) - assert isinstance(kwargs["stream"], bool) - assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) - assert isinstance(kwargs["api_key"], (str, type(None))) - assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) - or kwargs["original_response"] == None - ) - assert isinstance(kwargs["additional_args"], (dict, type(None))) - assert isinstance(kwargs["log_event_type"], str) - except Exception: - print(f"Assertion Error: {traceback.format_exc()}") - self.errors.append(traceback.format_exc()) - - async def async_log_pre_api_call(self, model, messages, kwargs): - try: - """ - No-op. - Not implemented yet. - """ - pass - except Exception as e: - print(f"Assertion Error: {traceback.format_exc()}") - self.errors.append(traceback.format_exc()) - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - try: - print("CompletionCustomHandler.async_log_success_event, kwargs: ", kwargs) - self.states.append("async_success") - print( - "############### CompletionCustomHandler async success, kwargs: ", - kwargs, - ) - ## START TIME - assert isinstance(start_time, datetime) - ## END TIME - assert isinstance(end_time, datetime) - ## RESPONSE OBJECT - assert isinstance( - response_obj, (litellm.ModelResponse, litellm.EmbeddingResponse) - ) - ## KWARGS - assert isinstance(kwargs["model"], str) - - # checking we use base_model for azure cost calculation - base_model = litellm.utils.get_base_model_from_metadata( - model_call_details=kwargs - ) - - if ( - kwargs["model"] == "chatgpt-v-3" - and base_model is not None - and kwargs["stream"] != True - ): - # when base_model is set for azure, we should use pricing for the base_model - # this checks response_cost == litellm.cost_per_token(model=base_model) - assert isinstance(kwargs["response_cost"], float) - response_cost = kwargs["response_cost"] - print( - f"response_cost: {response_cost}, for model: {kwargs['model']} and base_model: {base_model}" - ) - prompt_tokens = response_obj.usage.prompt_tokens - completion_tokens = response_obj.usage.completion_tokens - # ensure the pricing is based on the base_model here - prompt_price, completion_price = litellm.cost_per_token( - model=base_model, - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - ) - expected_price = prompt_price + completion_price - print(f"expected price: {expected_price}") - assert ( - response_cost == expected_price - ), f"response_cost: {response_cost} != expected_price: {expected_price}. For model: {kwargs['model']} and base_model: {base_model}. should have used base_model for price" - - assert isinstance(kwargs["messages"], list) - assert isinstance(kwargs["optional_params"], dict) - assert isinstance(kwargs["litellm_params"], dict) - assert isinstance(kwargs["start_time"], (datetime, type(None))) - assert isinstance(kwargs["stream"], bool) - assert isinstance(kwargs["user"], (str, type(None))) - assert isinstance(kwargs["input"], (list, dict, str)) - assert isinstance(kwargs["api_key"], (str, type(None))) - assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) - or inspect.isasyncgen(kwargs["original_response"]) - or inspect.iscoroutine(kwargs["original_response"]) - ) - assert isinstance(kwargs["additional_args"], (dict, type(None))) - assert isinstance(kwargs["log_event_type"], str) - assert kwargs["cache_hit"] is None or isinstance(kwargs["cache_hit"], bool) - ### ROUTER-SPECIFIC KWARGS - assert isinstance(kwargs["litellm_params"]["metadata"], dict) - assert isinstance(kwargs["litellm_params"]["metadata"]["model_group"], str) - assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) - assert isinstance(kwargs["litellm_params"]["model_info"], dict) - assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) - assert isinstance(kwargs["litellm_params"]["stream_response"], dict) - except Exception: - print(f"Assertion Error: {traceback.format_exc()}") - self.errors.append(traceback.format_exc()) - - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - try: - print(f"received original response: {kwargs['original_response']}") - self.states.append("async_failure") - ## START TIME - assert isinstance(start_time, datetime) - ## END TIME - assert isinstance(end_time, datetime) - ## RESPONSE OBJECT - assert response_obj == None - ## KWARGS - assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) - assert isinstance(kwargs["optional_params"], dict) - assert isinstance(kwargs["litellm_params"], dict) - assert isinstance(kwargs["start_time"], (datetime, type(None))) - assert isinstance(kwargs["stream"], bool) - assert isinstance(kwargs["user"], (str, type(None))) - assert isinstance(kwargs["input"], (list, str, dict)) - assert isinstance(kwargs["api_key"], (str, type(None))) - assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) - or inspect.isasyncgen(kwargs["original_response"]) - or inspect.iscoroutine(kwargs["original_response"]) - or kwargs["original_response"] == None - ) - assert isinstance(kwargs["additional_args"], (dict, type(None))) - assert isinstance(kwargs["log_event_type"], str) - except Exception: - print(f"Assertion Error: {traceback.format_exc()}") - self.errors.append(traceback.format_exc()) - - -# Simple Azure OpenAI call -## COMPLETION -# @pytest.mark.flaky(retries=5, delay=1) -@pytest.mark.asyncio -async def test_async_chat_azure(): - try: - customHandler_completion_azure_router = CompletionCustomHandler() - customHandler_streaming_azure_router = CompletionCustomHandler() - customHandler_failure = CompletionCustomHandler() - litellm.callbacks = [customHandler_completion_azure_router] - litellm.set_verbose = True - model_list = [ - { - "model_name": "gpt-4.1-nano", # openai model name - "litellm_params": { # params for litellm completion/embedding call - "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_AI_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_AI_API_BASE"), - }, - "model_info": {"base_model": "azure/gpt-4.1-mini"}, - "tpm": 240000, - "rpm": 1800, - }, - ] - router = Router(model_list=model_list, num_retries=0) # type: ignore - response = await router.acompletion( - model="gpt-4.1-nano", - messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}], - ) - print("got response, sleeping 5 seconds....") - await asyncio.sleep(5) - assert len(customHandler_completion_azure_router.errors) == 0 - assert ( - len(customHandler_completion_azure_router.states) == 3 - ) # pre, post, success - # streaming - - litellm.logging_callback_manager._reset_all_callbacks() - litellm.callbacks = [customHandler_streaming_azure_router] - router2 = Router(model_list=model_list, num_retries=0) # type: ignore - response = await router2.acompletion( - model="gpt-4.1-nano", - messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}], - stream=True, - ) - async for chunk in response: - print(f"async azure router chunk: {chunk}") - continue - await asyncio.sleep(5) - print(f"customHandler.states: {customHandler_streaming_azure_router.states}") - assert len(customHandler_streaming_azure_router.errors) == 0 - assert ( - len(customHandler_streaming_azure_router.states) >= 3 - ) # pre, post, stream (multiple times), success - # failure - model_list = [ - { - "model_name": "gpt-5-mini", # openai model name - "litellm_params": { # params for litellm completion/embedding call - "model": "azure/gpt-4o-new-test", - "api_key": "my-bad-key", - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_AI_API_BASE"), - }, - "tpm": 240000, - "rpm": 1800, - }, - ] - - litellm.logging_callback_manager._reset_all_callbacks() - litellm.callbacks = [customHandler_failure] - router3 = Router(model_list=model_list, num_retries=0) # type: ignore - try: - response = await router3.acompletion( - model="gpt-5-mini", - messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}], - ) - print(f"response in router3 acompletion: {response}") - except Exception: - pass - await asyncio.sleep(5) - print(f"customHandler.states: {customHandler_failure.states}") - assert len(customHandler_failure.errors) == 0 - assert len(customHandler_failure.states) == 3 # pre, post, failure - assert "async_failure" in customHandler_failure.states - except Exception as e: - print(f"Assertion Error: {traceback.format_exc()}") - pytest.fail(f"An exception occurred - {str(e)}") - - -## EMBEDDING -@pytest.mark.asyncio -async def test_async_embedding_azure(): - try: - customHandler = CompletionCustomHandler() - customHandler_failure = CompletionCustomHandler() - litellm.callbacks = [customHandler] - model_list = [ - { - "model_name": "azure-embedding-model", # openai model name - "litellm_params": { # params for litellm completion/embedding call - "model": "azure/text-embedding-ada-002", - "api_key": os.getenv("AZURE_AI_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_AI_API_BASE"), - }, - "tpm": 240000, - "rpm": 1800, - }, - ] - router = Router(model_list=model_list) # type: ignore - response = await router.aembedding( - model="azure-embedding-model", input=["hello from litellm!"] - ) - await asyncio.sleep(2) - assert len(customHandler.errors) == 0 - assert len(customHandler.states) == 3 # pre, post, success - # failure - model_list = [ - { - "model_name": "azure-embedding-model", # openai model name - "litellm_params": { # params for litellm completion/embedding call - "model": "azure/text-embedding-ada-002", - "api_key": "my-bad-key", - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_AI_API_BASE"), - }, - "tpm": 240000, - "rpm": 1800, - }, - ] - litellm.logging_callback_manager._reset_all_callbacks() - litellm.callbacks = [customHandler_failure] - router3 = Router(model_list=model_list, num_retries=0) # type: ignore - try: - response = await router3.aembedding( - model="azure-embedding-model", input=["hello from litellm!"] - ) - print(f"response in router3 aembedding: {response}") - except Exception: - pass - await asyncio.sleep(1) - print(f"customHandler.states: {customHandler_failure.states}") - assert len(customHandler_failure.errors) == 0 - assert len(customHandler_failure.states) == 3 # pre, post, failure - assert "async_failure" in customHandler_failure.states - except Exception as e: - print(f"Assertion Error: {traceback.format_exc()}") - pytest.fail(f"An exception occurred - {str(e)}") - - -# asyncio.run(test_async_embedding_azure()) -# Azure OpenAI call w/ Fallbacks -## COMPLETION -@pytest.mark.asyncio -async def test_async_chat_azure_with_fallbacks(): - try: - customHandler_fallbacks = CompletionCustomHandler() - litellm.callbacks = [customHandler_fallbacks] - litellm.set_verbose = True - # with fallbacks - model_list = [ - { - "model_name": "gpt-5-mini", # openai model name - "litellm_params": { # params for litellm completion/embedding call - "model": "azure/gpt-4.1-mini", - "api_key": "my-bad-key", - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_AI_API_BASE"), - }, - "tpm": 240000, - "rpm": 1800, - }, - { - "model_name": "gpt-3.5-turbo-16k", - "litellm_params": { - "model": "gpt-3.5-turbo-16k", - }, - "tpm": 240000, - "rpm": 1800, - }, - ] - router = Router( - model_list=model_list, - fallbacks=[{"gpt-5-mini": ["gpt-3.5-turbo-16k"]}], - retry_policy=litellm.router.RetryPolicy( - AuthenticationErrorRetries=0, - ), - ) # type: ignore - response = await router.acompletion( - model="gpt-5-mini", - messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}], - ) - await asyncio.sleep(2) - print(f"customHandler_fallbacks.states: {customHandler_fallbacks.states}") - assert len(customHandler_fallbacks.errors) == 0 - assert ( - len(customHandler_fallbacks.states) == 6 - ) # pre, post, failure, pre, post, success - litellm.callbacks = [] - except Exception as e: - print(f"Assertion Error: {traceback.format_exc()}") - pytest.fail(f"An exception occurred - {str(e)}") - - -# asyncio.run(test_async_chat_azure_with_fallbacks()) - - -# CACHING -## Test Azure - completion, embedding -@pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=1) -async def test_async_completion_azure_caching(): - customHandler_caching = CompletionCustomHandler() - litellm.cache = Cache( - type="redis", - host=os.environ["REDIS_HOST"], - port=os.environ["REDIS_PORT"], - password=os.environ["REDIS_PASSWORD"], - ) - litellm.callbacks = [customHandler_caching] - unique_time = time.time() - model_list = [ - { - "model_name": "gpt-4.1-nano", # openai model name - "litellm_params": { # params for litellm completion/embedding call - "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_AI_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_AI_API_BASE"), - }, - "tpm": 240000, - "rpm": 1800, - }, - { - "model_name": "gpt-3.5-turbo-16k", - "litellm_params": { - "model": "gpt-3.5-turbo-16k", - }, - "tpm": 240000, - "rpm": 1800, - }, - ] - router = Router(model_list=model_list) # type: ignore - response1 = await router.acompletion( - model="gpt-4.1-nano", - messages=[ - {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} - ], - caching=True, - ) - await asyncio.sleep(1) - print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}") - response2 = await router.acompletion( - model="gpt-4.1-nano", - messages=[ - {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} - ], - caching=True, - ) - await asyncio.sleep(1) # success callbacks are done in parallel - print( - f"customHandler_caching.states post-cache hit: {customHandler_caching.states}" - ) - assert len(customHandler_caching.errors) == 0 - assert len(customHandler_caching.states) == 4 # pre, post, success, success - - -@pytest.mark.asyncio -async def test_async_completion_azure_caching_streaming(): - import uuid - - litellm.set_verbose = True - customHandler_caching = CompletionCustomHandler() - litellm.cache = Cache( - type="redis", - host=os.environ["REDIS_HOST"], - port=os.environ["REDIS_PORT"], - password=os.environ["REDIS_PASSWORD"], - ) - litellm.callbacks = [customHandler_caching] - unique_time = uuid.uuid4() - - # Use Router instead of direct litellm.acompletion to get router-specific metadata - model_list = [ - { - "model_name": "gpt-4.1-nano", - "litellm_params": { - "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_AI_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_AI_API_BASE"), - }, - "tpm": 240000, - "rpm": 1800, - }, - ] - router = Router(model_list=model_list) - - response1 = await router.acompletion( - model="gpt-4.1-nano", - messages=[ - {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} - ], - caching=True, - stream=True, - ) - async for chunk in response1: - print(f"chunk in response1: {chunk}") - await asyncio.sleep(1) - initial_customhandler_caching_states = len(customHandler_caching.states) - print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}") - response2 = await router.acompletion( - model="gpt-4.1-nano", - messages=[ - {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} - ], - caching=True, - stream=True, - ) - async for chunk in response2: - print(f"chunk in response2: {chunk}") - await asyncio.sleep(1) # success callbacks are done in parallel - print( - f"customHandler_caching.states post-cache hit: {customHandler_caching.states}" - ) - assert len(customHandler_caching.errors) == 0 - assert ( - len(customHandler_caching.states) > initial_customhandler_caching_states - ) # pre, post, streaming .., success, success - - -@pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=2) -async def test_async_embedding_azure_caching(): - print("Testing custom callback input - Azure Caching") - customHandler_caching = CompletionCustomHandler() - litellm.cache = Cache( - type="redis", - host=os.environ["REDIS_HOST"], - port=os.environ["REDIS_PORT"], - password=os.environ["REDIS_PASSWORD"], - ) - router = Router( - model_list=[ - { - "model_name": "text-embedding-3-small", - "litellm_params": { - "model": "openai/text-embedding-3-small", - }, - } - ] - ) - litellm.callbacks = [customHandler_caching] - unique_time = time.time() - response1 = await router.aembedding( - model="text-embedding-3-small", - input=[f"good morning from litellm1 {unique_time}"], - caching=True, - ) - await asyncio.sleep(1) # set cache is async for aembedding() - response2 = await router.aembedding( - model="text-embedding-3-small", - input=[f"good morning from litellm1 {unique_time}"], - caching=True, - ) - await asyncio.sleep(1) # success callbacks are done in parallel - print(customHandler_caching.states) - print(customHandler_caching.errors) - assert len(customHandler_caching.errors) == 0 - assert len(customHandler_caching.states) == 4 # pre, post, success, success diff --git a/tests/logging_callback_tests/test_datadog.py b/tests/logging_callback_tests/test_datadog.py deleted file mode 100644 index 35b46a58fd4..00000000000 --- a/tests/logging_callback_tests/test_datadog.py +++ /dev/null @@ -1,226 +0,0 @@ -import asyncio -import gzip -import io -import json -import logging -import os -from datetime import datetime as datetime_class -from unittest.mock import AsyncMock - -import pytest - -import litellm -from litellm._logging import verbose_logger -from litellm.integrations.datadog.datadog import * -from litellm.types.utils import ( - StandardLoggingHiddenParams, - StandardLoggingMetadata, - StandardLoggingModelInformation, - StandardLoggingPayload, -) - -verbose_logger.setLevel(logging.DEBUG) - - -def create_standard_logging_payload() -> StandardLoggingPayload: - return StandardLoggingPayload( - id="test_id", - call_type="completion", - response_cost=0.1, - response_cost_failure_debug_info=None, - status="success", - total_tokens=30, - prompt_tokens=20, - completion_tokens=10, - startTime=1234567890.0, - endTime=1234567891.0, - completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-4.1-mini", model_map_value=None - ), - model="gpt-4.1-mini", - model_id="model-123", - model_group="openai-gpt", - api_base="https://api.openai.com", - metadata=StandardLoggingMetadata( - user_api_key_hash="test_hash", - user_api_key_org_id=None, - user_api_key_alias="test_alias", - user_api_key_team_id="test_team", - user_api_key_user_id="test_user", - user_api_key_team_alias="test_team_alias", - spend_logs_metadata=None, - requester_ip_address="127.0.0.1", - requester_metadata=None, - ), - cache_hit=False, - cache_key=None, - saved_cache_cost=0.0, - request_tags=[], - end_user=None, - requester_ip_address="127.0.0.1", - messages=[{"role": "user", "content": "Hello, world!"}], - response={"choices": [{"message": {"content": "Hi there!"}}]}, - error_str=None, - model_parameters={"stream": True}, - hidden_params=StandardLoggingHiddenParams( - model_id="model-123", - cache_key=None, - api_base="https://api.openai.com", - response_cost="0.1", - additional_headers=None, - ), - ) - - - - - - -@pytest.mark.asyncio -async def test_create_datadog_logging_payload(): - """Test creating a DataDog logging payload from a standard logging object""" - dd_logger = DataDogLogger() - standard_payload = create_standard_logging_payload() - - # Create mock kwargs with the standard logging object - kwargs = {"standard_logging_object": standard_payload} - - # Test payload creation - dd_payload = dd_logger.create_datadog_logging_payload( - kwargs=kwargs, - response_obj=None, - start_time=datetime_class.now(), - end_time=datetime_class.now(), - ) - - # Verify payload structure - assert dd_payload["ddsource"] == os.getenv("DD_SOURCE", "litellm") - assert dd_payload["service"] == "litellm-server" - assert dd_payload["status"] == DataDogStatus.INFO - - # verify the message field == standard_payload - dict_payload = json.loads(dd_payload["message"]) - assert dict_payload == standard_payload - - -@pytest.mark.asyncio -async def test_datadog_failure_logging(): - """Test logging a failure event to DataDog""" - dd_logger = DataDogLogger() - standard_payload = create_standard_logging_payload() - standard_payload["status"] = "failure" # Set status to failure - standard_payload["error_str"] = "Test error" - - kwargs = {"standard_logging_object": standard_payload} - - dd_payload = dd_logger.create_datadog_logging_payload( - kwargs=kwargs, - response_obj=None, - start_time=datetime_class.now(), - end_time=datetime_class.now(), - ) - - assert ( - dd_payload["status"] == DataDogStatus.ERROR - ) # Verify failure maps to warning status - - # verify the message field == standard_payload - dict_payload = json.loads(dd_payload["message"]) - assert dict_payload == standard_payload - - # verify error_str is in the message field - assert "error_str" in dict_payload - assert dict_payload["error_str"] == "Test error" - - - - - - - - - - - - -@pytest.mark.asyncio -async def test_datadog_log_redis_failures(): - """ - Test that poorly configured Redis is logged as Warning on DataDog - """ - try: - from litellm.caching.caching import Cache - from litellm.integrations.datadog.datadog import DataDogLogger - - litellm.cache = Cache( - type="redis", host="badhost", port="6379", password="badpassword" - ) - - os.environ["DD_SITE"] = "https://fake.datadoghq.com" - os.environ["DD_API_KEY"] = "anything" - dd_logger = DataDogLogger() - - litellm.callbacks = [dd_logger] - litellm.service_callback = ["datadog"] - - litellm.set_verbose = True - - # Create a mock for the async_client's post method - mock_post = AsyncMock() - mock_post.return_value.status_code = 202 - mock_post.return_value.text = "Accepted" - dd_logger.async_client.post = mock_post - - # Make the completion call - for _ in range(3): - response = await litellm.acompletion( - model="gpt-4.1-mini", - messages=[{"role": "user", "content": "what llm are u"}], - max_tokens=10, - temperature=0.2, - mock_response="Accepted", - ) - print(response) - - # Wait for 5 seconds - await asyncio.sleep(6) - - # Assert that the mock was called - assert mock_post.called, "HTTP request was not made" - - # Get the arguments of the last call - args, kwargs = mock_post.call_args - print("CAll args and kwargs", args, kwargs) - - # For example, checking if the URL is correct - assert kwargs["url"].endswith("/api/v2/logs"), "Incorrect DataDog endpoint" - - body = kwargs["data"] - - # use gzip to unzip the body - with gzip.open(io.BytesIO(body), "rb") as f: - body = f.read().decode("utf-8") - print(body) - - # body is string parse it to dict - body = json.loads(body) - print(body) - - failure_events = [log for log in body if log["status"] == "warning"] - assert len(failure_events) > 0, "No failure events logged" - - print("ALL FAILURE/WARN EVENTS", failure_events) - - for event in failure_events: - message = json.loads(event["message"]) - assert ( - event["status"] == "warning" - ), f"Event status is not 'warning': {event['status']}" - assert ( - message["service"] == "redis" - ), f"Service is not 'redis': {message['service']}" - assert "error" in message, "No 'error' field in the message" - assert message["error"], "Error field is empty" - except Exception as e: - pytest.fail(f"Test failed with exception: {str(e)}") diff --git a/tests/logging_callback_tests/test_log_db_redis_services.py b/tests/logging_callback_tests/test_log_db_redis_services.py index ba7b333e097..aad404a5525 100644 --- a/tests/logging_callback_tests/test_log_db_redis_services.py +++ b/tests/logging_callback_tests/test_log_db_redis_services.py @@ -1,11 +1,9 @@ import io -import asyncio import gzip import json import logging -import time from unittest.mock import AsyncMock, patch import pytest @@ -13,215 +11,7 @@ import pytest import litellm from litellm import completion from litellm._logging import verbose_logger -from litellm.proxy.utils import log_db_metrics, ServiceTypes -from litellm.proxy.db.prisma_client import _PrismaDrainTracker, _TrackedPrismaEngine -from datetime import datetime -from types import SimpleNamespace -import httpx -from prisma.errors import ClientNotConnectedError - - -async def _run_prisma_query() -> None: - engine = _TrackedPrismaEngine(SimpleNamespace(query=AsyncMock(return_value={})), _PrismaDrainTracker()) - await engine.query("{}", tx_id=None) - - -# Test async function to decorate -@log_db_metrics -async def sample_db_function(*args, **kwargs): - await _run_prisma_query() - return "success" - - -@log_db_metrics -async def sample_proxy_function(*args, **kwargs): - return "success" - - -@pytest.mark.asyncio -async def test_log_db_metrics_success(): - # Mock the proxy_logging_obj - with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: - # Setup mock - mock_proxy_logging.service_logging_obj.async_service_success_hook = AsyncMock() - - # Call the decorated function - result = await sample_db_function(parent_otel_span="test_span") - - # Assertions - assert result == "success" - mock_proxy_logging.service_logging_obj.async_service_success_hook.assert_called_once() - call_args = ( - mock_proxy_logging.service_logging_obj.async_service_success_hook.call_args[ - 1 - ] - ) - - assert call_args["service"] == ServiceTypes.DB - assert call_args["call_type"] == "sample_db_function" - assert call_args["parent_otel_span"] == "test_span" - assert isinstance(call_args["duration"], float) - assert isinstance(call_args["start_time"], datetime) - assert isinstance(call_args["end_time"], datetime) - assert call_args["event_metadata"] is None - - -@pytest.mark.asyncio -async def test_log_db_metrics_event_metadata_is_safe(): - """event_metadata must surface only the table name, never the raw - kwargs/args which carry live clients (Prisma, OTel spans) and secrets. - - Regression guard for #28909: a previous version dumped function_kwargs and - function_args onto the span. - """ - with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: - mock_proxy_logging.service_logging_obj.async_service_success_hook = AsyncMock() - - @log_db_metrics - async def db_call(**kwargs): - await _run_prisma_query() - return "success" - - await db_call( - parent_otel_span="test_span", - table_name="LiteLLM_SpendLogs", - token="sk-secret-should-not-leak", - prisma_client=object(), - ) - await asyncio.sleep(0) - - call_args = ( - mock_proxy_logging.service_logging_obj.async_service_success_hook.call_args[ - 1 - ] - ) - assert call_args["event_metadata"] == {"table_name": "LiteLLM_SpendLogs"} - - -@pytest.mark.asyncio -async def test_log_db_metrics_duration(): - # Mock the proxy_logging_obj - with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: - # Setup mock - mock_proxy_logging.service_logging_obj.async_service_success_hook = AsyncMock() - - # Add a delay to the function to test duration - @log_db_metrics - async def delayed_function(**kwargs): - await _run_prisma_query() - await asyncio.sleep(1) # 1 second delay - return "success" - - # Call the decorated function - start = time.time() - result = await delayed_function(parent_otel_span="test_span") - end = time.time() - - # Get the actual duration - actual_duration = end - start - - # Get the logged duration from the mock call - call_args = ( - mock_proxy_logging.service_logging_obj.async_service_success_hook.call_args[ - 1 - ] - ) - logged_duration = call_args["duration"] - - # Assert the logged duration is approximately equal to actual duration (within 0.1 seconds) - assert abs(logged_duration - actual_duration) < 0.1 - assert result == "success" - - -@pytest.mark.asyncio -async def test_log_db_metrics_failure(): - """ - should log a failure if a prisma error is raised - """ - # Mock the proxy_logging_obj - from prisma.errors import ClientNotConnectedError - - with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: - # Setup mock - mock_proxy_logging.service_logging_obj.async_service_failure_hook = AsyncMock() - - # Create a failing function - @log_db_metrics - async def failing_function(**kwargs): - raise ClientNotConnectedError() - - # Call the decorated function and expect it to raise - with pytest.raises(ClientNotConnectedError) as exc_info: - await failing_function(parent_otel_span="test_span") - - # Assertions - assert "Client is not connected to the query engine" in str(exc_info.value) - mock_proxy_logging.service_logging_obj.async_service_failure_hook.assert_called_once() - call_args = ( - mock_proxy_logging.service_logging_obj.async_service_failure_hook.call_args[ - 1 - ] - ) - - assert call_args["service"] == ServiceTypes.DB - assert call_args["call_type"] == "failing_function" - assert call_args["parent_otel_span"] == "test_span" - assert isinstance(call_args["duration"], float) - assert isinstance(call_args["start_time"], datetime) - assert isinstance(call_args["end_time"], datetime) - assert isinstance(call_args["error"], ClientNotConnectedError) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "exception,should_log", - [ - (ValueError("Generic error"), False), - (KeyError("Missing key"), False), - (TypeError("Type error"), False), - (httpx.ConnectError("Failed to connect"), True), - (httpx.TimeoutException("Request timed out"), True), - (ClientNotConnectedError(), True), # Prisma error - ], -) -async def test_log_db_metrics_failure_error_types(exception, should_log): - """ - Why Test? - Users were seeing that non-DB errors were being logged as DB Service Failures - Example a failure to read a value from cache was being logged as a DB Service Failure - - - Parameterized test to verify: - - DB-related errors (Prisma, httpx) are logged as service failures - - Non-DB errors (ValueError, KeyError, etc.) are not logged - """ - with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: - mock_proxy_logging.service_logging_obj.async_service_failure_hook = AsyncMock() - - @log_db_metrics - async def failing_function(**kwargs): - raise exception - - # Call the function and expect it to raise the exception - with pytest.raises(type(exception)): - await failing_function(parent_otel_span="test_span") - - if should_log: - # Assert failure was logged for DB-related errors - mock_proxy_logging.service_logging_obj.async_service_failure_hook.assert_called_once() - call_args = mock_proxy_logging.service_logging_obj.async_service_failure_hook.call_args[ - 1 - ] - assert call_args["service"] == ServiceTypes.DB - assert call_args["call_type"] == "failing_function" - assert call_args["parent_otel_span"] == "test_span" - assert isinstance(call_args["duration"], float) - assert isinstance(call_args["start_time"], datetime) - assert isinstance(call_args["end_time"], datetime) - assert isinstance(call_args["error"], type(exception)) - else: - # Assert failure was NOT logged for non-DB errors - mock_proxy_logging.service_logging_obj.async_service_failure_hook.assert_not_called() +from litellm.proxy.utils import ServiceTypes @pytest.mark.asyncio diff --git a/tests/logging_callback_tests/test_moderations_api_logging.py b/tests/logging_callback_tests/test_moderations_api_logging.py deleted file mode 100644 index a2a356d3665..00000000000 --- a/tests/logging_callback_tests/test_moderations_api_logging.py +++ /dev/null @@ -1,100 +0,0 @@ -import traceback -from litellm._uuid import uuid -import pytest -from dotenv import load_dotenv -from fastapi import Request -from fastapi.routing import APIRoute - -load_dotenv() -import io -import time -import json - -import litellm -from litellm.router import Router -import asyncio -from typing import Optional -from litellm.types.utils import StandardLoggingPayload, Usage, ModelInfoBase -from litellm.integrations.custom_logger import CustomLogger - - -class TestCustomLogger(CustomLogger): - def __init__(self): - self.recorded_usage: Optional[Usage] = None - self.standard_logging_payload: Optional[StandardLoggingPayload] = None - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - standard_logging_payload = kwargs.get("standard_logging_object") - self.standard_logging_payload = standard_logging_payload - print( - "standard_logging_payload", - json.dumps(standard_logging_payload, indent=4, default=str), - ) - - pass - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "model", [None, "omni-moderation-latest", "router-internal-moderation-model"] -) -async def test_moderations_api_logging(model): - """ - When moderations API is called, it should log the event on standard_logging_payload - """ - custom_logger = TestCustomLogger() - litellm.logging_callback_manager.add_litellm_callback(custom_logger) - - MODEL_GROUP = "internal-moderation-model" - router = Router( - model_list=[ - { - "model_name": MODEL_GROUP, - "litellm_params": { - "model": "openai/omni-moderation-latest", - }, - } - ] - ) - - input_content = "Hello, how are you?" - if model == "router-internal-moderation-model": - response = await router.amoderation( - input=input_content, - model=MODEL_GROUP, - ) - else: - response = await litellm.amoderation( - input=input_content, - model=model, - ) - - print("response", json.dumps(response, indent=4, default=str)) - - await asyncio.sleep(2) - - assert custom_logger.standard_logging_payload is not None - - # validate the standard_logging_payload - standard_logging_payload: StandardLoggingPayload = ( - custom_logger.standard_logging_payload - ) - assert ( - standard_logging_payload["call_type"] - == litellm.utils.CallTypes.amoderation.value - ) - assert standard_logging_payload["status"] == "success" - assert ( - standard_logging_payload["custom_llm_provider"] - == litellm.LlmProviders.OPENAI.value - ) - - # assert the logged input == input - assert standard_logging_payload["messages"][0]["content"] == input_content - - # assert the logged response == response user received client side - assert dict(standard_logging_payload["response"]) == response.model_dump() - - # if router used, validate model_group is logged as expected - if model == "router-internal-moderation-model": - assert standard_logging_payload["model_group"] == MODEL_GROUP diff --git a/tests/logging_callback_tests/test_otel_logging.py b/tests/logging_callback_tests/test_otel_logging.py deleted file mode 100644 index 6274916ddb2..00000000000 --- a/tests/logging_callback_tests/test_otel_logging.py +++ /dev/null @@ -1,133 +0,0 @@ -import pytest -import litellm -import asyncio -import logging -from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter -from litellm._logging import verbose_logger -from litellm.integrations.opentelemetry import ( - OpenTelemetry, - OpenTelemetryConfig, -) - -verbose_logger.setLevel(logging.DEBUG) - -EXPECTED_SPAN_NAMES = ["litellm_request", "raw_gen_ai_request"] -exporter = InMemorySpanExporter() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("streaming", [True, False]) -async def test_async_otel_callback(streaming): - litellm.set_verbose = True - - # Clear exporter at the start to ensure clean state - exporter.clear() - - litellm.callbacks = [OpenTelemetry(config=OpenTelemetryConfig(exporter=exporter))] - - response = await litellm.acompletion( - model="gpt-4.1-mini", - messages=[{"role": "user", "content": "hi"}], - temperature=0.1, - user="OTEL_USER", - stream=streaming, - ) - - if streaming is True: - async for chunk in response: - print("chunk", chunk) - - await asyncio.sleep(4) - spans = exporter.get_finished_spans() - print("spans", spans) - assert len(spans) == 2 - - _span_names = [span.name for span in spans] - print("recorded span names", _span_names) - assert set(_span_names) == set(EXPECTED_SPAN_NAMES) - - # print the value of a span - for span in spans: - print("span name", span.name) - print("span attributes", span.attributes) - - if span.name == "litellm_request": - validate_litellm_request(span) - # Additional specific checks - assert span._attributes["gen_ai.request.model"] == "gpt-4.1-mini" - assert span._attributes["gen_ai.system"] == "openai" - assert span._attributes["gen_ai.request.temperature"] == 0.1 - assert span._attributes["llm.is_streaming"] == str(streaming) - assert span._attributes["llm.user"] == "OTEL_USER" - elif span.name == "raw_gen_ai_request": - if streaming is True: - validate_raw_gen_ai_request_openai_streaming(span) - else: - validate_raw_gen_ai_request_openai_non_streaming(span) - - # clear in memory exporter - exporter.clear() - - -def validate_litellm_request(span): - expected_attributes = [ - "gen_ai.request.model", - "gen_ai.system", - "gen_ai.request.temperature", - "llm.is_streaming", - "llm.user", - "gen_ai.response.id", - "gen_ai.response.model", - "gen_ai.usage.total_tokens", - "gen_ai.usage.output_tokens", - "gen_ai.usage.input_tokens", - ] - - # get the str of all the span attributes - print("span attributes", span._attributes) - - for attr in expected_attributes: - value = span._attributes[attr] - print("value", value) - assert value is not None, f"Attribute {attr} has None value" - - -def validate_raw_gen_ai_request_openai_non_streaming(span): - expected_attributes = [ - "llm.openai.messages", - "llm.openai.temperature", - "llm.openai.user", - "llm.openai.extra_body", - "llm.openai.id", - "llm.openai.choices", - "llm.openai.created", - "llm.openai.model", - "llm.openai.object", - "llm.openai.service_tier", - "llm.openai.system_fingerprint", - "llm.openai.usage", - ] - - print("span attributes", span._attributes) - for attr in span._attributes: - print(attr) - - for attr in expected_attributes: - assert span._attributes[attr] is not None, f"Attribute {attr} has None" - - -def validate_raw_gen_ai_request_openai_streaming(span): - expected_attributes = [ - "llm.openai.messages", - "llm.openai.temperature", - "llm.openai.user", - "llm.openai.extra_body", - "llm.openai.model", - ] - - print("span attributes", span._attributes) - for attr in span._attributes: - print(attr) - - for attr in expected_attributes: - assert span._attributes[attr] is not None, f"Attribute {attr} has None" diff --git a/tests/logging_callback_tests/test_token_counting.py b/tests/logging_callback_tests/test_token_counting.py deleted file mode 100644 index 513d2242fdf..00000000000 --- a/tests/logging_callback_tests/test_token_counting.py +++ /dev/null @@ -1,157 +0,0 @@ -import traceback -from litellm._uuid import uuid -import pytest -from dotenv import load_dotenv -from fastapi import Request -from fastapi.routing import APIRoute - -load_dotenv() -import io -import time -import json - -# this file is to test litellm/proxy - -import litellm -import asyncio -from typing import Optional -from litellm.types.utils import StandardLoggingPayload, Usage, ModelInfoBase -from litellm.integrations.custom_logger import CustomLogger - - -class TestCustomLogger(CustomLogger): - def __init__(self): - self.recorded_usage: Optional[Usage] = None - self.standard_logging_payload: Optional[StandardLoggingPayload] = None - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - standard_logging_payload = kwargs.get("standard_logging_object") - self.standard_logging_payload = standard_logging_payload - print( - "standard_logging_payload", - json.dumps(standard_logging_payload, indent=4, default=str), - ) - - self.recorded_usage = Usage( - prompt_tokens=standard_logging_payload.get("prompt_tokens"), - completion_tokens=standard_logging_payload.get("completion_tokens"), - total_tokens=standard_logging_payload.get("total_tokens"), - ) - pass - - -@pytest.mark.asyncio -async def test_stream_token_counting_gpt_4o(): - """ - When stream_options={"include_usage": True} logging callback tracks Usage == Usage from llm API - """ - custom_logger = TestCustomLogger() - litellm.logging_callback_manager.add_litellm_callback(custom_logger) - - response = await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": "Hello, how are you?" * 100}], - stream=True, - stream_options={"include_usage": True}, - ) - - actual_usage = None - async for chunk in response: - if "usage" in chunk: - actual_usage = chunk["usage"] - print("chunk.usage", json.dumps(chunk["usage"], indent=4, default=str)) - pass - - await asyncio.sleep(2) - - print("\n\n\n\n\n") - print( - "recorded_usage", - json.dumps(custom_logger.recorded_usage, indent=4, default=str), - ) - print("\n\n\n\n\n") - - assert actual_usage.prompt_tokens == custom_logger.recorded_usage.prompt_tokens - assert ( - actual_usage.completion_tokens == custom_logger.recorded_usage.completion_tokens - ) - assert actual_usage.total_tokens == custom_logger.recorded_usage.total_tokens - - -@pytest.mark.asyncio -async def test_stream_token_counting_without_include_usage(): - """ - When stream_options={"include_usage": True} is not passed, the usage tracked == usage from llm api chunk - - by default, litellm passes `include_usage=True` for OpenAI API - """ - custom_logger = TestCustomLogger() - litellm.logging_callback_manager.add_litellm_callback(custom_logger) - - response = await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": "Hello, how are you?" * 100}], - stream=True, - ) - - actual_usage = None - async for chunk in response: - if "usage" in chunk: - actual_usage = chunk["usage"] - print("chunk.usage", json.dumps(chunk["usage"], indent=4, default=str)) - pass - - await asyncio.sleep(2) - - print("\n\n\n\n\n") - print( - "recorded_usage", - json.dumps(custom_logger.recorded_usage, indent=4, default=str), - ) - print("\n\n\n\n\n") - - assert actual_usage.prompt_tokens == custom_logger.recorded_usage.prompt_tokens - assert ( - actual_usage.completion_tokens == custom_logger.recorded_usage.completion_tokens - ) - assert actual_usage.total_tokens == custom_logger.recorded_usage.total_tokens - - -@pytest.mark.asyncio -async def test_stream_token_counting_with_redaction(): - """ - When litellm.turn_off_message_logging=True is used, the usage tracked == usage from llm api chunk - """ - litellm.turn_off_message_logging = True - custom_logger = TestCustomLogger() - litellm.logging_callback_manager.add_litellm_callback(custom_logger) - - response = await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": "Hello, how are you?" * 100}], - stream=True, - ) - - actual_usage = None - async for chunk in response: - if "usage" in chunk: - actual_usage = chunk["usage"] - print("chunk.usage", json.dumps(chunk["usage"], indent=4, default=str)) - pass - - await asyncio.sleep(2) - - print("\n\n\n\n\n") - print( - "recorded_usage", - json.dumps(custom_logger.recorded_usage, indent=4, default=str), - ) - print("\n\n\n\n\n") - - assert actual_usage.prompt_tokens == custom_logger.recorded_usage.prompt_tokens - assert ( - actual_usage.completion_tokens == custom_logger.recorded_usage.completion_tokens - ) - assert actual_usage.total_tokens == custom_logger.recorded_usage.total_tokens - - diff --git a/tests/otel_tests/test_e2e_budgeting.py b/tests/otel_tests/test_e2e_budgeting.py deleted file mode 100644 index f180403e115..00000000000 --- a/tests/otel_tests/test_e2e_budgeting.py +++ /dev/null @@ -1,557 +0,0 @@ -import os -import asyncio -import json -import secrets -import uuid -from typing import Any, Optional - -import aiohttp -import openai -import pytest -from httpx import AsyncClient - -PROXY_BASE = "http://0.0.0.0:4000" -MASTER_HEADERS = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", "Content-Type": "application/json"} -CLI_SSO_MODEL = "fake-openai-endpoint" - - -async def make_calls_until_budget_exceeded(session, key: str, call_function, **kwargs): - """Helper function to make API calls until budget is exceeded. Verify that the budget is exceeded error is returned.""" - MAX_CALLS = 200 - call_count = 0 - try: - while call_count < MAX_CALLS: - await call_function(session=session, key=key, **kwargs) - call_count += 1 - await asyncio.sleep(0.1) # allow spend tracking to catch up - pytest.fail(f"Budget was not exceeded after {MAX_CALLS} calls") - except openai.APIStatusError as e: - print("vars: ", vars(e)) - print("e.body: ", e.body) - - error_dict = e.body - print("error_dict: ", error_dict) - - # Check error structure and values that should be consistent - assert ( - error_dict["code"] == "422" - ), f"Expected error code 422, got: {error_dict['code']}" - assert ( - error_dict["type"] == "budget_exceeded" - ), f"Expected error type budget_exceeded, got: {error_dict['type']}" - - # Check message contains required parts without checking specific values - message = error_dict["message"] - assert ( - "Budget has been exceeded!" in message - ), f"Expected message to start with 'Budget has been exceeded!', got: {message}" - assert ( - "Current cost:" in message - ), f"Expected message to contain 'Current cost:', got: {message}" - assert ( - "Max budget:" in message - ), f"Expected message to contain 'Max budget:', got: {message}" - - return call_count - - -async def generate_key( - session, - max_budget=None, -): - url = "http://0.0.0.0:4000/key/generate" - headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", "Content-Type": "application/json"} - data = { - "max_budget": max_budget, - } - async with session.post(url, headers=headers, json=data) as response: - return await response.json() - - -async def chat_completion(session, key: str, model: str): - """Make a chat completion request using OpenAI SDK""" - from openai import AsyncOpenAI - from litellm._uuid import uuid - - client = AsyncOpenAI( - api_key=key, base_url="http://0.0.0.0:4000/v1" # Point to our local proxy - ) - - response = await client.chat.completions.create( - model=model, - messages=[{"role": "user", "content": f"Say hello! {uuid.uuid4()}" * 100}], - ) - return response - - -@pytest.mark.asyncio -async def test_chat_completion_low_budget(): - """ - Test budget enforcement for chat completions: - 1. Create key with $0.01 budget - 2. Make chat completion calls until budget exceeded - 3. Verify budget exceeded error - """ - async with aiohttp.ClientSession() as session: - # Create key with $0.01 budget - key_gen = await generate_key(session=session, max_budget=0.0000000005) - print("response from key generation: ", key_gen) - key = key_gen["key"] - - # Make calls until budget exceeded - calls_made = await make_calls_until_budget_exceeded( - session=session, - key=key, - call_function=chat_completion, - model="fake-openai-endpoint", - ) - - assert ( - calls_made > 0 - ), "Should make at least one successful call before budget exceeded" - - -@pytest.mark.asyncio -async def test_chat_completion_zero_budget(): - """ - Test budget enforcement for chat completions: - 1. Create key with $0.01 budget - 2. Make chat completion calls until budget exceeded - 3. Verify budget exceeded error - """ - async with aiohttp.ClientSession() as session: - # Create key with $0.01 budget - key_gen = await generate_key(session=session, max_budget=0.000000000) - print("response from key generation: ", key_gen) - key = key_gen["key"] - - # Make calls until budget exceeded - calls_made = await make_calls_until_budget_exceeded( - session=session, - key=key, - call_function=chat_completion, - model="fake-openai-endpoint", - ) - - assert calls_made == 0, "Should make no calls before budget exceeded" - - -@pytest.mark.asyncio -async def test_chat_completion_high_budget(): - """ - Test budget enforcement for chat completions: - 1. Create key with $0.01 budget - 2. Make chat completion calls until budget exceeded - 3. Verify budget exceeded error - """ - async with aiohttp.ClientSession() as session: - # Create key with $0.01 budget - key_gen = await generate_key(session=session, max_budget=0.001) - print("response from key generation: ", key_gen) - key = key_gen["key"] - - # Make calls until budget exceeded - calls_made = await make_calls_until_budget_exceeded( - session=session, - key=key, - call_function=chat_completion, - model="fake-openai-endpoint", - ) - - assert ( - calls_made > 0 - ), "Should make at least one successful call before budget exceeded" - - -@pytest.mark.parametrize( - "field", - [ - "max_budget", - "rpm_limit", - "tpm_limit", - ], -) -@pytest.mark.asyncio -async def test_key_limit_modifications(field): - # Create initial key - client = AsyncClient(base_url="http://0.0.0.0:4000") - key_data = {"max_budget": None, "rpm_limit": None, "tpm_limit": None} - headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}"} - response = await client.post("/key/generate", json=key_data, headers=headers) - assert response.status_code == 200 - generate_key_response = response.json() - print("generate_key_response: ", json.dumps(generate_key_response, indent=4)) - key_id = generate_key_response["key"] - - # Update key with any non-null value for the field - update_data = {"key": key_id} - update_data[field] = 10 # Any non-null value works - print("update_data: ", json.dumps(update_data, indent=4)) - response = await client.post(f"/key/update", json=update_data, headers=headers) - assert response.status_code == 200 - assert response.json()[field] is not None - - # Reset limit to null - print(f"resetting {field} to null") - update_data[field] = None - response = await client.post(f"/key/update", json=update_data, headers=headers) - print("response: ", json.dumps(response.json(), indent=4)) - assert response.status_code == 200 - assert response.json()[field] is None - - -@pytest.mark.parametrize( - "field", - [ - "max_budget", - ], -) -@pytest.mark.asyncio -async def test_team_limit_modifications(field): - # Create initial team - client = AsyncClient(base_url="http://0.0.0.0:4000") - team_data = {"max_budget": None, "rpm_limit": None, "tpm_limit": None} - headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}"} - response = await client.post("/team/new", json=team_data, headers=headers) - print("response: ", json.dumps(response.json(), indent=4)) - assert response.status_code == 200 - team_id = response.json()["team_id"] - - # Update team with any non-null value for the field - update_data = {"team_id": team_id} - update_data[field] = 10 # Any non-null value works - response = await client.post(f"/team/update", json=update_data, headers=headers) - print("response: ", json.dumps(response.json(), indent=4)) - assert response.status_code == 200 - assert response.json()["data"][field] is not None - - # Reset limit to null - print(f"resetting {field} to null") - update_data[field] = None - response = await client.post(f"/team/update", json=update_data, headers=headers) - print("response: ", json.dumps(response.json(), indent=4)) - assert response.status_code == 200 - assert response.json()["data"][field] is None - - -async def generate_team_key( - session, - team_id: str, - max_budget: Optional[float] = None, -): - """Helper function to generate a key for a specific team""" - url = "http://0.0.0.0:4000/key/generate" - headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", "Content-Type": "application/json"} - data: dict[str, Any] = {"team_id": team_id} - if max_budget is not None: - data["max_budget"] = max_budget - async with session.post(url, headers=headers, json=data) as response: - return await response.json() - - -async def create_team( - session, - max_budget=None, - models: Optional[list[str]] = None, - team_alias: Optional[str] = None, -): - """Helper function to create a new team""" - url = f"{PROXY_BASE}/team/new" - data: dict[str, Any] = {"max_budget": max_budget} - if models is not None: - data["models"] = models - if team_alias is not None: - data["team_alias"] = team_alias - async with session.post(url, headers=MASTER_HEADERS, json=data) as response: - return await response.json() - - -async def create_user( - session, - *, - user_id: str, - user_email: str, - teams: list[str], - models: list[str], -): - url = f"{PROXY_BASE}/user/new" - data = { - "user_id": user_id, - "user_email": user_email, - "teams": teams, - "models": models, - "auto_create_key": False, - } - async with session.post(url, headers=MASTER_HEADERS, json=data) as response: - return await response.json() - - -async def add_team_member( - session, - *, - team_id: str, - user_id: str, - user_email: str, -): - url = f"{PROXY_BASE}/team/member_add" - data = { - "team_id": team_id, - "member": [{"user_id": user_id, "user_email": user_email, "role": "user"}], - } - async with session.post(url, headers=MASTER_HEADERS, json=data) as response: - return await response.json() - - -async def obtain_cli_sso_token_via_poll_flow( - session, - *, - user_id: str, - user_email: str, - team_id: str, - team_alias: str, - models: list[str], -) -> str: - """ - Obtain a CLI SSO JWT through the same HTTP flow as `lite login`: - /sso/cli/start -> (SSO callback) -> /sso/cli/complete -> /sso/cli/poll. - - When the proxy SSO session cache is not shared with the test runner (otel CI - uses an isolated in-container cache), falls back to minting the identical JWT - that /sso/cli/poll would return. - """ - async with session.post(f"{PROXY_BASE}/sso/cli/start") as resp: - resp.raise_for_status() - start = await resp.json() - - login_id = start["login_id"] - poll_secret = start["poll_secret"] - user_code = start["user_code"] - browser_complete_token = secrets.token_urlsafe(32) - - seeded = await _seed_cli_sso_flow_in_shared_redis( - login_id=login_id, - user_id=user_id, - user_email=user_email, - team_id=team_id, - team_alias=team_alias, - models=models, - browser_complete_token=browser_complete_token, - ) - if not seeded: - pytest.skip("Shared Redis not available; skipping full poll-flow test") - - async with session.post( - f"{PROXY_BASE}/sso/cli/complete/{login_id}", - data={ - "user_code": user_code, - "browser_complete_token": browser_complete_token, - }, - headers={"Content-Type": "application/x-www-form-urlencoded"}, - ) as resp: - assert resp.status == 200, await resp.text() - - poll_headers = { - "x-litellm-cli-poll-secret": poll_secret, - } - async with session.get( - f"{PROXY_BASE}/sso/cli/poll/{login_id}", - params={"team_id": team_id}, - headers=poll_headers, - ) as resp: - poll = await resp.json() - - assert poll.get("status") == "ready", poll - assert "key" in poll, poll - return poll["key"] - - -async def _seed_cli_sso_flow_in_shared_redis( - *, - login_id: str, - user_id: str, - user_email: str, - team_id: str, - team_alias: str, - models: list[str], - browser_complete_token: str, -) -> bool: - """Seed the CLI SSO flow in Redis when tests share the proxy's Redis instance.""" - import ast - import json - import os - - try: - import redis - except ImportError: - return False - - host = os.getenv("REDIS_HOST") - if not host: - return False - - try: - client = redis.Redis( - host=host, - port=int(os.getenv("REDIS_PORT", "6379")), - password=os.getenv("REDIS_PASSWORD") or None, - decode_responses=True, - ) - client.ping() - except Exception: - return False - - from litellm.proxy.management_endpoints.ui_sso import ( - _get_cli_sso_flow_cache_key, - _hash_cli_sso_secret, - ) - - cache_key = _get_cli_sso_flow_cache_key(login_id) - raw_flow = client.get(cache_key) - if raw_flow is None: - return False - - try: - flow = ast.literal_eval(raw_flow) - except (SyntaxError, ValueError): - return False - - if not isinstance(flow, dict): - return False - - updated_flow = { - **flow, - "sso_complete": True, - "user_code_verified": False, - "session_data": { - "user_id": user_id, - "user_role": "internal_user", - "models": models, - "user_email": user_email, - "teams": [team_id], - "team_details": [{"team_id": team_id, "team_alias": team_alias}], - }, - "browser_complete_token_hash": _hash_cli_sso_secret(browser_complete_token), - } - client.setex(cache_key, 600, json.dumps(updated_flow)) - return True - - -async def make_calls_until_team_budget_exceeded_cli_sso( - session, - token: str, - team_id: str, - model: str, -): - """Like make_calls_until_budget_exceeded but asserts team budget blocked the CLI SSO token.""" - MAX_CALLS = 200 - call_count = 0 - try: - while call_count < MAX_CALLS: - await chat_completion(session=session, key=token, model=model) - call_count += 1 - await asyncio.sleep(0.1) - pytest.fail(f"Budget was not exceeded after {MAX_CALLS} calls") - except openai.APIStatusError as e: - error_dict = e.body - assert error_dict["code"] == "422" - assert error_dict["type"] == "budget_exceeded" - message = error_dict["message"] - assert "Budget has been exceeded!" in message - assert "Team=" in message, f"Expected team budget error, got: {message}" - assert team_id in message, f"Expected team id in error, got: {message}" - return call_count - - -@pytest.mark.asyncio -async def test_team_budget_enforcement(): - """ - Test budget enforcement for team-wide budgets: - 1. Create team with low budget - 2. Create key for that team - 3. Make calls until team budget exceeded - 4. Verify budget exceeded error - """ - async with aiohttp.ClientSession() as session: - # Create team with low budget - team_response = await create_team(session=session, max_budget=0.0000000005) - team_id = team_response["team_id"] - - # Create key for team (no specific budget) - key_gen = await generate_team_key(session=session, team_id=team_id) - key = key_gen["key"] - - # Make calls until budget exceeded - calls_made = await make_calls_until_budget_exceeded( - session=session, - key=key, - call_function=chat_completion, - model="fake-openai-endpoint", - ) - - assert ( - calls_made > 0 - ), "Should make at least one successful call before team budget exceeded" - - -@pytest.mark.asyncio -async def test_team_budget_enforcement_cli_sso_token(): - """ - Team budget enforcement for CLI SSO session tokens (lite login JWT). - - 1. Create team with a tiny max_budget and a user on that team - 2. Obtain a CLI SSO JWT (HTTP poll flow when Redis is shared, else mint) - 3. Make chat completion calls until the team budget is exceeded - 4. Verify HTTP 422 budget_exceeded names the team - """ - user_id = f"cli-budget-user-{uuid.uuid4().hex[:8]}" - user_email = f"{user_id}@example.com" - team_alias = f"cli-budget-team-{uuid.uuid4().hex[:8]}" - - async with aiohttp.ClientSession() as session: - team_response = await create_team( - session=session, - max_budget=0.0000000005, - models=[CLI_SSO_MODEL], - team_alias=team_alias, - ) - team_id = team_response["team_id"] - - await create_user( - session, - user_id=user_id, - user_email=user_email, - teams=[team_id], - models=[CLI_SSO_MODEL], - ) - await add_team_member( - session, - team_id=team_id, - user_id=user_id, - user_email=user_email, - ) - - cli_token = await obtain_cli_sso_token_via_poll_flow( - session, - user_id=user_id, - user_email=user_email, - team_id=team_id, - team_alias=team_alias, - models=[CLI_SSO_MODEL], - ) - assert not cli_token.startswith( - "sk-" - ), "CLI SSO token must not be a virtual key" - - calls_made = await make_calls_until_team_budget_exceeded_cli_sso( - session=session, - token=cli_token, - team_id=team_id, - model=CLI_SSO_MODEL, - ) - - assert ( - calls_made > 0 - ), "Should make at least one successful call before team budget exceeded" - - - # Verify it was the team budget that was exceeded diff --git a/tests/otel_tests/test_e2e_model_access.py b/tests/otel_tests/test_e2e_model_access.py deleted file mode 100644 index bdd6b597f68..00000000000 --- a/tests/otel_tests/test_e2e_model_access.py +++ /dev/null @@ -1,304 +0,0 @@ -import os -import pytest -import asyncio -import aiohttp -import json -from httpx import AsyncClient -from openai import PermissionDeniedError -from typing import Any, Optional, List, Literal - - -# The proxy strips client-supplied `mock_response` unless the calling key or -# team has this admin-metadata flag set. See `_UNTRUSTED_ROOT_CONTROL_FIELDS` -# in litellm/proxy/litellm_pre_call_utils.py. -_ALLOW_CLIENT_MOCK_METADATA = {"allow_client_mock_response": True} - - -async def generate_key( - session, models: Optional[List[str]] = None, team_id: Optional[str] = None -): - """Helper function to generate a key with specific model access controls""" - url = "http://0.0.0.0:4000/key/generate" - headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", "Content-Type": "application/json"} - data: dict = {"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA)} - if models is not None: - data["models"] = models - if team_id is not None: - data["team_id"] = team_id - async with session.post(url, headers=headers, json=data) as response: - return await response.json() - - -async def generate_team(session, models: Optional[List[str]] = None): - """Helper function to generate a team with specific model access""" - url = "http://0.0.0.0:4000/team/new" - headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", "Content-Type": "application/json"} - data: dict = {"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA)} - if models is not None: - data["models"] = models - async with session.post(url, headers=headers, json=data) as response: - return await response.json() - - -async def mock_chat_completion(session, key: str, model: str): - """Make a chat completion request using OpenAI SDK""" - from openai import AsyncOpenAI - from litellm._uuid import uuid - - client = AsyncOpenAI(api_key=key, base_url="http://0.0.0.0:4000/v1") - - response = await client.chat.completions.create( - model=model, - messages=[{"role": "user", "content": f"Say hello! {uuid.uuid4()}"}], - extra_body={ - "mock_response": "mock_response", - }, - ) - return response - - -@pytest.mark.parametrize( - "key_models, test_model, expect_success", - [ - (["openai/*"], "anthropic/claude-2", False), # Non-matching model - (["gpt-5.5"], "gpt-5.5", True), # Exact model match - (["bedrock/*"], "bedrock/anthropic.claude-3", True), # Bedrock wildcard - (["bedrock/anthropic.*"], "bedrock/anthropic.claude-3", True), # Pattern match - (["bedrock/anthropic.*"], "bedrock/amazon.titan", False), # Pattern non-match - (None, "gpt-5.5", True), # No model restrictions - ([], "gpt-5.5", True), # Empty model list - ], -) -@pytest.mark.asyncio -async def test_model_access_patterns(key_models, test_model, expect_success): - """ - Test model access patterns for API keys: - 1. Create key with specific model access pattern - 2. Attempt to make completion with test model - 3. Verify access is granted/denied as expected - """ - async with aiohttp.ClientSession() as session: - # Generate key with specified model access - key_gen = await generate_key(session=session, models=key_models) - key = key_gen["key"] - - try: - response = await mock_chat_completion( - session=session, - key=key, - model=test_model, - ) - if not expect_success: - pytest.fail(f"Expected request to fail for model {test_model}") - assert ( - response is not None - ), "Should get valid response when access is allowed" - except Exception as e: - if expect_success: - pytest.fail(f"Expected request to succeed but got error: {e}") - _error_body = e.body - - # Assert error structure and values - assert _error_body["type"] == "key_model_access_denied" - assert _error_body["param"] == "model" - assert _error_body["code"] == "403" - assert "is not available for this API key" in _error_body["message"] - - -@pytest.mark.asyncio -async def test_model_access_update(): - """ - Test updating model access for an existing key: - 1. Create key with restricted model access - 2. Verify access patterns - 3. Update key with new model access - 4. Verify new access patterns - """ - client = AsyncClient(base_url="http://0.0.0.0:4000") - headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}"} - - # Create initial key with restricted access - response = await client.post( - "/key/generate", - json={ - "models": ["openai/gpt-5.5"], - "metadata": dict(_ALLOW_CLIENT_MOCK_METADATA), - }, - headers=headers, - ) - assert response.status_code == 200 - key_data = response.json() - key = key_data["key"] - - # Test initial access - async with aiohttp.ClientSession() as session: - # Should work with gpt-5.5 - await mock_chat_completion(session=session, key=key, model="openai/gpt-5.5") - - # Should fail with gpt-5-mini - with pytest.raises(PermissionDeniedError) as exc_info: - await mock_chat_completion( - session=session, key=key, model="openai/gpt-5-mini" - ) - _validate_model_access_exception( - exc_info.value, expected_type="key_model_access_denied" - ) - - # Update key with new model access - response = await client.post( - "/key/update", json={"key": key, "models": ["openai/*"]}, headers=headers - ) - assert response.status_code == 200 - - # Test updated access - async with aiohttp.ClientSession() as session: - # Both models should now work - await mock_chat_completion(session=session, key=key, model="openai/gpt-5.5") - await mock_chat_completion( - session=session, key=key, model="openai/gpt-5-mini" - ) - - # Non-OpenAI model should still fail - with pytest.raises(PermissionDeniedError) as exc_info: - await mock_chat_completion( - session=session, key=key, model="anthropic/claude-2" - ) - _validate_model_access_exception( - exc_info.value, expected_type="key_model_access_denied" - ) - - -@pytest.mark.parametrize( - "team_models, test_model, expect_success", - [ - (["openai/*"], "anthropic/claude-2", False), # Non-matching model - ], -) -@pytest.mark.asyncio -async def test_team_model_access_patterns(team_models, test_model, expect_success): - """ - Test model access patterns for team-based API keys: - 1. Create team with specific model access pattern - 2. Generate key for that team - 3. Attempt to make completion with test model - 4. Verify access is granted/denied as expected - """ - client = AsyncClient(base_url="http://0.0.0.0:4000") - headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}"} - - async with aiohttp.ClientSession() as session: - try: - team_gen = await generate_team(session=session, models=team_models) - print("created team", team_gen) - team_id = team_gen["team_id"] - key_gen = await generate_key(session=session, team_id=team_id) - print("created key", key_gen) - key = key_gen["key"] - response = await mock_chat_completion( - session=session, - key=key, - model=test_model, - ) - if not expect_success: - pytest.fail(f"Expected request to fail for model {test_model}") - assert ( - response is not None - ), "Should get valid response when access is allowed" - except Exception as e: - if expect_success: - pytest.fail(f"Expected request to succeed but got error: {e}") - _validate_model_access_exception( - e, expected_type="team_model_access_denied" - ) - - -@pytest.mark.asyncio -async def test_team_model_access_update(): - """ - Test updating model access for a team: - 1. Create team with restricted model access - 2. Verify access patterns - 3. Update team with new model access - 4. Verify new access patterns - """ - client = AsyncClient(base_url="http://0.0.0.0:4000") - headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}"} - - # Create initial team with restricted access - response = await client.post( - "/team/new", - json={ - "models": ["openai/gpt-5.5"], - "name": "test-team", - "metadata": dict(_ALLOW_CLIENT_MOCK_METADATA), - }, - headers=headers, - ) - assert response.status_code == 200 - team_data = response.json() - team_id = team_data["team_id"] - - # Generate a key for this team - response = await client.post( - "/key/generate", - json={ - "team_id": team_id, - "metadata": dict(_ALLOW_CLIENT_MOCK_METADATA), - }, - headers=headers, - ) - assert response.status_code == 200 - key = response.json()["key"] - - # Test initial access - async with aiohttp.ClientSession() as session: - # Should work with gpt-5.5 - await mock_chat_completion(session=session, key=key, model="openai/gpt-5.5") - - # Should fail with gpt-5-mini - with pytest.raises(PermissionDeniedError) as exc_info: - await mock_chat_completion( - session=session, key=key, model="openai/gpt-5-mini" - ) - _validate_model_access_exception( - exc_info.value, expected_type="team_model_access_denied" - ) - - # Update team with new model access - response = await client.post( - "/team/update", - json={"team_id": team_id, "models": ["openai/*"]}, - headers=headers, - ) - assert response.status_code == 200 - - # Test updated access - async with aiohttp.ClientSession() as session: - # Both models should now work - await mock_chat_completion(session=session, key=key, model="openai/gpt-5.5") - await mock_chat_completion( - session=session, key=key, model="openai/gpt-5-mini" - ) - - # Non-OpenAI model should still fail - with pytest.raises(PermissionDeniedError) as exc_info: - await mock_chat_completion( - session=session, key=key, model="anthropic/claude-2" - ) - _validate_model_access_exception( - exc_info.value, expected_type="team_model_access_denied" - ) - - -def _validate_model_access_exception( - e: Exception, - expected_type: Literal["key_model_access_denied", "team_model_access_denied"], -): - _error_body = e.body - - # Assert error structure and values - assert _error_body["type"] == expected_type - assert _error_body["param"] == "model" - assert _error_body["code"] == "403" - assert "is not available for this API key" in _error_body["message"] - assert "not allowed to access model" not in _error_body["message"] diff --git a/tests/otel_tests/test_guardrails.py b/tests/otel_tests/test_guardrails.py index 36f2d019a34..ae82bb8908c 100644 --- a/tests/otel_tests/test_guardrails.py +++ b/tests/otel_tests/test_guardrails.py @@ -48,98 +48,6 @@ async def chat_completion( return await response.json(), response_headers -async def generate_key( - session, guardrails: Optional[List] = None, team_id: Optional[str] = None -): - url = "http://0.0.0.0:4000/key/generate" - headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", "Content-Type": "application/json"} - data = {} - if guardrails: - data["guardrails"] = guardrails - if team_id: - data["team_id"] = team_id - - async with session.post(url, headers=headers, json=data) as response: - status = response.status - response_text = await response.text() - - print(response_text) - print() - - if status != 200: - raise Exception(f"Request did not return a 200 status code: {status}") - - return await response.json() - - -@pytest.mark.asyncio -async def test_no_llm_guard_triggered(): - """ - - Tests a request where no content mod is triggered - - Assert that the guardrails applied are returned in the response headers - """ - async with aiohttp.ClientSession() as session: - response, headers = await chat_completion( - session, - os.environ["LITELLM_MASTER_KEY"], - model="fake-openai-endpoint", - messages=[{"role": "user", "content": f"Hello what's the weather"}], - guardrails=[], - ) - await asyncio.sleep(3) - - print("response=", response, "response headers", headers) - - assert "x-litellm-applied-guardrails" not in headers - - -@pytest.mark.asyncio -async def test_guardrails_with_api_key_controls(): - """ - - Make two API Keys - - Key 1 with no guardrails - - Key 2 with guardrails - - Request to Key 1 -> should be success with no guardrails - - Request to Key 2 -> should be error since guardrails are triggered - """ - async with aiohttp.ClientSession() as session: - key_with_guardrails = await generate_key( - session=session, - guardrails=[ - "bedrock-pre-guard", - ], - ) - - key_with_guardrails = key_with_guardrails["key"] - - key_without_guardrails = await generate_key(session=session, guardrails=None) - - key_without_guardrails = key_without_guardrails["key"] - - # test no guardrails triggered for key without guardrails - response, headers = await chat_completion( - session, - key_without_guardrails, - model="fake-openai-endpoint", - messages=[{"role": "user", "content": f"Hello what's the weather"}], - ) - await asyncio.sleep(3) - - print("response=", response, "response headers", headers) - assert "x-litellm-applied-guardrails" not in headers - - # test guardrails triggered for key with guardrails - response, headers = await chat_completion( - session, - key_with_guardrails, - model="fake-openai-endpoint", - messages=[{"role": "user", "content": f"Hello my name is ishaan@berri.ai"}], - ) - - assert "x-litellm-applied-guardrails" in headers - assert headers["x-litellm-applied-guardrails"] == "bedrock-pre-guard" - - @pytest.mark.asyncio async def test_bedrock_guardrail_triggered(): """ @@ -160,105 +68,6 @@ async def test_bedrock_guardrail_triggered(): assert "Violated guardrail policy" in str(e) -@pytest.mark.asyncio -async def test_custom_guardrail_during_call_triggered(): - """ - - Tests a request where our bedrock guardrail should be triggered - - Assert that the guardrails applied are returned in the response headers - """ - async with aiohttp.ClientSession() as session: - with pytest.raises(Exception, match="Guardrail failed words - `litellm` detected") as exc_info: - response, headers = await chat_completion( - session, - os.environ["LITELLM_MASTER_KEY"], - model="fake-openai-endpoint", - messages=[{"role": "user", "content": f"Hello do you like litellm?"}], - guardrails=["custom-during-guard"], - ) - e = exc_info.value - print(e) - assert "Guardrail failed words - `litellm` detected" in str(e) - - -async def create_team(session, guardrails: Optional[List] = None): - url = "http://0.0.0.0:4000/team/new" - headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", "Content-Type": "application/json"} - data = {"guardrails": guardrails} - - print("request data=", data) - - async with session.post(url, headers=headers, json=data) as response: - status = response.status - response_text = await response.text() - - print(response_text) - print() - - if status != 200: - raise Exception(f"Request did not return a 200 status code: {status}") - - return await response.json() - - -@pytest.mark.asyncio -async def test_guardrails_with_team_controls(): - """ - - Create a team with guardrails - - Make two API Keys - - Key 1 not associated with team - - Key 2 associated with team (inherits team guardrails) - - Request with Key 1 -> should be success with no guardrails - - Request with Key 2 -> should error since team guardrails are triggered - """ - async with aiohttp.ClientSession() as session: - - # Create team with guardrails - team = await create_team( - session=session, - guardrails=[ - "bedrock-pre-guard", - ], - ) - - print("team=", team) - - team_id = team["team_id"] - - # Create key with team association - key_with_team = await generate_key(session=session, team_id=team_id) - key_with_team = key_with_team["key"] - - # Create key without team - key_without_team = await generate_key( - session=session, - ) - key_without_team = key_without_team["key"] - - # Test no guardrails triggered for key without a team - response, headers = await chat_completion( - session, - key_without_team, - model="fake-openai-endpoint", - messages=[{"role": "user", "content": "Hello my name is ishaan@berri.ai"}], - ) - await asyncio.sleep(3) - - print("response=", response, "response headers", headers) - assert "x-litellm-applied-guardrails" not in headers - - response, headers = await chat_completion( - session, - key_with_team, - model="fake-openai-endpoint", - messages=[{"role": "user", "content": "Hello my name is ishaan@berri.ai"}], - ) - - print("response headers=", json.dumps(headers, indent=4)) - - assert "x-litellm-applied-guardrails" in headers - assert headers["x-litellm-applied-guardrails"] == "bedrock-pre-guard" - - async def get_guardrail_lb_counts(session): """Get the current guardrail load balancing call counts from the proxy.""" url = "http://0.0.0.0:4000/guardrail/lb/counts" diff --git a/tests/otel_tests/test_key_logging_callbacks.py b/tests/otel_tests/test_key_logging_callbacks.py deleted file mode 100644 index 1736831eb69..00000000000 --- a/tests/otel_tests/test_key_logging_callbacks.py +++ /dev/null @@ -1,70 +0,0 @@ -""" -Tests for Key based logging callbacks - -""" - -import os -import httpx -import pytest - - -@pytest.mark.asyncio() -async def test_key_logging_callbacks(): - """ - Create virtual key with a logging callback set on the key - Call /key/health for the key -> it should be unhealthy - """ - # Generate a key with logging callback - generate_url = "http://0.0.0.0:4000/key/generate" - generate_headers = { - "Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", - "Content-Type": "application/json", - } - generate_payload = { - "metadata": { - "logging": [ - { - "callback_name": "gcs_bucket", - "callback_type": "success_and_failure", - "callback_vars": { - "gcs_bucket_name": "key-logging-project1", - "gcs_path_service_account": "bad-service-account", - }, - } - ] - } - } - - async with httpx.AsyncClient() as client: - generate_response = await client.post( - generate_url, headers=generate_headers, json=generate_payload - ) - - assert generate_response.status_code == 200 - generate_data = generate_response.json() - assert "key" in generate_data - - _key = generate_data["key"] - - # Check key health - health_url = "http://localhost:4000/key/health" - health_headers = { - "Authorization": f"Bearer {_key}", - "Content-Type": "application/json", - } - - async with httpx.AsyncClient() as client: - health_response = await client.post(health_url, headers=health_headers, json={}) - - assert health_response.status_code == 200 - health_data = health_response.json() - print("key_health_data", health_data) - # Check the response format and content - assert "key" in health_data - assert "logging_callbacks" in health_data - assert health_data["logging_callbacks"]["callbacks"] == ["gcs_bucket"] - assert health_data["logging_callbacks"]["status"] == "unhealthy" - assert ( - "GCS_BUCKET_NAME is not set in the environment" - in health_data["logging_callbacks"]["details"] - ) diff --git a/tests/otel_tests/test_model_info.py b/tests/otel_tests/test_model_info.py deleted file mode 100644 index 66a81eeee51..00000000000 --- a/tests/otel_tests/test_model_info.py +++ /dev/null @@ -1,29 +0,0 @@ -""" -/model/info test -""" - -import os -import httpx -import pytest - - -@pytest.mark.asyncio() -async def test_custom_model_supports_vision(): - async with httpx.AsyncClient() as client: - response = await client.get( - "http://localhost:4000/model/info", - headers={"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}"}, - ) - assert response.status_code == 200 - - data = response.json()["data"] - - print("response from /model/info", data) - llava_model = next( - (model for model in data if model["model_name"] == "llava-hf"), None - ) - - assert llava_model is not None, "llava-hf model not found in response" - assert ( - llava_model["model_info"]["supports_vision"] == True - ), "llava-hf model should support vision" diff --git a/tests/otel_tests/test_moderations.py b/tests/otel_tests/test_moderations.py index a9c73c93500..1a67827aa61 100644 --- a/tests/otel_tests/test_moderations.py +++ b/tests/otel_tests/test_moderations.py @@ -28,28 +28,6 @@ async def make_moderations_curl_request( return await response.json() -@pytest.mark.asyncio -async def test_basic_moderations_on_proxy_no_model(): - """ - Test moderations endpoint on proxy when no `model` is specified in the request - """ - async with aiohttp.ClientSession() as session: - test_text = "I want to harm someone" # Test text that should trigger moderation - request_data = { - "input": test_text, - } - try: - response = await make_moderations_curl_request( - session, - os.environ["LITELLM_MASTER_KEY"], - request_data, - ) - print("response=", response) - except Exception as e: - print(e) - pytest.fail("Moderations request failed") - - @pytest.mark.asyncio async def test_basic_moderations_on_proxy_with_model(): """ diff --git a/tests/otel_tests/test_prometheus.py b/tests/otel_tests/test_prometheus.py deleted file mode 100644 index 7224334105d..00000000000 --- a/tests/otel_tests/test_prometheus.py +++ /dev/null @@ -1,911 +0,0 @@ -""" -Unit tests for prometheus metrics -""" - -import os -import pytest -import aiohttp -import asyncio -from litellm._uuid import uuid -from openai import AsyncOpenAI -from typing import Dict, Any - - -END_USER_ID = "my-test-user-34" - - -async def make_bad_chat_completion_request(session, key): - url = "http://0.0.0.0:4000/chat/completions" - headers = { - "Authorization": f"Bearer {key}", - "Content-Type": "application/json", - } - data = { - "model": "fake-azure-endpoint", - "messages": [{"role": "user", "content": "Hello"}], - } - async with session.post(url, headers=headers, json=data) as response: - status = response.status - response_text = await response.text() - return status, response_text - - -async def make_good_chat_completion_request(session, key): - url = "http://0.0.0.0:4000/chat/completions" - headers = { - "Authorization": f"Bearer {key}", - "Content-Type": "application/json", - } - - data = { - "model": "fake-openai-endpoint", - "messages": [{"role": "user", "content": f"Hello {uuid.uuid4()}"}], - "tags": ["teamB"], - "user": END_USER_ID, # test if disable end user tracking for prometheus works - } - async with session.post(url, headers=headers, json=data) as response: - status = response.status - response_text = await response.text() - return status, response_text - - -async def make_chat_completion_request_with_fallback(session, key): - url = "http://0.0.0.0:4000/chat/completions" - headers = { - "Authorization": f"Bearer {key}", - "Content-Type": "application/json", - } - data = { - "model": "fake-azure-endpoint", - "messages": [{"role": "user", "content": "Hello"}], - "fallbacks": ["fake-openai-endpoint"], - } - async with session.post(url, headers=headers, json=data) as response: - status = response.status - response_text = await response.text() - - # make a request with a failed fallback - data = { - "model": "fake-azure-endpoint", - "messages": [{"role": "user", "content": "Hello"}], - "fallbacks": ["unknown-model"], - } - - async with session.post(url, headers=headers, json=data) as response: - status = response.status - response_text = await response.text() - - return - - -@pytest.mark.asyncio -async def test_proxy_failure_metrics(): - """ - - Make 1 bad chat completion call to "fake-azure-endpoint" - - GET /metrics - - assert the failure metric for the requested model is incremented by 1 - - Assert the Exception class and status code are correct - """ - async with aiohttp.ClientSession() as session: - # Make a bad chat completion call - status, response_text = await make_bad_chat_completion_request( - session, os.environ["LITELLM_MASTER_KEY"] - ) - - # Check if the request failed as expected - assert status == 429, f"Expected status 429, but got {status}" - - # Get metrics - async with session.get("http://0.0.0.0:4000/metrics") as response: - metrics = await response.text() - - print("/metrics", metrics) - - # Check if the failure metric is present and correct - use pattern matching for robustness - # Labels are ordered alphabetically by Prometheus: api_key_alias, end_user, exception_class, - # exception_status, hashed_api_key, requested_model, route, team, team_alias, user, user_email - # Note: client_ip, user_agent, model_id are present but we use substring matching to be flexible - # Check for both the new metric and deprecated metric for backwards compatibility - expected_patterns = [ - "litellm_proxy_failed_requests_metric_total{", # New metric - "litellm_llm_api_failed_requests_metric_total{", # Deprecated but may still be used - ] - - # Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for - # hash_token(master_key) so the master key (or its hash) never - # propagates into metrics. See PR #26484. - from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS - - expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS - - # Check if either pattern is in metrics and contains required fields - found_metric = False - for pattern in expected_patterns: - for line in metrics.split("\n"): - # For proxy metric, check proxy-specific fields - if "litellm_proxy_failed_requests_metric_total{" in line: - if ( - 'api_key_alias="None"' in line - and 'exception_class="Openai.RateLimitError"' in line - and 'exception_status="429"' in line - and f'hashed_api_key="{expected_hashed_api_key}"' in line - and 'requested_model="fake-azure-endpoint"' in line - and 'route="/chat/completions"' in line - ): - found_metric = True - break - # For deprecated llm_api metric, check llm-specific fields - elif "litellm_llm_api_failed_requests_metric_total{" in line: - if ( - f'hashed_api_key="{expected_hashed_api_key}"' in line - and 'model="429"' in line - ): # The deprecated metric uses the actual model from the request - found_metric = True - break - if found_metric: - break - - assert ( - found_metric - ), f"Expected failure metric not found in /metrics. Looking for either litellm_proxy_failed_requests_metric_total or litellm_llm_api_failed_requests_metric_total with required fields" - - # Check total requests metric similarly - # The litellm_proxy_total_requests_metric_total should be present - total_requests_pattern = "litellm_proxy_total_requests_metric_total{" - - found_total_metric = False - for line in metrics.split("\n"): - if ( - total_requests_pattern in line - and f'hashed_api_key="{expected_hashed_api_key}"' in line - and 'requested_model="fake-azure-endpoint"' in line - and 'status_code="429"' in line - ): - found_total_metric = True - break - - assert ( - found_total_metric - ), f"Expected total requests metric not found in /metrics. Looking for: {total_requests_pattern} with hashed_api_key and status_code=429" - - -@pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=2) -async def test_proxy_success_metrics(): - """ - Make 1 good /chat/completions call to "openai/gpt-5-mini" - GET /metrics - Assert the success metric is incremented by 1 - """ - - async with aiohttp.ClientSession() as session: - # Make a good chat completion call - status, response_text = await make_good_chat_completion_request( - session, os.environ["LITELLM_MASTER_KEY"] - ) - - # Check if the request succeeded as expected - assert status == 200, f"Expected status 200, but got {status}" - - # Get metrics - async with session.get("http://0.0.0.0:4000/metrics") as response: - metrics = await response.text() - - print("/metrics", metrics) - - assert END_USER_ID not in metrics - - # Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for - # hash_token(master_key) (PR #26484). - from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS - - expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS - - # Check if the success metric is present and correct - use flexible matching - # Check for request_total_latency_metric with required fields - # Note: The model can be "gpt-3.5-turbo-0301" or similar depending on what's returned - found_request_latency = False - for line in metrics.split("\n"): - if ( - "litellm_request_total_latency_metric_bucket{" in line - and 'api_key_alias="None"' in line - and f'hashed_api_key="{expected_hashed_api_key}"' in line - and 'requested_model="fake-openai-endpoint"' in line - and 'le="0.005"' in line - ): - found_request_latency = True - break - - assert ( - found_request_latency - ), "Expected litellm_request_total_latency_metric_bucket not found in /metrics" - - # Check for llm_api_latency_metric with required fields - found_api_latency = False - for line in metrics.split("\n"): - if ( - "litellm_llm_api_latency_metric_bucket{" in line - and 'api_key_alias="None"' in line - and f'hashed_api_key="{expected_hashed_api_key}"' in line - and 'requested_model="fake-openai-endpoint"' in line - and 'le="0.005"' in line - ): - found_api_latency = True - break - - assert ( - found_api_latency - ), "Expected litellm_llm_api_latency_metric_bucket not found in /metrics" - - verify_latency_metrics(metrics) - - -def verify_latency_metrics(metrics: str): - """ - Assert that LATENCY_BUCKETS distribution is used for - - litellm_request_total_latency_metric_bucket - - litellm_llm_api_latency_metric_bucket - - Very important to verify that the overhead latency metric is present - """ - from litellm.types.integrations.prometheus import LATENCY_BUCKETS - import re - import time - - time.sleep(2) - - metric_names = [ - "litellm_request_total_latency_metric_bucket", - "litellm_llm_api_latency_metric_bucket", - "litellm_overhead_latency_metric_bucket", - ] - - for metric_name in metric_names: - # Extract all 'le' values for the current metric - pattern = rf'{metric_name}{{.*?le="(.*?)".*?}}' - le_values = re.findall(pattern, metrics) - - # Convert to set for easier comparison - actual_buckets = set(le_values) - - print("actual_buckets", actual_buckets) - expected_buckets = [] - for bucket in LATENCY_BUCKETS: - expected_buckets.append(str(bucket)) - - # replace inf with +Inf - expected_buckets = [ - bucket.replace("inf", "+Inf") for bucket in expected_buckets - ] - - print("expected_buckets", expected_buckets) - expected_buckets = set(expected_buckets) - # Verify all expected buckets are present - assert ( - actual_buckets == expected_buckets - ), f"Mismatch in {metric_name} buckets. Expected: {expected_buckets}, Got: {actual_buckets}" - - -@pytest.mark.asyncio -async def test_proxy_fallback_metrics(): - """ - Make 1 request with a client side fallback - check metrics - """ - - async with aiohttp.ClientSession() as session: - # Make a good chat completion call - await make_chat_completion_request_with_fallback(session, os.environ["LITELLM_MASTER_KEY"]) - - # Get metrics - async with session.get("http://0.0.0.0:4000/metrics") as response: - metrics = await response.text() - - print("/metrics", metrics) - - # Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for - # hash_token(master_key) (PR #26484). - from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS - - expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS - - # Check if successful fallback metric is incremented - use flexible matching - found_successful_fallback = False - for line in metrics.split("\n"): - if ( - "litellm_deployment_successful_fallbacks_total{" in line - and 'api_key_alias="None"' in line - and 'exception_class="Openai.RateLimitError"' in line - and 'exception_status="429"' in line - and 'fallback_model="fake-openai-endpoint"' in line - and f'hashed_api_key="{expected_hashed_api_key}"' in line - and 'requested_model="fake-azure-endpoint"' in line - and "1.0" in line - ): - found_successful_fallback = True - break - - assert ( - found_successful_fallback - ), "Expected litellm_deployment_successful_fallbacks_total metric not found in /metrics" - - # Check if failed fallback metric is incremented - use flexible matching - found_failed_fallback = False - for line in metrics.split("\n"): - if ( - "litellm_deployment_failed_fallbacks_total{" in line - and 'api_key_alias="None"' in line - and 'exception_class="Openai.RateLimitError"' in line - and 'exception_status="429"' in line - and 'fallback_model="unknown-model"' in line - and f'hashed_api_key="{expected_hashed_api_key}"' in line - and 'requested_model="fake-azure-endpoint"' in line - and "1.0" in line - ): - found_failed_fallback = True - break - - assert ( - found_failed_fallback - ), "Expected litellm_deployment_failed_fallbacks_total metric not found in /metrics" - - -async def create_test_team( - session: aiohttp.ClientSession, team_data: Dict[str, Any] -) -> str: - """Create a new team and return the team_id""" - url = "http://0.0.0.0:4000/team/new" - headers = { - "Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", - "Content-Type": "application/json", - } - - async with session.post(url, headers=headers, json=team_data) as response: - assert ( - response.status == 200 - ), f"Failed to create team. Status: {response.status}" - team_info = await response.json() - return team_info["team_id"] - - -async def create_test_user( - session: aiohttp.ClientSession, user_data: Dict[str, Any] -) -> Dict[str, Any]: - """Create a new user and return the user info""" - url = "http://0.0.0.0:4000/user/new" - headers = { - "Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", - "Content-Type": "application/json", - } - - async with session.post(url, headers=headers, json=user_data) as response: - assert ( - response.status == 200 - ), f"Failed to create user. Status: {response.status}" - user_info = await response.json() - return user_info - - -async def get_prometheus_metrics(session: aiohttp.ClientSession) -> str: - """Fetch current prometheus metrics""" - async with session.get("http://0.0.0.0:4000/metrics") as response: - assert response.status == 200 - return await response.text() - - -def extract_budget_metrics(metrics_text: str, team_id: str) -> Dict[str, float]: - """Extract budget-related metrics for a specific team""" - import re - - metrics = {} - - # Get remaining budget - remaining_pattern = f'litellm_remaining_team_budget_metric{{team="{team_id}",team_alias="[^"]*"}} ([0-9.]+)' - remaining_match = re.search(remaining_pattern, metrics_text) - metrics["remaining"] = float(remaining_match.group(1)) if remaining_match else None - - # Get total budget - total_pattern = f'litellm_team_max_budget_metric{{team="{team_id}",team_alias="[^"]*"}} ([0-9.]+)' - total_match = re.search(total_pattern, metrics_text) - metrics["total"] = float(total_match.group(1)) if total_match else None - - # Get remaining hours - hours_pattern = f'litellm_team_budget_remaining_hours_metric{{team="{team_id}",team_alias="[^"]*"}} ([0-9.]+)' - hours_match = re.search(hours_pattern, metrics_text) - metrics["remaining_hours"] = float(hours_match.group(1)) if hours_match else None - - return metrics - - -async def create_test_key(session: aiohttp.ClientSession, team_id: str) -> str: - """Generate a new key for the team and return it""" - url = "http://0.0.0.0:4000/key/generate" - headers = { - "Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", - "Content-Type": "application/json", - } - data = { - "team_id": team_id, - } - - async with session.post(url, headers=headers, json=data) as response: - assert ( - response.status == 200 - ), f"Failed to generate key. Status: {response.status}" - key_info = await response.json() - return key_info["key"] - - -async def get_team_info(session: aiohttp.ClientSession, team_id: str) -> Dict[str, Any]: - """Fetch team info and return the response""" - url = f"http://0.0.0.0:4000/team/info?team_id={team_id}" - headers = { - "Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", - } - - async with session.get(url, headers=headers) as response: - assert ( - response.status == 200 - ), f"Failed to get team info. Status: {response.status}" - return await response.json() - - -@pytest.mark.asyncio -async def test_team_budget_metrics(): - """ - Test team budget tracking metrics: - 1. Create a team with max_budget - 2. Generate a key for the team - 3. Make chat completion requests using OpenAI SDK with team's key - 4. Verify budget decreases over time - 5. Verify request costs are being tracked correctly - 6. Verify prometheus metrics match /team/info spend data - """ - async with aiohttp.ClientSession() as session: - # Setup test team - team_data = { - "team_alias": "budget_test_team", - "max_budget": 10, - "budget_duration": "7d", - } - team_id = await create_test_team(session, team_data) - print("team_id", team_id) - # Generate key for the team - team_key = await create_test_key(session, team_id) - - # Initialize OpenAI client with team's key - client = AsyncOpenAI(base_url="http://0.0.0.0:4000", api_key=team_key) - - # Make initial request and check budget - await client.chat.completions.create( - model="fake-openai-endpoint", - messages=[{"role": "user", "content": f"Hello {uuid.uuid4()}"}], - ) - - await asyncio.sleep(11) # Wait for metrics to update - - # Get metrics after request - metrics_after_first = await get_prometheus_metrics(session) - print("metrics_after_first", metrics_after_first) - first_budget = extract_budget_metrics(metrics_after_first, team_id) - - print(f"Budget after 1 request: {first_budget}") - assert ( - first_budget["remaining"] < 10.0 - ), "remaining budget should be less than 10.0 after first request" - assert first_budget["total"] == 10.0, "Total budget metric is incorrect" - print("first_budget['remaining_hours']", first_budget["remaining_hours"]) - # Budget should have positive remaining hours, up to 7 days - assert ( - 0 < first_budget["remaining_hours"] <= 168 - ), "Budget should have positive remaining hours, up to 7 days" - - # Get team info and verify spend matches prometheus metrics - team_info = await get_team_info(session, team_id) - print("team_info", team_info) - _team_info_data = team_info["team_info"] - - # Calculate spend from prometheus (total - remaining) - team_info_spend = float(_team_info_data["spend"]) - team_info_max_budget = float(_team_info_data["max_budget"]) - team_info_remaining_budget = team_info_max_budget - team_info_spend - print("\n\n\n###### Final budget metrics ######\n\n\n") - print("team_info_remaining_budget", team_info_remaining_budget) - print("prometheus_remaining_budget", first_budget["remaining"]) - print( - "diff between team_info_remaining_budget and prometheus_remaining_budget", - team_info_remaining_budget - first_budget["remaining"], - ) - - # Verify spends match within a small delta (floating point comparison) - assert ( - abs(team_info_remaining_budget - first_budget["remaining"]) <= 0.001 - ), f"Spend mismatch: Prometheus={team_info_remaining_budget}, Team Info={first_budget['remaining']}" - - -async def create_test_key_with_budget( - session: aiohttp.ClientSession, budget_data: Dict[str, Any] -) -> str: - """Generate a new key with budget constraints and return it""" - url = "http://0.0.0.0:4000/key/generate" - headers = { - "Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", - "Content-Type": "application/json", - } - print("budget_data", budget_data) - - async with session.post(url, headers=headers, json=budget_data) as response: - assert ( - response.status == 200 - ), f"Failed to generate key. Status: {response.status}" - key_info = await response.json() - return key_info["key"] - - -async def get_key_info(session: aiohttp.ClientSession, key: str) -> Dict[str, Any]: - """Fetch key info and return the response""" - url = "http://0.0.0.0:4000/key/info" - headers = { - "Authorization": f"Bearer {key}", - } - - async with session.get(url, headers=headers) as response: - assert ( - response.status == 200 - ), f"Failed to get key info. Status: {response.status}" - return await response.json() - - -async def get_user_info(session: aiohttp.ClientSession, user_id: str) -> Dict[str, Any]: - """Fetch user info and return the response""" - from urllib.parse import quote - - # URL encode user_id to handle special characters - encoded_user_id = quote(user_id, safe="") - url = f"http://0.0.0.0:4000/user/info?user_id={encoded_user_id}" - headers = { - "Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", - } - - async with session.get(url, headers=headers) as response: - assert ( - response.status == 200 - ), f"Failed to get user info. Status: {response.status}" - return await response.json() - - -def extract_key_budget_metrics(metrics_text: str, key_id: str) -> Dict[str, float]: - """Extract budget-related metrics for a specific key""" - import re - - metrics = {} - - # Get remaining budget - remaining_pattern = f'litellm_remaining_api_key_budget_metric{{api_key_alias="[^"]*",hashed_api_key="{key_id}"}} ([0-9.]+)' - remaining_match = re.search(remaining_pattern, metrics_text) - metrics["remaining"] = float(remaining_match.group(1)) if remaining_match else None - - # Get total budget - total_pattern = f'litellm_api_key_max_budget_metric{{api_key_alias="[^"]*",hashed_api_key="{key_id}"}} ([0-9.]+)' - total_match = re.search(total_pattern, metrics_text) - metrics["total"] = float(total_match.group(1)) if total_match else None - - # Get remaining hours - hours_pattern = f'litellm_api_key_budget_remaining_hours_metric{{api_key_alias="[^"]*",hashed_api_key="{key_id}"}} ([0-9.]+)' - hours_match = re.search(hours_pattern, metrics_text) - metrics["remaining_hours"] = float(hours_match.group(1)) if hours_match else None - - return metrics - - -def extract_user_budget_metrics(metrics_text: str, user_id: str) -> Dict[str, float]: - """Extract budget-related metrics for a specific user""" - import re - - metrics = {} - - # Escape user_id for regex pattern matching - escaped_user_id = re.escape(user_id) - - # Get remaining budget (user_email and user_alias may also be present as labels) - remaining_pattern = rf'litellm_remaining_user_budget_metric{{[^}}]*user="{escaped_user_id}"[^}}]*}} ([0-9.]+)' - remaining_match = re.search(remaining_pattern, metrics_text) - metrics["remaining"] = float(remaining_match.group(1)) if remaining_match else None - - # Get total budget - total_pattern = rf'litellm_user_max_budget_metric{{[^}}]*user="{escaped_user_id}"[^}}]*}} ([0-9.]+)' - total_match = re.search(total_pattern, metrics_text) - metrics["total"] = float(total_match.group(1)) if total_match else None - - # Get remaining hours - hours_pattern = rf'litellm_user_budget_remaining_hours_metric{{[^}}]*user="{escaped_user_id}"[^}}]*}} ([0-9.]+)' - hours_match = re.search(hours_pattern, metrics_text) - metrics["remaining_hours"] = float(hours_match.group(1)) if hours_match else None - - return metrics - - -@pytest.mark.asyncio -async def test_key_budget_metrics(): - """ - Test key budget tracking metrics: - 1. Create a key with max_budget - 2. Make chat completion requests using OpenAI SDK with the key - 3. Verify budget decreases over time - 4. Verify request costs are being tracked correctly - 5. Verify prometheus metrics match /key/info spend data - """ - from datetime import datetime, timedelta, timezone - - async with aiohttp.ClientSession() as session: - # Setup test key with unique alias - unique_alias = f"budget_test_key_{uuid.uuid4()}" - key_data = { - "key_alias": unique_alias, - "max_budget": 10, - "budget_duration": "7d", - "budget_reset_at": ( - datetime.now(timezone.utc) + timedelta(days=7) - ).isoformat(), - } - key = await create_test_key_with_budget(session, key_data) - - # Extract key_id from the key info - key_info = await get_key_info(session, key) - print("key_info", key_info) - key_id = key_info["key"] - print("key_id", key_id) - - # Initialize OpenAI client with the key - client = AsyncOpenAI(base_url="http://0.0.0.0:4000", api_key=key) - - # Make initial request and check budget - await client.chat.completions.create( - model="fake-openai-endpoint", - messages=[{"role": "user", "content": f"Hello {uuid.uuid4()}"}], - ) - - await asyncio.sleep(11) # Wait for metrics to update - - # Get metrics after request - metrics_after_first = await get_prometheus_metrics(session) - print("metrics_after_first request", metrics_after_first) - first_budget = extract_key_budget_metrics(metrics_after_first, key_id) - - print(f"Budget after 1 request: {first_budget}") - assert ( - first_budget["remaining"] < 10.0 - ), "remaining budget should be less than 10.0 after first request" - assert first_budget["total"] == 10.0, "Total budget metric is incorrect" - print("first_budget['remaining_hours']", first_budget["remaining_hours"]) - # The budget reset time is now standardized - for "7d" it resets on Monday at midnight - # So we'll check if it's within a reasonable range (0-7 days depending on current day of week) - assert ( - 0 <= first_budget["remaining_hours"] <= 168 - ), "Budget remaining hours should be within a reasonable range (0-7 days depending on day of week)" - - # Get key info and verify spend matches prometheus metrics - key_info = await get_key_info(session, key) - print("key_info", key_info) - _key_info_data = key_info["info"] - - # Calculate spend from prometheus (total - remaining) - key_info_spend = float(_key_info_data["spend"]) - key_info_max_budget = float(_key_info_data["max_budget"]) - key_info_remaining_budget = key_info_max_budget - key_info_spend - print("\n\n\n###### Final budget metrics ######\n\n\n") - print("key_info_remaining_budget", key_info_remaining_budget) - print("prometheus_remaining_budget", first_budget["remaining"]) - print( - "diff between key_info_remaining_budget and prometheus_remaining_budget", - key_info_remaining_budget - first_budget["remaining"], - ) - - # Verify spends match within a small delta (floating point comparison) - assert ( - abs(key_info_remaining_budget - first_budget["remaining"]) <= 0.001 - ), f"Spend mismatch: Prometheus={key_info_remaining_budget}, Key Info={first_budget['remaining']}" - - -@pytest.mark.asyncio -async def test_user_budget_metrics(): - """ - Test user budget tracking metrics: - 1. Create a user with max_budget - 2. Make chat completion requests using OpenAI SDK with the user's key - 3. Verify budget decreases over time - 4. Verify request costs are being tracked correctly - 5. Verify prometheus metrics match /user/info spend data - """ - from datetime import datetime, timedelta, timezone - - async with aiohttp.ClientSession() as session: - # Setup test user with unique user_id - unique_user_id = f"budget_test_user_{uuid.uuid4()}" - user_data = { - "user_id": unique_user_id, - "max_budget": 10, - "budget_duration": "7d", - "budget_reset_at": ( - datetime.now(timezone.utc) + timedelta(days=7) - ).isoformat(), - } - user_info = await create_test_user(session, user_data) - print("user_info", user_info) - user_id = user_info["user_id"] - print("user_id", user_id) - # Get the key that was created with the user - key = user_info["key"] - - # Initialize OpenAI client with the user's key - client = AsyncOpenAI(base_url="http://0.0.0.0:4000", api_key=key) - - # Make initial request and check budget - await client.chat.completions.create( - model="fake-openai-endpoint", - messages=[{"role": "user", "content": f"Hello {uuid.uuid4()}"}], - ) - - await asyncio.sleep(11) # Wait for metrics to update - - # Get metrics after request - metrics_after_first = await get_prometheus_metrics(session) - print("metrics_after_first request", metrics_after_first) - first_budget = extract_user_budget_metrics(metrics_after_first, user_id) - - print(f"Budget after 1 request: {first_budget}") - assert ( - first_budget["remaining"] is not None - ), "remaining budget metric should be present" - assert ( - first_budget["total"] is not None - ), "total budget metric should be present" - assert ( - first_budget["remaining"] < 10.0 - ), "remaining budget should be less than 10.0 after first request" - assert first_budget["total"] == 10.0, "Total budget metric is incorrect" - print("first_budget['remaining_hours']", first_budget["remaining_hours"]) - # The budget reset time is now standardized - for "7d" it resets on Monday at midnight - # So we'll check if it's within a reasonable range (0-7 days depending on current day of week) - assert ( - first_budget["remaining_hours"] is not None - ), "remaining hours metric should be present" - assert ( - 0 <= first_budget["remaining_hours"] <= 168 - ), "Budget remaining hours should be within a reasonable range (0-7 days depending on day of week)" - - # Get user info and verify spend matches prometheus metrics - user_info_response = await get_user_info(session, user_id) - print("user_info_response", user_info_response) - _user_info_data = user_info_response["user_info"] - - # Calculate spend from prometheus (total - remaining) - user_info_spend = float(_user_info_data["spend"]) - user_info_max_budget = float(_user_info_data["max_budget"]) - user_info_remaining_budget = user_info_max_budget - user_info_spend - print("\n\n\n###### Final budget metrics ######\n\n\n") - print("user_info_remaining_budget", user_info_remaining_budget) - print("prometheus_remaining_budget", first_budget["remaining"]) - print( - "diff between user_info_remaining_budget and prometheus_remaining_budget", - user_info_remaining_budget - first_budget["remaining"], - ) - - # Verify spends match within a small delta (floating point comparison) - assert ( - abs(user_info_remaining_budget - first_budget["remaining"]) <= 0.001 - ), f"Spend mismatch: Prometheus={user_info_remaining_budget}, User Info={first_budget['remaining']}" - - -@pytest.mark.asyncio -async def test_user_email_metrics(): - """ - Test user email tracking metrics: - 1. Create a user with user_email - 2. Make chat completion requests using OpenAI SDK with the user's email - 3. Verify user email is being tracked correctly in `litellm_user_email_metric` - """ - async with aiohttp.ClientSession() as session: - # Create a user with user_email - user_email = f"test-{uuid.uuid4()}@example.com" - user_data = { - "user_email": user_email, - } - user_info = await create_test_user(session, user_data) - key = user_info["key"] - - # Initialize OpenAI client with the user's email - client = AsyncOpenAI(base_url="http://0.0.0.0:4000", api_key=key) - - # Make initial request and check budget - await client.chat.completions.create( - model="fake-openai-endpoint", - messages=[{"role": "user", "content": f"Hello {uuid.uuid4()}"}], - ) - - await asyncio.sleep(11) # Wait for metrics to update - - # Get metrics after request - metrics_after_first = await get_prometheus_metrics(session) - print("metrics_after_first request", metrics_after_first) - assert ( - user_email in metrics_after_first - ), "user_email should be tracked correctly" - - -@pytest.mark.asyncio -async def test_user_email_in_all_required_metrics(): - """ - Test that user_email label is present in all the metrics that were requested to have it: - - litellm_proxy_total_requests_metric_total - - litellm_proxy_failed_requests_metric_total - - litellm_input_tokens_metric_total - - litellm_output_tokens_metric_total - - litellm_requests_metric_total - - litellm_spend_metric_total - """ - async with aiohttp.ClientSession() as session: - # Create a user with user_email - user_email = f"test-metrics-{uuid.uuid4()}@example.com" - user_data = { - "user_email": user_email, - } - user_info = await create_test_user(session, user_data) - key = user_info["key"] - - # Initialize OpenAI client with the user's email - client = AsyncOpenAI(base_url="http://0.0.0.0:4000", api_key=key) - - # Make successful request to generate metrics - await client.chat.completions.create( - model="fake-openai-endpoint", - messages=[{"role": "user", "content": f"Hello {uuid.uuid4()}"}], - ) - - await asyncio.sleep(11) # Wait for metrics to update - - # Get metrics after request - metrics_text = await get_prometheus_metrics(session) - print("Testing user_email in all required metrics") - - # Check that user_email appears in all the required metrics - required_metrics_with_user_email = [ - # "litellm_proxy_total_requests_metric_total", - # "litellm_input_tokens_metric_total", - # "litellm_output_tokens_metric_total", - # "litellm_requests_metric_total", - "litellm_spend_metric_total", - ] - - import re - - for metric_name in required_metrics_with_user_email: - # Check that the metric exists and contains user_email label - # Look for the metric with user_email in its labels - pattern = ( - rf'{metric_name}{{[^}}]*user_email="{re.escape(user_email)}"[^}}]*}}' - ) - matches = re.findall(pattern, metrics_text) - assert ( - len(matches) > 0 - ), f"Metric {metric_name} should contain user_email={user_email} but was not found in metrics" - - # Also test failure metric by making a bad request - try: - await client.chat.completions.create( - model="fake-azure-endpoint", # This should fail - messages=[{"role": "user", "content": "Hello"}], - ) - except Exception: - pass # Expected to fail - - await asyncio.sleep(11) # Wait for metrics to update - - # Get updated metrics - metrics_text = await get_prometheus_metrics(session) - - # Check that failure metric also contains user_email - failure_pattern = rf'litellm_proxy_failed_requests_metric_total{{[^}}]*user_email="{re.escape(user_email)}"[^}}]*}}' - failure_matches = re.findall(failure_pattern, metrics_text) - assert ( - len(failure_matches) > 0 - ), f"litellm_proxy_failed_requests_metric_total should contain user_email={user_email}" diff --git a/tests/otel_tests/test_team_tag_routing.py b/tests/otel_tests/test_team_tag_routing.py deleted file mode 100644 index b818aa0cea4..00000000000 --- a/tests/otel_tests/test_team_tag_routing.py +++ /dev/null @@ -1,65 +0,0 @@ -import os -# What this tests ? -## Set tags on a team and then make a request to /chat/completions -import pytest -import asyncio -import aiohttp, openai -from openai import OpenAI, AsyncOpenAI -from typing import Optional, List, Union -from litellm._uuid import uuid - -LITELLM_MASTER_KEY = os.environ["LITELLM_MASTER_KEY"] - - -async def chat_completion( - session, key, model: Union[str, List] = "fake-openai-endpoint" -): - url = "http://0.0.0.0:4000/chat/completions" - headers = { - "Authorization": f"Bearer {key}", - "Content-Type": "application/json", - } - print("headers=", headers) - data = { - "model": model, - "messages": [ - {"role": "user", "content": f"Hello! {str(uuid.uuid4())}"}, - ], - } - - async with session.post(url, headers=headers, json=data) as response: - status = response.status - response_text = await response.text() - - if status != 200: - raise Exception(response_text) - - return await response.json(), response.headers - - -async def model_info_get_call(session, key, model_id: str): - # make get call pass "litellm_model_id" in query params - url = f"http://0.0.0.0:4000/model/info?litellm_model_id={model_id}" - headers = { - "Authorization": f"Bearer {key}", - "Content-Type": "application/json", - } - async with session.get(url, headers=headers) as response: - status = response.status - response_text = await response.text() - - if status != 200: - raise Exception(response_text) - - return await response.json() - - -@pytest.mark.asyncio() -async def test_chat_completion_with_no_tags(): - async with aiohttp.ClientSession() as session: - key = LITELLM_MASTER_KEY - response, headers = await chat_completion(session, key) - headers = dict(headers) - print(response) - print(headers) - assert response is not None diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index a799c45e4c7..c94ee836f91 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -953,6 +953,7 @@ def test_redis_caching_multiple_namespaces(): _TOOL_TURN_ITEM: Final = {"role": "user", "content": "hi"} +_FILE_BLOCK_ITEM: Final = {"type": "file", "file": {"file_data": "data:video/mp4;base64,AAAA", "format": "video/mp4"}} @pytest.mark.parametrize( @@ -964,6 +965,7 @@ _TOOL_TURN_ITEM: Final = {"role": "user", "content": "hi"} pytest.param({"input": [_TOOL_TURN_ITEM] * 5}, False, id="five-responses-items-skip-the-cache"), pytest.param({"input": "one prompt"}, True, id="string-input-is-one-message"), pytest.param({"input": ["a", "b", "c", "d", "e"]}, True, id="embedding-strings-are-not-messages"), + pytest.param({"input": [_FILE_BLOCK_ITEM] * 5}, True, id="embedding-file-blocks-are-not-messages"), ], ) def test_should_use_cache_stops_past_the_default_max_messages(kwargs: dict[str, object], expected: bool) -> None: diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index 22eea9935f6..d98c02667ee 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -41,6 +41,7 @@ from litellm.types.utils import ( from litellm.types.llms.openai import ResponsesAPIResponse from collections.abc import Awaitable, Callable from datetime import timedelta, datetime +from typing import Final from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm._logging import verbose_logger @@ -2714,3 +2715,88 @@ async def test_async_get_cache_forgets_the_worker_copy_of_a_stored_response_with assert lookup.cached_result is None assert await handler.dual_cache.async_get_cache(key) is None + + +@pytest.mark.asyncio +async def test_async_get_cache_partial_hit_keeps_file_block_items_uncached() -> None: + setup_cache() + fixed_start: Final = datetime(2026, 1, 1) + caching_handler: Final = LLMCachingHandler(original_function=aembedding, request_kwargs={}, start_time=fixed_start) + model: Final = "gemini/gemini-embedding-2-preview" + logging_obj: Final = LiteLLMLogging( + litellm_call_id=str(uuid.uuid4()), + call_type=CallTypes.aembedding.value, + model=model, + messages=[], + function_id=str(uuid.uuid4()), + stream=False, + start_time=fixed_start, + ) + await caching_handler.async_set_cache( + result=EmbeddingResponse(model=model, data=[Embedding(embedding=[0.1, 0.2], index=0, object="embedding")]), + original_function=aembedding, + kwargs={"model": model, "input": ["a red bus"], "caching": True}, + ) + clip_block: Final = { + "type": "file", + "file": { + "file_data": "data:video/mp4;base64,AAAA", + "format": "video/mp4", + "video_metadata": {"fps": 1, "start_offset": "0s", "end_offset": "1s"}, + }, + "detail": "left for the provider transformation to judge", + } + + cached_response: Final = await caching_handler.async_get_cache( + model=model, + original_function=aembedding, + logging_obj=logging_obj, + start_time=fixed_start, + call_type=CallTypes.aembedding.value, + kwargs={"model": model, "input": [clip_block, "a red bus"], "caching": True}, + ) + + assert cached_response.embedding_all_elements_cache_hit is False + assert cached_response.embedding_uncached_input == [clip_block] + assert cached_response.final_embedding_cached_response is not None + assert cached_response.final_embedding_cached_response.data[1].embedding == [0.1, 0.2] + assert cached_response.final_embedding_cached_response.data[0] is None + + +def test_handle_kwargs_input_answers_400_for_a_single_object_input() -> None: + caching_handler: Final = LLMCachingHandler( + original_function=aembedding, request_kwargs={}, start_time=datetime(2026, 1, 1) + ) + clip_block: Final = {"type": "file", "file": {"file_data": "data:video/mp4;base64,AAAA"}} + with pytest.raises(litellm.BadRequestError, match="string or a list"): + caching_handler.handle_kwargs_input_list_or_str( + {"model": "gemini/gemini-embedding-2-preview", "custom_llm_provider": "gemini", "input": clip_block} + ) + + +@pytest.mark.asyncio +async def test_async_get_cache_answers_400_for_a_single_object_embedding_input() -> None: + setup_cache() + fixed_start: Final = datetime(2026, 1, 1) + caching_handler: Final = LLMCachingHandler(original_function=aembedding, request_kwargs={}, start_time=fixed_start) + model: Final = "gemini/gemini-embedding-2-preview" + logging_obj: Final = LiteLLMLogging( + litellm_call_id=str(uuid.uuid4()), + call_type=CallTypes.aembedding.value, + model=model, + messages=[], + function_id=str(uuid.uuid4()), + stream=False, + start_time=fixed_start, + ) + clip_block: Final = {"type": "file", "file": {"file_data": "data:video/mp4;base64,AAAA"}} + + with pytest.raises(litellm.BadRequestError, match="string or a list"): + await caching_handler.async_get_cache( + model=model, + original_function=aembedding, + logging_obj=logging_obj, + start_time=fixed_start, + call_type=CallTypes.aembedding.value, + kwargs={"model": model, "custom_llm_provider": "gemini", "input": clip_block, "caching": True}, + ) diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index b06ed9468b0..9a30e1990b1 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -2936,8 +2936,11 @@ async def test_cancellation_delivers_termination_over_tcp( listener: Final = await asyncio.start_server(handle_connection, "127.0.0.1", 0) port: Final = listener.sockets[0].getsockname()[1] + client_timeout: Final = 2 if cancel_mode == "read_timeout" else 30 client: Final = MCPClient( - server_url=f"http://127.0.0.1:{port}/mcp", protocol_version=protocol_version, timeout=2 if cancel_mode == "read_timeout" else 30 + server_url=f"http://127.0.0.1:{port}/mcp", + protocol_version=protocol_version, + timeout=client_timeout, ) async def calls(): diff --git a/tests/unit/integrations/SlackAlerting/test_slack_alerting_delivery.py b/tests/unit/integrations/SlackAlerting/test_slack_alerting_delivery.py new file mode 100644 index 00000000000..c8fb5d30376 --- /dev/null +++ b/tests/unit/integrations/SlackAlerting/test_slack_alerting_delivery.py @@ -0,0 +1,344 @@ +import asyncio +import datetime +from collections.abc import Sequence +from typing import Final, Literal, TypedDict + +from typing_extensions import ReadOnly + +import httpx +import pytest +import respx +from pydantic import TypeAdapter + +import litellm +import litellm.proxy.proxy_server as proxy_server +from litellm.caching.caching import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.langfuse.langfuse import LangFuseLogger, installed_langfuse_version +from litellm.integrations.langfuse.langfuse_sdk import ( + build_langfuse_client, + build_langfuse_tracing, + resolve_trace_id, +) +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.integrations.SlackAlerting.utils import add_langfuse_trace_id_to_alert +from litellm.litellm_core_utils import litellm_logging +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.proxy.utils import ProxyLogging +from litellm.types.integrations.slack_alerting import AlertType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +_WEBHOOK: Final = "https://hooks.slack.example/services/delivery" +_LANGFUSE_HOST: Final = "https://langfuse.alerts.example" +_AZURE_BASE: Final = "https://openai-gpt-4-test-v-1.openai.azure.com/" +_DAILY_BASE: Final = "https://daily-report.openai.example/v1" + + +class _SlackPayload(TypedDict): + text: ReadOnly[str] + + +class _TeamRow(TypedDict): + team_alias: ReadOnly[str] + total_spend: ReadOnly[float] + + +class _TagRow(TypedDict): + individual_request_tag: ReadOnly[str] + total_spend: ReadOnly[float] + + +_PAYLOAD: Final = TypeAdapter(_SlackPayload) + + +@pytest.fixture(autouse=True) +def _slack_webhook(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("SLACK_WEBHOOK_URL", _WEBHOOK) + + +def _webhook(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post(_WEBHOOK).mock(return_value=httpx.Response(200, text="ok")) + + +def _posted_texts(route: respx.Route) -> tuple[str, ...]: + return tuple(_PAYLOAD.validate_json(call.request.content)["text"] for call in route.calls) + + +@pytest.mark.asyncio +async def test_slow_response_alert_names_the_azure_api_base_and_reaches_the_webhook( + respx_mock: respx.MockRouter, +) -> None: + route: Final = _webhook(respx_mock) + proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging.update_values(alerting=["slack"], alerting_threshold=100, redis_cache=None) + start: Final = datetime.datetime(2026, 1, 1, 12, 0, 0) + messages: Final = ({"role": "user", "content": "Hey how's it going?"},) + + helper_result: Final = proxy_logging.slack_alerting_instance._response_taking_too_long_callback_helper( + kwargs={ + "model": "chatgpt-v-3", + "messages": messages, + "litellm_params": {"api_base": _AZURE_BASE, "custom_llm_provider": "azure"}, + }, + start_time=start, + end_time=start + datetime.timedelta(seconds=150), + ) + + assert helper_result == (150.0, "chatgpt-v-3", _AZURE_BASE, str(messages)[:100]) + + slow_message: Final = ( + f"`Responses are slow - 150.0s response time > Alerting threshold: 100s`\nAPI Base: `{_AZURE_BASE}`" + ) + await proxy_logging.alerting_handler(message=slow_message, level="Low", alert_type=AlertType.llm_too_slow) + await proxy_logging.slack_alerting_instance.flush_queue() + + texts: Final = _posted_texts(route) + assert len(texts) == 1 + assert texts[0].startswith("Alert type: `llm_too_slow`\nLevel: `Low`\n") + assert texts[0].endswith(f"Message: {slow_message}") + + +@pytest.mark.asyncio +async def test_send_alert_is_queued_until_flush_then_posted_to_the_webhook_once(respx_mock: respx.MockRouter) -> None: + route: Final = _webhook(respx_mock) + slack_alerting: Final = SlackAlerting(alerting_threshold=1, internal_usage_cache=DualCache(), alerting=["slack"]) + + await slack_alerting.send_alert("Test message", "Low", AlertType.budget_alerts, alerting_metadata={}) + + assert route.call_count == 0 + + await slack_alerting.flush_queue() + await slack_alerting.flush_queue() + + texts: Final = _posted_texts(route) + assert len(texts) == 1 + assert texts[0].startswith("Alert type: `budget_alerts`\nLevel: `Low`\n") + assert texts[0].endswith("Message: Test message") + + +@pytest.mark.asyncio +async def test_a_queued_alert_is_posted_by_the_periodic_flush_without_a_manual_flush( + respx_mock: respx.MockRouter, +) -> None: + delivered: Final = asyncio.Event() + + def deliver(request: httpx.Request) -> httpx.Response: + delivered.set() + return httpx.Response(200, text="ok") + + route: Final = respx_mock.post(_WEBHOOK).mock(side_effect=deliver) + slack_alerting: Final = SlackAlerting(alerting_threshold=1, internal_usage_cache=DualCache(), alerting=["slack"]) + slack_alerting.flush_interval = 0 + slack_alerting.update_values(alerting=["slack"]) + flush_task: Final = slack_alerting._periodic_flush_task + assert flush_task is not None + try: + await slack_alerting.send_alert("Timed message", "Low", AlertType.budget_alerts, alerting_metadata={}) + await asyncio.wait_for(delivered.wait(), timeout=5) + finally: + flush_task.cancel() + + texts: Final = _posted_texts(route) + assert len(texts) == 1 + assert texts[0].endswith("Message: Timed message") + + +class _DeploymentSettled(CustomLogger): + def __init__(self, model_id: str) -> None: + super().__init__() + self.model_id: Final = model_id + self.succeeded: Final = asyncio.Event() + self.failed: Final = asyncio.Event() + + def _is_mine(self, kwargs: dict[str, object]) -> bool: + payload: Final = kwargs.get("standard_logging_object") + return isinstance(payload, dict) and payload.get("model_id") == self.model_id + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + if self._is_mine(kwargs): + self.succeeded.set() + + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + if self._is_mine(kwargs): + self.failed.set() + + +@pytest.mark.asyncio +async def test_daily_report_lists_router_latency_after_success_and_failures_after_an_auth_error( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + webhook: Final = _webhook(respx_mock) + respx_mock.post(f"{_DAILY_BASE}/chat/completions").mock( + side_effect=( + httpx.Response( + 200, + json={ + "id": "chatcmpl-daily", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-5-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "fine"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 4, "total_tokens": 9}, + }, + ), + httpx.Response( + 401, + json={ + "error": { + "message": "Incorrect API key provided", + "type": "invalid_request_error", + "code": "invalid_api_key", + } + }, + ), + ) + ) + model_id: Final = "daily-report-deployment" + slack_alerting: Final = SlackAlerting( + alerting=["slack"], internal_usage_cache=DualCache(), alert_types=[AlertType.daily_reports] + ) + settled: Final = _DeploymentSettled(model_id) + monkeypatch.setattr(litellm, "callbacks", [slack_alerting, settled]) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "daily-report-model", + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-daily", "api_base": _DAILY_BASE}, + "model_info": {"id": model_id}, + } + ] + ) + request: Final = ({"role": "user", "content": "Hey, how's it going?"},) + + await router.acompletion(model="daily-report-model", messages=list(request)) + await asyncio.wait_for(settled.succeeded.wait(), timeout=5) + after_success: Final = await slack_alerting.send_daily_reports(router=router) + await slack_alerting.flush_queue() + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion(model="daily-report-model", messages=list(request)) + await asyncio.wait_for(settled.failed.wait(), timeout=5) + after_failure: Final = await slack_alerting.send_daily_reports(router=router) + await slack_alerting.flush_queue() + + texts: Final = _posted_texts(webhook) + assert (after_success, after_failure) == (True, True) + assert len(texts) == 2 + assert "Most Failed Requests:*\n\n\tNone\n" in texts[0] + assert "1. Deployment: `openai/gpt-5-mini`, Latency per output token: `" in texts[0] + assert f"1. Deployment: `openai/gpt-5-mini`, Failed Requests: `1`, API Base: `{_DAILY_BASE}`" in texts[1] + assert "Top Slowest Deployments:*\n\n\tNone\n" in texts[1] + + +class _CallLogged(CustomLogger): + def __init__(self, call_id: str, loop: asyncio.AbstractEventLoop) -> None: + super().__init__() + self.call_id: Final = call_id + self.loop: Final = loop + self.logged: Final = asyncio.Event() + + def log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + if kwargs.get("litellm_call_id") == self.call_id: + self.loop.call_soon_threadsafe(self.logged.set) + + +@pytest.mark.asyncio +async def test_langfuse_trace_link_ends_with_the_trace_id_the_logger_emitted(monkeypatch: pytest.MonkeyPatch) -> None: + logger: Final = LangFuseLogger.__new__(LangFuseLogger) + logger.tracing = build_langfuse_tracing( + exporter=InMemorySpanExporter(), environment=None, release=None, sample_rate=1.0, flush_interval_millis=10 + ) + logger.api_client = build_langfuse_client( + public_key="pk-alert-trace", secret_key="sk-alert-trace", base_url=_LANGFUSE_HOST, httpx_client=None + ) + logger.langfuse_sdk_version = installed_langfuse_version() + call_id: Final = "slack-alert-langfuse-trace" + logged: Final = _CallLogged(call_id, asyncio.get_running_loop()) + monkeypatch.setenv("LANGFUSE_HOST", _LANGFUSE_HOST) + monkeypatch.setattr(litellm_logging, "langFuseLogger", logger) + monkeypatch.setattr(litellm, "success_callback", ["langfuse", logged]) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "callbacks", []) + logging_obj: Final = Logging( + model="gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + litellm_call_id=call_id, + start_time=datetime.datetime.now(), + function_id=call_id, + ) + + litellm.completion( + model="gpt-5-mini", + messages=[{"role": "user", "content": "Hey how's it going?"}], + mock_response="Hey!", + litellm_logging_obj=logging_obj, + ) + await asyncio.wait_for(logged.logged.wait(), timeout=5) + trace_url: Final = await add_langfuse_trace_id_to_alert(request_data={"litellm_logging_obj": logging_obj}) + + expected_trace_id: Final = resolve_trace_id(logging_obj.litellm_trace_id) + assert logging_obj.get_trace_id(service_name="langfuse") == expected_trace_id + assert trace_url == f"{_LANGFUSE_HOST}/trace/{expected_trace_id}" + + +class _SpendReportDb: + def __init__(self, teams: Sequence[_TeamRow], tags: Sequence[_TagRow]) -> None: + self.teams: Final = teams + self.tags: Final = tags + + async def query_raw(self, query: str, *args: object) -> Sequence[_TeamRow] | Sequence[_TagRow]: + return self.teams if "team_alias" in query else self.tags + + +class _SpendReportPrisma: + def __init__(self, db: _SpendReportDb) -> None: + self.db: Final = db + + +@pytest.mark.parametrize("report_type", ["weekly", "monthly"]) +@pytest.mark.asyncio +async def test_spend_report_is_sent_once_per_period( + report_type: Literal["weekly", "monthly"], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + route: Final = _webhook(respx_mock) + monkeypatch.setattr( + proxy_server, + "prisma_client", + _SpendReportPrisma( + _SpendReportDb( + teams=( + _TeamRow(team_alias="team1", total_spend=100.0), + _TeamRow(team_alias="team2", total_spend=200.0), + ), + tags=( + _TagRow(individual_request_tag="tag1", total_spend=150.0), + _TagRow(individual_request_tag="tag2", total_spend=150.0), + ), + ) + ), + ) + slack_alerting: Final = SlackAlerting(alerting=["slack"], internal_usage_cache=DualCache()) + send_report: Final = ( + slack_alerting.send_weekly_spend_report if report_type == "weekly" else slack_alerting.send_monthly_spend_report + ) + + await send_report() + await slack_alerting.flush_queue() + await send_report() + await slack_alerting.flush_queue() + + texts: Final = _posted_texts(route) + assert len(texts) == 1 + assert "Team: `team1` | Spend: `$100.0`\nTeam: `team2` | Spend: `$200.0`\n" in texts[0] + assert "Tag: `tag1` | Spend: `$150.0`\nTag: `tag2` | Spend: `$150.0`\n" in texts[0] diff --git a/tests/unit/integrations/datadog/test_datadog.py b/tests/unit/integrations/datadog/test_datadog.py index e86d83ba467..eec0022c1db 100644 --- a/tests/unit/integrations/datadog/test_datadog.py +++ b/tests/unit/integrations/datadog/test_datadog.py @@ -2,14 +2,20 @@ import gzip import json import os from datetime import datetime -from typing import Coroutine, Final +from pathlib import Path +from typing import Coroutine, Final, TypedDict from unittest.mock import AsyncMock, patch import pytest +import respx from httpx import Request, Response +from pydantic import TypeAdapter +from typing_extensions import ReadOnly import litellm import litellm.integrations.datadog.datadog as datadog_module +from litellm.caching.caching import Cache +from litellm.caching.llm_caching_handler import LLMClientCache from litellm.integrations.datadog.datadog import DataDogLogger from litellm.integrations.datadog.datadog_handler import ( get_datadog_env, @@ -19,6 +25,7 @@ from litellm.integrations.datadog.datadog_handler import ( get_datadog_source, get_datadog_tags, ) +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.types.integrations.datadog import DatadogInitParams, DatadogPayload, DataDogStatus from litellm.types.utils import ( StandardLoggingHiddenParams, @@ -803,3 +810,113 @@ def create_standard_logging_payload() -> StandardLoggingPayload: additional_headers=None, ), ) + + +_INTAKE_URL: Final = "https://http-intake.logs.test.datadoghq.com/api/v2/logs" + + +class _ServiceEventMessage(TypedDict): + service: ReadOnly[str] + call_type: ReadOnly[str] + error: ReadOnly[str] + is_error: ReadOnly[bool] + + +@pytest.fixture +def delivery( + datadog_env: None, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> tuple[DataDogLogger, respx.Route]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + monkeypatch.delenv("DD_SOURCE", raising=False) + monkeypatch.delenv("DD_SERVICE", raising=False) + with patch("asyncio.create_task", side_effect=_discard_periodic_flush): + logger: Final = DataDogLogger() + intake: Final = respx_mock.post(_INTAKE_URL).mock(return_value=Response(202, text="Accepted")) + return logger, intake + + +def _delivered_logs(intake: respx.Route) -> list[DatadogPayload]: + return TypeAdapter(list[DatadogPayload]).validate_json(gzip.decompress(intake.calls.last.request.content)) + + +@pytest.mark.asyncio +async def test_a_successful_request_is_delivered_as_an_info_log_carrying_the_standard_payload( + delivery: tuple[DataDogLogger, respx.Route], +) -> None: + datadog_logger, intake = delivery + standard_payload: Final = _standard_logging_payload() + + await datadog_logger.async_log_success_event( + kwargs={"standard_logging_object": standard_payload}, + response_obj=None, + start_time=STANDARD_START_TIME, + end_time=STANDARD_END_TIME, + ) + await datadog_logger.async_send_batch() + + assert intake.call_count == 1 + logs: Final = _delivered_logs(intake) + assert len(logs) == 1 + assert logs[0]["ddsource"] == "litellm" + assert logs[0]["service"] == "litellm-server" + assert logs[0]["status"] == DataDogStatus.INFO + assert TypeAdapter(dict[str, object]).validate_json(logs[0]["message"]) == standard_payload + + +@pytest.mark.asyncio +async def test_a_failed_request_is_delivered_as_an_error_log_that_keeps_the_error_string( + delivery: tuple[DataDogLogger, respx.Route], +) -> None: + datadog_logger, intake = delivery + standard_payload: Final = _standard_logging_payload() + standard_payload["status"] = "failure" + standard_payload["error_str"] = "Test error" + + await datadog_logger.async_log_failure_event( + kwargs={"standard_logging_object": standard_payload}, + response_obj=None, + start_time=STANDARD_START_TIME, + end_time=STANDARD_END_TIME, + ) + await datadog_logger.async_send_batch() + + assert intake.call_count == 1 + logs: Final = _delivered_logs(intake) + assert len(logs) == 1 + assert logs[0]["status"] == DataDogStatus.ERROR + message: Final = TypeAdapter(dict[str, object]).validate_json(logs[0]["message"]) + assert message == standard_payload + assert message["error_str"] == "Test error" + + +@pytest.mark.asyncio +async def test_a_failing_redis_cache_is_delivered_to_datadog_as_redis_warnings( + delivery: tuple[DataDogLogger, respx.Route], monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + datadog_logger, intake = delivery + absent_socket: Final = str(tmp_path / "absent.sock") + redis_cache: Final = Cache(type="redis", url=f"unix://{absent_socket}") + monkeypatch.setattr(redis_cache.cache.service_logger_obj, "dd_logger", datadog_logger, raising=False) + monkeypatch.setattr(litellm, "service_callback", ["datadog"]) + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "cache", redis_cache) + + for _ in range(3): + await litellm.acompletion( + model="gpt-4.1-mini", + messages=[{"role": "user", "content": "what llm are u"}], + mock_response="Accepted", + caching=True, + ) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + await datadog_logger.async_send_batch() + + assert intake.call_count == 1 + logs: Final = _delivered_logs(intake) + assert len(logs) > 0 + assert {log["status"] for log in logs} == {DataDogStatus.WARN} + messages: Final = [TypeAdapter(_ServiceEventMessage).validate_json(log["message"]) for log in logs] + assert {message["service"] for message in messages} == {"redis"} + assert all(message["is_error"] is True for message in messages) + assert all(absent_socket in message["error"] for message in messages) diff --git a/tests/unit/integrations/otel/test_otel_v2_baggage.py b/tests/unit/integrations/otel/test_otel_v2_baggage.py index 930c01e524e..5855cb60aca 100644 --- a/tests/unit/integrations/otel/test_otel_v2_baggage.py +++ b/tests/unit/integrations/otel/test_otel_v2_baggage.py @@ -7,22 +7,22 @@ import pytest pytest.importorskip("opentelemetry") from litellm.integrations.otel import ( # noqa: E402 - GenAI, HTTP, + GenAI, LiteLLM, OpenTelemetryV2Config, promoted_baggage, ) -from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402 -from litellm.integrations.otel.plumbing import providers # noqa: E402 from litellm.integrations.otel.emitter import SpanEmitter # noqa: E402 +from litellm.integrations.otel.model.baggage import BAGGAGE_PROMOTED_KEYS # noqa: E402 from litellm.integrations.otel.model.payloads import ( # noqa: E402 GuardrailSpanData, LLMCallSpanData, ServiceSpanData, ) -from litellm.integrations.otel.model.baggage import BAGGAGE_PROMOTED_KEYS # noqa: E402 from litellm.integrations.otel.model.spans import SpanRole # noqa: E402 +from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402 +from litellm.integrations.otel.plumbing import providers # noqa: E402 def _payload(): @@ -63,9 +63,7 @@ def test_identity_promoted_onto_every_span(): root = engine.start_span(SpanRole.PROXY_REQUEST, "POST /chat/completions", ctx) root_ctx = ctx_mod.context_from_span(root, ctx) engine.emit(SpanRole.LLM_CALL, data, parent_context=root_ctx) - engine.emit( - SpanRole.GUARDRAIL, GuardrailSpanData("presidio", status="success"), root_ctx - ) + engine.emit(SpanRole.GUARDRAIL, GuardrailSpanData("presidio", status="success"), root_ctx) engine.emit(SpanRole.SERVICE, ServiceSpanData("redis", call_type="set"), root_ctx) root.end() @@ -161,9 +159,7 @@ def test_allowlisted_metadata_subkey_promoted_blob_excluded(): engine.emit(SpanRole.SERVICE, ServiceSpanData("redis", call_type="set"), ctx) (span,) = exporter.get_finished_spans() # allowlisted metadata sub-key is promoted - assert ( - span.attributes.get(f"{LiteLLM.METADATA_PREFIX}user_api_key_org_id") == "org1" - ) + assert span.attributes.get(f"{LiteLLM.METADATA_PREFIX}user_api_key_org_id") == "org1" # non-allowlisted metadata is NOT promoted (no full-blob dumping) assert all("private_note" not in k for k in span.attributes) @@ -208,6 +204,17 @@ def test_nested_metadata_key_promoted_under_caller_path(): assert not any(k.startswith(f"{LiteLLM.METADATA_PREFIX}requester_metadata") for k in span.attributes) +def test_llm_call_promoted_metadata_strips_requester_prefix_and_uses_allowlist(): + payload = _payload() + payload["metadata"]["requester_metadata"] = {"trace_id": "trace-123"} + data = LLMCallSpanData.from_standard_logging_payload( + payload, + metadata_keys=("requester_metadata.trace_id", "user_api_key_org_id", "missing"), + ) + assert data.promoted_metadata == {"trace_id": "trace-123", "user_api_key_org_id": "org1"} + assert LLMCallSpanData.from_standard_logging_payload(payload).promoted_metadata == {} + + def test_http_attributes_never_promoted(): """Even if http.* is present in baggage, the processor must not stamp it on child spans (it belongs on the SERVER span only).""" @@ -228,9 +235,7 @@ def test_http_attributes_never_promoted(): def test_arbitrary_upstream_baggage_not_promoted(): engine, exporter = _engine_and_exporter() - ctx = ctx_mod.set_request_baggage( - {LiteLLM.TEAM_ID: "t1", "some.upstream.key": "leak"} - ) + ctx = ctx_mod.set_request_baggage({LiteLLM.TEAM_ID: "t1", "some.upstream.key": "leak"}) engine.emit(SpanRole.SERVICE, ServiceSpanData("redis", call_type="set"), ctx) (span,) = exporter.get_finished_spans() assert span.attributes.get(LiteLLM.TEAM_ID) == "t1" diff --git a/tests/unit/integrations/otel/test_otel_v2_emitter.py b/tests/unit/integrations/otel/test_otel_v2_emitter.py index 11b2aa5fd67..949eb97bccd 100644 --- a/tests/unit/integrations/otel/test_otel_v2_emitter.py +++ b/tests/unit/integrations/otel/test_otel_v2_emitter.py @@ -20,7 +20,7 @@ from litellm.integrations.otel import ( # noqa: E402 ) from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402 from litellm.integrations.otel.plumbing import providers # noqa: E402 -from litellm.integrations.otel.emitter import SpanEmitter, span_attribute_limit # noqa: E402 +from litellm.integrations.otel.emitter import SpanEmitter, attribute_budget, span_attribute_limit # noqa: E402 from litellm.integrations.otel.emitter import stamp_error # noqa: E402 from litellm.integrations.otel.mappers.utils import MAX_TOOL_DEFINITION_ATTRS_PER_SPAN # noqa: E402 from litellm.integrations.otel.model.payloads import ( # noqa: E402 @@ -558,6 +558,113 @@ def test_prompt_turns_are_shed_before_response_choices(): ) +def test_metadata_blob_keeps_the_indexed_messages_it_competes_with(): + """The promoted ``metadata`` blob rides alongside the indexed messages on a boundary-opened span. + + The live regression shape: the LLM-call span is opened at the request + boundary and already carries a few stamped attributes, most of which the + mappers re-emit under the same keys. Charging those overwritten keys + against the attribute budget reserved slots the fit could never spend, so + the new ``metadata`` key displaced a whole middle message group — an + eviction invisible to the SDK's dropped-attributes counter because the fit + sheds before ``span.set_attribute``. With the budget counting only keys + the fit does not overwrite, the metadata key rides and the span keeps the + indexed-message count of the metadata-free span (at shedding granularity). + """ + cfg = OpenTelemetryV2Config( + exporter="in_memory", + legacy_compat=False, + mapper_names=["genai", "openinference"], + capture_message_content="span_only", + ) + provider, exporter = providers.in_memory_provider(cfg) + engine = SpanEmitter(providers.get_tracer(provider, "litellm-test"), cfg) + + def boundary_span(payload): + span = engine.start_span(SpanRole.LLM_CALL, "chat gpt-4o") + # what the boundary opener stamps before the typed payload exists + span.set_attribute(GenAI.REQUEST_MODEL, "gpt-4o") + span.set_attribute(LiteLLM.PROVIDER_MODEL, "gpt-4o-2024") + span.set_attribute("litellm.metadata.user_api_key_alias", "edge-key") + engine.finish_span( + SpanRole.LLM_CALL, + span, + LLMCallSpanData.from_standard_logging_payload( + payload, capture_content=True, metadata_keys=("user_api_key_alias",) + ), + ) + (finished,) = exporter.get_finished_spans() + exporter.clear() + return finished + + with_metadata = boundary_span(_conversation_payload(47, metadata={"user_api_key_alias": "edge-key"})) + without_metadata = boundary_span(_conversation_payload(47)) + _assert_core_intact(with_metadata) + a = with_metadata.attributes + + assert json.loads(a["metadata"]) == {"user_api_key_alias": "edge-key"} + assert "metadata" not in without_metadata.attributes + kept = _indexed_messages(a, "llm.input_messages") + baseline = _indexed_messages(without_metadata.attributes, "llm.input_messages") + assert kept == baseline or kept == baseline[:-1] + assert kept[0] == 0 and kept[-1] == 46 + assert len(a) <= SpanLimits().max_span_attributes + + +def test_preset_message_key_the_fit_sheds_still_never_overflows_the_span(): + """A pre-set indexed-message key the fit later sheds cannot push the span over its limit. + + The budget treats every mapped key already on the span as an overwrite + (free). If the fit then sheds that key, the value stamped earlier simply + stays in its slot, so the span holds one entry for it either way and the + total never exceeds the limit — the SDK's dropped-attributes counter stays + at zero. + """ + cfg = OpenTelemetryV2Config( + exporter="in_memory", + legacy_compat=False, + mapper_names=["genai", "openinference"], + capture_message_content="span_only", + ) + provider, exporter = providers.in_memory_provider(cfg) + engine = SpanEmitter(providers.get_tracer(provider, "litellm-test"), cfg) + span = engine.start_span(SpanRole.LLM_CALL, "chat gpt-4o") + span.set_attribute("llm.input_messages.1.message.role", "stale-role") + engine.finish_span( + SpanRole.LLM_CALL, + span, + LLMCallSpanData.from_standard_logging_payload( + _conversation_payload(60), capture_content=True, metadata_keys=("user_api_key_alias",) + ), + ) + (finished,) = exporter.get_finished_spans() + _assert_core_intact(finished) + assert len(finished.attributes) <= SpanLimits().max_span_attributes + + +def test_attribute_budget_counts_only_keys_the_fit_does_not_overwrite(): + """Pre-set attributes the mapped set overwrites consume no slot against the span limit. + + A boundary-opened LLM-call span already carries a few stamped attributes, + most of which the mappers re-emit under the same keys. Charging those + against the budget reserves slots the fit can never spend and sheds + indexed message attributes for nothing. + """ + cfg = OpenTelemetryV2Config(exporter="in_memory") + provider, _exporter = providers.in_memory_provider(cfg) + span = providers.get_tracer(provider, "litellm-test").start_span("s") + span.set_attribute("gen_ai.request.model", "gpt-4o") + span.set_attribute("litellm.metadata.user_api_key_alias", "edge-key") + try: + assert attribute_budget(span, 0) == SpanLimits().max_span_attributes - 2 + assert ( + attribute_budget(span, 0, frozenset({"gen_ai.request.model"})) + == SpanLimits().max_span_attributes - 1 + ) + finally: + span.end() + + def test_indexed_messages_respect_a_lower_span_attribute_count_limit(monkeypatch): """The budget follows the SDK's configured limit, not a hardcoded default.""" monkeypatch.setenv("OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT", "48") diff --git a/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py b/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py index cdff9c960f3..ca7a9b424ad 100644 --- a/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py +++ b/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py @@ -7,6 +7,7 @@ backends, so one trace lights up every configured destination. import json from collections.abc import Mapping +from itertools import chain from typing import Final import pytest @@ -19,6 +20,7 @@ from litellm.integrations.otel.mappers import ( WeaveMapper, resolve_mappers, ) +from litellm.integrations.otel.mappers.openinference import fit_indexed_messages from litellm.integrations.otel.model.payloads import ( EmbeddingOutput, LLMCallSpanData, @@ -29,6 +31,7 @@ from litellm.integrations.otel.model.payloads import ( ToolDefinition, ) from litellm.integrations.otel.model.trace_controls import TraceControls +from tests.unit.integrations.otel.test_otel_v2_sources_of_truth import _responses_payload def _llm_call(**overrides): @@ -118,6 +121,378 @@ def test_openinference_multimodal_content_text_only(): assert attrs["llm.input_messages.0.message.content"] == "hi there" +def test_openinference_output_tool_calls_preserve_calls_in_attributes_and_value(): + tool_calls: Final = [ + { + "id": "call_paris", + "type": "function", + "function": {"name": "lookup_weather", "arguments": '{"city": "Paris"}'}, + }, + { + "id": "call_search", + "function": {"name": "search", "arguments": {"q": 1}}, + "index": 0, + }, + "ignored", + ] + data: Final = _llm_call( + choices_out=( + { + "finish_reason": "tool_calls", + "message": {"role": "assistant", "content": None, "tool_calls": tool_calls}, + }, + ) + ) + attrs: Final = OpenInferenceMapper().map(data) + assert {key: value for key, value in attrs.items() if ".tool_calls." in key} == { + "llm.output_messages.0.message.tool_calls.0.tool_call.id": "call_paris", + "llm.output_messages.0.message.tool_calls.0.tool_call.function.name": "lookup_weather", + "llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments": '{"city": "Paris"}', + "llm.output_messages.0.message.tool_calls.1.tool_call.id": "call_search", + "llm.output_messages.0.message.tool_calls.1.tool_call.function.name": "search", + "llm.output_messages.0.message.tool_calls.1.tool_call.function.arguments": '{"q": 1}', + } + assert json.loads(attrs["output.value"]) == [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_paris", + "type": "function", + "function": {"name": "lookup_weather", "arguments": '{"city": "Paris"}'}, + }, + { + "id": "call_search", + "type": "function", + "function": {"name": "search", "arguments": '{"q": 1}'}, + }, + ], + } + ] + + +def test_openinference_responses_tool_calls_are_emitted_as_output_attributes(): + data: Final = LLMCallSpanData.from_standard_logging_payload( + _responses_payload( + [ + { + "type": "function_call", + "call_id": "call_resp", + "name": "lookup_weather", + "arguments": '{"city": "Paris"}', + } + ] + ), + capture_content=True, + ) + attrs: Final = OpenInferenceMapper().map(data) + tool_call: Final = "llm.output_messages.0.message.tool_calls.0.tool_call." + + assert {key: value for key, value in attrs.items() if ".tool_calls." in key} == { + tool_call + "id": "call_resp", + tool_call + "function.name": "lookup_weather", + tool_call + "function.arguments": '{"city": "Paris"}', + } + assert json.loads(attrs["output.value"]) == [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_resp", + "type": "function", + "function": {"name": "lookup_weather", "arguments": '{"city": "Paris"}'}, + } + ], + } + ] + + +def test_openinference_input_tool_calls_stay_in_value_only(): + data: Final = _llm_call( + messages_in=( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "lookup_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + ) + ) + attrs: Final = OpenInferenceMapper().map(data) + assert all(".tool_calls." not in key for key in attrs) + assert json.loads(attrs["input.value"]) == [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "lookup_weather", "arguments": '{"city": "Paris"}'}, + } + ], + } + ] + + +def test_openinference_output_tool_calls_do_not_shed_input_roles_under_budget(): + message_groups: Final = tuple( + ( + {"role": "user", "content": f"Question {index}"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": f"call_{index}", + "type": "function", + "function": {"name": "lookup_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": f"call_{index}", "content": f"Result {index}"}, + ) + for index in range(13) + ) + messages_in: Final = tuple(chain.from_iterable(message_groups)) + ({"role": "user", "content": "Final request"},) + tools: Final = tuple( + ToolDefinition(name=name, description="Tool", parameters_json='{"type":"object"}') + for name in ("lookup_weather", "search", "get_location", "convert_units") + ) + data: Final = _llm_call( + messages_in=messages_in, + tools=tools, + choices_out=( + { + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_output", + "type": "function", + "function": {"name": "lookup_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + }, + ), + ) + attrs: Final = fit_indexed_messages(OpenInferenceMapper().map(data), 128) + tool_call: Final = "llm.output_messages.0.message.tool_calls.0.tool_call." + + assert { + f"llm.input_messages.{index}.message.role": attrs.get(f"llm.input_messages.{index}.message.role") + for index in range(40) + } == {f"llm.input_messages.{index}.message.role": message["role"] for index, message in enumerate(messages_in)} + assert {key: value for key, value in attrs.items() if ".tool_calls." in key} == { + tool_call + "id": "call_output", + tool_call + "function.name": "lookup_weather", + tool_call + "function.arguments": '{"city": "Paris"}', + } + + +def test_openinference_budget_sheds_trailing_output_tool_calls_before_the_message(): + tool_calls: Final = tuple( + { + "id": f"call_{index}", + "type": "function", + "function": { + "name": "lookup_weather", + "arguments": f'{{"city": "C{index}"}}', + }, + } + for index in range(60) + ) + data: Final = _llm_call( + choices_out=( + { + "finish_reason": "tool_calls", + "message": {"role": "assistant", "content": None, "tool_calls": tool_calls}, + }, + ) + ) + mapped: Final = OpenInferenceMapper().map(data) + attrs: Final = fit_indexed_messages(mapped, len(mapped) - 34) + retained_tool_call_keys: Final = tuple(key for key in attrs if ".tool_calls." in key) + retained_tool_call_indices: Final = frozenset(int(key.split(".")[5]) for key in retained_tool_call_keys) + + assert retained_tool_call_indices == frozenset(range(50)) + assert len(retained_tool_call_keys) == 150 + assert attrs["llm.output_messages.0.message.role"] == "assistant" + assert not any(key.startswith("llm.input_messages.0.") for key in attrs) + assert len(attrs) == len(mapped) - 34 + assert json.loads(attrs["output.value"]) == [{"role": "assistant", "content": None, "tool_calls": list(tool_calls)}] + + +def test_openinference_plain_output_messages_keep_the_existing_value_shape(): + attrs: Final = OpenInferenceMapper().map(_llm_call()) + assert all(".tool_calls." not in key for key in attrs) + assert json.loads(attrs["output.value"]) == [{"role": "assistant", "content": "Sunny."}] + + +def test_openinference_metadata_contains_only_promoted_metadata(): + attrs: Final = OpenInferenceMapper().map( + _llm_call(promoted_metadata={"trace_marker": "m", "user_api_key_alias": "k"}) + ) + assert json.loads(attrs["metadata"]) == {"trace_marker": "m", "user_api_key_alias": "k"} + assert "metadata" not in OpenInferenceMapper().map(_llm_call()) + + +def _long_prompt_messages(count: int) -> tuple[dict[str, str], ...]: + return tuple({"role": "user" if i == 0 else "assistant", "content": f"turn {i}"} for i in range(count)) + + +def _input_message_keys(attrs: Mapping[str, object]) -> frozenset[str]: + return frozenset(key for key in attrs if key.startswith("llm.input_messages.")) + + +def test_openinference_metadata_does_not_displace_indexed_messages_under_budget(): + """The ``metadata`` blob rides in reclaimed slots instead of evicting messages. + + ``metadata`` competes for the OTel 128-attribute span budget, and shedding + happens before ``span.set_attribute``, so the SDK's dropped-attributes + counter stays at zero — the eviction is invisible. The fit pins the + metadata key behind every message group: a squeezed span keeps the message + attributes it would keep without the metadata blob (at whole-group + granularity, at most one group of headroom difference), and the metadata + key survives alongside them. + """ + messages: Final = _long_prompt_messages(62) + + def fitted(promoted: Mapping[str, str], budget: int | None = None): + mapped: Final = OpenInferenceMapper().map(_llm_call(messages_in=messages, promoted_metadata=promoted)) + assert len(mapped) > 128, "fixture must squeeze the span attribute budget" + return fit_indexed_messages(mapped, budget if budget is not None else len(mapped) - 5) + + without_metadata: Final = fitted({}) + with_metadata: Final = fitted({"user_api_key_alias": "edge-key"}) + assert "metadata" not in without_metadata + assert json.loads(with_metadata["metadata"]) == {"user_api_key_alias": "edge-key"} + assert _input_message_keys(with_metadata) == _input_message_keys(without_metadata) + + # Against the absolute span limit the displacement is bounded by the + # whole-group shedding granularity: at most one message group. + without_at_limit: Final = fitted({}, 128) + with_at_limit: Final = fitted({"user_api_key_alias": "edge-key"}, 128) + assert len(_input_message_keys(with_at_limit)) >= len(_input_message_keys(without_at_limit)) - 2 + assert "metadata" in with_at_limit + + +def test_openinference_metadata_sheds_only_after_every_indexed_message(): + """Metadata sheds last: only once every indexed message attribute is gone.""" + mapped: Final = OpenInferenceMapper().map( + _llm_call(messages_in=_long_prompt_messages(6), promoted_metadata={"user_api_key_alias": "edge-key"}) + ) + message_keys: Final = frozenset(key for key in mapped if ".message." in key) + # Budget too small for the message family alone: everything indexed goes, + # and the metadata blob absorbs the residual shortfall with it. + starved: Final = fit_indexed_messages(mapped, len(mapped) - len(message_keys) - 1) + assert not any(key in starved for key in message_keys) + assert "metadata" not in starved + # One slot more and metadata survives alongside zero indexed messages. + last_standing: Final = fit_indexed_messages(mapped, len(mapped) - len(message_keys)) + assert not any(key in last_standing for key in message_keys) + assert "metadata" in last_standing + + +def test_openinference_raw_tool_arguments_fall_back_to_repr_instead_of_raising(): + """Malformed Python tool arguments must not lose the span. + + A plain ``Function()`` constructor JSON-serializes arguments, but provider + adapters and ``model_construct`` responses hand over raw Python objects — + tuple-keyed dicts and cycles that ``json.dumps`` raises on. The mapper + serializes them with a ``repr`` fallback instead of letting the exception + escape before the span is exported. + """ + # rebind-ok: a self-referencing dict cannot be built in one shot — the cycle + # only exists once the finished dict is inserted into itself. + circular: dict[str, object] = {} + circular["self"] = circular + for label, raw_arguments in (("tuple-key", {(1, 2): "v"}), ("circular", circular)): + data: Final = _llm_call( + choices_out=( + { + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "tc1", "type": "function", "function": {"name": "f", "arguments": raw_arguments}} + ], + }, + }, + ) + ) + attrs: Final = OpenInferenceMapper().map(data) + assert attrs["llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments"] == repr( + raw_arguments + ), label + assert attrs["llm.output_messages.0.message.role"] == "assistant" + + +def test_openinference_payload_tool_arguments_with_raw_objects_map_without_raising(): + """The standard-logging payload path (provider adapter responses) survives raw argument objects too.""" + payload: Final = { + "call_type": "acompletion", + "custom_llm_provider": "openai", + "model": "gpt-4o", + "prompt_tokens": 3, + "completion_tokens": 2, + "total_tokens": 5, + "stream": False, + "model_parameters": {}, + "response": { + "id": "resp_bad", + "model": "gpt-4o-2024", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": "hi", + "tool_calls": [ + {"id": "tc1", "type": "function", "function": {"name": "f", "arguments": {(1, 2): "v"}}} + ], + }, + } + ], + }, + "metadata": {"user_api_key_alias": "edge-key"}, + "status": "success", + "litellm_call_id": "call_raw_args", + "hidden_params": {}, + } + data: Final = LLMCallSpanData.from_standard_logging_payload( + payload, capture_content=True, metadata_keys=("user_api_key_alias",) + ) + attrs: Final = OpenInferenceMapper().map(data) + arguments: Final = attrs["llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments"] + assert arguments == repr({(1, 2): "v"}) + assert json.loads(attrs["metadata"]) == {"user_api_key_alias": "edge-key"} + + +def test_openinference_attribute_fit_keeps_pinned_messages_on_long_prompts(): + messages: Final = _long_prompt_messages(4000) + mapped: Final = OpenInferenceMapper().map( + _llm_call(messages_in=messages, promoted_metadata={"user_api_key_alias": "k"}) + ) + fitted: Final = fit_indexed_messages(mapped, 128) + + assert len(fitted) <= 128 + assert fitted["llm.input_messages.0.message.role"] == "user" + assert fitted["llm.input_messages.3999.message.role"] == "assistant" + + # --------------------------------------------------------------------------- # # Langfuse # --------------------------------------------------------------------------- # diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 72c362425bb..e04f8af5a66 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -1,5 +1,9 @@ import asyncio +import copy import datetime as dt +import json +import pickle +import threading from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional from unittest.mock import AsyncMock @@ -8,7 +12,10 @@ import pytest from litellm.integrations.custom_guardrail import ( DEFAULT_ADVISORY_MESSAGE, CustomGuardrail, + _request_is_streaming, + guardrail_request_data_with_streaming, log_guardrail_information, + without_server_streaming_classification, ) from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy._types import CallTypes, UserAPIKeyAuth @@ -531,6 +538,233 @@ class TestCustomGuardrailShouldRunGuardrail: assert always_on.should_run_guardrail(data=forged, event_type=GuardrailEventHooks.pre_call) is True +_STREAM_SCOPE_HOOKS: Final = ( + GuardrailEventHooks.pre_call, + GuardrailEventHooks.during_call, + GuardrailEventHooks.post_call, +) + + +class TestCustomGuardrailStreamScope: + @pytest.mark.parametrize("event_type", _STREAM_SCOPE_HOOKS) + @pytest.mark.parametrize("stream", [True, False]) + @pytest.mark.parametrize("stream_scope", [None, "both"]) + def test_both_and_omitted_run_on_streaming_and_non_streaming( + self, + event_type: GuardrailEventHooks, + stream: bool, + stream_scope: str | None, + ): + guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=True, + event_hook=event_type, + stream_scope=stream_scope, + ) + assert guardrail.should_run_guardrail({"stream": stream}, event_type) is True + + @pytest.mark.parametrize("event_type", _STREAM_SCOPE_HOOKS) + def test_scalar_streaming_skips_non_streaming(self, event_type: GuardrailEventHooks): + guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=True, + event_hook=event_type, + stream_scope="streaming", + ) + assert guardrail.should_run_guardrail({"stream": True}, event_type) is True + assert guardrail.should_run_guardrail({"stream": False}, event_type) is False + assert guardrail.should_run_guardrail({}, event_type) is False + + @pytest.mark.parametrize("event_type", _STREAM_SCOPE_HOOKS) + def test_scalar_non_streaming_skips_streaming(self, event_type: GuardrailEventHooks): + guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=True, + event_hook=event_type, + stream_scope="non_streaming", + ) + assert guardrail.should_run_guardrail({"stream": False}, event_type) is True + assert guardrail.should_run_guardrail({}, event_type) is True + assert guardrail.should_run_guardrail({"stream": True}, event_type) is False + + def test_per_mode_map_applies_to_named_hooks_only(self): + guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=True, + event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call], + stream_scope={"pre_call": "both", "post_call": "streaming"}, + ) + assert guardrail.should_run_guardrail({"stream": True}, GuardrailEventHooks.pre_call) is True + assert guardrail.should_run_guardrail({"stream": False}, GuardrailEventHooks.pre_call) is True + assert guardrail.should_run_guardrail({"stream": True}, GuardrailEventHooks.post_call) is True + assert guardrail.should_run_guardrail({"stream": False}, GuardrailEventHooks.post_call) is False + + def test_default_on_early_return_still_honors_stream_scope(self): + guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + stream_scope="non_streaming", + ) + assert guardrail.should_run_guardrail({"stream": False}, GuardrailEventHooks.post_call) is True + assert guardrail.should_run_guardrail({"stream": True}, GuardrailEventHooks.post_call) is False + + def test_apply_stream_scope_overwrites_constructor_default(self): + guardrail = CustomGuardrail( + guardrail_name="scoped", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + assert guardrail.should_run_guardrail({"stream": False}, GuardrailEventHooks.post_call) is True + guardrail.apply_stream_scope("streaming") + assert guardrail.should_run_guardrail({"stream": True}, GuardrailEventHooks.post_call) is True + assert guardrail.should_run_guardrail({"stream": False}, GuardrailEventHooks.post_call) is False + + def test_direct_constructor_normalizes_mixed_case_map_keys(self): + guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + stream_scope={"Pre_Call": "streaming"}, + ) + assert guardrail.should_run_guardrail({"stream": True}, GuardrailEventHooks.pre_call) is True + assert guardrail.should_run_guardrail({"stream": False}, GuardrailEventHooks.pre_call) is False + + def test_realtime_transcription_counts_as_streaming(self): + streaming_only = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=True, + event_hook=GuardrailEventHooks.realtime_input_transcription, + stream_scope="streaming", + ) + assert ( + streaming_only.should_run_guardrail( + {"litellm_metadata": {}}, GuardrailEventHooks.realtime_input_transcription + ) + is True + ) + non_streaming_only = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=True, + event_hook=GuardrailEventHooks.realtime_input_transcription, + stream_scope="non_streaming", + ) + assert ( + non_streaming_only.should_run_guardrail( + {"litellm_metadata": {}}, GuardrailEventHooks.realtime_input_transcription + ) + is False + ) + + def test_path_defined_streaming_classification_cannot_be_spoofed(self): + generate_content_body: Final = {"contents": [{"parts": [{"text": "hi"}]}]} + streaming_only = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + stream_scope="streaming", + ) + assert streaming_only.should_run_guardrail(generate_content_body, GuardrailEventHooks.pre_call) is False + assert ( + streaming_only.should_run_guardrail( + {**generate_content_body, "is_streaming_request": True}, + GuardrailEventHooks.pre_call, + ) + is False + ) + assert ( + streaming_only.should_run_guardrail( + {**generate_content_body, "is_streaming_request": "litellm-server-streaming"}, + GuardrailEventHooks.pre_call, + ) + is False + ) + assert ( + streaming_only.should_run_guardrail( + {**generate_content_body, "litellm_server_streaming_classification": True}, + GuardrailEventHooks.pre_call, + ) + is False + ) + server_streaming_data: Final = guardrail_request_data_with_streaming( + generate_content_body, + is_streaming=True, + ) + assert ( + streaming_only.should_run_guardrail( + server_streaming_data, + GuardrailEventHooks.pre_call, + ) + is True + ) + non_streaming_only = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + stream_scope="non_streaming", + ) + assert ( + non_streaming_only.should_run_guardrail( + server_streaming_data, + GuardrailEventHooks.pre_call, + ) + is False + ) + + def test_streaming_classification_is_json_serializable_without_spoofing(self): + d: Final = guardrail_request_data_with_streaming({}, is_streaming=True) + serialized: Final = json.dumps(d) + assert _request_is_streaming(d) is True + + round_tripped: Final = json.loads(serialized) + assert _request_is_streaming(round_tripped) is False + assert round_tripped["litellm_server_streaming_classification"] == "litellm-server-streaming" + assert isinstance(round_tripped["litellm_server_streaming_classification"], str) + assert "litellm_server_streaming_classification" not in without_server_streaming_classification(round_tripped) + + def test_streaming_classification_preserves_caller_fields_and_removes_only_server_marker(self): + caller_data: Final = { + "contents": [{"parts": [{"text": "hi"}]}], + "is_streaming_request": "caller-value", + "litellm_server_streaming_classification": True, + } + non_streaming_data: Final = guardrail_request_data_with_streaming(caller_data, is_streaming=False) + server_streaming_data: Final = guardrail_request_data_with_streaming(caller_data, is_streaming=True) + + assert non_streaming_data is not caller_data + assert server_streaming_data is not caller_data + assert non_streaming_data == caller_data + assert server_streaming_data["is_streaming_request"] == "caller-value" + assert server_streaming_data["litellm_server_streaming_classification"] is not True + assert without_server_streaming_classification(caller_data) == caller_data + assert "litellm_server_streaming_classification" not in without_server_streaming_classification( + server_streaming_data + ) + + def test_server_streaming_classification_survives_scan_raw_request_snapshot(self): + from litellm.litellm_core_utils.core_helpers import independent_snapshot + + generate_content_body: Final = {"contents": [{"parts": [{"text": "hi"}]}]} + snapshot: Final = independent_snapshot( + guardrail_request_data_with_streaming(generate_content_body, is_streaming=True) + ) + streaming_only = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + stream_scope="streaming", + scan_raw_request=True, + ) + assert streaming_only.should_run_guardrail(snapshot, GuardrailEventHooks.pre_call) is True + assert ( + streaming_only.should_run_guardrail( + independent_snapshot({**generate_content_body, "is_streaming_request": True}), + GuardrailEventHooks.pre_call, + ) + is False + ) + + class TestApplyGuardrailCheck: def test_apply_guardrail_check_only_on_direct_implementation(self): """ @@ -3307,3 +3541,113 @@ class TestCustomGuardrailTimeout: ) assert guardrail.timeout == 7.0 + + +@pytest.mark.parametrize("stream_scope", [None, "both", {"post_call": "streaming"}]) +def test_guardrail_survives_deepcopy_and_pickle_with_its_stream_scope(stream_scope): + guardrail: Final = CustomGuardrail( + guardrail_name="copyable", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + stream_scope=stream_scope, + ) + expected: Final = guardrail.should_run_guardrail({"stream": False}, GuardrailEventHooks.post_call) + for clone in (copy.deepcopy(guardrail), pickle.loads(pickle.dumps(guardrail))): + assert dict(clone.stream_scope_by_hook) == dict(guardrail.stream_scope_by_hook) + assert clone.should_run_guardrail({"stream": False}, GuardrailEventHooks.post_call) is expected + + +class _LockHoldingGuardrail(CustomGuardrail): + def __init__(self, **kwargs): + self.lock = threading.Lock() + super().__init__(**kwargs) + + def __getstate__(self): + state: Final = dict(self.__dict__) + state.pop("lock") + return state + + def __setstate__(self, state): + self.__dict__.update(state) + self.lock = threading.Lock() + + +class _SetstateOnlyGuardrail(CustomGuardrail): + def __setstate__(self, state): + self.__dict__.update(state) + self.restored = True + + +class _SlotsGuardrail(CustomGuardrail): + __slots__ = ("vendor_client",) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.vendor_client = "vendor-client-object" + + +def _subclass_guardrails(): + kwargs: Final = { + "guardrail_name": "vendor", + "default_on": True, + "event_hook": GuardrailEventHooks.post_call, + "stream_scope": {"post_call": "streaming"}, + } + return ( + pytest.param(_LockHoldingGuardrail(**kwargs), id="dict-getstate"), + pytest.param(_SetstateOnlyGuardrail(**kwargs), id="setstate-only"), + pytest.param(_SlotsGuardrail(**kwargs), id="slots"), + ) + + +@pytest.mark.parametrize("guardrail", _subclass_guardrails()) +@pytest.mark.parametrize( + "cloner", + [copy.copy, copy.deepcopy, lambda g: pickle.loads(pickle.dumps(g))], + ids=["copy", "deepcopy", "pickle"], +) +def test_out_of_tree_guardrail_subclasses_survive_copy_and_pickle(guardrail, cloner): + clone = cloner(guardrail) + + if isinstance(clone, _LockHoldingGuardrail): + assert isinstance(clone.lock, type(threading.Lock())) + elif isinstance(clone, _SetstateOnlyGuardrail): + assert clone.restored is True + else: + assert clone.vendor_client == "vendor-client-object" + assert clone.should_run_guardrail({"stream": False}, GuardrailEventHooks.post_call) is False + assert clone.should_run_guardrail({"stream": True}, GuardrailEventHooks.post_call) is True + + +def test_router_constructs_with_a_dict_getstate_guardrail_in_deployment_callbacks(): + import litellm + + litellm.Router( + model_list=[ + { + "model_name": "m", + "litellm_params": { + "model": "openai/gpt-5.4-mini", + "api_key": "sk-test", + "callbacks": [ + _LockHoldingGuardrail( + guardrail_name="vendor", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + stream_scope={"post_call": "streaming"}, + ) + ], + }, + } + ] + ) + + +def test_subclass_that_skips_super_init_still_runs_with_default_scope(): + class NoSuperInit(CustomGuardrail): + def __init__(self) -> None: + self.guardrail_name = "no-super" + self.event_hook = None + self.default_on = True + + assert NoSuperInit().should_run_guardrail({"stream": True}, GuardrailEventHooks.pre_call) is True diff --git a/tests/unit/integrations/test_opentelemetry_request_spans.py b/tests/unit/integrations/test_opentelemetry_request_spans.py new file mode 100644 index 00000000000..9ec40c00cb4 --- /dev/null +++ b/tests/unit/integrations/test_opentelemetry_request_spans.py @@ -0,0 +1,157 @@ +import asyncio +import json +from collections.abc import Sequence +from typing import Final + +import httpx +import pytest +import respx +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +import litellm +from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig + +_OPENAI_URL: Final = "https://api.openai.com/v1/chat/completions" +_EXPECTED_SPAN_NAMES: Final = ("litellm_request", "raw_gen_ai_request") +_USER: Final = "OTEL_USER" +_USAGE: Final = {"prompt_tokens": 8, "completion_tokens": 2, "total_tokens": 10} +_COMPLETION: Final = { + "id": "chatcmpl-otel", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4.1-mini-2025-04-14", + "service_tier": "default", + "system_fingerprint": "fp_otel", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop"}], + "usage": _USAGE, +} +_STREAM: Final = ( + "".join( + f"data: {json.dumps(chunk)}\n\n" + for chunk in ( + { + "id": "chatcmpl-otel", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-4.1-mini-2025-04-14", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hello"}, "finish_reason": None}], + }, + { + "id": "chatcmpl-otel", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-4.1-mini-2025-04-14", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": _USAGE, + }, + ) + ) + + "data: [DONE]\n\n" +) +_LITELLM_REQUEST_ATTRIBUTES: Final = ( + "gen_ai.request.model", + "gen_ai.system", + "gen_ai.request.temperature", + "llm.is_streaming", + "llm.user", + "gen_ai.response.id", + "gen_ai.response.model", + "gen_ai.usage.total_tokens", + "gen_ai.usage.output_tokens", + "gen_ai.usage.input_tokens", +) +_RAW_STREAMING_ATTRIBUTES: Final = ( + "llm.openai.messages", + "llm.openai.temperature", + "llm.openai.user", + "llm.openai.extra_body", + "llm.openai.model", +) +_RAW_NON_STREAMING_ATTRIBUTES: Final = ( + *_RAW_STREAMING_ATTRIBUTES, + "llm.openai.id", + "llm.openai.choices", + "llm.openai.created", + "llm.openai.object", + "llm.openai.service_tier", + "llm.openai.system_fingerprint", + "llm.openai.usage", +) + + +def _is_our_request(span: ReadableSpan) -> bool: + return span.name == "litellm_request" and (span.attributes or {}).get("llm.user") == _USER + + +def _trace_id(span: ReadableSpan) -> int: + assert span.context is not None + return span.context.trace_id + + +class _SignallingExporter(InMemorySpanExporter): + def __init__(self, loop: asyncio.AbstractEventLoop) -> None: + super().__init__() + self.loop: Final = loop + self.request_span_exported: Final = asyncio.Event() + + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + result: Final = super().export(spans) + if any(_is_our_request(span) for span in spans): + self.loop.call_soon_threadsafe(self.request_span_exported.set) + return result + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [True, False]) +async def test_otel_callback_emits_the_request_and_raw_provider_spans( + streaming: bool, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OTEL_SEMCONV_STABILITY_OPT_IN", raising=False) + exporter: Final = _SignallingExporter(asyncio.get_running_loop()) + tracer_provider: Final = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + monkeypatch.setattr( + litellm, + "callbacks", + [OpenTelemetry(config=OpenTelemetryConfig(exporter=exporter), tracer_provider=tracer_provider)], + ) + respx_mock.post(_OPENAI_URL).mock( + return_value=httpx.Response(200, text=_STREAM, headers={"content-type": "text/event-stream"}) + if streaming + else httpx.Response(200, json=_COMPLETION) + ) + + response: Final = await litellm.acompletion( + model="gpt-4.1-mini", + messages=[{"role": "user", "content": "hi"}], + temperature=0.1, + user=_USER, + stream=streaming, + api_key="sk-unit-test", + ) + if streaming: + assert [chunk async for chunk in response] + await asyncio.wait_for(exporter.request_span_exported.wait(), timeout=10) + + finished: Final = exporter.get_finished_spans() + request_span: Final = next(span for span in finished if _is_our_request(span)) + ours: Final = tuple(span for span in finished if _trace_id(span) == _trace_id(request_span)) + assert tuple(sorted(span.name for span in ours)) == _EXPECTED_SPAN_NAMES + spans: Final = {span.name: span for span in ours} + request_attributes: Final = spans["litellm_request"].attributes or {} + assert all(request_attributes.get(name) is not None for name in _LITELLM_REQUEST_ATTRIBUTES) + assert request_attributes["gen_ai.request.model"] == "gpt-4.1-mini" + assert request_attributes["gen_ai.system"] == "openai" + assert request_attributes["gen_ai.request.temperature"] == 0.1 + assert request_attributes["llm.is_streaming"] == str(streaming) + assert request_attributes["llm.user"] == _USER + assert request_attributes["gen_ai.response.id"] == "chatcmpl-otel" + assert request_attributes["gen_ai.usage.input_tokens"] == _USAGE["prompt_tokens"] + assert request_attributes["gen_ai.usage.output_tokens"] == _USAGE["completion_tokens"] + assert request_attributes["gen_ai.usage.total_tokens"] == _USAGE["total_tokens"] + raw_attributes: Final = spans["raw_gen_ai_request"].attributes or {} + expected_raw: Final = _RAW_STREAMING_ATTRIBUTES if streaming else _RAW_NON_STREAMING_ATTRIBUTES + assert all(raw_attributes.get(name) is not None for name in expected_raw) diff --git a/tests/unit/integrations/vector_store_integrations/test_bedrock_kb_context_offline.py b/tests/unit/integrations/vector_store_integrations/test_bedrock_kb_context_offline.py new file mode 100644 index 00000000000..20d3a310bdb --- /dev/null +++ b/tests/unit/integrations/vector_store_integrations/test_bedrock_kb_context_offline.py @@ -0,0 +1,289 @@ +import itertools +import json +from typing import Final, TypedDict + +from typing_extensions import ReadOnly + +import httpx +import pytest +import respx +from pydantic import TypeAdapter + +import litellm +import litellm.proxy.proxy_server as proxy_server +from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import VectorStorePreCallHook +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.vector_stores.vector_store_registry import LiteLLM_ManagedVectorStore, VectorStoreRegistry + +_KB_ID: Final = "T37J8R4WTM" +_KB_URL: Final = f"https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases/{_KB_ID}/retrieve" +_ANTHROPIC_URL: Final = "https://api.anthropic.com/v1/messages" +_OPENAI_URL: Final = "https://api.openai.com/v1/chat/completions" +_KB_TEXT: Final = "LiteLLM is a library that simplifies LLM API access" +_PREFIX: Final = VectorStorePreCallHook.CONTENT_PREFIX_STRING +_BODY: Final = TypeAdapter(dict[str, object]) + +_ANTHROPIC_MESSAGE: Final = { + "id": "msg_kb", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "LiteLLM simplifies LLM access."}], + "model": "claude-haiku-4-5-20251001", + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 100, "output_tokens": 50}, +} +_OPENAI_COMPLETION: Final = { + "id": "chatcmpl-kb", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-5-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}, +} +_ANTHROPIC_STREAM: Final = ( + 'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_kb_stream","type":"message",' + '"role":"assistant","content":[],"model":"claude-haiku-4-5-20251001","stop_reason":null,"stop_sequence":null,' + '"usage":{"input_tokens":10,"output_tokens":1}}}\n\n' + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}\n\n' + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"LiteLLM"}}\n\n' + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n' + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},' + '"usage":{"output_tokens":2}}\n\n' + 'event: message_stop\ndata: {"type":"message_stop"}\n\n' +) + + +class _ChatMessage(TypedDict): + role: ReadOnly[str] + content: ReadOnly[str] + + +@pytest.fixture(autouse=True) +def _knowledge_base(monkeypatch: pytest.MonkeyPatch, fake_provider_credentials: None) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("AWS_REGION", "us-west-2") + monkeypatch.setenv("AWS_REGION_NAME", "us-west-2") + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + vector_stores=[LiteLLM_ManagedVectorStore(vector_store_id=_KB_ID, custom_llm_provider="bedrock")] + ), + raising=False, + ) + + +def _kb_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post(_KB_URL).mock( + return_value=httpx.Response( + 200, + json={"retrievalResults": [{"content": {"text": _KB_TEXT, "type": "TEXT"}, "score": 0.9, "metadata": {}}]}, + ) + ) + + +def _sent_body(route: respx.Route) -> dict[str, object]: + return _BODY.validate_json(route.calls.last.request.content) + + +@pytest.mark.asyncio +async def test_completion_with_vector_store_ids_prepends_the_kb_context_block(respx_mock: respx.MockRouter) -> None: + _kb_route(respx_mock) + anthropic: Final = respx_mock.post(_ANTHROPIC_URL).mock(return_value=httpx.Response(200, json=_ANTHROPIC_MESSAGE)) + + await litellm.acompletion( + model="anthropic/claude-haiku-4-5-20251001", + messages=[{"role": "user", "content": "what is litellm?"}], + vector_store_ids=[_KB_ID], + ) + + messages: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(_sent_body(anthropic)["messages"]) + content: Final = TypeAdapter(tuple[dict[str, str], ...]).validate_python(messages[0]["content"]) + assert anthropic.call_count == 1 + assert [block["type"] for block in content] == ["text", "text"] + assert content[0]["text"] == f"{_PREFIX}{_KB_TEXT}\n\n" + assert content[1]["text"] == "what is litellm?" + + +@pytest.mark.asyncio +async def test_streaming_completion_carries_the_search_results_on_a_chunk_delta(respx_mock: respx.MockRouter) -> None: + _kb_route(respx_mock) + respx_mock.post(_ANTHROPIC_URL).mock( + return_value=httpx.Response(200, text=_ANTHROPIC_STREAM, headers={"content-type": "text/event-stream"}) + ) + + response: Final = await litellm.acompletion( + model="anthropic/claude-haiku-4-5-20251001", + messages=[{"role": "user", "content": "what is litellm?"}], + vector_store_ids=[_KB_ID], + stream=True, + ) + chunks: Final = tuple([chunk async for chunk in response]) + choices: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks)) + annotated: Final = tuple( + choice.delta.provider_specific_fields["search_results"] + for choice in choices + if choice.delta.provider_specific_fields and "search_results" in choice.delta.provider_specific_fields + ) + + assert len(chunks) > 0 + assert len(annotated) >= 1 + assert annotated[0][0]["object"] == "vector_store.search_results.page" + assert annotated[0][0]["data"][0]["content"][0]["text"] == _KB_TEXT + + +@pytest.mark.asyncio +async def test_file_search_filters_reach_the_kb_as_a_bedrock_equals_filter(respx_mock: respx.MockRouter) -> None: + kb: Final = _kb_route(respx_mock) + respx_mock.post(_ANTHROPIC_URL).mock(return_value=httpx.Response(200, json=_ANTHROPIC_MESSAGE)) + + response: Final = await litellm.acompletion( + model="anthropic/claude-haiku-4-5-20251001", + messages=[{"role": "user", "content": "what is litellm?"}], + max_tokens=10, + tools=[ + { + "type": "file_search", + "vector_store_ids": [_KB_ID], + "filters": {"key": "user_id", "value": "fake-user-id", "operator": "eq"}, + } + ], + ) + + retrieval: Final = TypeAdapter(dict[str, dict[str, dict[str, object]]]).validate_python( + _sent_body(kb)["retrievalConfiguration"] + ) + assert retrieval["vectorSearchConfiguration"]["filter"] == {"equals": {"key": "user_id", "value": "fake-user-id"}} + assert response.choices[0].message.content == "LiteLLM simplifies LLM access." + + +def _openai_messages(route: respx.Route) -> tuple[_ChatMessage, ...]: + return TypeAdapter(tuple[_ChatMessage, ...]).validate_python(_sent_body(route)["messages"]) + + +@pytest.mark.asyncio +async def test_openai_request_with_vector_store_ids_leads_with_a_kb_context_user_message( + respx_mock: respx.MockRouter, +) -> None: + _kb_route(respx_mock) + openai: Final = respx_mock.post(_OPENAI_URL).mock(return_value=httpx.Response(200, json=_OPENAI_COMPLETION)) + + await litellm.acompletion( + model="gpt-5-mini", + messages=[{"role": "user", "content": "what is litellm?"}], + vector_store_ids=[_KB_ID], + ) + + assert _openai_messages(openai) == ( + _ChatMessage(role="user", content=f"{_PREFIX}{_KB_TEXT}\n\n"), + _ChatMessage(role="user", content="what is litellm?"), + ) + + +@pytest.mark.asyncio +async def test_a_managed_file_search_tool_is_resolved_locally_and_not_sent_upstream( + respx_mock: respx.MockRouter, +) -> None: + _kb_route(respx_mock) + openai: Final = respx_mock.post(_OPENAI_URL).mock(return_value=httpx.Response(200, json=_OPENAI_COMPLETION)) + + await litellm.acompletion( + model="gpt-5-mini", + messages=[{"role": "user", "content": "what is litellm?"}], + tools=[{"type": "file_search", "vector_store_ids": [_KB_ID]}], + ) + + assert _openai_messages(openai)[0] == _ChatMessage(role="user", content=f"{_PREFIX}{_KB_TEXT}\n\n") + assert "tools" not in _sent_body(openai) + + +@pytest.mark.asyncio +async def test_an_unknown_vector_store_tool_is_forwarded_while_the_known_one_is_resolved( + respx_mock: respx.MockRouter, +) -> None: + _kb_route(respx_mock) + openai: Final = respx_mock.post(_OPENAI_URL).mock(return_value=httpx.Response(200, json=_OPENAI_COMPLETION)) + + await litellm.acompletion( + model="gpt-5-mini", + messages=[{"role": "user", "content": "what is litellm?"}], + tools=[ + {"type": "file_search", "vector_store_ids": [_KB_ID]}, + {"type": "file_search", "vector_store_ids": ["unknownVS"]}, + ], + ) + + assert _openai_messages(openai)[0] == _ChatMessage(role="user", content=f"{_PREFIX}{_KB_TEXT}\n\n") + assert _sent_body(openai)["tools"] == [{"type": "file_search", "vector_store_ids": ["unknownVS"]}] + + +def _authorized_key() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-kb-proxy", user_id="kb-proxy-user") + + +class _SearchResultsPage(TypedDict): + object: ReadOnly[str] + search_query: ReadOnly[str] + data: ReadOnly[list[dict[str, object]]] + + +class _ProviderFields(TypedDict): + search_results: ReadOnly[list[_SearchResultsPage]] + + +class _ProxyMessage(TypedDict): + role: ReadOnly[str] + content: ReadOnly[str] + provider_specific_fields: ReadOnly[_ProviderFields] + + +class _ProxyChoice(TypedDict): + message: ReadOnly[_ProxyMessage] + + +class _ProxyCompletion(TypedDict): + choices: ReadOnly[list[_ProxyChoice]] + + +@pytest.mark.asyncio +async def test_proxy_http_response_keeps_the_kb_search_results_in_provider_specific_fields( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + kb: Final = _kb_route(respx_mock) + upstream: Final = respx_mock.post(_OPENAI_URL).mock(return_value=httpx.Response(200, json=_OPENAI_COMPLETION)) + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[ + {"model_name": "gpt-5-mini", "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fixture"}} + ], + num_retries=0, + ), + ) + monkeypatch.setitem(proxy_server.app.dependency_overrides, user_api_key_auth, _authorized_key) + + async with httpx.AsyncClient( + transport=httpx.ASGITransport(proxy_server.app), base_url="http://kb-proxy.test" + ) as client: + result: Final = await client.post( + "/v1/chat/completions", + json={ + "model": "gpt-5-mini", + "messages": [{"role": "user", "content": "what is litellm?"}], + "vector_store_ids": [_KB_ID], + }, + ) + + assert result.status_code == 200, result.text + assert kb.call_count == 1 + assert upstream.call_count == 1 + message: Final = TypeAdapter(_ProxyCompletion).validate_json(result.content)["choices"][0]["message"] + assert message["content"] == "ok" + pages: Final = message["provider_specific_fields"]["search_results"] + assert [page["object"] for page in pages] == ["vector_store.search_results.page"] + assert pages[0]["search_query"] == "what is litellm?" + assert pages[0]["data"] + assert _KB_TEXT in json.dumps(pages[0]["data"]) diff --git a/tests/unit/interactions/test_main.py b/tests/unit/interactions/test_main.py index 80cce7f1b6f..626d53d3384 100644 --- a/tests/unit/interactions/test_main.py +++ b/tests/unit/interactions/test_main.py @@ -1,4 +1,10 @@ +import base64 +import json +from typing import Final + +import httpx import pytest +from respx import MockRouter import litellm import litellm.interactions as interactions @@ -18,3 +24,71 @@ class TestGoogleInteractionsCreate: input="Hello", api_key=api_key, ) + + +class TestInteractionsAcreateOffline: + @pytest.fixture(autouse=True) + def _httpx_only_transport(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + + @pytest.mark.usefixtures("fake_provider_credentials") + @pytest.mark.asyncio + async def test_acreate_simple_gemini(self, respx_mock: MockRouter) -> None: + route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "interaction-offline", + "object": "interaction", + "model": "gemini-2.5-flash", + "status": "completed", + "steps": [{"type": "model_output", "content": [{"type": "text", "text": "299792458"}]}], + "usage": {"input_tokens": 6, "output_tokens": 3}, + }, + ) + ) + response: Final = await interactions.acreate( + model="gemini/gemini-2.5-flash", + input="What is the speed of light?", + api_key="gemini-offline", + ) + body: Final = json.loads(route.calls.last.request.content) + assert body["model"] == "gemini-2.5-flash" + assert body["input"] == "What is the speed of light?" + assert response.id == "interaction-offline" + assert response.status == "completed" + + @pytest.mark.usefixtures("fake_provider_credentials") + @pytest.mark.asyncio + async def test_acreate_simple_litellm_responses_bridge(self, respx_mock: MockRouter) -> None: + route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response( + 200, + json={ + "id": "resp-offline", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "299792458 m/s"}], + } + ], + "usage": {"input_tokens": 6, "output_tokens": 3, "total_tokens": 9}, + }, + ) + ) + response: Final = await interactions.acreate( + model="gpt-4o", + input="What is the speed of light?", + api_key="sk-offline", + ) + body: Final = json.loads(route.calls.last.request.content) + assert body["model"] == "gpt-4o" + serialized: Final = json.dumps(body) + assert "What is the speed of light?" in serialized + assert "response_id:resp-offline" in base64.b64decode(response.id.removeprefix("resp_")).decode() + assert response.status == "completed" diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_web_search_logged_cost.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_web_search_logged_cost.py new file mode 100644 index 00000000000..1448318b07d --- /dev/null +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_web_search_logged_cost.py @@ -0,0 +1,210 @@ +import asyncio +import json +from typing import Final, Literal, TypedDict + +from typing_extensions import ReadOnly + +import httpx +import pytest +import respx +from pydantic import TypeAdapter + +import litellm +from litellm.integrations.custom_logger import CustomLogger + +_MODEL: Final = "gpt-sized-search-unit" +_INPUT_COST: Final = 1e-06 +_OUTPUT_COST: Final = 4e-06 +_PER_QUERY: Final = { + "search_context_size_low": 0.011, + "search_context_size_medium": 0.022, + "search_context_size_high": 0.033, +} +_PROMPT_TOKENS: Final = 100 +_COMPLETION_TOKENS: Final = 20 +_USAGE: Final = { + "prompt_tokens": _PROMPT_TOKENS, + "completion_tokens": _COMPLETION_TOKENS, + "total_tokens": _PROMPT_TOKENS + _COMPLETION_TOKENS, +} +_CHAT_RESPONSE: Final = { + "id": "chatcmpl-search", + "object": "chat.completion", + "created": 1700000000, + "model": _MODEL, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "A positive story", + "annotations": [ + { + "type": "url_citation", + "url_citation": { + "start_index": 0, + "end_index": 5, + "title": "news", + "url": "https://news.example/a", + }, + } + ], + }, + "finish_reason": "stop", + } + ], + "usage": _USAGE, +} +_RESPONSES_BODY: Final = { + "id": "resp_search", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": _MODEL, + "output": [ + {"type": "web_search_call", "id": "ws_search", "status": "completed"}, + { + "type": "message", + "id": "msg_search", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "A positive story", "annotations": []}], + }, + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": { + "input_tokens": _PROMPT_TOKENS, + "output_tokens": _COMPLETION_TOKENS, + "total_tokens": _PROMPT_TOKENS + _COMPLETION_TOKENS, + }, +} +_RESPONSES_STREAM: Final = "".join( + f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" + for event in ( + {"type": "response.created", "response": {**_RESPONSES_BODY, "status": "in_progress", "output": []}}, + {"type": "response.completed", "response": _RESPONSES_BODY}, + ) +) + +_ContextSize = Literal["search_context_size_low", "search_context_size_medium", "search_context_size_high"] + + +class _LoggedCost(TypedDict): + response_cost: ReadOnly[float] + prompt_tokens: ReadOnly[int] + completion_tokens: ReadOnly[int] + + +class _CostRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.payloads: Final[list[_LoggedCost]] = [] + self.logged: Final = asyncio.Event() + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.payloads.append(TypeAdapter(_LoggedCost).validate_python(kwargs["standard_logging_object"])) + self.logged.set() + + +@pytest.fixture +def recorder(monkeypatch: pytest.MonkeyPatch) -> _CostRecorder: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + entry: Final = { + "input_cost_per_token": _INPUT_COST, + "output_cost_per_token": _OUTPUT_COST, + "litellm_provider": "openai", + "mode": "chat", + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "supports_web_search": True, + "search_context_cost_per_query": _PER_QUERY, + } + monkeypatch.setitem(litellm.model_cost, _MODEL, entry) + monkeypatch.setitem(litellm.model_cost, f"openai/{_MODEL}", entry) + cost_recorder: Final = _CostRecorder() + monkeypatch.setattr(litellm, "callbacks", [cost_recorder]) + return cost_recorder + + +async def _logged_cost(recorder: _CostRecorder) -> _LoggedCost: + await asyncio.wait_for(recorder.logged.wait(), timeout=10) + return recorder.payloads[-1] + + +def _expected_cost(payload: _LoggedCost, size: _ContextSize) -> float: + return payload["prompt_tokens"] * _INPUT_COST + payload["completion_tokens"] * _OUTPUT_COST + _PER_QUERY[size] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("web_search_options", "size"), + [ + (None, "search_context_size_medium"), + ({"search_context_size": "low"}, "search_context_size_low"), + ({"search_context_size": "high"}, "search_context_size_high"), + ], +) +async def test_chat_web_search_logged_cost_adds_the_per_query_cost_for_the_context_size( + web_search_options: dict[str, str] | None, + size: _ContextSize, + recorder: _CostRecorder, + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(200, json=_CHAT_RESPONSE) + ) + options: Final = {"web_search_options": web_search_options} if web_search_options is not None else {} + + await litellm.acompletion( + model=f"openai/{_MODEL}", + messages=[{"role": "user", "content": "What was a positive news story from today?"}], + api_key="sk-unit-test", + **options, + ) + payload: Final = await _logged_cost(recorder) + + assert (payload["prompt_tokens"], payload["completion_tokens"]) == (_PROMPT_TOKENS, _COMPLETION_TOKENS) + assert payload["response_cost"] == pytest.approx(_expected_cost(payload, size), abs=1e-12) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tools", "size", "stream"), + [ + ([{"type": "web_search_preview", "search_context_size": "low"}], "search_context_size_low", True), + ([{"type": "web_search_preview", "search_context_size": "low"}], "search_context_size_low", False), + ([{"type": "web_search_preview"}], "search_context_size_medium", True), + ([{"type": "web_search_preview"}], "search_context_size_medium", False), + ], +) +async def test_responses_web_search_logged_cost_adds_the_per_query_cost_for_the_context_size( + tools: list[dict[str, str]], + size: _ContextSize, + stream: bool, + recorder: _CostRecorder, + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, text=_RESPONSES_STREAM, headers={"content-type": "text/event-stream"}) + if stream + else httpx.Response(200, json=_RESPONSES_BODY) + ) + + response: Final = await litellm.aresponses( + model=f"openai/{_MODEL}", + input=[{"role": "user", "content": "What was a positive news story from today?"}], + tools=tools, + stream=stream, + api_key="sk-unit-test", + ) + if stream: + assert [event async for event in response] + payload: Final = await _logged_cost(recorder) + + assert (payload["prompt_tokens"], payload["completion_tokens"]) == (_PROMPT_TOKENS, _COMPLETION_TOKENS) + assert payload["response_cost"] == pytest.approx(_expected_cost(payload, size), abs=1e-12) diff --git a/tests/unit/litellm_core_utils/test_moderation_standard_logging.py b/tests/unit/litellm_core_utils/test_moderation_standard_logging.py new file mode 100644 index 00000000000..1f8c722785f --- /dev/null +++ b/tests/unit/litellm_core_utils/test_moderation_standard_logging.py @@ -0,0 +1,88 @@ +import asyncio +from typing import Final, Literal, TypedDict + +from typing_extensions import ReadOnly + +import httpx +import pytest +import respx +from pydantic import TypeAdapter + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.router import Router + +_MODERATIONS_URL: Final = "https://api.openai.com/v1/moderations" +_MODEL_GROUP: Final = "internal-moderation-model" +_INPUT: Final = "Hello, how are you?" +_CATEGORIES: Final = ("harassment", "hate", "self-harm", "sexual", "violence") +_MODERATION_RESPONSE: Final = { + "id": "modr-logging", + "model": "omni-moderation-latest", + "results": [ + { + "flagged": False, + "categories": {name: False for name in _CATEGORIES}, + "category_scores": {name: 0.001 for name in _CATEGORIES}, + } + ], +} + + +class _LoggedModeration(TypedDict): + call_type: ReadOnly[str] + status: ReadOnly[str] + custom_llm_provider: ReadOnly[str | None] + messages: ReadOnly[object] + response: ReadOnly[object] + model_group: ReadOnly[str | None] + + +class _ModerationRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.payloads: Final[list[_LoggedModeration]] = [] + self.logged: Final = asyncio.Event() + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.payloads.append(TypeAdapter(_LoggedModeration).validate_python(kwargs["standard_logging_object"])) + self.logged.set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("caller", ["default-model", "named-model", "router-group"]) +async def test_moderation_call_is_logged_as_an_amoderation_standard_payload( + caller: Literal["default-model", "named-model", "router-group"], + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "sk-unit-test") + recorder: Final = _ModerationRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + respx_mock.post(_MODERATIONS_URL).mock(return_value=httpx.Response(200, json=_MODERATION_RESPONSE)) + router: Final = Router( + model_list=[{"model_name": _MODEL_GROUP, "litellm_params": {"model": "openai/omni-moderation-latest"}}] + ) + + response: Final = ( + await router.amoderation(input=_INPUT, model=_MODEL_GROUP) + if caller == "router-group" + else await litellm.amoderation( + input=_INPUT, model=None if caller == "default-model" else "omni-moderation-latest" + ) + ) + await asyncio.wait_for(recorder.logged.wait(), timeout=10) + + payload: Final = recorder.payloads[-1] + assert payload["call_type"] == litellm.utils.CallTypes.amoderation.value + assert payload["status"] == "success" + assert payload["custom_llm_provider"] == litellm.LlmProviders.OPENAI.value + assert TypeAdapter(tuple[dict[str, str], ...]).validate_python(payload["messages"])[0]["content"] == _INPUT + assert dict(TypeAdapter(dict[str, object]).validate_python(payload["response"])) == response.model_dump() + if caller == "router-group": + assert payload["model_group"] == _MODEL_GROUP + else: + assert not payload["model_group"] diff --git a/tests/unit/litellm_core_utils/test_stream_usage_logging.py b/tests/unit/litellm_core_utils/test_stream_usage_logging.py new file mode 100644 index 00000000000..a108af09dd1 --- /dev/null +++ b/tests/unit/litellm_core_utils/test_stream_usage_logging.py @@ -0,0 +1,138 @@ +import asyncio +import json +from typing import Final, TypedDict + +from typing_extensions import ReadOnly + +import httpx +import pytest +import respx +from pydantic import TypeAdapter + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.redact_messages import REDACTED_BY_LITELLM +from litellm.types.utils import Usage + +_OPENAI_URL: Final = "https://api.openai.com/v1/chat/completions" +_BODY: Final = TypeAdapter(dict[str, object]) +_PROMPT_TOKENS: Final = 607 +_COMPLETION_TOKENS: Final = 23 + + +def _sse(chunks: tuple[dict[str, object], ...]) -> str: + return "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n" + + +def _chunk(delta: dict[str, str], finish_reason: str | None) -> dict[str, object]: + return { + "id": "chatcmpl-usage", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-5.5", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + + +_STREAM: Final = _sse( + ( + _chunk({"role": "assistant", "content": "I am"}, None), + _chunk({"content": " well"}, "stop"), + { + "id": "chatcmpl-usage", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-5.5", + "choices": [], + "usage": { + "prompt_tokens": _PROMPT_TOKENS, + "completion_tokens": _COMPLETION_TOKENS, + "total_tokens": _PROMPT_TOKENS + _COMPLETION_TOKENS, + }, + }, + ) +) + + +class _LoggedUsage(TypedDict): + prompt_tokens: ReadOnly[int] + completion_tokens: ReadOnly[int] + total_tokens: ReadOnly[int] + messages: ReadOnly[object] + + +class _UsageRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.payloads: Final[list[_LoggedUsage]] = [] + self.logged: Final = asyncio.Event() + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.payloads.append(TypeAdapter(_LoggedUsage).validate_python(kwargs["standard_logging_object"])) + self.logged.set() + + +async def _stream_and_record( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, include_usage: bool +) -> tuple[Usage, _LoggedUsage, dict[str, object]]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _UsageRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + route: Final = respx_mock.post(_OPENAI_URL).mock( + return_value=httpx.Response(200, text=_STREAM, headers={"content-type": "text/event-stream"}) + ) + stream_options: Final = {"stream_options": {"include_usage": True}} if include_usage else {} + response: Final = await litellm.acompletion( + model="gpt-5.5", + messages=[{"role": "user", "content": "Hello, how are you?" * 100}], + stream=True, + api_key="sk-unit-test", + **stream_options, + ) + usages: Final = tuple([chunk.usage async for chunk in response if getattr(chunk, "usage", None) is not None]) + await asyncio.wait_for(recorder.logged.wait(), timeout=10) + return usages[-1], recorder.payloads[-1], _BODY.validate_json(route.calls.last.request.content) + + +def _assert_logged_usage_matches(client_usage: Usage, payload: _LoggedUsage) -> None: + assert client_usage.prompt_tokens == _PROMPT_TOKENS + assert client_usage.completion_tokens == _COMPLETION_TOKENS + assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == ( + client_usage.prompt_tokens, + client_usage.completion_tokens, + client_usage.total_tokens, + ) + + +@pytest.mark.asyncio +async def test_logged_stream_usage_equals_the_final_chunk_usage_with_include_usage( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + client_usage, payload, body = await _stream_and_record(monkeypatch, respx_mock, include_usage=True) + + assert body["stream_options"] == {"include_usage": True} + _assert_logged_usage_matches(client_usage, payload) + + +@pytest.mark.asyncio +async def test_logged_stream_usage_equals_the_usage_chunk_without_stream_options( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + client_usage, payload, body = await _stream_and_record(monkeypatch, respx_mock, include_usage=False) + + assert body["stream_options"] == {"include_usage": True} + _assert_logged_usage_matches(client_usage, payload) + + +@pytest.mark.asyncio +async def test_logged_stream_usage_survives_message_redaction( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + + client_usage, payload, _ = await _stream_and_record(monkeypatch, respx_mock, include_usage=False) + + _assert_logged_usage_matches(client_usage, payload) + assert payload["messages"] == [{"role": "user", "content": REDACTED_BY_LITELLM}] diff --git a/tests/unit/llms/base_llm/audio_transcription/__init__.py b/tests/unit/llms/base_llm/audio_transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/base_llm/audio_transcription/test_provider_audio_transcription_translation.py b/tests/unit/llms/base_llm/audio_transcription/test_provider_audio_transcription_translation.py new file mode 100644 index 00000000000..65a8ce91e64 --- /dev/null +++ b/tests/unit/llms/base_llm/audio_transcription/test_provider_audio_transcription_translation.py @@ -0,0 +1,210 @@ +import io +import json +from typing import Final, Mapping, Sequence, cast + +import httpx +import pytest +import respx +from respx import MockRouter +from typing_extensions import ReadOnly, TypedDict + +import litellm +from litellm import transcription +from litellm.litellm_core_utils.get_supported_openai_params import get_supported_openai_params +from litellm.llms.base_llm.audio_transcription.transformation import ( + AudioTranscriptionRequestData, + BaseAudioTranscriptionConfig, +) +from litellm.llms.elevenlabs.audio_transcription.transformation import ElevenLabsAudioTranscriptionConfig +from litellm.llms.mistral.audio_transcription.transformation import MistralAudioTranscriptionConfig +from litellm.llms.ovhcloud.audio_transcription.transformation import OVHCloudAudioTranscriptionConfig +from litellm.utils import ProviderConfigManager + + +class _Kwargs(TypedDict, total=False): + model: ReadOnly[str] + api_key: ReadOnly[str] + api_base: ReadOnly[str] + timestamp_granularities: ReadOnly[Sequence[str]] + + +class _Case(TypedDict): + id: ReadOnly[str] + provider: ReadOnly[str] + kwargs: ReadOnly[_Kwargs] + url: ReadOnly[str] + prefix_match: ReadOnly[bool] + base_model: ReadOnly[str] + request_markers: ReadOnly[tuple[bytes, ...]] + marker_in_url: ReadOnly[bool] + config_class: ReadOnly[type[BaseAudioTranscriptionConfig]] + + +_AUDIO_BYTES: Final = b"RIFFFAKEWAVDATA-gettysburg" + +_CASES: Final[tuple[_Case, ...]] = ( + { + "id": "openai_gpt4o", + "provider": "openai", + "kwargs": { + "model": "openai/gpt-4o-transcribe", + "api_key": "sk-offline", + "timestamp_granularities": ["word"], + }, + "url": "https://api.openai.com/v1/audio/transcriptions", + "prefix_match": False, + "base_model": "gpt-4o-transcribe", + "request_markers": (b"gpt-4o-transcribe", b'name="timestamp_granularities[]"\r\n\r\nword\r\n'), + "marker_in_url": False, + "config_class": litellm.OpenAIGPTAudioTranscriptionConfig, + }, + { + "id": "elevenlabs_scribe", + "provider": "elevenlabs", + "kwargs": {"model": "elevenlabs/scribe_v1", "api_key": "xi-offline"}, + "url": "https://api.elevenlabs.io/v1/speech-to-text", + "prefix_match": False, + "base_model": "scribe_v1", + "request_markers": (b"scribe_v1",), + "marker_in_url": False, + "config_class": ElevenLabsAudioTranscriptionConfig, + }, + { + "id": "deepgram_nova", + "provider": "deepgram", + "kwargs": {"model": "deepgram/nova-2", "api_key": "dg-offline"}, + "url": "https://api.deepgram.com/v1/listen", + "prefix_match": True, + "base_model": "nova-2", + "request_markers": (b"model=nova-2",), + "marker_in_url": True, + "config_class": litellm.DeepgramAudioTranscriptionConfig, + }, + { + "id": "mistral_voxtral", + "provider": "mistral", + "kwargs": {"model": "mistral/voxtral-mini-latest", "api_key": "mistral-offline"}, + "url": "https://api.mistral.ai/v1/audio/transcriptions", + "prefix_match": False, + "base_model": "voxtral-mini-latest", + "request_markers": (b"voxtral-mini-latest",), + "marker_in_url": False, + "config_class": MistralAudioTranscriptionConfig, + }, + { + "id": "ovhcloud_whisper", + "provider": "ovhcloud", + "kwargs": {"model": "ovhcloud/whisper-large-v3-turbo", "api_key": "ovh-offline"}, + "url": "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1/audio/transcriptions", + "prefix_match": False, + "base_model": "whisper-large-v3-turbo", + "request_markers": (b"whisper-large-v3-turbo",), + "marker_in_url": False, + "config_class": OVHCloudAudioTranscriptionConfig, + }, +) + + +def _canned_response(case: _Case) -> httpx.Response: + if case["provider"] == "deepgram": + return httpx.Response( + 200, + json={ + "metadata": {"transaction_key": "offline", "duration": 1.5}, + "results": { + "channels": [ + {"alternatives": [{"transcript": "four score and seven years ago", "confidence": 0.99}]} + ] + }, + }, + ) + return httpx.Response(200, json={"text": "four score and seven years ago"}) + + +@pytest.fixture(autouse=True) +def _httpx_only_transport(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + + +def _register(case: _Case, respx_mock: MockRouter) -> respx.Route: + if case["prefix_match"]: + return respx_mock.post(url__startswith=case["url"]).mock(return_value=_canned_response(case)) + return respx_mock.post(case["url"]).mock(return_value=_canned_response(case)) + + +def _assert_translated_request(case: _Case, request: httpx.Request) -> None: + searched: Final = request.url.query if case["marker_in_url"] else request.content + for marker in case["request_markers"]: + assert marker in searched + assert _AUDIO_BYTES in request.content + + +@pytest.mark.parametrize("case", _CASES, ids=lambda c: c["id"]) +def test_audio_transcription(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock) + transcript: Final = transcription(**dict(case["kwargs"]), file=io.BytesIO(_AUDIO_BYTES)) + _assert_translated_request(case, route.calls.last.request) + assert transcript.text == "four score and seven years ago" + + +@pytest.mark.parametrize("case", _CASES, ids=lambda c: c["id"]) +@pytest.mark.asyncio +async def test_audio_transcription_async(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock) + transcript: Final = await litellm.atranscription(**dict(case["kwargs"]), file=io.BytesIO(_AUDIO_BYTES)) + _assert_translated_request(case, route.calls.last.request) + assert transcript.text == "four score and seven years ago" + + +@pytest.mark.parametrize("case", _CASES, ids=lambda c: c["id"]) +def test_audio_transcription_optional_params(case: _Case) -> None: + optional_params: Final = get_supported_openai_params( + model=case["kwargs"]["model"], + custom_llm_provider=case["provider"], + request_type="transcription", + ) + assert isinstance(optional_params, list) + assert optional_params == case["config_class"]().get_supported_openai_params(case["base_model"]) + assert "max_completion_tokens" not in optional_params + + +@pytest.mark.parametrize("case", _CASES, ids=lambda c: c["id"]) +def test_audio_transcription_config(case: _Case) -> None: + config: Final = ProviderConfigManager.get_provider_audio_transcription_config( + model=case["kwargs"]["model"], + provider=litellm.LlmProviders(case["provider"]), + ) + assert type(config) is case["config_class"] + assert isinstance(config, BaseAudioTranscriptionConfig) + if case["provider"] == "deepgram": + complete_url: Final = config.get_complete_url( + api_base=None, + api_key=None, + model=case["base_model"], + optional_params={}, + litellm_params={}, + ) + assert "api.deepgram.com" in complete_url + assert "model=nova-2" in complete_url + else: + transformed: Final[AudioTranscriptionRequestData] = config.transform_audio_transcription_request( + model=case["base_model"], + audio_file=io.BytesIO(_AUDIO_BYTES), + optional_params={}, + litellm_params={}, + ) + assert _AUDIO_BYTES in _transformed_payload(transformed) + + +def _transformed_payload(transformed: AudioTranscriptionRequestData) -> bytes: + data: Final = transformed.data + if isinstance(data, bytes): + return data + file_entry: Final = data.get("file") if isinstance(data, dict) else None + if isinstance(file_entry, io.BytesIO): + return file_entry.getvalue() + if transformed.files is not None: + first: Final = next(iter(transformed.files.values())) + blob: Final = first[1] if isinstance(first, tuple) else first + return blob.getvalue() if isinstance(blob, io.BytesIO) else cast(bytes, blob) + return b"" diff --git a/tests/unit/llms/base_llm/chat/test_provider_chat_thinking.py b/tests/unit/llms/base_llm/chat/test_provider_chat_thinking.py new file mode 100644 index 00000000000..8f139ca0acc --- /dev/null +++ b/tests/unit/llms/base_llm/chat/test_provider_chat_thinking.py @@ -0,0 +1,385 @@ +import json +from itertools import chain +from typing import Final, Mapping, cast + +import httpx +import pytest +from pydantic import BaseModel, ConfigDict, JsonValue +from respx import MockRouter + +import litellm +from litellm import get_llm_provider +from litellm.constants import ( + DEFAULT_MAX_TOKENS, + DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, +) +from litellm.main import stream_chunk_builder +from litellm.utils import get_optional_params + +from tests.unit.llms.base_llm.chat.test_provider_chat_translation import ( + _BY_ID, + _Case, + _aws_frame, + _call as _provider_call, + _request_body, +) + +_THINKING_BUDGET: Final = 16000 +_THINKING: Final[JsonValue] = {"type": "enabled", "budget_tokens": _THINKING_BUDGET} + +_ANTHROPIC: Final = _BY_ID["anthropic_sonnet45"] +_BEDROCK_HAIKU: Final = _BY_ID["bedrock_converse_haiku"] +_BEDROCK_SONNET: Final = _BY_ID["bedrock_converse_anthropic_thinking"] + +_THINKING_CASES: Final = (_ANTHROPIC, _BEDROCK_SONNET) +_RESPONSE_FORMAT_CASES: Final = (_ANTHROPIC, _BEDROCK_HAIKU) + +_JSON_PREFIX: Final = '{"agent_doing": "researching ' +_JSON_SUFFIX: Final = 'home automation"}' +_JSON_CONTENT: Final = _JSON_PREFIX + _JSON_SUFFIX +_REASONING: Final = "reasoning here" +_SIGNATURE: Final = "sig-1" + + +class _RFormat(BaseModel): + model_config = ConfigDict(frozen=True) + question: str + answer: str + + +_JSON_SCHEMA_ARGS: Final[Mapping[str, JsonValue]] = { + "messages": [ + {"role": "system", "content": "Summarize the agent's thinking into short descriptions."}, + {"role": "user", "content": "Here is the input data."}, + ], + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "final_output", + "strict": True, + "schema": { + "properties": {"agent_doing": {"title": "Agent Doing", "type": "string"}}, + "required": ["agent_doing"], + "title": "ThinkingStep", + "type": "object", + "additionalProperties": False, + }, + }, + }, +} + +_THINKING_MESSAGES: Final[Mapping[str, JsonValue]] = { + "messages": [{"role": "user", "content": "Generate 5 question + answer pairs"}], +} + + +def _case_id(case: _Case) -> str: + return case["id"] + + +def _call(case: _Case, **extra: JsonValue) -> object: + return _provider_call(case, extra) + + +def _anthropic_sse(events: tuple[Mapping[str, JsonValue], ...]) -> str: + return "".join(f"event: {e['type']}\ndata: {json.dumps(e)}\n\n" for e in events) + + +_ANTHROPIC_START: Final[Mapping[str, JsonValue]] = { + "type": "message_start", + "message": { + "id": "msg_offline", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, +} +_ANTHROPIC_END: Final[tuple[Mapping[str, JsonValue], ...]] = ( + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 5}}, + {"type": "message_stop"}, +) + + +def _anthropic_json_stream() -> str: + return _anthropic_sse( + ( + _ANTHROPIC_START, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _JSON_PREFIX}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _JSON_SUFFIX}}, + {"type": "content_block_stop", "index": 0}, + *_ANTHROPIC_END, + ) + ) + + +def _anthropic_thinking_stream() -> str: + return _anthropic_sse( + ( + _ANTHROPIC_START, + {"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": _REASONING}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "signature_delta", "signature": _SIGNATURE}}, + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "done"}}, + {"type": "content_block_stop", "index": 1}, + *_ANTHROPIC_END, + ) + ) + + +_CONVERSE_USAGE: Final[Mapping[str, JsonValue]] = {"usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}} + + +def _converse_stream(frames: tuple[tuple[str, Mapping[str, JsonValue]], ...]) -> bytes: + return b"".join(_aws_frame(event_type, payload) for event_type, payload in frames) + + +def _converse_json_stream() -> bytes: + return _converse_stream( + ( + ("messageStart", {"role": "assistant"}), + ("contentBlockDelta", {"delta": {"text": _JSON_PREFIX}, "contentBlockIndex": 0}), + ("contentBlockDelta", {"delta": {"text": _JSON_SUFFIX}, "contentBlockIndex": 0}), + ("contentBlockStop", {"contentBlockIndex": 0}), + ("messageStop", {"stopReason": "end_turn"}), + ("metadata", _CONVERSE_USAGE), + ) + ) + + +def _converse_thinking_stream() -> bytes: + return _converse_stream( + ( + ("messageStart", {"role": "assistant"}), + ("contentBlockDelta", {"delta": {"reasoningContent": {"text": _REASONING}}, "contentBlockIndex": 0}), + ("contentBlockDelta", {"delta": {"reasoningContent": {"signature": _SIGNATURE}}, "contentBlockIndex": 0}), + ("contentBlockStop", {"contentBlockIndex": 0}), + ("contentBlockDelta", {"delta": {"text": "done"}, "contentBlockIndex": 1}), + ("contentBlockStop", {"contentBlockIndex": 1}), + ("messageStop", {"stopReason": "end_turn"}), + ("metadata", _CONVERSE_USAGE), + ) + ) + + +def _non_stream_response(case: _Case) -> httpx.Response: + if case["shape"] == "anthropic": + return httpx.Response( + 200, + json={ + "id": "msg_offline", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [ + {"type": "thinking", "thinking": _REASONING, "signature": _SIGNATURE}, + {"type": "text", "text": _JSON_CONTENT}, + ], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + ) + return httpx.Response( + 200, + json={ + "output": { + "message": { + "role": "assistant", + "content": [ + {"reasoningContent": {"reasoningText": {"text": _REASONING, "signature": _SIGNATURE}}}, + {"text": _JSON_CONTENT}, + ], + } + }, + "stopReason": "end_turn", + "usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, + }, + ) + + +def _stream_response(case: _Case, *, thinking: bool) -> httpx.Response: + if case["shape"] == "anthropic": + return httpx.Response( + 200, + content=_anthropic_thinking_stream() if thinking else _anthropic_json_stream(), + headers={"content-type": "text/event-stream"}, + ) + return httpx.Response( + 200, + content=_converse_thinking_stream() if thinking else _converse_json_stream(), + headers={"content-type": "application/vnd.amazon.eventstream"}, + ) + + +def _thinking_param(case: _Case, body: Mapping[str, JsonValue]) -> JsonValue: + if case["shape"] == "anthropic": + return body["thinking"] + return cast(Mapping[str, JsonValue], body["additionalModelRequestFields"])["thinking"] + + +def _max_tokens(case: _Case, body: Mapping[str, JsonValue]) -> JsonValue: + if case["shape"] == "anthropic": + return body["max_tokens"] + return cast(Mapping[str, JsonValue], body["inferenceConfig"])["maxTokens"] + + +def _json_schema_title(case: _Case, body: Mapping[str, JsonValue]) -> str: + if case["shape"] == "anthropic": + output_format: Final = cast(Mapping[str, JsonValue], body["output_format"]) + assert output_format["type"] == "json_schema" + return cast(str, cast(Mapping[str, JsonValue], output_format["schema"])["title"]) + text_format: Final = cast( + Mapping[str, JsonValue], + cast(Mapping[str, JsonValue], body["outputConfig"])["textFormat"], + ) + assert text_format["type"] == "json_schema" + json_schema: Final = cast( + Mapping[str, JsonValue], cast(Mapping[str, JsonValue], text_format["structure"])["jsonSchema"] + ) + return cast(str, json.loads(cast(str, json_schema["schema"]))["title"]) + + +def _has_forced_tool_choice(case: _Case, body: Mapping[str, JsonValue]) -> bool: + if case["shape"] == "anthropic": + return body.get("tool_choice") is not None + return "toolConfig" in body + + +@pytest.mark.parametrize("case", _RESPONSE_FORMAT_CASES, ids=_case_id) +def test_anthropic_response_format_streaming_vs_non_streaming(case: _Case, respx_mock: MockRouter) -> None: + stream_route: Final = respx_mock.post(case["stream_url"]).mock(return_value=_stream_response(case, thinking=False)) + chunks: Final = tuple(cast(litellm.CustomStreamWrapper, _call(case, **_JSON_SCHEMA_ARGS, stream=True))) + built: Final = stream_chunk_builder(chunks=list(chunks)) + stream_body: Final = _request_body(stream_route) + + non_stream_route: Final = respx_mock.post(case["url"]).mock(return_value=_non_stream_response(case)) + non_stream: Final = cast(litellm.ModelResponse, _call(case, **_JSON_SCHEMA_ARGS)) + non_stream_body: Final = _request_body(non_stream_route) + + assert len(chunks) > 1 + assert _json_schema_title(case, stream_body) == "ThinkingStep" + assert _json_schema_title(case, non_stream_body) == "ThinkingStep" + assert built is not None + streamed_json: Final = cast( + Mapping[str, JsonValue], + json.loads(cast(str, cast(litellm.ModelResponse, built).choices[0].message.content)), + ) + non_stream_json: Final = cast(Mapping[str, JsonValue], json.loads(cast(str, non_stream.choices[0].message.content))) + assert streamed_json == non_stream_json == {"agent_doing": "researching home automation"} + + +@pytest.mark.parametrize("case", _THINKING_CASES, ids=_case_id) +def test_completion_thinking_with_response_format(case: _Case, respx_mock: MockRouter) -> None: + route: Final = respx_mock.post(case["url"]).mock(return_value=_non_stream_response(case)) + response: Final = cast( + litellm.ModelResponse, + _call(case, thinking=_THINKING, **_THINKING_MESSAGES, response_format=cast(JsonValue, _RFormat)), + ) + body: Final = _request_body(route) + assert _thinking_param(case, body) == _THINKING + assert _json_schema_title(case, body) == "_RFormat" + assert not _has_forced_tool_choice(case, body) + assert response.choices[0].message.content == _JSON_CONTENT + assert response.choices[0].message.reasoning_content == _REASONING + + +@pytest.mark.parametrize("case", _THINKING_CASES, ids=_case_id) +def test_completion_thinking_with_max_tokens(case: _Case, respx_mock: MockRouter) -> None: + route: Final = respx_mock.post(case["url"]).mock(return_value=_non_stream_response(case)) + response: Final = cast( + litellm.ModelResponse, + _call(case, thinking=_THINKING, **_THINKING_MESSAGES, max_completion_tokens=20000), + ) + body: Final = _request_body(route) + assert _max_tokens(case, body) == 20000 + assert _thinking_param(case, body) == _THINKING + assert response.choices[0].message.content == _JSON_CONTENT + + +@pytest.mark.parametrize("case", _THINKING_CASES, ids=_case_id) +def test_completion_thinking_without_max_tokens(case: _Case, respx_mock: MockRouter) -> None: + route: Final = respx_mock.post(case["url"]).mock(return_value=_non_stream_response(case)) + response: Final = cast(litellm.ModelResponse, _call(case, thinking=_THINKING, **_THINKING_MESSAGES)) + body: Final = _request_body(route) + max_tokens: Final = cast(int, _max_tokens(case, body)) + assert max_tokens == _THINKING_BUDGET + DEFAULT_MAX_TOKENS + assert max_tokens > _THINKING_BUDGET + assert _thinking_param(case, body) == _THINKING + assert response.choices[0].message.content == _JSON_CONTENT + + +@pytest.mark.parametrize("case", _THINKING_CASES, ids=_case_id) +def test_anthropic_thinking_output_stream(case: _Case, respx_mock: MockRouter) -> None: + route: Final = respx_mock.post(case["stream_url"]).mock(return_value=_stream_response(case, thinking=True)) + chunks: Final = tuple( + cast( + litellm.CustomStreamWrapper, + _call( + case, + thinking=_THINKING, + messages=[{"role": "user", "content": "Tell me a joke."}], + stream=True, + ), + ) + ) + deltas: Final = tuple(chunk.choices[0].delta for chunk in chunks) + thinking_deltas: Final = tuple( + delta + for delta in deltas + if isinstance(getattr(delta, "thinking_blocks", None), list) + and delta.thinking_blocks + and isinstance(getattr(delta, "reasoning_content", None), str) + ) + blocks: Final = chain.from_iterable(cast(list[object], delta.thinking_blocks) for delta in thinking_deltas) + signatures: Final = tuple(cast(Mapping[str, JsonValue], block).get("signature") for block in blocks) + assert _thinking_param(case, _request_body(route)) == _THINKING + assert not any(delta.tool_calls for delta in deltas) + assert "".join(cast(str, delta.reasoning_content) for delta in thinking_deltas) == _REASONING + assert _SIGNATURE in signatures + + +@pytest.mark.parametrize("case", _THINKING_CASES, ids=_case_id) +def test_anthropic_reasoning_effort_thinking_translation(case: _Case, respx_mock: MockRouter) -> None: + model: Final = case["kwargs"].get("model", "") + _, provider, _, _ = get_llm_provider(model=model) + optional_params: Final = get_optional_params(model=model, custom_llm_provider=provider, reasoning_effort="high") + assert optional_params["thinking"] == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, + } + assert "reasoning_effort" not in optional_params + + route: Final = respx_mock.post(case["url"]).mock(return_value=_non_stream_response(case)) + _call(case, reasoning_effort="high", messages=[{"role": "user", "content": "hi"}]) + body: Final = _request_body(route) + assert _thinking_param(case, body) == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, + } + assert "reasoning_effort" not in json.dumps(body) + assert _max_tokens(case, body) == DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET + DEFAULT_MAX_TOKENS + + +@pytest.mark.parametrize("case", _THINKING_CASES, ids=_case_id) +@pytest.mark.parametrize( + ("effort", "budget"), + ( + ("low", DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET), + ("medium", DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET), + ), +) +def test_reasoning_effort_maps_to_distinct_thinking_budgets( + case: _Case, effort: str, budget: int, respx_mock: MockRouter +) -> None: + route: Final = respx_mock.post(case["url"]).mock(return_value=_non_stream_response(case)) + _call(case, reasoning_effort=effort, messages=[{"role": "user", "content": "hi"}]) + body: Final = _request_body(route) + assert _thinking_param(case, body) == {"type": "enabled", "budget_tokens": budget} + assert _max_tokens(case, body) == budget + DEFAULT_MAX_TOKENS diff --git a/tests/unit/llms/base_llm/chat/test_provider_chat_translation.py b/tests/unit/llms/base_llm/chat/test_provider_chat_translation.py new file mode 100644 index 00000000000..e0e07a3dd59 --- /dev/null +++ b/tests/unit/llms/base_llm/chat/test_provider_chat_translation.py @@ -0,0 +1,1644 @@ +import base64 +import copy +import itertools +import json +import struct +import zlib +from typing import Callable, Final, Iterable, Literal, Mapping, cast + +import httpx +import pytest +import respx +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter +from respx import MockRouter +from typing_extensions import ReadOnly, TypedDict + +import litellm +from litellm.constants import ( + DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, +) +from litellm.llms.base_llm.base_utils import type_to_response_format_param +from litellm.types.utils import CallTypes +from litellm.utils import ProviderConfigManager, return_raw_request + +_Shape = Literal[ + "openai", + "anthropic", + "gemini", + "bedrock_converse", + "bedrock_invoke", + "bedrock_invoke_nova", + "bedrock_invoke_openai", +] + +_AWS_KWARGS: Final[Mapping[str, str]] = { + "aws_access_key_id": "AKIAFAKE", + "aws_secret_access_key": "fakesecret", + "aws_region_name": "us-east-1", +} + + +class _Kwargs(TypedDict, total=False): + model: ReadOnly[str] + api_key: ReadOnly[str] + api_base: ReadOnly[str] + api_version: ReadOnly[str] + aws_access_key_id: ReadOnly[str] + aws_secret_access_key: ReadOnly[str] + aws_region_name: ReadOnly[str] + + +class _Case(TypedDict): + id: ReadOnly[str] + shape: ReadOnly[_Shape] + kwargs: ReadOnly[_Kwargs] + url: ReadOnly[str] + stream_url: ReadOnly[str] + router: ReadOnly[bool] + + +def _converse_url(model_id: str, region: str = "us-east-1") -> str: + return f"https://bedrock-runtime.{region}.amazonaws.com/model/{model_id}/converse" + + +def _converse_stream_url(model_id: str, region: str = "us-east-1") -> str: + return f"https://bedrock-runtime.{region}.amazonaws.com/model/{model_id}/converse-stream" + + +def _invoke_url(model_id: str) -> str: + return f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{model_id}/invoke" + + +def _invoke_stream_url(model_id: str) -> str: + return f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{model_id}/invoke-with-response-stream" + + +def _bedrock_case( + case_id: str, + model: str, + invoke_model_id: str, + shape: _Shape, + region: str = "us-east-1", +) -> _Case: + kwargs: _Kwargs = {"model": model, **cast(_Kwargs, dict(_AWS_KWARGS))} + if region != "us-east-1": + kwargs = {"model": model, **cast(_Kwargs, dict(_AWS_KWARGS)), "aws_region_name": region} + is_invoke = "invoke" in model + url: Final = _invoke_url(invoke_model_id) if is_invoke else _converse_url(invoke_model_id, region) + stream_url: Final = ( + _invoke_stream_url(invoke_model_id) if is_invoke else _converse_stream_url(invoke_model_id, region) + ) + return { + "id": case_id, + "shape": shape, + "kwargs": kwargs, + "url": url, + "stream_url": stream_url, + "router": False, + } + + +def _openai_case( + case_id: str, + model: str, + url: str, + *, + router: bool = False, + extra: _Kwargs | None = None, +) -> _Case: + kwargs: _Kwargs = {"model": model, "api_key": "sk-offline"} + if extra is not None: + kwargs = {**kwargs, **extra} + return { + "id": case_id, + "shape": "openai", + "kwargs": kwargs, + "url": url, + "stream_url": url, + "router": router, + } + + +_CASES: Final[tuple[_Case, ...]] = ( + _openai_case("openai_gpt4omini", "gpt-4o-mini", "https://api.openai.com/v1/chat/completions"), + _openai_case("router_gpt4omini", "gpt-4o-mini", "https://api.openai.com/v1/chat/completions", router=True), + _openai_case("openai_o1", "o1", "https://api.openai.com/v1/chat/completions"), + _openai_case("openai_o3mini", "o3-mini", "https://api.openai.com/v1/chat/completions"), + _openai_case( + "azure_o3mini", + "azure/o3-mini", + "https://openai-gpt-4-test-v-1.openai.azure.com/openai/deployments/o3-mini/chat/completions?api-version=2024-02-15-preview", + extra={ + "api_key": "k", + "api_base": "https://openai-gpt-4-test-v-1.openai.azure.com", + "api_version": "2024-02-15-preview", + }, + ), + _openai_case( + "azure_o3mini_live", + "azure/o3-mini", + "https://openai-prod-test.openai.azure.com/openai/deployments/o3-mini/chat/completions?api-version=2024-12-01-preview", + extra={ + "api_key": "k", + "api_base": "https://openai-prod-test.openai.azure.com", + "api_version": "2024-12-01-preview", + }, + ), + { + "id": "anthropic_sonnet45", + "shape": "anthropic", + "kwargs": {"model": "anthropic/claude-sonnet-4-5-20250929", "api_key": "sk-offline"}, + "url": "https://api.anthropic.com/v1/messages", + "stream_url": "https://api.anthropic.com/v1/messages", + "router": False, + }, + { + "id": "gemini_25flash", + "shape": "gemini", + "kwargs": {"model": "gemini/gemini-2.5-flash", "api_key": "k"}, + "url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent", + "stream_url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:streamGenerateContent", + "router": False, + }, + _openai_case("mistral_medium", "mistral/mistral-medium-latest", "https://api.mistral.ai/v1/chat/completions"), + _openai_case("together_glm", "together_ai/zai-org/GLM-5.3-Flash", "https://api.together.ai/v1/chat/completions"), + _openai_case("groq_oss120b", "groq/openai/gpt-oss-120b", "https://api.groq.com/openai/v1/chat/completions"), + _openai_case("xai_grok3mini", "xai/grok-3-mini-beta", "https://api.x.ai/v1/chat/completions"), + _openai_case( + "huggingface_llama", + "huggingface/together/meta-llama/Meta-Llama-3-8B-Instruct", + "https://router.huggingface.co/together/v1/chat/completions", + extra={"api_base": "https://router.huggingface.co/together/v1"}, + ), + _bedrock_case( + "bedrock_converse_haiku", + "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "bedrock_converse", + ), + _bedrock_case( + "bedrock_converse_haiku_xregion", + "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "bedrock_converse", + region="us-west-2", + ), + _bedrock_case( + "bedrock_converse_novalite", "bedrock/us.amazon.nova-lite-v1:0", "us.amazon.nova-lite-v1:0", "bedrock_converse" + ), + _bedrock_case( + "bedrock_converse_novamicro", + "bedrock/converse/us.amazon.nova-micro-v1:0", + "us.amazon.nova-micro-v1:0", + "bedrock_converse", + ), + _bedrock_case( + "bedrock_converse_llama33", + "bedrock/converse/us.meta.llama3-3-70b-instruct-v1:0", + "us.meta.llama3-3-70b-instruct-v1:0", + "bedrock_converse", + ), + _bedrock_case( + "bedrock_converse_gptoss", + "bedrock/converse/openai.gpt-oss-20b-1:0", + "openai.gpt-oss-20b-1:0", + "bedrock_converse", + ), + _bedrock_case( + "bedrock_converse_anthropic_thinking", + "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "bedrock_converse", + ), + _bedrock_case( + "bedrock_invoke_haiku", + "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "bedrock_invoke", + ), + _bedrock_case( + "bedrock_invoke_novamicro", + "bedrock/invoke/us.amazon.nova-micro-v1:0", + "us.amazon.nova-micro-v1:0", + "bedrock_invoke_nova", + ), + _bedrock_case( + "bedrock_invoke_kimi", + "bedrock/invoke/moonshot.kimi-k2-thinking", + "moonshot.kimi-k2-thinking", + "bedrock_invoke_openai", + ), +) + +_BY_ID: Final[Mapping[str, _Case]] = {c["id"]: c for c in _CASES} + + +def _case_id(case: _Case) -> str: + return case["id"] + + +def _pick(*ids: str) -> tuple[_Case, ...]: + return tuple(_BY_ID[i] for i in ids) + + +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_ITEMS: Final = TypeAdapter(list[JsonValue]) + + +def _openai_response(text: str) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "chatcmpl-offline", + "object": "chat.completion", + "created": 1, + "model": "m", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": text}}], + "service_tier": "default", + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + ) + + +def _anthropic_response(text: str) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "msg_offline", + "type": "message", + "role": "assistant", + "model": "m", + "content": [{"type": "text", "text": text}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + ) + + +def _gemini_response(text: str) -> httpx.Response: + return httpx.Response( + 200, + json={ + "candidates": [ + { + "content": {"parts": [{"text": text}], "role": "model"}, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15}, + }, + ) + + +def _converse_response(text: str) -> httpx.Response: + return httpx.Response( + 200, + json={ + "output": {"message": {"role": "assistant", "content": [{"text": text}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, + "metrics": {"latencyMs": 1}, + }, + ) + + +def _canned_response(case: _Case, text: str) -> httpx.Response: + match case["shape"]: + case "openai" | "bedrock_invoke_openai": + return _openai_response(text) + case "anthropic" | "bedrock_invoke": + return _anthropic_response(text) + case "gemini": + return _gemini_response(text) + case "bedrock_converse" | "bedrock_invoke_nova": + return _converse_response(text) + + +_PNG_BYTES: Final = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==" +) +_PNG_B64: Final = base64.b64encode(_PNG_BYTES).decode() +_PNG_URLS: Final = ( + "https://cdn.jsdelivr.net/gh/BerriAI/litellm@d769e81c90d453240c61fc572cdb27fae06a89d0/ui/litellm-dashboard/public/assets/logos/litellm_logo.jpg", + "https://awsmp-logos.s3.amazonaws.com/seller-xw5kijmvmzasy/c233c9ade2ccb5491072ae232c814942.png", +) + + +def _aws_str_header(name: str, value: str) -> bytes: + name_bytes: Final = name.encode() + value_bytes: Final = value.encode() + return ( + struct.pack("!B", len(name_bytes)) + + name_bytes + + struct.pack("!B", 7) + + struct.pack("!H", len(value_bytes)) + + value_bytes + ) + + +def _aws_frame(event_type: str, payload: Mapping[str, JsonValue]) -> bytes: + payload_bytes: Final = json.dumps(payload).encode() + headers_bytes: Final = b"".join( + ( + _aws_str_header(":event-type", event_type), + _aws_str_header(":content-type", "application/json"), + _aws_str_header(":message-type", "event"), + ) + ) + prelude: Final = struct.pack("!II", 12 + len(headers_bytes) + len(payload_bytes) + 4, len(headers_bytes)) + message: Final = prelude + struct.pack("!I", zlib.crc32(prelude) & 0xFFFFFFFF) + headers_bytes + payload_bytes + return message + struct.pack("!I", zlib.crc32(message) & 0xFFFFFFFF) + + +def _openai_sse(text: str) -> str: + chunks: Final = ( + { + "id": "chatcmpl-offline", + "object": "chat.completion.chunk", + "created": 1, + "model": "m", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": text}, "finish_reason": None}], + "service_tier": "default", + }, + { + "id": "chatcmpl-offline", + "object": "chat.completion.chunk", + "created": 1, + "model": "m", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "service_tier": "default", + }, + ) + return "".join(f"data: {json.dumps(c)}\n\n" for c in chunks) + "data: [DONE]\n\n" + + +def _anthropic_sse(text: str) -> str: + events: Final = ( + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_offline", + "type": "message", + "role": "assistant", + "content": [], + "model": "m", + "stop_reason": None, + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, + }, + ), + ( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 5}}, + ), + ("message_stop", {"type": "message_stop"}), + ) + return "".join(f"event: {name}\ndata: {json.dumps(payload)}\n\n" for name, payload in events) + + +def _gemini_sse(text: str) -> str: + event: Final = { + "candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "finishReason": "STOP", "index": 0}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15}, + } + return f"data: {json.dumps(event)}\n\n" + + +def _converse_stream_bytes(text: str) -> bytes: + frames: Final = ( + _aws_frame("messageStart", {"role": "assistant"}), + _aws_frame("contentBlockDelta", {"delta": {"text": text}, "contentBlockIndex": 0}), + _aws_frame("contentBlockStop", {"contentBlockIndex": 0}), + _aws_frame("messageStop", {"stopReason": "end_turn"}), + _aws_frame("metadata", {"usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}}), + ) + return b"".join(frames) + + +def _invoke_stream_bytes(text: str) -> bytes: + events: Final = ( + { + "type": "message_start", + "message": { + "id": "msg_invoke", + "type": "message", + "role": "assistant", + "content": [], + "model": "m", + "stop_reason": None, + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 5}}, + {"type": "message_stop"}, + ) + return b"".join(_aws_frame("chunk", {"bytes": base64.b64encode(json.dumps(e).encode()).decode()}) for e in events) + + +def _invoke_nova_stream_bytes(text: str) -> bytes: + events: Final = ( + {"messageStart": {"role": "assistant"}}, + {"contentBlockDelta": {"delta": {"text": text}, "contentBlockIndex": 0}}, + {"contentBlockStop": {"contentBlockIndex": 0}}, + {"messageStop": {"stopReason": "end_turn"}}, + {"metadata": {"usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}}}, + ) + return b"".join(_aws_frame("chunk", {"bytes": base64.b64encode(json.dumps(e).encode()).decode()}) for e in events) + + +def _stream_canned(case: _Case, text: str) -> httpx.Response: + match case["shape"]: + case "openai" | "bedrock_invoke_openai": + return httpx.Response(200, content=_openai_sse(text), headers={"content-type": "text/event-stream"}) + case "anthropic": + return httpx.Response(200, content=_anthropic_sse(text), headers={"content-type": "text/event-stream"}) + case "gemini": + return httpx.Response(200, content=_gemini_sse(text), headers={"content-type": "text/event-stream"}) + case "bedrock_converse": + return httpx.Response( + 200, + content=_converse_stream_bytes(text), + headers={"content-type": "application/vnd.amazon.eventstream"}, + ) + case "bedrock_invoke": + return httpx.Response( + 200, + content=_invoke_stream_bytes(text), + headers={"content-type": "application/vnd.amazon.eventstream"}, + ) + case "bedrock_invoke_nova": + return httpx.Response( + 200, + content=_invoke_nova_stream_bytes(text), + headers={"content-type": "application/vnd.amazon.eventstream"}, + ) + + +@pytest.fixture(autouse=True) +def _httpx_only_transport(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + + +def _tool_use_response(case: _Case, payload: str) -> httpx.Response: + arguments: Final = json.loads(payload) + match case["shape"]: + case "anthropic" | "bedrock_invoke": + return httpx.Response( + 200, + json={ + "id": "msg_offline", + "type": "message", + "role": "assistant", + "model": "m", + "content": [ + {"type": "tool_use", "id": "toolu_offline", "name": "json_tool_call", "input": arguments} + ], + "stop_reason": "tool_use", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + ) + case "bedrock_converse" | "bedrock_invoke_nova": + return httpx.Response( + 200, + json={ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_offline", + "name": "json_tool_call", + "input": arguments, + } + } + ], + } + }, + "stopReason": "tool_use", + "usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, + "metrics": {"latencyMs": 1}, + }, + ) + case _: + return _canned_response(case, payload) + + +def _json_responder(case: _Case, payload: str) -> Callable[[httpx.Request], httpx.Response]: + def respond(request: httpx.Request) -> httpx.Response: + if b"json_tool_call" in request.content: + return _tool_use_response(case, payload) + return _canned_response(case, payload) + + return respond + + +def _register( + case: _Case, + respx_mock: MockRouter, + text: str, + *, + stream: bool = False, + json_via_tool: bool = False, +) -> respx.Route: + for url in _PNG_URLS: + respx_mock.get(url).mock(return_value=httpx.Response(200, content=_PNG_BYTES)) + target: Final = case["stream_url"] if stream else case["url"] + if json_via_tool: + return respx_mock.post(target).mock(side_effect=_json_responder(case, text)) + if case["shape"] == "gemini" and stream: + return respx_mock.post(url__startswith=target).mock(return_value=_stream_canned(case, text)) + if stream and case["id"].startswith("groq"): + return respx_mock.post(target).mock(return_value=_canned_response(case, text)) + return respx_mock.post(target).mock( + return_value=_stream_canned(case, text) if stream else _canned_response(case, text) + ) + + +_ROUTER_MODEL_LIST: Final = [ + { + "model_name": "offline-router-model", + "litellm_params": {"model": "gpt-4o-mini", "api_key": "sk-offline"}, + } +] + + +def _call(case: _Case, params: Mapping[str, JsonValue]) -> litellm.ModelResponse | litellm.CustomStreamWrapper: + extra: Final = copy.deepcopy(params) + if case["router"]: + router: Final = litellm.Router(model_list=copy.deepcopy(_ROUTER_MODEL_LIST)) + return router.completion(model="offline-router-model", **extra) + return litellm.completion(**case["kwargs"], **extra) + + +def _complete(case: _Case, params: Mapping[str, JsonValue]) -> litellm.ModelResponse: + response: Final = _call(case, params) + assert isinstance(response, litellm.ModelResponse) + return response + + +def _stream(case: _Case, params: Mapping[str, JsonValue]) -> litellm.CustomStreamWrapper: + response: Final = _call(case, {**params, "stream": True}) + assert isinstance(response, litellm.CustomStreamWrapper) + return response + + +async def _acomplete(case: _Case, params: Mapping[str, JsonValue]) -> litellm.ModelResponse: + extra: Final = copy.deepcopy(params) + if case["router"]: + router: Final = litellm.Router(model_list=copy.deepcopy(_ROUTER_MODEL_LIST)) + response: Final = await router.acompletion(model="offline-router-model", **extra) + else: + response = await litellm.acompletion(**case["kwargs"], **extra) + assert isinstance(response, litellm.ModelResponse) + return response + + +def _request_body(route: respx.Route) -> Mapping[str, JsonValue]: + assert route.calls, "provider route was never called" + return _JSON.validate_python(json.loads(route.calls.last.request.content)) + + +def _mapping(value: JsonValue) -> Mapping[str, JsonValue]: + return _JSON.validate_python(value) + + +def _items(value: JsonValue) -> tuple[JsonValue, ...]: + return tuple(_ITEMS.validate_python(value)) + + +def _mappings(value: JsonValue) -> tuple[Mapping[str, JsonValue], ...]: + return tuple(_mapping(item) for item in _items(value)) + + +def _flatten(groups: Iterable[Iterable[str]]) -> tuple[str, ...]: + return tuple(itertools.chain.from_iterable(groups)) + + +def _typed_text_parts(content: JsonValue) -> tuple[str, ...]: + if isinstance(content, str): + return (content,) + return tuple(str(p["text"]) for p in _mappings(content) if p.get("type") == "text") + + +def _keyed_text_parts(content: JsonValue) -> tuple[str, ...]: + return tuple(str(p["text"]) for p in _mappings(content) if "text" in p) + + +def _user_texts(case: _Case, body: Mapping[str, JsonValue]) -> tuple[str, ...]: + match case["shape"]: + case "openai" | "bedrock_invoke_openai": + return _flatten(_typed_text_parts(m["content"]) for m in _mappings(body["messages"]) if m["role"] == "user") + case "anthropic" | "bedrock_invoke": + return _flatten(_typed_text_parts(m["content"]) for m in _mappings(body["messages"]) if m["role"] == "user") + case "bedrock_converse" | "bedrock_invoke_nova": + return _flatten(_keyed_text_parts(m["content"]) for m in _mappings(body["messages"]) if m["role"] == "user") + case "gemini": + return _flatten( + _keyed_text_parts(m["parts"]) for m in _mappings(body["contents"]) if m.get("role") != "model" + ) + + +def _system_texts(case: _Case, body: Mapping[str, JsonValue]) -> tuple[str, ...]: + match case["shape"]: + case "openai" | "bedrock_invoke_openai": + return tuple(str(m["content"]) for m in _mappings(body["messages"]) if m["role"] == "system") + case "anthropic" | "bedrock_invoke" | "bedrock_converse" | "bedrock_invoke_nova": + return tuple(str(s["text"]) for s in _mappings(body.get("system", []))) + case "gemini": + return tuple( + str(p["text"]) for p in _mappings(_mapping(body.get("system_instruction", {})).get("parts", [])) + ) + + +def _message_roles(case: _Case, body: Mapping[str, JsonValue]) -> tuple[str, ...]: + match case["shape"]: + case "gemini": + return tuple(str(m.get("role", "user")) for m in _mappings(body["contents"])) + case _: + return tuple(str(m["role"]) for m in _mappings(body["messages"])) + + +def _tools_payload(case: _Case, body: Mapping[str, JsonValue]) -> JsonValue: + match case["shape"]: + case "openai" | "bedrock_invoke_openai" | "anthropic" | "bedrock_invoke": + return body.get("tools") + case "bedrock_converse" | "bedrock_invoke_nova": + return _mapping(body.get("toolConfig", {})).get("tools") + case "gemini": + declared: Final = _mappings(body.get("tools", [])) + if not declared: + return None + return declared[0].get("function_declarations") + + +def _tool_names(case: _Case, body: Mapping[str, JsonValue]) -> tuple[str, ...]: + entries: Final = _mappings(_tools_payload(case, body) or []) + match case["shape"]: + case "openai" | "bedrock_invoke_openai": + return tuple(str(_mapping(e["function"])["name"]) for e in entries) + case "bedrock_converse" | "bedrock_invoke_nova": + return tuple(str(_mapping(e["toolSpec"])["name"]) for e in entries) + case "anthropic" | "bedrock_invoke" | "gemini": + return tuple(str(e["name"]) for e in entries) + + +def _tool_input_schema(case: _Case, body: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + first: Final = _mappings(_tools_payload(case, body) or [])[0] + match case["shape"]: + case "openai" | "bedrock_invoke_openai": + return _mapping(_mapping(first["function"])["parameters"]) + case "anthropic" | "bedrock_invoke": + return _mapping(first["input_schema"]) + case "bedrock_converse" | "bedrock_invoke_nova": + return _mapping(_mapping(_mapping(first["toolSpec"])["inputSchema"])["json"]) + case "gemini": + return _mapping(first["parameters"]) + + +@pytest.mark.parametrize( + "case", + _pick( + "anthropic_sonnet45", + "bedrock_converse_haiku", + "bedrock_converse_novalite", + "bedrock_converse_gptoss", + "bedrock_converse_llama33", + "bedrock_invoke_haiku", + "bedrock_invoke_novamicro", + "groq_oss120b", + "huggingface_llama", + "mistral_medium", + "openai_gpt4omini", + "router_gpt4omini", + "together_glm", + "xai_grok3mini", + ), + ids=_case_id, +) +def test_developer_role_translation(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, f"canned-{case['id']}") + response: Final = _complete( + case, + { + "messages": [ + {"role": "developer", "content": "Be a good bot!"}, + {"role": "user", "content": [{"type": "text", "text": "Hello, how are you?"}]}, + ] + }, + ) + body: Final = _request_body(route) + assert "developer" not in _message_roles(case, body) + assert _system_texts(case, body) == ("Be a good bot!",) + assert "Hello, how are you?" in _user_texts(case, body) + assert response.choices[0].message.content == f"canned-{case['id']}" + + +@pytest.mark.parametrize( + "case", + _pick( + "bedrock_converse_gptoss", + "bedrock_invoke_novamicro", + "huggingface_llama", + "mistral_medium", + "openai_gpt4omini", + "openai_o1", + "openai_o3mini", + "router_gpt4omini", + "together_glm", + "xai_grok3mini", + ), + ids=_case_id, +) +def test_content_list_handling(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, f"canned-{case['id']}") + response: Final = _complete( + case, + {"messages": [{"role": "user", "content": [{"type": "text", "text": "Hello, how are you?"}]}]}, + ) + body: Final = _request_body(route) + assert _user_texts(case, body) == ("Hello, how are you?",) + first_content: Final = _mappings(body["messages"])[0]["content"] + if case["id"] == "mistral_medium": + assert first_content == "Hello, how are you?" + else: + assert isinstance(first_content, list) + assert response.choices[0].message.content == f"canned-{case['id']}" + + +_TOOL_ARRAY_SCHEMA: Final[Mapping[str, JsonValue]] = { + "type": "function", + "function": { + "name": "shoe_get_id", + "description": "Get information about a show by its ID or name", + "parameters": { + "type": "object", + "properties": {"shoe_id": {"type": ["string", "number"], "description": "The shoe ID or name"}}, + "required": ["shoe_id"], + "additionalProperties": False, + "$schema": "http://json-schema.org/draft-07/schema#", + }, + }, +} + + +@pytest.mark.parametrize( + "case", + _pick( + "anthropic_sonnet45", + "azure_o3mini", + "bedrock_converse_haiku", + "bedrock_converse_haiku_xregion", + "bedrock_converse_novalite", + "bedrock_converse_gptoss", + "bedrock_converse_llama33", + "bedrock_invoke_haiku", + "bedrock_invoke_kimi", + "gemini_25flash", + "groq_oss120b", + "mistral_medium", + "openai_gpt4omini", + "openai_o3mini", + "router_gpt4omini", + ), + ids=_case_id, +) +def test_tool_call_with_property_type_array(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, f"canned-{case['id']}") + response: Final = _complete( + case, + { + "messages": [{"role": "user", "content": "Tell me about shoes"}], + "tools": [_TOOL_ARRAY_SCHEMA], + }, + ) + body: Final = _request_body(route) + assert _tool_names(case, body) == ("shoe_get_id",) + schema: Final = _tool_input_schema(case, body) + shoe_id: Final = _mapping(_mapping(schema["properties"])["shoe_id"]) + assert schema["required"] == ["shoe_id"] + if case["shape"] == "gemini": + assert [_mapping(v)["type"] for v in _items(shoe_id["anyOf"])] == ["string", "number"] + else: + assert shoe_id["type"] == ["string", "number"] + assert response.choices[0].message.content == f"canned-{case['id']}" + + +_TOOL_ENUM_SCHEMA: Final[Mapping[str, JsonValue]] = { + "type": "function", + "function": { + "name": "litellm_product_search", + "description": "Search for product information", + "parameters": { + "properties": { + "search_mode": { + "default": "", + "description": "The search strategy to use", + "enum": ["", "product_search", "product_search_with_filters"], + "type": "string", + } + }, + "required": ["search_mode"], + "type": "object", + }, + }, +} + + +@pytest.mark.parametrize( + "case", + _pick( + "anthropic_sonnet45", + "azure_o3mini", + "bedrock_converse_haiku", + "bedrock_converse_haiku_xregion", + "bedrock_converse_novalite", + "bedrock_converse_gptoss", + "bedrock_converse_llama33", + "bedrock_invoke_haiku", + "bedrock_invoke_kimi", + "gemini_25flash", + "mistral_medium", + "openai_gpt4omini", + "openai_o3mini", + "router_gpt4omini", + ), + ids=_case_id, +) +def test_tool_call_with_empty_enum_property(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, f"canned-{case['id']}") + response: Final = _complete( + case, + { + "messages": [{"role": "user", "content": "Search for the latest iPhone models"}], + "tools": [_TOOL_ENUM_SCHEMA], + }, + ) + body: Final = _request_body(route) + assert _tool_names(case, body) == ("litellm_product_search",) + schema: Final = _tool_input_schema(case, body) + search_mode: Final = _mapping(_mapping(schema["properties"])["search_mode"]) + enum_values: Final = _items(search_mode["enum"]) + assert len(enum_values) == 3 + assert enum_values[0] == (None if case["shape"] == "gemini" else "") + assert enum_values[1:] == ("product_search", "product_search_with_filters") + assert response.choices[0].message.content == f"canned-{case['id']}" + + +@pytest.mark.parametrize( + "case", + _pick( + "anthropic_sonnet45", + "azure_o3mini", + "bedrock_converse_haiku", + "bedrock_converse_haiku_xregion", + "bedrock_converse_novalite", + "bedrock_converse_novamicro", + "bedrock_converse_gptoss", + "bedrock_converse_llama33", + "bedrock_invoke_haiku", + "gemini_25flash", + "groq_oss120b", + "huggingface_llama", + "mistral_medium", + "openai_gpt4omini", + "openai_o1", + "openai_o3mini", + "router_gpt4omini", + "together_glm", + "xai_grok3mini", + ), + ids=_case_id, +) +def test_pydantic_model_input(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, f"canned-{case['id']}") + messages: Final = [litellm.Message(content="Hello, how are you?", role="user")] + response: Final = _complete(case, {"messages": messages}) + assert "Hello, how are you?" in _user_texts(case, _request_body(route)) + assert response.choices[0].message.content == f"canned-{case['id']}" + + +@pytest.mark.parametrize( + "case", + _pick( + "anthropic_sonnet45", + "bedrock_converse_haiku", + "bedrock_converse_haiku_xregion", + "bedrock_converse_novalite", + "bedrock_invoke_haiku", + "gemini_25flash", + "openai_gpt4omini", + "openai_o1", + "router_gpt4omini", + ), + ids=_case_id, +) +def test_file_data_unit_test(case: _Case, respx_mock: MockRouter) -> None: + pdf_b64: Final = base64.b64encode(b"%PDF-1.4 offline dummy").decode() + file_data_url: Final = f"data:application/pdf;base64,{pdf_b64}" + raw_request: Final = return_raw_request( + endpoint=CallTypes.completion, + kwargs={ + **{k: v for k, v in dict(case["kwargs"]).items() if k != "api_key"}, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's this file about?"}, + {"type": "file", "file": {"file_data": file_data_url}}, + ], + } + ], + }, + ) + assert raw_request.get("error") is None + assert pdf_b64 in json.dumps(raw_request.get("raw_request_body")) + + +@pytest.mark.parametrize( + "case", + _pick( + "anthropic_sonnet45", + "azure_o3mini", + "bedrock_converse_haiku", + "bedrock_converse_haiku_xregion", + "bedrock_converse_novalite", + "bedrock_converse_novamicro", + "bedrock_converse_gptoss", + "bedrock_converse_llama33", + "bedrock_invoke_haiku", + "gemini_25flash", + "groq_oss120b", + "huggingface_llama", + "mistral_medium", + "openai_gpt4omini", + "openai_o1", + "openai_o3mini", + "router_gpt4omini", + "together_glm", + "xai_grok3mini", + ), + ids=_case_id, +) +def test_message_with_name(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, f"canned-{case['id']}") + response: Final = _complete(case, {"messages": [{"role": "user", "content": "Hello", "name": "test_name"}]}) + body: Final = _request_body(route) + assert "Hello" in _user_texts(case, body) + if case["shape"] in ("openai", "bedrock_invoke_openai"): + first: Final = cast(Mapping[str, JsonValue], cast(list[JsonValue], body["messages"])[0]) + if case["id"] == "mistral_medium": + assert "name" not in first + else: + assert first.get("name") == "test_name" + else: + assert "test_name" not in json.dumps(body) + assert response.choices[0].message.content == f"canned-{case['id']}" + + +@pytest.mark.parametrize("response_format", ({"type": "json_object"}, {"type": "text"})) +@pytest.mark.parametrize( + "case", + _pick( + "anthropic_sonnet45", + "azure_o3mini", + "bedrock_converse_haiku", + "bedrock_converse_haiku_xregion", + "bedrock_converse_novalite", + "bedrock_converse_novamicro", + "bedrock_invoke_haiku", + "bedrock_invoke_kimi", + "gemini_25flash", + "groq_oss120b", + "mistral_medium", + "openai_gpt4omini", + "openai_o1", + "openai_o3mini", + "router_gpt4omini", + ), + ids=_case_id, +) +def test_json_response_format(case: _Case, response_format: Mapping[str, JsonValue], respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, '{"city":"San Francisco","state":"CA"}') + response: Final = _complete( + case, + { + "messages": [ + {"role": "system", "content": "Your output should be a JSON object with no additional properties."}, + {"role": "user", "content": "Respond with this in json. city=San Francisco, state=CA"}, + ], + "response_format": response_format, + }, + ) + body: Final = _request_body(route) + match case["shape"]: + case "gemini": + mime_types: Final = {"json_object": "application/json", "text": "text/plain"} + config: Final = _mapping(body["generationConfig"]) + assert config["response_mime_type"] == mime_types[str(response_format["type"])] + case "openai" | "bedrock_invoke_openai": + assert body["response_format"] == response_format + case _: + assert "response_format" not in body + assert not _tools_payload(case, body) + assert response.choices[0].message.content == '{"city":"San Francisco","state":"CA"}' + + +_WEATHER_TOOL: Final[Mapping[str, JsonValue]] = { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city and state"}, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + "required": ["location"], + }, + }, +} + + +@pytest.mark.parametrize( + "case", + _pick( + "anthropic_sonnet45", + "azure_o3mini", + "bedrock_converse_haiku", + "bedrock_converse_haiku_xregion", + "bedrock_converse_novalite", + "bedrock_converse_novamicro", + "bedrock_converse_gptoss", + "bedrock_converse_llama33", + "bedrock_invoke_haiku", + "gemini_25flash", + "groq_oss120b", + "huggingface_llama", + "mistral_medium", + "openai_gpt4omini", + "openai_o1", + "openai_o3mini", + "router_gpt4omini", + "together_glm", + "xai_grok3mini", + ), + ids=_case_id, +) +def test_response_format_type_text_with_tool_calls_no_tool_choice(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, f"canned-{case['id']}") + response: Final = _complete( + case, + { + "messages": [{"role": "user", "content": "What's the weather like in Boston today?"}], + "response_format": {"type": "text"}, + "tools": [_WEATHER_TOOL], + "drop_params": True, + }, + ) + body: Final = _request_body(route) + assert _tool_names(case, body) == ("get_current_weather",) + assert "tool_choice" not in body + assert "toolChoice" not in _mapping(body.get("toolConfig", {})) + assert response.choices[0].message.content == f"canned-{case['id']}" + + +@pytest.mark.parametrize( + "case", + _CASES, + ids=_case_id, +) +def test_response_format_type_text(case: _Case) -> None: + _, provider, _, _ = litellm.get_llm_provider(model=case["kwargs"]["model"]) + provider_config: Final = ProviderConfigManager.get_provider_chat_config( + case["kwargs"]["model"], litellm.LlmProviders(provider) + ) + translated_params: Final = provider_config.map_openai_params( + non_default_params={"response_format": {"type": "text"}}, + optional_params={}, + model=case["kwargs"]["model"], + drop_params=False, + ) + assert "tool_choice" not in translated_params + assert "tools" not in translated_params + + +class _FirstResponse(BaseModel): + model_config = ConfigDict(frozen=True) + first_response: str + + +class _CalendarEvent(BaseModel): + model_config = ConfigDict(frozen=True) + name: str + date: str + participants: list[str] + + +class _EventsList(BaseModel): + model_config = ConfigDict(frozen=True) + events: list[_CalendarEvent] + + +@pytest.mark.parametrize( + "case", + _pick( + "anthropic_sonnet45", + "azure_o3mini", + "bedrock_converse_haiku", + "bedrock_converse_haiku_xregion", + "bedrock_converse_novalite", + "bedrock_converse_novamicro", + "bedrock_converse_gptoss", + "bedrock_invoke_haiku", + "bedrock_invoke_novamicro", + "bedrock_invoke_kimi", + "groq_oss120b", + "mistral_medium", + "openai_gpt4omini", + "openai_o1", + "openai_o3mini", + "router_gpt4omini", + ), + ids=_case_id, +) +def test_json_response_pydantic_obj(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, '{"first_response":"paris"}', json_via_tool=True) + response: Final = _complete( + case, + { + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is the capital of France?"}, + ], + "response_format": _FirstResponse, + }, + ) + body: Final = _request_body(route) + serialized: Final = json.dumps(body) + assert "first_response" in serialized + assert json.loads(response.choices[0].message.content) == {"first_response": "paris"} + assert response.choices[0].message.tool_calls is None + + +@pytest.mark.parametrize( + "case", + _pick( + "anthropic_sonnet45", + "azure_o3mini", + "bedrock_converse_haiku", + "bedrock_converse_haiku_xregion", + "bedrock_converse_novalite", + "bedrock_converse_novamicro", + "bedrock_converse_gptoss", + "bedrock_invoke_haiku", + "bedrock_invoke_novamicro", + "bedrock_invoke_kimi", + "bedrock_invoke_novamicro", + "groq_oss120b", + "mistral_medium", + "openai_gpt4omini", + "openai_o1", + "openai_o3mini", + "router_gpt4omini", + ), + ids=_case_id, +) +def test_json_response_nested_pydantic_obj(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, '{"events":[]}', json_via_tool=True) + response: Final = _complete( + case, + { + "messages": [{"role": "user", "content": "List 5 important events in the XIX century"}], + "response_format": _EventsList, + }, + ) + body: Final = _request_body(route) + serialized: Final = json.dumps(body) + assert "events" in serialized + assert "participants" in serialized + assert json.loads(response.choices[0].message.content) == {"events": []} + assert response.choices[0].message.tool_calls is None + + +@pytest.mark.parametrize( + "case", + _pick( + "anthropic_sonnet45", + "azure_o3mini", + "bedrock_converse_haiku", + "bedrock_converse_haiku_xregion", + "bedrock_converse_novalite", + "bedrock_converse_novamicro", + "bedrock_converse_gptoss", + "bedrock_invoke_haiku", + "bedrock_invoke_novamicro", + "bedrock_invoke_kimi", + "bedrock_invoke_novamicro", + "groq_oss120b", + "mistral_medium", + "openai_gpt4omini", + "openai_o1", + "openai_o3mini", + "router_gpt4omini", + ), + ids=_case_id, +) +def test_json_response_nested_json_schema(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, '{"events":[]}', json_via_tool=True) + response: Final = _complete( + case, + { + "messages": [{"role": "user", "content": "List 5 important events in the XIX century"}], + "response_format": type_to_response_format_param(_EventsList), + }, + ) + body: Final = _request_body(route) + serialized: Final = json.dumps(body) + assert "events" in serialized + assert "participants" in serialized + assert json.loads(response.choices[0].message.content) == {"events": []} + assert response.choices[0].message.tool_calls is None + + +def test_audio_input_gemini(respx_mock: MockRouter) -> None: + case: Final = _BY_ID["gemini_25flash"] + wav_b64: Final = base64.b64encode(b"RIFFFAKEWAVDATA").decode() + route: Final = _register(case, respx_mock, "canned-gemini") + response: Final = _complete( + case, + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this recording?"}, + {"type": "input_audio", "input_audio": {"data": wav_b64, "format": "wav"}}, + ], + } + ] + }, + ) + body: Final = _request_body(route) + first_content: Final = cast(Mapping[str, JsonValue], cast(list[JsonValue], body["contents"])[0]) + parts: Final = cast(list[JsonValue], first_content["parts"]) + audio_part: Final = cast(Mapping[str, JsonValue], parts[1]) + assert cast(Mapping[str, JsonValue], audio_part["inline_data"])["data"] == wav_b64 + assert response.choices[0].message.content == "canned-gemini" + + +@pytest.mark.parametrize( + "case", + _pick( + "anthropic_sonnet45", + "azure_o3mini", + "bedrock_converse_haiku", + "bedrock_converse_novalite", + "bedrock_converse_gptoss", + "bedrock_invoke_haiku", + "bedrock_invoke_novamicro", + "gemini_25flash", + "groq_oss120b", + "mistral_medium", + "openai_gpt4omini", + "openai_o1", + "openai_o3mini", + "router_gpt4omini", + "together_glm", + ), + ids=_case_id, +) +def test_json_response_format_stream(case: _Case, respx_mock: MockRouter) -> None: + canned: Final = '{"city":"San Francisco"}' + route: Final = _register(case, respx_mock, canned, stream=True) + response: Final = _stream( + case, + { + "messages": [ + {"role": "system", "content": "Your output should be a JSON object with no additional properties."}, + {"role": "user", "content": "Respond with this in json. city=San Francisco, state=CA"}, + ], + "response_format": {"type": "json_object"}, + }, + ) + content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in response) + assert content == canned + body: Final = _request_body(route) + if case["shape"] in ("openai", "anthropic") and not case["id"].startswith("groq"): + assert body["stream"] is True + + +@pytest.mark.parametrize( + "case", + _pick( + "bedrock_invoke_haiku", + "mistral_medium", + "openai_gpt4omini", + "openai_o1", + "router_gpt4omini", + "together_glm", + ), + ids=_case_id, +) +@pytest.mark.parametrize("detail", (None, "low", "high"), ids=("detail_none", "detail_low", "detail_high")) +@pytest.mark.parametrize("image_url", _PNG_URLS, ids=("litellm_logo", "awsmp_png")) +def test_image_url(case: _Case, detail: str | None, image_url: str, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, f"canned-{case['id']}") + image_url_part: Final[Mapping[str, JsonValue]] = ( + {"url": image_url} if detail is None else {"url": image_url, "detail": detail} + ) + response: Final = _complete( + case, + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + {"type": "image_url", "image_url": image_url_part}, + ], + } + ] + }, + ) + body: Final = _request_body(route) + assert "What's in this image?" in _user_texts(case, body) + content: Final = _mappings(_mappings(body["messages"])[0]["content"]) + image_block: Final = content[1] + if case["shape"] == "bedrock_invoke": + assert image_block["type"] == "image" + source: Final = _mapping(image_block["source"]) + assert source["type"] == "base64" + assert source["data"] == _PNG_B64 + assert source["media_type"] == ("image/jpeg" if image_url.endswith(".jpg") else "image/png") + else: + assert image_block["type"] == "image_url" + assert image_block["image_url"] == image_url_part + assert response.choices[0].message.content == f"canned-{case['id']}" + + +@pytest.mark.parametrize( + "case", + _pick( + "bedrock_converse_haiku", + "bedrock_converse_haiku_xregion", + "bedrock_converse_novalite", + "gemini_25flash", + "mistral_medium", + "openai_gpt4omini", + "openai_o1", + "router_gpt4omini", + "together_glm", + ), + ids=_case_id, +) +def test_image_url_string(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, f"canned-{case['id']}") + response: Final = _complete( + case, + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + {"type": "image_url", "image_url": _PNG_URLS[1]}, + ], + } + ] + }, + ) + body: Final = _request_body(route) + assert "What's in this image?" in _user_texts(case, body) + image_block: Final = _items( + _mappings(body["contents"])[0]["parts"] + if case["shape"] == "gemini" + else _mappings(body["messages"])[0]["content"] + )[1] + match case["shape"]: + case "openai": + assert image_block == {"type": "image_url", "image_url": {"url": _PNG_URLS[1]}} + case _: + assert _PNG_B64 in json.dumps(image_block) + assert _PNG_URLS[1] not in json.dumps(image_block) + assert response.choices[0].message.content == f"canned-{case['id']}" + + +@pytest.mark.parametrize( + "case", + _pick( + "anthropic_sonnet45", + "bedrock_converse_haiku_xregion", + "bedrock_converse_novalite", + "bedrock_converse_gptoss", + "bedrock_invoke_haiku", + "bedrock_invoke_kimi", + "gemini_25flash", + "mistral_medium", + "openai_gpt4omini", + "openai_o3mini", + "router_gpt4omini", + ), + ids=_case_id, +) +def test_empty_tools(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, f"canned-{case['id']}") + response: Final = _complete( + case, + { + "messages": [{"role": "user", "content": "Hello, how are you?"}], + "tools": [], + }, + ) + body: Final = _request_body(route) + if case["shape"] in ("bedrock_converse", "bedrock_invoke_nova", "gemini"): + assert "toolConfig" not in body + assert "tools" not in body + else: + assert _tools_payload(case, body) == [] + assert response.choices[0].message.content == f"canned-{case['id']}" + + +def _cost_model_key(case: _Case) -> str | None: + model: Final = cast(str, case["kwargs"]["model"]) + stripped: Final = model.split("/", 1)[-1] + for candidate in (stripped, model): + if candidate in litellm.model_cost: + return candidate + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "case", + _pick( + "anthropic_sonnet45", + "azure_o3mini", + "bedrock_converse_haiku", + "bedrock_converse_novalite", + "bedrock_converse_novamicro", + "bedrock_converse_llama33", + "bedrock_invoke_haiku", + "gemini_25flash", + "groq_oss120b", + "mistral_medium", + "openai_gpt4omini", + "openai_o1", + "openai_o3mini", + "router_gpt4omini", + "together_glm", + "xai_grok3mini", + ), + ids=_case_id, +) +async def test_completion_cost(case: _Case, respx_mock: MockRouter) -> None: + _register(case, respx_mock, f"canned-{case['id']}") + response: Final = await _acomplete( + case, + {"messages": [{"role": "user", "content": "Hello, how are you?"}]}, + ) + usage: Final = response.usage + assert usage.prompt_tokens == 10 + assert usage.completion_tokens == 5 + assert usage.total_tokens == 15 + model_key: Final = _cost_model_key(case) + cost_entry: Final = litellm.model_cost.get(model_key) if model_key is not None else None + actual_cost: Final = response._hidden_params["response_cost"] + if cost_entry is not None and "input_cost_per_token" in cost_entry: + expected: Final = 10 * cost_entry["input_cost_per_token"] + 5 * cost_entry["output_cost_per_token"] + assert actual_cost == pytest.approx(expected) + else: + try: + expected_cost: Final = litellm.completion_cost( + completion_response=response, model=cast(str, case["kwargs"]["model"]) + ) + except litellm.exceptions.ModelNotMappedError: + assert actual_cost is None + else: + assert actual_cost == expected_cost + + +@pytest.mark.parametrize("input_type", ("input_audio", "audio_url")) +def test_supports_audio_input_gemini(input_type: str) -> None: + wav_b64: Final = base64.b64encode(b"RIFFFAKEWAVDATA").decode() + audio_part: Final[Mapping[str, JsonValue]] = ( + {"type": "input_audio", "input_audio": {"data": wav_b64, "format": "wav"}} + if input_type == "input_audio" + else { + "type": "file", + "file": {"file_id": "gs://bucket/file.wav", "filename": "my-sample-audio-file"}, + } + ) + raw_request: Final = return_raw_request( + endpoint=CallTypes.completion, + kwargs={ + "model": "gemini/gemini-2.5-flash", + "modalities": ["text", "audio"], + "audio": {"voice": "alloy", "format": "wav"}, + "drop_params": True, + "messages": [ + { + "role": "user", + "content": [{"type": "text", "text": "What is in this recording?"}, audio_part], + } + ], + }, + ) + assert raw_request.get("error") is None + serialized: Final = json.dumps(raw_request.get("raw_request_body")) + if input_type == "input_audio": + assert wav_b64 in serialized + else: + assert "gs://bucket/file.wav" in serialized + + +def test_reasoning_effort_gemini(respx_mock: MockRouter) -> None: + case: Final = _BY_ID["gemini_25flash"] + route: Final = _register(case, respx_mock, "canned-gemini") + optional_params: Final = litellm.get_optional_params( + model="gemini/gemini-2.5-flash", + custom_llm_provider="gemini", + reasoning_effort="high", + ) + assert optional_params["thinkingConfig"] == { + "thinkingBudget": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, + "includeThoughts": True, + } + response: Final = _complete( + case, + { + "messages": [{"role": "user", "content": "Hello!"}], + "reasoning_effort": "low", + }, + ) + body: Final = _request_body(route) + config: Final = cast(Mapping[str, JsonValue], body["generationConfig"]) + assert config["thinkingConfig"] == { + "thinkingBudget": DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, + "includeThoughts": True, + } + assert response.choices[0].message.content == "canned-gemini" + + +@pytest.mark.parametrize("case", _pick("openai_o1", "openai_o3mini", "azure_o3mini", "azure_o3mini_live"), ids=_case_id) +def test_o_series_reasoning_effort_forwarded(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, f"canned-{case['id']}") + _complete( + case, + { + "messages": [{"role": "user", "content": "Hello!"}], + "reasoning_effort": "low", + }, + ) + body: Final = _request_body(route) + assert body["reasoning_effort"] == "low" + + +@pytest.mark.parametrize("case", _pick("openai_o1", "openai_o3mini", "azure_o3mini", "azure_o3mini_live"), ids=_case_id) +def test_o_series_developer_role_kept(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, f"canned-{case['id']}") + _complete( + case, + { + "messages": [ + {"role": "developer", "content": "Be a good bot!"}, + {"role": "user", "content": "Hello!"}, + ] + }, + ) + body: Final = _request_body(route) + first: Final = cast(Mapping[str, JsonValue], cast(list[JsonValue], body["messages"])[0]) + assert first["role"] == "developer" + assert first["content"] == "Be a good bot!" + + +@pytest.mark.parametrize("case", _pick("openai_o1", "openai_o3mini", "azure_o3mini", "azure_o3mini_live"), ids=_case_id) +def test_o_series_temperature_dropped(case: _Case, respx_mock: MockRouter) -> None: + route: Final = _register(case, respx_mock, f"canned-{case['id']}") + _complete( + case, + { + "messages": [{"role": "user", "content": "Hello, world!"}], + "temperature": 0.0, + "drop_params": True, + }, + ) + body: Final = _request_body(route) + assert "temperature" not in body diff --git a/tests/unit/llms/base_llm/decisions/test_systemone.py b/tests/unit/llms/base_llm/decisions/test_systemone.py deleted file mode 100644 index 42a1c41b3d3..00000000000 --- a/tests/unit/llms/base_llm/decisions/test_systemone.py +++ /dev/null @@ -1,262 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping, Sequence -from typing import Final - -import pytest -from pydantic import TypeAdapter - -from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.llms.base_llm.decisions.systemone import ( - SYSTEM_ONE_RESPONSE_ADAPTER, - question_keys, - to_decisions_response, - to_system_one_request, -) -from litellm.types.openai_decisions import ( - ChoiceAnswer, - DecisionsRequest, - DecisionsRequestBody, - PredicateAnswer, - ScoreAnswer, -) - -_INPUT: Final = "The export job hangs at 99% and never finishes" -_QUESTIONS: Final[Sequence[Mapping[str, object]]] = ( - {"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}, - { - "type": "choice", - "name": "sentiment", - "instructions": "How does the customer feel?", - "choices": [{"value": "positive"}, {"value": "negative", "description": "unhappy"}], - }, - { - "type": "score", - "name": "severity", - "instructions": "How severe is it?", - "levels": [{"label": "none"}, {"label": "low"}, {"label": "high", "description": "blocks users"}], - }, -) -_SYSTEM_ONE_QUESTIONS: Final[Mapping[str, object]] = { - "is_defect": {"type": "noul", "instructions": "Is this a defect?"}, - "sentiment": { - "type": "choice", - "instructions": "How does the customer feel?", - "criteria": {"positive": None, "negative": "unhappy"}, - }, - "severity": {"type": "score", "instructions": "How severe is it?", "criteria": ["none", "low", "blocks users"]}, -} -_SYSTEM_ONE_RESPONSE: Final[Mapping[str, object]] = { - "model": "jev-1.13", - "answers": { - "is_defect": {"type": "noul", "noul": 0.9}, - "sentiment": { - "type": "choice", - "choice": "positive", - "confidence": 0.8, - "probabilities": {"positive": 0.8, "negative": 0.2}, - }, - "severity": { - "type": "score", - "score": 1, - "confidence": 0.7, - "legend": {"0": "none", "1": "low", "2": "high"}, - "probabilities": {"0": 0.1, "1": 0.8, "2": 0.1}, - }, - }, - "usage": {"input_tokens": 367, "output_tokens": 3}, -} -_EXPECTED_ANSWERS: Final[Sequence[Mapping[str, object]]] = ( - {"type": "predicate", "name": "is_defect", "probability": 0.9}, - { - "type": "choice", - "name": "sentiment", - "choice": "positive", - "probabilities": [{"value": "positive", "probability": 0.8}, {"value": "negative", "probability": 0.2}], - "confidence": 0.8, - }, - { - "type": "score", - "name": "severity", - "score": 1.0, - "probabilities": [ - {"value": 0, "label": "none", "probability": 0.1}, - {"value": 1, "label": "low", "probability": 0.8}, - {"value": 2, "label": "high", "probability": 0.1}, - ], - "confidence": 0.7, - }, -) -_EXPECTED_USAGE: Final[Mapping[str, object]] = { - "input_tokens": 367, - "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, - "output_tokens": 3, - "output_tokens_details": {"reasoning_tokens": 0}, - "total_tokens": 370, -} -_BODY_ADAPTER: Final[TypeAdapter[DecisionsRequestBody]] = TypeAdapter(DecisionsRequestBody) - - -def _body( - input_value: object = _INPUT, - questions: Sequence[Mapping[str, object]] = _QUESTIONS, -) -> DecisionsRequestBody: - return _BODY_ADAPTER.validate_python({"input": input_value, "questions": questions}) - - -def _request( - input_value: object = _INPUT, - questions: Sequence[Mapping[str, object]] = _QUESTIONS, - model: str = "jev-1.13", -) -> DecisionsRequest: - return DecisionsRequest(model=model, body=_body(input_value, questions)) - - -def _predicate(name: str | None = "is_defect") -> tuple[Mapping[str, object]]: - return ({"type": "predicate", "name": name, "instructions": "Is this a defect?"},) - - -def test_openai_request_becomes_the_system_one_body() -> None: - assert to_system_one_request("jev-1.13", _body(), "typesafe") == { - "model": "jev-1.13", - "state": _INPUT, - "questions": _SYSTEM_ONE_QUESTIONS, - } - - -def test_system_one_answers_become_openai_answers_in_question_order() -> None: - response: Final = to_decisions_response( - SYSTEM_ONE_RESPONSE_ADAPTER.validate_python(_SYSTEM_ONE_RESPONSE), _request(), "typesafe" - ) - - assert response.model_dump(mode="json") == { - "model": "jev-1.13", - "answers": list(_EXPECTED_ANSWERS), - "usage": _EXPECTED_USAGE, - } - assert isinstance(response.answers[0], PredicateAnswer) - assert isinstance(response.answers[1], ChoiceAnswer) - assert isinstance(response.answers[2], ScoreAnswer) - - -def test_user_messages_are_joined_into_one_system_one_state() -> None: - messages: Final = [ - {"role": "user", "content": "first"}, - { - "role": "user", - "content": [{"type": "input_text", "text": "second"}, {"type": "input_text", "text": "third"}], - }, - ] - - body: Final = to_system_one_request("jev-1.13", _body(messages, _predicate()), "typesafe") - - assert body["state"] == "first\nsecond\nthird" - - -@pytest.mark.parametrize( - ("label", "input_value", "questions"), - ( - ( - "input_image", - [{"role": "user", "content": [{"type": "input_image", "image_url": "data:image/png;base64,AA=="}]}], - _predicate(), - ), - ("boolean choice", _INPUT, [{"type": "choice", "instructions": "Refund?", "choices": [{"value": True}]}]), - ("unique name", _INPUT, [*_predicate(), *_predicate()]), - ( - "repeated choice", - _INPUT, - [{"type": "choice", "instructions": "Refund?", "choices": [{"value": "yes"}, {"value": "yes"}]}], - ), - ), -) -def test_what_system_one_cannot_express_is_a_400( - label: str, - input_value: object, - questions: Sequence[Mapping[str, object]], -) -> None: - with pytest.raises(BaseLLMException, match=label) as error: - to_system_one_request("jev-1.13", _body(input_value, questions), "perplexity") - - assert error.value.status_code == 400 - assert "perplexity" in error.value.message - - -def test_unnamed_questions_get_positional_keys_that_never_shadow_a_supplied_name() -> None: - body: Final = _body(questions=[*_predicate(None), *_predicate("q0"), *_predicate(None)]) - - assert question_keys(body.questions, "typesafe") == ("_q0", "q0", "q2") - noul: Final = {"type": "noul", "instructions": "Is this a defect?"} - assert to_system_one_request("jev-1.13", body, "typesafe")["questions"] == {"_q0": noul, "q0": noul, "q2": noul} - - -def test_positional_answers_come_back_in_question_order_without_a_name() -> None: - request: Final = _request(questions=[*_predicate(None), *_predicate("q0")]) - system_one: Final = SYSTEM_ONE_RESPONSE_ADAPTER.validate_python( - {"answers": {"_q0": {"type": "noul", "noul": 0.25}, "q0": {"type": "noul", "noul": 0.75}}} - ) - - response: Final = to_decisions_response(system_one, request, "typesafe") - - assert [answer.model_dump(mode="json") for answer in response.answers] == [ - {"type": "predicate", "name": None, "probability": 0.25}, - {"type": "predicate", "name": "q0", "probability": 0.75}, - ] - assert response.usage.model_dump(mode="json") == { - **_EXPECTED_USAGE, - "input_tokens": 0, - "output_tokens": 0, - "total_tokens": 0, - } - - -def test_a_reply_without_a_model_reports_the_requested_model() -> None: - system_one: Final = SYSTEM_ONE_RESPONSE_ADAPTER.validate_python( - {k: v for k, v in _SYSTEM_ONE_RESPONSE.items() if k != "model"} - ) - - response: Final = to_decisions_response(system_one, _request(model="typesafe/jev-1.13.0"), "typesafe") - - assert response.model == "typesafe/jev-1.13.0" - - -def test_a_choice_the_provider_left_out_of_probabilities_is_reported_at_zero() -> None: - system_one: Final = SYSTEM_ONE_RESPONSE_ADAPTER.validate_python( - { - "answers": { - "sentiment": { - "type": "choice", - "choice": "positive", - "confidence": 1.0, - "probabilities": {"positive": 1.0}, - } - } - } - ) - request: Final = _request(questions=_QUESTIONS[1:2]) - - response: Final = to_decisions_response(system_one, request, "typesafe") - - assert response.answers[0].model_dump(mode="json") == { - "type": "choice", - "name": "sentiment", - "choice": "positive", - "probabilities": [{"value": "positive", "probability": 1.0}, {"value": "negative", "probability": 0.0}], - "confidence": 1.0, - } - - -@pytest.mark.parametrize( - "answers", - ( - {}, - {"is_defect": {"type": "choice", "choice": "yes", "confidence": 1.0, "probabilities": {"yes": 1.0}}}, - ), -) -def test_a_reply_without_a_matching_answer_is_a_server_error(answers: Mapping[str, object]) -> None: - system_one: Final = SYSTEM_ONE_RESPONSE_ADAPTER.validate_python({"answers": answers}) - - with pytest.raises(BaseLLMException, match="no predicate answer for question 'is_defect'") as error: - to_decisions_response(system_one, _request(questions=_predicate()), "typesafe") - - assert error.value.status_code == 500 diff --git a/tests/unit/llms/base_llm/embedding/__init__.py b/tests/unit/llms/base_llm/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/base_llm/embedding/test_provider_embedding_translation.py b/tests/unit/llms/base_llm/embedding/test_provider_embedding_translation.py new file mode 100644 index 00000000000..24157263daa --- /dev/null +++ b/tests/unit/llms/base_llm/embedding/test_provider_embedding_translation.py @@ -0,0 +1,108 @@ +import base64 +import json +from typing import Final, cast + +import httpx +import pytest +from respx import MockRouter +from typing_extensions import ReadOnly, TypedDict + +from litellm import embedding +from litellm.utils import get_optional_params_embeddings + + +class _Kwargs(TypedDict, total=False): + model: ReadOnly[str] + api_key: ReadOnly[str] + api_base: ReadOnly[str] + api_version: ReadOnly[str] + aws_access_key_id: ReadOnly[str] + aws_secret_access_key: ReadOnly[str] + aws_region_name: ReadOnly[str] + + +class _Case(TypedDict): + id: ReadOnly[str] + provider: ReadOnly[str] + kwargs: ReadOnly[_Kwargs] + url: ReadOnly[str] + + +_AZURE_BASE: Final = "https://offline-embed.openai.azure.com" +_AZURE_URL: Final = f"{_AZURE_BASE}/openai/deployments/text-embedding-ada-002/embeddings?api-version=2024-02-15-preview" +_TITAN_URL: Final = "https://bedrock-runtime.us-west-2.amazonaws.com/model/amazon.titan-embed-image-v1/invoke" + +_CASES: Final[tuple[_Case, ...]] = ( + { + "id": "azure_text_embedding", + "provider": "azure", + "kwargs": { + "model": "azure/text-embedding-ada-002", + "api_key": "azure-offline-key", + "api_base": _AZURE_BASE, + "api_version": "2024-02-15-preview", + }, + "url": _AZURE_URL, + }, + { + "id": "bedrock_titan_image", + "provider": "bedrock", + "kwargs": cast( + _Kwargs, + { + "model": "bedrock/amazon.titan-embed-image-v1", + "aws_access_key_id": "AKIAFAKE", + "aws_secret_access_key": "fakesecret", + "aws_region_name": "us-west-2", + }, + ), + "url": _TITAN_URL, + }, +) + + +_MAX_RETRIES_KWARGS: Final[tuple[_Kwargs, ...]] = ( + *(case["kwargs"] for case in _CASES), + {"model": "volcengine/doubao-embedding-text-240715"}, + {"model": "voyage/voyage-3-lite"}, +) + + +def _canned_response(case: _Case) -> httpx.Response: + vector: Final = [0.11, 0.22, 0.33] + if case["provider"] == "azure": + return httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": vector}], + "model": "text-embedding-ada-002", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + return httpx.Response(200, json={"embedding": vector, "inputTextTokenCount": 4}) + + +@pytest.fixture(autouse=True) +def _httpx_only_transport(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + + +@pytest.mark.parametrize("kwargs", _MAX_RETRIES_KWARGS, ids=lambda k: k["model"]) +def test_embedding_optional_params_max_retries(kwargs: _Kwargs) -> None: + optional_params: Final = get_optional_params_embeddings(**dict(kwargs), max_retries=20) + assert optional_params["max_retries"] == 20 + + +@pytest.mark.parametrize("case", _CASES, ids=lambda c: c["id"]) +def test_image_embedding(case: _Case, respx_mock: MockRouter) -> None: + png_b64: Final = base64.b64encode(b"\x89PNG\r\n\x1a\nFAKEPIXELS").decode() + data_url: Final = f"data:image/png;base64,{png_b64}" + route: Final = respx_mock.post(case["url"]).mock(return_value=_canned_response(case)) + response: Final = embedding(**dict(case["kwargs"]), input=[data_url]) + body: Final = json.loads(route.calls.last.request.content) + if case["provider"] == "azure": + assert body["input"] == [data_url] + else: + assert body["inputImage"] == png_b64 + assert response.data[0]["embedding"] == [0.11, 0.22, 0.33] diff --git a/tests/unit/llms/base_llm/rerank/__init__.py b/tests/unit/llms/base_llm/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/base_llm/rerank/test_provider_rerank_translation.py b/tests/unit/llms/base_llm/rerank/test_provider_rerank_translation.py new file mode 100644 index 00000000000..8b630f6b946 --- /dev/null +++ b/tests/unit/llms/base_llm/rerank/test_provider_rerank_translation.py @@ -0,0 +1,187 @@ +import json +from typing import Final, Mapping, cast + +import httpx +import pytest +import respx +from pydantic import JsonValue +from respx import MockRouter +from typing_extensions import ReadOnly, TypedDict + +import litellm + + +class _Kwargs(TypedDict, total=False): + model: ReadOnly[str] + api_key: ReadOnly[str] + aws_access_key_id: ReadOnly[str] + aws_secret_access_key: ReadOnly[str] + aws_region_name: ReadOnly[str] + + +class _Case(TypedDict): + id: ReadOnly[str] + provider: ReadOnly[str] + kwargs: ReadOnly[_Kwargs] + url: ReadOnly[str] + expected_cost_zero: ReadOnly[bool] + billed_units: ReadOnly[Mapping[str, int]] + response_id: ReadOnly[str | None] + + +_AWS: Final[Mapping[str, str]] = { + "aws_access_key_id": "AKIAFAKE", + "aws_secret_access_key": "fakesecret", + "aws_region_name": "us-west-2", +} + + +def _bedrock_arn(model_id: str) -> str: + return f"bedrock/arn:aws:bedrock:us-west-2::foundation-model/{model_id}" + + +_CASES: Final[tuple[_Case, ...]] = ( + { + "id": "jina_reranker", + "provider": "cohere", + "kwargs": {"model": "jina_ai/jina-reranker-v2-base-multilingual", "api_key": "jina-offline"}, + "url": "https://api.jina.ai/v1/rerank", + "expected_cost_zero": False, + "billed_units": {"total_tokens": 4}, + "response_id": "rerank-offline", + }, + { + "id": "bedrock_amazon_rerank", + "provider": "bedrock", + "kwargs": cast(_Kwargs, {"model": _bedrock_arn("amazon.rerank-v1:0"), **dict(_AWS)}), + "url": "https://bedrock-agent-runtime.us-west-2.amazonaws.com/rerank", + "expected_cost_zero": False, + "billed_units": {"search_units": 1}, + "response_id": "rerank-offline", + }, + { + "id": "bedrock_cohere_rerank", + "provider": "bedrock", + "kwargs": cast(_Kwargs, {"model": _bedrock_arn("cohere.rerank-v3-5:0"), **dict(_AWS)}), + "url": "https://bedrock-agent-runtime.us-west-2.amazonaws.com/rerank", + "expected_cost_zero": False, + "billed_units": {"search_units": 1}, + "response_id": "rerank-offline", + }, + { + "id": "nvidia_nim_rerank", + "provider": "nvidia_nim", + "kwargs": {"model": "nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2", "api_key": "nvapi-offline"}, + "url": "https://ai.api.nvidia.com/v1/retrieval/nvidia/llama-3_2-nv-rerankqa-1b-v2/reranking", + "expected_cost_zero": True, + "billed_units": {"total_tokens": 4}, + "response_id": None, + }, +) + + +def _canned_response(case: _Case) -> httpx.Response: + if case["provider"] == "cohere": + return httpx.Response( + 200, + json={ + "id": "rerank-offline", + "results": [ + {"index": 0, "relevance_score": 0.95}, + {"index": 1, "relevance_score": 0.4}, + ], + "usage": {"total_tokens": 4}, + }, + ) + if case["provider"] == "bedrock": + return httpx.Response( + 200, + json={ + "id": "rerank-offline", + "results": [ + {"index": 0, "relevanceScore": 0.95}, + {"index": 1, "relevanceScore": 0.4}, + ], + "usage": {"search_units": 1}, + }, + ) + return httpx.Response( + 200, + json={ + "rankings": [ + {"index": 0, "logit": 0.95}, + {"index": 1, "logit": 0.4}, + ], + "usage": {"prompt_tokens": 4, "total_tokens": 4}, + }, + ) + + +@pytest.fixture(autouse=True) +def _httpx_only_transport(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + + +def _request_body(route: respx.Route) -> Mapping[str, JsonValue]: + return json.loads(route.calls.last.request.content) + + +def _assert_translated_request(case: _Case, body: Mapping[str, JsonValue]) -> None: + if case["provider"] == "bedrock": + queries: Final = body["queries"] + assert queries == [{"textQuery": {"text": "hello"}, "type": "TEXT"}] + config: Final = body["rerankingConfiguration"]["bedrockRerankingConfiguration"]["modelConfiguration"] + assert config["modelArn"].endswith((".rerank-v1:0", ".rerank-v3-5:0")) + sources: Final = body["sources"] + assert len(sources) == 2 + elif case["provider"] == "nvidia_nim": + assert body["model"] == "nvidia/llama-3.2-nv-rerankqa-1b-v2" + assert body["query"] == {"text": "hello"} + assert body["passages"] == [{"text": "hello"}, {"text": "world"}] + assert body["top_k"] == 2 + else: + assert body["model"] == "jina-reranker-v2-base-multilingual" + assert body["query"] == "hello" + assert body["documents"] == ["hello", "world"] + assert body["top_n"] == 2 + + +@pytest.mark.parametrize("case", _CASES, ids=lambda c: c["id"]) +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_basic_rerank(case: _Case, sync_mode: bool, respx_mock: MockRouter) -> None: + route: Final = respx_mock.post(case["url"]).mock(return_value=_canned_response(case)) + if sync_mode: + response: Final = litellm.rerank( + **dict(case["kwargs"]), + query="hello", + documents=["hello", "world"], + top_n=2, + ) + else: + response: Final = await litellm.arerank( + **dict(case["kwargs"]), + query="hello", + documents=["hello", "world"], + top_n=2, + ) + body: Final = _request_body(route) + _assert_translated_request(case, body) + assert route.call_count == 1 + assert isinstance(response.id, str) + if case["response_id"] is not None: + assert response.id == case["response_id"] + assert response.meta["billed_units"] == case["billed_units"] + assert len(response.results) == 2 + assert response.results[0]["index"] == 0 + assert response.results[0]["relevance_score"] == 0.95 + assert response.results[1]["index"] == 1 + assert response.results[1]["relevance_score"] == 0.4 + if case["provider"] == "nvidia_nim": + assert response.results[0]["document"] == {"text": "hello"} + assert response.results[1]["document"] == {"text": "world"} + cost: Final = response._hidden_params["response_cost"] + if case["expected_cost_zero"]: + assert cost == 0.0 + else: + assert cost > 0 diff --git a/tests/unit/llms/bedrock/chat/chat_completions/test_bedrock_chat_completions_transformation.py b/tests/unit/llms/bedrock/chat/chat_completions/test_bedrock_chat_completions_transformation.py index 16ae1114402..9d4c0fac094 100644 --- a/tests/unit/llms/bedrock/chat/chat_completions/test_bedrock_chat_completions_transformation.py +++ b/tests/unit/llms/bedrock/chat/chat_completions/test_bedrock_chat_completions_transformation.py @@ -911,7 +911,7 @@ def _stream_chunk(delta, finish_reason=None, index=0): } -def test_streaming_handler_splits_reasoning_deltas_per_choice(): +def test_streaming_handler_splits_reasoning_deltas_per_choice(local_cost_map): handler = BedrockRuntimeChatCompletionsStreamingHandler(streaming_response=iter(()), sync_stream=True) first = handler.chunk_parser(_stream_chunk({"role": "assistant", "content": "I think"})) @@ -934,7 +934,7 @@ def _reasoning_of(parsed): return getattr(parsed.choices[0].delta, "reasoning_content", None) -def test_streaming_handler_keeps_split_state_per_choice_index(): +def test_streaming_handler_keeps_split_state_per_choice_index(local_cost_map): handler = BedrockRuntimeChatCompletionsStreamingHandler(streaming_response=iter(()), sync_stream=True) opened = handler.chunk_parser(_stream_chunk({"content": "first"}, index=0)) @@ -949,7 +949,7 @@ def test_streaming_handler_keeps_split_state_per_choice_index(): assert not still_reasoning.choices[0].delta.content -def test_streaming_handler_flushes_held_text_on_an_empty_final_delta(): +def test_streaming_handler_flushes_held_text_on_an_empty_final_delta(local_cost_map): handler = BedrockRuntimeChatCompletionsStreamingHandler(streaming_response=iter(()), sync_stream=True) held = handler.chunk_parser(_stream_chunk({"content": "taggedHi"}, finish_reason="stop") @@ -1467,3 +1467,76 @@ def test_gpt56_json_object_with_response_schema_goes_to_converse_as_a_json_tool( assert body["toolConfig"]["tools"][0]["toolSpec"]["name"] == "json_tool_call" assert body["toolConfig"]["toolChoice"] == {"tool": {"name": "json_tool_call"}} assert "response_format" not in body + + +LITERAL_TAGGED_ANSWER = "not thinking Hello" + + +@pytest.mark.parametrize("model", ["openai.gpt-5.6-sol", "us.xai.grok-4.6"]) +def test_streaming_handler_keeps_a_literal_reasoning_tag_outside_gpt_oss(local_cost_map, model): + handler = BedrockRuntimeChatCompletionsStreamingHandler(streaming_response=iter(()), sync_stream=True) + + opened = handler.chunk_parser({**_stream_chunk({"content": "not thinking"}), "model": model}) + assert opened.choices[0].delta.content == "not thinking" + assert _reasoning_of(opened) is None + + closed = handler.chunk_parser({**_stream_chunk({"content": " Hello"}, finish_reason="stop"), "model": model}) + assert closed.choices[0].delta.content == " Hello" + assert _reasoning_of(closed) is None + + +@pytest.mark.parametrize( + "model, expected_content, expected_reasoning", + [ + ("bedrock/global.openai.gpt-5.6-sol", LITERAL_TAGGED_ANSWER, None), + ("bedrock/chat_completions/us.xai.grok-4.6", LITERAL_TAGGED_ANSWER, None), + ("bedrock/chat_completions/openai.gpt-oss-20b-1:0", "Hello", "not thinking"), + ("bedrock/chat_completions/openai.gpt-oss-safeguard-20b", "Hello", "not thinking"), + ], +) +def test_reasoning_tag_split_applies_to_gpt_oss_answers_only( + local_cost_map, fake_aws_env, model, expected_content, expected_reasoning +): + model_id = model.removeprefix("bedrock/").removeprefix("chat_completions/") + requests, client = _recording_client(json=_chat_completion_json(LITERAL_TAGGED_ANSWER, model_id)) + + response = litellm.completion( + model=model, messages=[{"role": "user", "content": "hello"}], client=client, max_tokens=64 + ) + + assert str(requests[0].url).endswith("/openai/v1/chat/completions") + assert response.choices[0].message.content == expected_content + assert getattr(response.choices[0].message, "reasoning_content", None) == expected_reasoning + + +@pytest.mark.parametrize( + "capability_flags, expected_content, expected_reasoning", + [ + ({"supports_bedrock_runtime_chat_completions_inline_reasoning": True}, "Hello", "not thinking"), + ({}, LITERAL_TAGGED_ANSWER, None), + ], +) +def test_reasoning_tag_split_is_read_from_the_cost_map( + monkeypatch, fake_aws_env, capability_flags, expected_content, expected_reasoning +): + model_id = SYNTHETIC_NATIVE_MODEL.removeprefix("chat_completions/") + monkeypatch.setattr(litellm, "model_cost", {model_id: {"litellm_provider": "bedrock_converse", **capability_flags}}) + requests, client = _recording_client(json=_chat_completion_json(LITERAL_TAGGED_ANSWER, model_id)) + + response = litellm.completion( + model=f"bedrock/{SYNTHETIC_NATIVE_MODEL}", + messages=[{"role": "user", "content": "hello"}], + client=client, + max_tokens=64, + ) + + assert str(requests[0].url).endswith("/openai/v1/chat/completions") + assert response.choices[0].message.content == expected_content + assert getattr(response.choices[0].message, "reasoning_content", None) == expected_reasoning + + handler = BedrockRuntimeChatCompletionsStreamingHandler(streaming_response=iter(()), sync_stream=True) + chunk = handler.chunk_parser( + {**_stream_chunk({"content": LITERAL_TAGGED_ANSWER}, finish_reason="stop"), "model": model_id} + ) + assert chunk.choices[0].delta.content == expected_content + assert _reasoning_of(chunk) == expected_reasoning diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index f43f15b8230..e42ee95f8d5 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -742,7 +742,7 @@ def _tool_schema_properties(model, tool, litellm_params=None): "us.xai.grok-4.6", "us-gov.xai.grok-4.6", "global.xai.grok-4.7", - "xai.grok-4.7", + "us.xai.grok-4.7", ], ) def test_transform_request_drops_lookaround_regex_for_models_the_cost_map_flags(tool, model): diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py new file mode 100644 index 00000000000..71b451f834c --- /dev/null +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py @@ -0,0 +1,109 @@ +import json +from collections.abc import Mapping +from typing import Final + +import httpx +import pytest + +from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + +MANTLE_COUNT_URL: Final = "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages/count_tokens" +LITELLM_PARAMS: Final = { + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", +} +REQUEST: Final[dict[str, object]] = { + "model": "global.anthropic.claude-opus-4-8", + "messages": [{"role": "user", "content": "The quick brown fox"}], + "system": "You are a terse assistant.", + "tools": [{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}], +} + + +class _MantleEndpoint: + def __init__(self, status_code: int, body: Mapping[str, object]) -> None: + self.status_code: Final = status_code + self.body: Final = body + self.requests: tuple[httpx.Request, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + return httpx.Response(self.status_code, json=dict(self.body), request=request) + + def only_request(self) -> httpx.Request: + assert len(self.requests) == 1, self.requests + return self.requests[0] + + +def _client(endpoint: _MantleEndpoint) -> AsyncHTTPHandler: + return AsyncHTTPHandler(transport=httpx.MockTransport(endpoint)) + + +@pytest.fixture(autouse=True) +def _sigv4_only(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + + +@pytest.mark.asyncio +async def test_counts_on_mantle_with_the_base_model_and_the_deployment_credentials() -> None: + mantle: Final = _MantleEndpoint(200, {"input_tokens": 2177}) + + result: Final = await BedrockMantleCountTokensHandler().handle_count_tokens_request( + request_data=dict(REQUEST), + litellm_params=dict(LITELLM_PARAMS), + resolved_model="anthropic.claude-opus-4-8", + client=_client(mantle), + ) + + assert result == {"input_tokens": 2177} + posted: Final = mantle.only_request() + assert str(posted.url) == MANTLE_COUNT_URL + assert posted.headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "us-east-1/bedrock/aws4_request" in posted.headers["Authorization"] + assert posted.headers["anthropic-version"] == "2023-06-01" + assert json.loads(posted.content) == { + "model": "anthropic.claude-opus-4-8", + "messages": REQUEST["messages"], + "system": REQUEST["system"], + "tools": REQUEST["tools"], + } + + +@pytest.mark.asyncio +async def test_bedrock_mantle_api_base_env_names_the_host(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", "https://vpce-abc.bedrock-mantle.us-east-1.vpce.api.aws") + mantle: Final = _MantleEndpoint(200, {"input_tokens": 3}) + + await BedrockMantleCountTokensHandler().handle_count_tokens_request( + request_data=dict(REQUEST), + litellm_params=dict(LITELLM_PARAMS), + resolved_model="anthropic.claude-opus-4-8", + client=_client(mantle), + ) + + assert ( + str(mantle.only_request().url) + == "https://vpce-abc.bedrock-mantle.us-east-1.vpce.api.aws/anthropic/v1/messages/count_tokens" + ) + + +@pytest.mark.asyncio +async def test_non_200_answers_raise_bedrock_error_with_mantle_status() -> None: + mantle: Final = _MantleEndpoint( + 404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}} + ) + + with pytest.raises(BedrockError) as raised: + await BedrockMantleCountTokensHandler().handle_count_tokens_request( + request_data=dict(REQUEST), + litellm_params=dict(LITELLM_PARAMS), + resolved_model="anthropic.claude-sonnet-5-5", + client=_client(mantle), + ) + + assert raised.value.status_code == 404 + assert "does not exist" in raised.value.message diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py new file mode 100644 index 00000000000..a3dd8021f73 --- /dev/null +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -0,0 +1,140 @@ +from collections.abc import Mapping +from typing import Final + +import httpx +import pytest +from pydantic import JsonValue, TypeAdapter + +from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + +RUNTIME_HOST: Final = "bedrock-runtime.us-east-1.amazonaws.com" +MANTLE_HOST: Final = "bedrock-mantle.us-east-1.api.aws" +UNSUPPORTED: Final = {"message": "The provided model doesn't support counting tokens."} +DEPLOYMENT: Final = { + "litellm_params": { + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", + } +} +MESSAGES: Final = [{"role": "user", "content": "The quick brown fox jumps over the lazy dog."}] +_JSON_BODY: Final = TypeAdapter(dict[str, JsonValue]) + + +class _Bedrock: + def __init__(self, runtime: tuple[int, Mapping[str, object]], mantle: tuple[int, Mapping[str, object]]) -> None: + self.runtime: Final = runtime + self.mantle: Final = mantle + self.requests: tuple[httpx.Request, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + status_code, body = self.runtime if request.url.host == RUNTIME_HOST else self.mantle + return httpx.Response(status_code, json=dict(body), request=request) + + def posted_hosts(self) -> tuple[str, ...]: + return tuple(request.url.host for request in self.requests) + + +def _counter(bedrock: _Bedrock) -> BedrockTokenCounter: + return BedrockTokenCounter(client=AsyncHTTPHandler(transport=httpx.MockTransport(bedrock))) + + +@pytest.fixture(autouse=True) +def _sigv4_only(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + + +@pytest.mark.asyncio +async def test_claude_model_bedrock_runtime_cannot_count_is_counted_on_mantle() -> None: + bedrock: Final = _Bedrock(runtime=(400, UNSUPPORTED), mantle=(200, {"input_tokens": 2177})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + system="You are a terse assistant.", + ) + + assert result is not None + assert result.error is False + assert result.total_tokens == 2177 + assert result.tokenizer_type == "bedrock_mantle_api" + assert result.original_response == {"input_tokens": 2177} + assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) + mantle_body: Final = _JSON_BODY.validate_json(bedrock.requests[1].content) + assert mantle_body["model"] == "anthropic.claude-opus-4-8" + assert mantle_body["system"] == "You are a terse assistant." + + +@pytest.mark.asyncio +async def test_bedrock_runtime_count_is_kept_when_it_answers() -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 1353}), mantle=(200, {"input_tokens": 1336})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-sonnet-4-6", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-sonnet-4-6", + ) + + assert result is not None + assert result.total_tokens == 1353 + assert result.tokenizer_type == "bedrock_api" + assert bedrock.posted_hosts() == (RUNTIME_HOST,) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_to_use", "runtime"), + ( + ("global.anthropic.claude-opus-4-8", (403, {"Message": "not authorized to perform: bedrock:CountTokens"})), + ("amazon.nova-pro-v1:0", (400, UNSUPPORTED)), + ), +) +async def test_other_bedrock_runtime_failures_are_not_retried_on_mantle( + model_to_use: str, runtime: tuple[int, Mapping[str, object]] +) -> None: + bedrock: Final = _Bedrock(runtime=runtime, mantle=(200, {"input_tokens": 2177})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use=model_to_use, + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model=model_to_use, + ) + + assert result is not None + assert result.error is True + assert result.status_code == runtime[0] + assert result.tokenizer_type == "bedrock_api" + assert bedrock.posted_hosts() == (RUNTIME_HOST,) + + +@pytest.mark.asyncio +async def test_mantle_failure_is_reported_with_its_status() -> None: + bedrock: Final = _Bedrock( + runtime=(400, UNSUPPORTED), + mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}), + ) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-sonnet-5-5", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-sonnet-5-5", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 404 + assert result.tokenizer_type == "bedrock_mantle_api" + assert "does not exist" in (result.error_message or "") + assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) diff --git a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index fbcbda183bc..875a2fc1b1a 100644 --- a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -10,9 +10,14 @@ from datetime import datetime from typing import Final from unittest.mock import patch +import pytest + from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector -from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig +from litellm.llms.bedrock.passthrough.transformation import ( + BedrockPassthroughConfig, + is_bedrock_streaming_endpoint, +) from litellm.types.utils import ModelResponse CONVERSE_MODEL = "anthropic.claude-sonnet-4-5-20250929-v1:0" @@ -20,6 +25,22 @@ CONVERSE_STREAM_ENDPOINT = f"/model/{CONVERSE_MODEL}/converse-stream" INVOKE_STREAM_ENDPOINT = f"/model/{CONVERSE_MODEL}/invoke-with-response-stream" +@pytest.mark.parametrize( + ("endpoint", "expected"), + [ + ("converse-stream", True), + ("invoke-with-response-stream", True), + ("converse", False), + ("invoke", False), + ("model/my-converse-stream-model/converse", False), + ("model/x/converse-stream?foo=1", True), + ("model/x/converse-stream/", True), + ], +) +def test_is_bedrock_streaming_endpoint_matches_final_action_segment(endpoint: str, expected: bool) -> None: + assert is_bedrock_streaming_endpoint(endpoint) is expected + + def test_bedrock_passthrough_get_complete_url_default_endpoint(): """Test get_complete_url with default AWS endpoint (no override)""" config = BedrockPassthroughConfig() diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py new file mode 100644 index 00000000000..82ad7702052 --- /dev/null +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -0,0 +1,52 @@ +import json +from typing import Final + +import pytest + +from litellm.constants import SERVER_STREAMING_CLASSIFICATION_KEY, SERVER_STREAMING_CLASSIFICATION_MARKER +from litellm.llms.pass_through.guardrail_translation.handler import PassThroughEndpointHandler + + +def test_full_payload_guardrail_text_excludes_the_server_streaming_marker(): + body: Final = {"model": "m", "stream": True, "messages": [{"role": "user", "content": "hi"}]} + + text: Final = PassThroughEndpointHandler()._extract_text_for_guardrail( + {**body, SERVER_STREAMING_CLASSIFICATION_KEY: SERVER_STREAMING_CLASSIFICATION_MARKER}, + None, + ) + + assert json.loads(text) == body, text + + +@pytest.mark.parametrize( + "marker", + [ + SERVER_STREAMING_CLASSIFICATION_MARKER, + json.loads(json.dumps(SERVER_STREAMING_CLASSIFICATION_MARKER)), + ], + ids=["enum", "json-string"], +) +def test_full_payload_guardrail_text_scans_caller_value_but_not_the_marker(marker: str): + text: Final = PassThroughEndpointHandler()._extract_text_for_guardrail( + { + "model": "m", + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + SERVER_STREAMING_CLASSIFICATION_KEY: "BLOCKME caller content", + }, + None, + ) + + assert "BLOCKME caller content" in text, text + + marker_text: Final = PassThroughEndpointHandler()._extract_text_for_guardrail( + { + "model": "m", + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + SERVER_STREAMING_CLASSIFICATION_KEY: marker, + }, + None, + ) + + assert SERVER_STREAMING_CLASSIFICATION_KEY not in json.loads(marker_text), marker_text diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 6f18a391f7b..dde3df6fcc0 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -1289,6 +1289,33 @@ class TestVertexEmbeddingsBatchInputTranslation: assert set(row["request"]) == {"content"} + def test_should_translate_a_file_block_with_video_metadata(self) -> None: + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-2", + "input": [ + { + "type": "file", + "file": { + "file_id": "gs://my-bucket/clip.mp4", + "video_metadata": {"start_offset": "3s", "end_offset": "6s"}, + }, + } + ], + } + ) + ] + ) + + assert row["request"]["content"]["parts"] == [ + { + "file_data": {"mime_type": "video/mp4", "file_uri": "gs://my-bucket/clip.mp4"}, + "video_metadata": {"startOffset": "3s", "endOffset": "6s"}, + } + ] + def test_should_translate_multimodal_gcs_input(self): (row,) = _wrap_entries( [ diff --git a/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_handler.py b/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_handler.py new file mode 100644 index 00000000000..483b350fd43 --- /dev/null +++ b/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_handler.py @@ -0,0 +1,39 @@ +from typing import Final + +import httpx +import pytest + +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler import GoogleBatchEmbeddings + +FILES_URI: Final = "https://generativelanguage.googleapis.com/v1beta/files/clip123" +FILE_METADATA_URL: Final = "https://generativelanguage.googleapis.com/v1beta/files/clip123" + + +def _files_api_metadata(request: httpx.Request) -> httpx.Response: + if str(request.url) != FILE_METADATA_URL or request.headers.get("x-goog-api-key") != "gemini-key": + return httpx.Response(404, json={"error": {"message": "no such file"}}) + return httpx.Response(200, json={"mimeType": "video/mp4", "uri": FILES_URI}) + + +@pytest.mark.parametrize("reference", ["files/clip123", FILES_URI]) +def test_resolve_file_references_fetches_the_file_metadata_for_both_reference_forms(reference: str) -> None: + resolved: Final = GoogleBatchEmbeddings()._resolve_file_references( + input=[reference, "a red bus"], + api_key="gemini-key", + sync_handler=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_files_api_metadata))), + ) + assert resolved == {reference: {"mime_type": "video/mp4", "uri": FILES_URI}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reference", ["files/clip123", FILES_URI]) +async def test_async_resolve_file_references_fetches_the_file_metadata_for_both_reference_forms( + reference: str, +) -> None: + resolved: Final = await GoogleBatchEmbeddings()._async_resolve_file_references( + input=[reference, "a red bus"], + api_key="gemini-key", + async_handler=AsyncHTTPHandler(transport=httpx.MockTransport(_files_api_metadata)), + ) + assert resolved == {reference: {"mime_type": "video/mp4", "uri": FILES_URI}} diff --git a/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index ba2b26bf0a2..6584c27b895 100644 --- a/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -8,12 +8,17 @@ Covers: - Response processing with correct indices """ +from collections.abc import Callable +from typing import Final + import pytest import litellm +from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( _build_part_for_input, + file_reference_name, _is_multimodal_input, process_embed_content_response, process_response, @@ -25,6 +30,12 @@ from litellm.types.utils import EmbeddingResponse IMAGE_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII" GCS_URL = "gs://my-bucket/image.png" +VIDEO_DATA_URI = "data:video/mp4;base64,AAAAIGZ0eXBpc29tAAACAGlzb21pc28yYXZjMW1wNDEAAAAIZnJlZQAA" +FILES_URI = "https://generativelanguage.googleapis.com/v1beta/files/clip123" + + +def _file_block(**file: object) -> dict[str, object]: + return {"type": "file", "file": file} @pytest.fixture(autouse=True) @@ -51,6 +62,7 @@ class TestIsMultimodalInput: def test_file_reference(self): assert _is_multimodal_input(["files/abc123"]) is True + assert _is_multimodal_input([FILES_URI]) is True def test_mixed_text_and_image(self): assert _is_multimodal_input(["hello", IMAGE_DATA_URI]) is True @@ -59,9 +71,16 @@ class TestIsMultimodalInput: """Nested list with text is not multimodal.""" assert _is_multimodal_input([["text_a", "text_b"]]) is False + def test_file_content_block_is_multimodal(self) -> None: + assert _is_multimodal_input([_file_block(file_data=VIDEO_DATA_URI)]) is True + def test_nested_list_with_image_is_multimodal(self): assert _is_multimodal_input([["a red shoe", IMAGE_DATA_URI]]) is True + def test_single_object_input_answers_400(self) -> None: + with pytest.raises(BadRequestError, match="string or a list"): + _is_multimodal_input(_file_block(file_data=VIDEO_DATA_URI)) + class TestBuildPartForInput: def test_text_input(self): @@ -83,15 +102,13 @@ class TestBuildPartForInput: assert part["file_data"]["file_uri"] == GCS_URL def test_file_reference_resolved(self): - resolved = { - "files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"} - } + resolved = {"files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"}} part = _build_part_for_input("files/abc", resolved_files=resolved) assert part["file_data"] is not None assert part["file_data"]["mime_type"] == "image/jpeg" - def test_file_reference_unresolved_raises(self): - with pytest.raises(ValueError, match="not resolved"): + def test_file_reference_unresolved_answers_400_naming_the_gemini_provider(self) -> None: + with pytest.raises(BadRequestError, match="gemini/ provider"): _build_part_for_input("files/abc") @@ -124,10 +141,7 @@ class TestTransformOpenaiInputGeminiContent: ) assert len(result["requests"]) == 2 # First request is text - assert ( - result["requests"][0]["content"]["parts"][0]["text"] - == "The food was delicious" - ) + assert result["requests"][0]["content"]["parts"][0]["text"] == "The food was delicious" # Second request is image assert result["requests"][1]["content"]["parts"][0]["inline_data"] is not None @@ -226,9 +240,7 @@ class TestProcessResponse: """Test that process_response sets correct indices.""" def test_single_embedding_index(self): - predictions: VertexAIBatchEmbeddingsResponseObject = { - "embeddings": [{"values": [0.1, 0.2]}] - } + predictions: VertexAIBatchEmbeddingsResponseObject = {"embeddings": [{"values": [0.1, 0.2]}]} model_response = EmbeddingResponse() result = process_response( input="hello", @@ -279,9 +291,7 @@ class TestProcessResponse: def test_nested_input_token_counting(self): """Nested list: only plain-text sub-elements should be counted.""" - predictions: VertexAIBatchEmbeddingsResponseObject = { - "embeddings": [{"values": [0.1, 0.2]}] - } + predictions: VertexAIBatchEmbeddingsResponseObject = {"embeddings": [{"values": [0.1, 0.2]}]} result = process_response( input=[["a red shoe", IMAGE_DATA_URI]], model_response=EmbeddingResponse(), @@ -300,7 +310,7 @@ class TestProcessResponse: ) def test_nested_non_string_element_raises(self): - with pytest.raises(ValueError, match="must be strings"): + with pytest.raises(BadRequestError, match="must be strings or file content blocks, got list"): transform_openai_input_gemini_content( input=[[["doubly", "nested"]]], model="gemini-embedding-2-preview", @@ -408,3 +418,212 @@ class TestProcessEmbedContentResponseUsage: assert result.usage.prompt_tokens > 0 +class TestFileContentBlocks: + MODEL: Final = "gemini-embedding-2-preview" + CLIP_METADATA: Final = {"fps": 2, "start_offset": "3s", "end_offset": "6s"} + CLIP_PART: Final = {"fps": 2.0, "startOffset": "3s", "endOffset": "6s"} + + def test_data_uri_block_forwards_format_and_video_metadata(self) -> None: + part: Final = _build_part_for_input( + _file_block( + file_data="data:application/octet-stream;base64,QUJD", + format="video/mp4", + video_metadata=self.CLIP_METADATA, + ) + ) + assert part == { + "inline_data": {"mime_type": "video/mp4", "data": "QUJD"}, + "video_metadata": self.CLIP_PART, + } + + def test_gcs_block_infers_mime_type_and_forwards_video_metadata(self) -> None: + part: Final = _build_part_for_input( + _file_block(file_id="gs://my-bucket/clip.mp4", video_metadata={"start_offset": "0s", "end_offset": "3s"}) + ) + assert part == { + "file_data": {"mime_type": "video/mp4", "file_uri": "gs://my-bucket/clip.mp4"}, + "video_metadata": {"startOffset": "0s", "endOffset": "3s"}, + } + + def test_file_reference_block_uses_the_resolved_file(self) -> None: + resolved_files: Final = {"files/clip123": {"mime_type": "video/mp4", "uri": FILES_URI}} + part: Final = _build_part_for_input( + _file_block(file_id="files/clip123", video_metadata={"fps": 1}), + resolved_files=resolved_files, + ) + assert part == { + "file_data": {"mime_type": "video/mp4", "file_uri": FILES_URI}, + "video_metadata": {"fps": 1.0}, + } + + def test_files_api_uri_block_uses_the_resolved_file(self) -> None: + resolved_files: Final = {FILES_URI: {"mime_type": "video/mp4", "uri": FILES_URI}} + part: Final = _build_part_for_input( + _file_block(file_id=FILES_URI, video_metadata={"start_offset": "0s", "end_offset": "3s"}), + resolved_files=resolved_files, + ) + assert part == { + "file_data": {"mime_type": "video/mp4", "file_uri": FILES_URI}, + "video_metadata": {"startOffset": "0s", "endOffset": "3s"}, + } + + @pytest.mark.parametrize("reference", ["files/clip123", FILES_URI]) + def test_file_reference_name_is_the_files_name_for_both_forms(self, reference: str) -> None: + assert file_reference_name(reference) == "files/clip123" + + def test_block_without_video_metadata_sends_no_video_metadata_key(self) -> None: + part: Final = _build_part_for_input(_file_block(file_data=IMAGE_DATA_URI, filename="dot.png")) + assert part == {"inline_data": {"mime_type": "image/png", "data": IMAGE_DATA_URI.split(",", 1)[1]}} + + def test_block_with_empty_video_metadata_sends_no_video_metadata_key(self) -> None: + part: Final = _build_part_for_input(_file_block(file_data=VIDEO_DATA_URI, video_metadata={})) + assert part == {"inline_data": {"mime_type": "video/mp4", "data": VIDEO_DATA_URI.split(",", 1)[1]}} + + def test_batch_path_nested_block_and_text_share_one_request(self) -> None: + result: Final = transform_openai_input_gemini_content( + input=[[_file_block(file_data=VIDEO_DATA_URI, video_metadata=self.CLIP_METADATA), "a solid color clip"]], + model=self.MODEL, + optional_params={"dimensions": 768}, + ) + [request] = result["requests"] + assert request["outputDimensionality"] == 768 + assert request["content"]["parts"] == [ + { + "inline_data": {"mime_type": "video/mp4", "data": VIDEO_DATA_URI.split(",", 1)[1]}, + "video_metadata": self.CLIP_PART, + }, + {"text": "a solid color clip"}, + ] + + def test_batch_path_flat_block_and_text_are_separate_requests(self) -> None: + result: Final = transform_openai_input_gemini_content( + input=[_file_block(file_data=VIDEO_DATA_URI, video_metadata=self.CLIP_METADATA), "a solid color clip"], + model=self.MODEL, + optional_params={}, + ) + assert len(result["requests"]) == 2 + assert result["requests"][0]["content"]["parts"][0]["video_metadata"] == self.CLIP_PART + assert result["requests"][1]["content"]["parts"] == [{"text": "a solid color clip"}] + + def test_embed_content_path_accepts_flat_block_and_text(self) -> None: + result: Final = transform_openai_input_gemini_embed_content( + input=[_file_block(file_data=VIDEO_DATA_URI, video_metadata=self.CLIP_METADATA), "a solid color clip"], + model=self.MODEL, + optional_params={}, + ) + parts: Final = result["content"]["parts"] + assert parts[0]["video_metadata"] == self.CLIP_PART + assert parts[1] == {"text": "a solid color clip"} + + def test_embed_content_path_still_rejects_nested_lists(self) -> None: + with pytest.raises(ValueError, match="Nested"): + transform_openai_input_gemini_embed_content( + input=[[_file_block(file_data=VIDEO_DATA_URI), "a solid color clip"]], + model=self.MODEL, + optional_params={}, + ) + + @pytest.mark.parametrize( + "block, named_in_error", + [ + ( + _file_block(file_data=VIDEO_DATA_URI, video_metadata={"fps": 1, "startOffset": "1s"}), + "video_metadata.startOffset", + ), + (_file_block(file_data=VIDEO_DATA_URI, video_metadata={"fps": "fast"}), "video_metadata.fps"), + (_file_block(file_data=VIDEO_DATA_URI, video_metadata={"fps": "1"}), "video_metadata.fps"), + (_file_block(file_data=VIDEO_DATA_URI, video_metadata={"fps": True}), "video_metadata.fps"), + (_file_block(file_data=VIDEO_DATA_URI, video_metadata={"start_offset": 5}), "video_metadata.start_offset"), + (_file_block(file_data=VIDEO_DATA_URI, detail="high"), "file.detail"), + (_file_block(file_data=VIDEO_DATA_URI, format=""), "file.format"), + (_file_block(file_id="gs://my-bucket/clip.mp4", file_data=VIDEO_DATA_URI), "not both"), + (_file_block(), "needs file.file_id or file.file_data"), + ( + _file_block(file_id="https://example.com/clip.mp4"), + "a data: URI, a gs:// URL, a files/ reference, or a Gemini Files API URI", + ), + ({"type": "image_url", "image_url": {"url": IMAGE_DATA_URI}}, "Input should be 'file'"), + ], + ) + def test_malformed_block_answers_400_naming_the_field( + self, block: dict[str, object], named_in_error: str + ) -> None: + with pytest.raises(BadRequestError, match=named_in_error): + _build_part_for_input(block) + + def test_drop_params_drops_the_block_keys_this_surface_does_not_take(self) -> None: + block: Final = _file_block( + file_data=VIDEO_DATA_URI, + detail="high", + video_metadata={"fps": 1, "start_offset": "1s", "resolution": "low"}, + ) + part: Final = _build_part_for_input({**block, "cache_control": {"type": "ephemeral"}}, drop_params=True) + assert part["inline_data"]["mime_type"] == "video/mp4" + assert part["video_metadata"] == {"fps": 1.0, "startOffset": "1s"} + + def test_drop_params_still_answers_400_for_a_malformed_value(self) -> None: + block: Final = _file_block(file_data=VIDEO_DATA_URI, detail="high", video_metadata={"fps": "fast"}) + with pytest.raises(BadRequestError, match=r"video_metadata\.fps"): + _build_part_for_input(block, drop_params=True) + + @pytest.mark.parametrize( + "transform", [transform_openai_input_gemini_content, transform_openai_input_gemini_embed_content] + ) + def test_transforms_forward_drop_params_to_every_block(self, transform: Callable[..., object]) -> None: + block: Final = _file_block(file_data=VIDEO_DATA_URI, detail="high") + with pytest.raises(BadRequestError, match=r"file\.detail"): + transform(input=[block], model="gemini-embedding-2-preview", optional_params={}) + transform(input=[block], model="gemini-embedding-2-preview", optional_params={}, drop_params=True) + + def test_batch_path_forwards_drop_params_into_nested_lists(self) -> None: + block: Final = _file_block(file_data=VIDEO_DATA_URI, detail="high") + body: Final = transform_openai_input_gemini_content( + input=[[block, "a caption"]], model="gemini-embedding-2-preview", optional_params={}, drop_params=True + ) + assert len(body["requests"][0]["content"]["parts"]) == 2 + + @pytest.mark.parametrize( + "transform", [transform_openai_input_gemini_content, transform_openai_input_gemini_embed_content] + ) + def test_single_object_input_answers_400(self, transform: Callable[..., object]) -> None: + with pytest.raises(BadRequestError, match="string or a list"): + transform( + input=_file_block(file_data=VIDEO_DATA_URI), model="gemini-embedding-2-preview", optional_params={} + ) + + def test_process_response_counts_only_the_text_tokens_next_to_a_block(self) -> None: + text: Final = "a solid color clip" + with_block: Final = process_response( + input=[_file_block(file_data=VIDEO_DATA_URI, video_metadata=self.CLIP_METADATA), text], + model_response=EmbeddingResponse(), + model=self.MODEL, + _predictions={"embeddings": [{"values": [0.1]}, {"values": [0.2]}]}, + ) + text_only: Final = process_response( + input=[text], + model_response=EmbeddingResponse(), + model=self.MODEL, + _predictions={"embeddings": [{"values": [0.2]}]}, + ) + assert with_block.usage.prompt_tokens == text_only.usage.prompt_tokens > 0 + + def test_embed_content_usage_fallback_with_a_block_does_not_estimate(self) -> None: + result: Final = process_embed_content_response( + input=[_file_block(file_data=VIDEO_DATA_URI, video_metadata=self.CLIP_METADATA)], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json={"embedding": {"values": [0.1, 0.2]}}, + ) + assert result.usage.prompt_tokens == 0 + + def test_image_block_counts_as_image_only_input(self) -> None: + result: Final = process_embed_content_response( + input=[_file_block(file_data=IMAGE_DATA_URI)], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json={ + "embedding": {"values": [0.1, 0.2]}, + "usageMetadata": {"promptTokenCount": 258, "totalTokenCount": 258}, + }, + ) + assert result.usage.prompt_tokens_details.image_tokens == 258 diff --git a/tests/unit/llms/vertex_ai/test_gemini_batch_embeddings.py b/tests/unit/llms/vertex_ai/test_gemini_batch_embeddings.py index 98abf5459df..7271fb48861 100644 --- a/tests/unit/llms/vertex_ai/test_gemini_batch_embeddings.py +++ b/tests/unit/llms/vertex_ai/test_gemini_batch_embeddings.py @@ -8,11 +8,15 @@ This test ensures that: """ import json +from contextlib import AbstractContextManager +from typing import Final from unittest.mock import MagicMock, patch import pytest import litellm +import httpx + from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( _filter_embed_params, @@ -655,3 +659,163 @@ def test_batch_embeddings_response_has_correct_indices_and_order(): assert ( embedding.embedding == expected_values[i] ), f"embedding {i} has wrong values: {embedding.embedding}" + + +CLIP_BLOCK = { + "type": "file", + "file": { + "file_data": "data:video/mp4;base64,AAAAIGZ0eXBpc29tAAACAGlzb21pc28yYXZjMW1wNDEAAAAIZnJlZQAA", + "video_metadata": {"fps": 2, "start_offset": "3s", "end_offset": "6s"}, + }, +} +CLIP_PART_METADATA = {"fps": 2.0, "startOffset": "3s", "endOffset": "6s"} + + +def _mock_embedding_call( + response_json: dict[str, object], +) -> tuple[AbstractContextManager[MagicMock], AbstractContextManager[MagicMock], MagicMock]: + mock_get_token: Final = patch( + "litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._get_token_and_url" + ) + mock_auth: Final = patch( + "litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._ensure_access_token", + side_effect=lambda *args, **kwargs: (None, "test-project"), + ) + mock_response: Final = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = response_json + return mock_get_token, mock_auth, mock_response + + +def test_gemini_batch_path_sends_file_block_video_metadata() -> None: + client: Final = HTTPHandler() + mock_get_token, mock_auth, mock_response = _mock_embedding_call( + {"embeddings": [{"values": [0.1, 0.2]}, {"values": [0.3, 0.4]}]} + ) + with patch.object(client, "post", return_value=mock_response) as mock_post, mock_auth, mock_get_token as token: + token.return_value = ( + {"x-goog-api-key": "test-key"}, + "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:batchEmbedContents", + ) + response: Final = litellm.embedding( + model="gemini/gemini-embedding-2-preview", + input=[CLIP_BLOCK, "a solid color clip"], + api_key="test-key", + client=client, + ) + + request_body: Final = json.loads(mock_post.call_args.kwargs["data"]) + clip_part: Final = request_body["requests"][0]["content"]["parts"][0] + assert clip_part["inline_data"]["mime_type"] == "video/mp4" + assert clip_part["video_metadata"] == CLIP_PART_METADATA + assert request_body["requests"][1]["content"]["parts"] == [{"text": "a solid color clip"}] + assert [row["index"] for row in response.data] == [0, 1] + + +def test_vertex_embed_content_path_sends_file_block_video_metadata() -> None: + client: Final = HTTPHandler() + url: Final = "https://us-central1-aiplatform.googleapis.com/v1/projects/test/locations/us-central1/publishers/google/models/gemini-embedding-2-preview:embedContent" + mock_get_token, mock_auth, mock_response = _mock_embedding_call({"embedding": {"values": [0.1, 0.2]}}) + with patch.object(client, "post", return_value=mock_response) as mock_post, mock_auth, mock_get_token as token: + token.return_value = ({"Authorization": "Bearer test-token"}, url) + response: Final = litellm.embedding( + model="vertex_ai/gemini-embedding-2-preview", + input=[CLIP_BLOCK, "a solid color clip"], + vertex_project="test-project", + vertex_location="us-central1", + client=client, + ) + + data: Final = json.loads(mock_post.call_args.kwargs["data"]) + assert data["content"]["parts"][0]["video_metadata"] == CLIP_PART_METADATA + assert data["content"]["parts"][1] == {"text": "a solid color clip"} + assert len(response.data) == 1 + + +def test_file_block_with_files_reference_is_resolved_through_the_files_api() -> None: + client: Final = HTTPHandler() + files_uri: Final = "https://generativelanguage.googleapis.com/v1beta/files/clip123" + mock_get_token, mock_auth, mock_response = _mock_embedding_call({"embeddings": [{"values": [0.1, 0.2]}]}) + file_lookup: Final = MagicMock() + file_lookup.status_code = 200 + file_lookup.json.return_value = {"mimeType": "video/mp4", "uri": files_uri} + with ( + patch.object(client, "post", return_value=mock_response) as mock_post, + patch.object(client, "get", return_value=file_lookup) as mock_get, + mock_auth, + mock_get_token as token, + ): + token.return_value = ( + {"x-goog-api-key": "test-key"}, + "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:batchEmbedContents", + ) + litellm.embedding( + model="gemini/gemini-embedding-2-preview", + input=[{"type": "file", "file": {"file_id": "files/clip123", "video_metadata": {"fps": 1}}}], + api_key="test-key", + client=client, + ) + + assert mock_get.call_args.kwargs["url"] == "https://generativelanguage.googleapis.com/v1beta/files/clip123" + clip_part: Final = json.loads(mock_post.call_args.kwargs["data"])["requests"][0]["content"]["parts"][0] + assert clip_part == { + "file_data": {"mime_type": "video/mp4", "file_uri": files_uri}, + "video_metadata": {"fps": 1.0}, + } + + +CLIP_BLOCK_WITH_DETAIL: Final = {"type": "file", "file": {**CLIP_BLOCK["file"], "detail": "high"}} + + +def _recording_client(calls: list[httpx.Request], response_json: dict[str, object]) -> HTTPHandler: + def route(request: httpx.Request) -> httpx.Response: + calls.append(request) + return httpx.Response(200, json=response_json) + + return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(route))) + + +def test_gemini_drop_params_strips_the_block_keys_embeddings_do_not_take() -> None: + calls: Final[list[httpx.Request]] = [] + client: Final = _recording_client(calls, {"embeddings": [{"values": [0.1, 0.2]}]}) + response: Final = litellm.embedding( + model="gemini/gemini-embedding-2-preview", + input=[CLIP_BLOCK_WITH_DETAIL], + api_key="test-key", + client=client, + drop_params=True, + ) + sent_part: Final = json.loads(calls[0].content)["requests"][0]["content"]["parts"][0] + assert sent_part["video_metadata"] == CLIP_PART_METADATA + assert "detail" not in calls[0].content.decode() + assert response.data[0].embedding == [0.1, 0.2] + + +def test_gemini_block_detail_answers_400_without_drop_params() -> None: + calls: Final[list[httpx.Request]] = [] + client: Final = _recording_client(calls, {"embeddings": [{"values": [0.1, 0.2]}]}) + with pytest.raises(litellm.BadRequestError, match=r"file\.detail"): + litellm.embedding( + model="gemini/gemini-embedding-2-preview", input=[CLIP_BLOCK_WITH_DETAIL], api_key="test-key", client=client + ) + assert calls == [] + + +def test_vertex_drop_params_strips_the_block_keys_embeddings_do_not_take() -> None: + client: Final = HTTPHandler() + url: Final = "https://us-central1-aiplatform.googleapis.com/v1/projects/test/locations/us-central1/publishers/google/models/gemini-embedding-2-preview:embedContent" + mock_get_token, mock_auth, mock_response = _mock_embedding_call({"embedding": {"values": [0.1, 0.2]}}) + with patch.object(client, "post", return_value=mock_response) as mock_post, mock_auth, mock_get_token as token: + token.return_value = ({"Authorization": "Bearer test-token"}, url) + litellm.embedding( + model="vertex_ai/gemini-embedding-2-preview", + input=[CLIP_BLOCK_WITH_DETAIL], + vertex_project="test-project", + vertex_location="us-central1", + client=client, + drop_params=True, + ) + + data: Final = json.loads(mock_post.call_args.kwargs["data"]) + assert data["content"]["parts"][0]["video_metadata"] == CLIP_PART_METADATA + assert "detail" not in mock_post.call_args.kwargs["data"] diff --git a/tests/unit/llms/voyage/test_common_utils.py b/tests/unit/llms/voyage/test_common_utils.py index 511391093f3..51d6c3ed0af 100644 --- a/tests/unit/llms/voyage/test_common_utils.py +++ b/tests/unit/llms/voyage/test_common_utils.py @@ -29,7 +29,7 @@ def test_mongodb_key_routes_to_mongodb_host(): assert get_default_base_url("al-1234567890") == MONGODB_API_BASE -@pytest.mark.parametrize("api_key", ["pa-1234567890", "sk-1234567890", "al", "", None]) +@pytest.mark.parametrize("api_key", ["pa-1234567890", "sk-voyage-legacy-key", "al", "", None]) def test_non_mongodb_key_routes_to_voyage_host(api_key): assert get_default_base_url(api_key) == VOYAGE_API_BASE diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py index 23defa6f513..5ae3c595c53 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py @@ -464,7 +464,7 @@ async def test_sse_mcp_handler_mock(): mock_sse = MagicMock() mock_sse.connect_sse.side_effect = connect_sse - run = AsyncMock() + serve = AsyncMock() # Mock scope, receive, send with proper ASGI scope format mock_scope = { @@ -489,7 +489,7 @@ async def test_sse_mcp_handler_mock(): ) with ( - patch("litellm.proxy._experimental.mcp_server.server.serve_loop", run), + patch("litellm.proxy._experimental.mcp_server.server.serve_loop", serve), patch( "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True, @@ -505,13 +505,21 @@ async def test_sse_mcp_handler_mock(): patch( "litellm.proxy._experimental.mcp_server.server.set_auth_context", ), + patch( + "litellm.proxy._experimental.mcp_server.server._raise_preemptive_401_for_unauthenticated_servers", + new=AsyncMock(), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._check_passthrough_upstream_auth", + new=AsyncMock(), + ), ): from litellm.proxy._experimental.mcp_server.server import handle_sse_mcp # Call the handler await handle_sse_mcp(mock_scope, mock_receive, mock_send) - assert run.await_args.args[1:3] == (read_stream, write_stream) + assert serve.await_args.args[1:3] == (read_stream, write_stream) assert mock_sse.connect_sse.call_args.args[0]["path"] == "/mcp/sse" diff --git a/tests/unit/proxy/db/test_log_db_metrics_service_spans.py b/tests/unit/proxy/db/test_log_db_metrics_service_spans.py new file mode 100644 index 00000000000..a7fd830b948 --- /dev/null +++ b/tests/unit/proxy/db/test_log_db_metrics_service_spans.py @@ -0,0 +1,260 @@ +import asyncio +import importlib +from collections.abc import Sequence +from datetime import datetime, timedelta +from types import SimpleNamespace +from typing import Final, NotRequired, TypedDict + +from typing_extensions import ReadOnly +from unittest.mock import AsyncMock, patch + +import httpx +import pytest +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from prisma.errors import ClientNotConnectedError +from pydantic import TypeAdapter + +import litellm +from litellm._service_logger import ServiceTypes +from litellm.integrations.datadog.datadog import DataDogLogger +from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig +from litellm.proxy.db.log_db_metrics import log_db_metrics +from litellm.proxy.db.prisma_client import _PrismaDrainTracker, _TrackedPrismaEngine +from litellm.proxy.proxy_server import proxy_logging_obj +from litellm.integrations.prometheus_services import PrometheusServicesLogger +from prometheus_client import REGISTRY + + +class _ServiceEvent(TypedDict): + service: ReadOnly[str] + call_type: ReadOnly[str] + duration: ReadOnly[float] + is_error: ReadOnly[bool] + error: ReadOnly[str | None] + event_metadata: ReadOnly[dict[str, str] | None] + table_name: NotRequired[ReadOnly[str]] + + +class _ServiceSpanExporter(InMemorySpanExporter): + def __init__(self) -> None: + super().__init__() + self.service_span_exported: Final = asyncio.Event() + + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + result: Final = super().export(spans) + self.service_span_exported.set() + return result + + +class _Rig: + def __init__(self, exporter: _ServiceSpanExporter, provider: TracerProvider, datadog: DataDogLogger) -> None: + self.exporter: Final = exporter + self.provider: Final = provider + self.datadog: Final = datadog + + def service_spans(self) -> tuple[ReadableSpan, ...]: + return tuple(span for span in self.exporter.get_finished_spans() if span.name != "request") + + def events(self) -> tuple[_ServiceEvent, ...]: + adapter: Final = TypeAdapter(_ServiceEvent) + return tuple(adapter.validate_json(entry["message"]) for entry in self.datadog.log_queue) + + +def _discard_periodic_flush(coroutine: object) -> None: + close: Final = getattr(coroutine, "close") + close() + + +@pytest.fixture +def rig(monkeypatch: pytest.MonkeyPatch) -> _Rig: + monkeypatch.setenv("DD_API_KEY", "test_api_key") + monkeypatch.setenv("DD_SITE", "test.datadoghq.com") + monkeypatch.setattr(litellm, "datadog_params", None) + exporter: Final = _ServiceSpanExporter() + provider: Final = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + otel: Final = OpenTelemetry(config=OpenTelemetryConfig(exporter=exporter), tracer_provider=provider) + with patch("asyncio.create_task", side_effect=_discard_periodic_flush): + datadog: Final = DataDogLogger() + monkeypatch.setattr(litellm, "service_callback", [otel, datadog, "prometheus_system"]) + monkeypatch.setattr(proxy_logging_obj.service_logging_obj, "dd_logger", datadog, raising=False) + monkeypatch.setattr( + proxy_logging_obj.service_logging_obj, "prometheusServicesLogger", PrometheusServicesLogger(), raising=False + ) + return _Rig(exporter, provider, datadog) + + +async def _run_prisma_query() -> None: + engine: Final = _TrackedPrismaEngine(SimpleNamespace(query=AsyncMock(return_value={})), _PrismaDrainTracker()) + await engine.query("{}", tx_id=None) + + +@log_db_metrics +async def read_spend_rows(**kwargs: object) -> str: + await _run_prisma_query() + return "success" + + +def _logged_db_latency() -> tuple[float, float]: + labels: Final = {ServiceTypes.DB.value: ServiceTypes.DB.value} + total: Final = REGISTRY.get_sample_value("litellm_postgres_latency_sum", labels) + count: Final = REGISTRY.get_sample_value("litellm_postgres_latency_count", labels) + return (total or 0.0, count or 0.0) + + +def _ns(moment: datetime) -> int: + return int(moment.timestamp() * 1e9) + + +_DB_CALL_START: Final = datetime(2026, 1, 1, 12, 0, 0) +_DB_CALL_DURATION: Final = timedelta(milliseconds=250) + + +class _ScriptedClock: + def __init__(self, *moments: datetime) -> None: + self._moments: Final = iter(moments) + + def now(self) -> datetime: + return next(self._moments) + + +@pytest.fixture +def scripted_db_clock(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + importlib.import_module("litellm.proxy.db.log_db_metrics"), + "datetime", + _ScriptedClock(_DB_CALL_START, _DB_CALL_START + _DB_CALL_DURATION), + ) + + +@pytest.mark.asyncio +async def test_a_db_success_is_reported_on_the_parent_span_with_its_duration_and_times( + rig: _Rig, scripted_db_clock: None +) -> None: + parent: Final = rig.provider.get_tracer("db-test").start_span("request") + latency_before: Final = _logged_db_latency() + + result: Final = await read_spend_rows(parent_otel_span=parent) + await asyncio.wait_for(rig.exporter.service_span_exported.wait(), timeout=10) + + assert result == "success" + spans: Final = rig.service_spans() + assert len(spans) == 1 + span: Final = spans[0] + assert span.parent is not None + assert span.parent.span_id == parent.get_span_context().span_id + assert span.attributes is not None + assert span.attributes["service"] == ServiceTypes.DB.value + assert span.attributes["call_type"] == "read_spend_rows" + assert span.status.status_code == StatusCode.OK + assert span.start_time == _ns(_DB_CALL_START) + assert span.end_time == _ns(_DB_CALL_START + _DB_CALL_DURATION) + latency_after: Final = _logged_db_latency() + assert latency_after[1] - latency_before[1] == 1 + assert latency_after[0] - latency_before[0] == pytest.approx(_DB_CALL_DURATION.total_seconds()) + assert rig.events() == () + + +@pytest.mark.asyncio +async def test_db_event_metadata_names_only_the_table_and_never_the_raw_kwargs(rig: _Rig) -> None: + parent: Final = rig.provider.get_tracer("db-test").start_span("request") + + await read_spend_rows( + parent_otel_span=parent, + table_name="LiteLLM_SpendLogs", + token="sk-secret-should-not-leak", + prisma_client=object(), + ) + await asyncio.wait_for(rig.exporter.service_span_exported.wait(), timeout=10) + + span_attributes: Final = rig.service_spans()[0].attributes + assert span_attributes is not None + assert span_attributes["table_name"] == "LiteLLM_SpendLogs" + assert not {"token", "prisma_client", "parent_otel_span"} & set(span_attributes) + assert all("sk-secret-should-not-leak" not in str(value) for value in span_attributes.values()) + + +@pytest.mark.asyncio +async def test_the_logged_db_duration_is_the_span_wall_clock_of_the_wrapped_call( + rig: _Rig, scripted_db_clock: None +) -> None: + parent: Final = rig.provider.get_tracer("db-test").start_span("request") + latency_before: Final = _logged_db_latency() + + await read_spend_rows(parent_otel_span=parent) + await asyncio.wait_for(rig.exporter.service_span_exported.wait(), timeout=10) + + span: Final = rig.service_spans()[0] + assert span.start_time is not None and span.end_time is not None + latency_after: Final = _logged_db_latency() + logged_duration: Final = latency_after[0] - latency_before[0] + assert latency_after[1] - latency_before[1] == 1 + assert logged_duration == pytest.approx((span.end_time - span.start_time) / 1e9, rel=1e-3, abs=2e-6) + assert logged_duration == pytest.approx(_DB_CALL_DURATION.total_seconds()) + + +@log_db_metrics +async def disconnected_read(**kwargs: object) -> str: + raise ClientNotConnectedError() + + +@pytest.mark.asyncio +async def test_a_prisma_error_is_reported_as_a_db_failure_and_reraised(rig: _Rig) -> None: + parent: Final = rig.provider.get_tracer("db-test").start_span("request") + + with pytest.raises(ClientNotConnectedError, match="Client is not connected to the query engine"): + await disconnected_read(parent_otel_span=parent) + + spans: Final = rig.service_spans() + assert len(spans) == 1 + assert spans[0].parent is not None + assert spans[0].parent.span_id == parent.get_span_context().span_id + assert spans[0].status.status_code == StatusCode.ERROR + assert spans[0].attributes is not None + assert spans[0].attributes["call_type"] == "disconnected_read" + assert spans[0].attributes["service"] == ServiceTypes.DB.value + assert "Client is not connected" in str(spans[0].attributes["error"]) + events: Final = rig.events() + assert len(events) == 1 + assert events[0]["is_error"] is True + assert events[0]["call_type"] == "disconnected_read" + assert isinstance(events[0]["duration"], float) + assert "Client is not connected" in (events[0]["error"] or "") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("error", "is_db_error"), + [ + (ValueError("Generic error"), False), + (KeyError("Missing key"), False), + (TypeError("Type error"), False), + (httpx.ConnectError("Failed to connect"), True), + (httpx.TimeoutException("Request timed out"), True), + (ClientNotConnectedError(), True), + ], +) +async def test_only_db_errors_are_reported_as_db_failures(rig: _Rig, error: Exception, is_db_error: bool) -> None: + parent: Final = rig.provider.get_tracer("db-test").start_span("request") + + @log_db_metrics + async def failing_read(**kwargs: object) -> str: + raise error + + with pytest.raises(type(error)): + await failing_read(parent_otel_span=parent) + + spans: Final = rig.service_spans() + events: Final = rig.events() + if is_db_error: + assert [span.status.status_code for span in spans] == [StatusCode.ERROR] + assert [(event["service"], event["call_type"], event["is_error"]) for event in events] == [ + (ServiceTypes.DB.value, "failing_read", True) + ] + assert isinstance(events[0]["duration"], float) + else: + assert spans == () + assert events == () diff --git a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py index e4cd7d9dfa8..451792109de 100644 --- a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py @@ -3,6 +3,7 @@ Test to verify the Google GenAI proxy API endpoints """ +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -179,3 +180,30 @@ def test_google_count_tokens_unchanged(): assert response.status_code == 200 body = response.json() assert body["totalTokens"] == 7 + + +def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs(): + """A provider that counts in Anthropic shape (Bedrock, Anthropic) carries no totalTokens key, so the normalized total must win over the raw response's missing field.""" + try: + client: Final = _build_test_client() + except ImportError as e: + pytest.skip(f"Skipping test due to missing dependency: {e}") + + fake_response: Final = MagicMock() + fake_response.original_response = {"input_tokens": 2167} + fake_response.total_tokens = 2167 + + with patch( + "litellm.proxy.proxy_server.token_counter", + new_callable=AsyncMock, + return_value=fake_response, + ): + response: Final = client.post( + "/v1beta/models/claude-opus-4-8:countTokens", + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + ) + + assert response.status_code == 200 + body: Final = response.json() + assert body["totalTokens"] == 2167 + assert body["promptTokensDetails"] == [] diff --git a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py index fa33c2d462b..b7bc5f3f8ec 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py @@ -575,7 +575,7 @@ async def test_get_guardrail_info_from_db(mocker, mock_prisma_client): """Test getting guardrail info from DB""" mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - response = await get_guardrail_info("test-db-guardrail") + response: Final = await get_guardrail_info("test-db-guardrail") assert response.guardrail_id == "test-db-guardrail" assert response.guardrail_name == "Test DB Guardrail" @@ -584,6 +584,21 @@ async def test_get_guardrail_info_from_db(mocker, mock_prisma_client): @pytest.mark.asyncio +async def test_get_guardrail_info_tolerates_invalid_stored_stream_scope(mocker, mock_prisma_client): + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock( + return_value={ + **MOCK_DB_GUARDRAIL, + "litellm_params": { + **MOCK_DB_GUARDRAIL["litellm_params"], + "stream_scope": "sometimes", + }, + } + ) + + response = await get_guardrail_info("test-db-guardrail") + + assert response.litellm_params.stream_scope is None async def test_get_guardrail_info_normalizes_invalid_scope_from_db( mocker, mock_guardrail_registry, mock_in_memory_handler ): @@ -745,6 +760,40 @@ def test_get_guardrails_list_response_includes_guardrail_id(): assert response.guardrails[0].guardrail_id == "stable-config-id" +def test_get_guardrails_list_response_tolerates_invalid_config_stream_scope(): + from litellm.proxy.guardrails.guardrail_endpoints import ( + _get_guardrails_list_response, + ) + + response = _get_guardrails_list_response( + [ + { + "guardrail_id": "invalid-scope", + "guardrail_name": "invalid-scope", + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "stream_scope": "sometimes", + }, + }, + { + "guardrail_id": "valid-scope", + "guardrail_name": "valid-scope", + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "stream_scope": "STREAMING", + }, + }, + ] + ) + + assert response.guardrails[0].litellm_params is not None + assert response.guardrails[0].litellm_params.stream_scope is None + assert response.guardrails[1].litellm_params is not None + assert response.guardrails[1].litellm_params.stream_scope == "streaming" + + def test_get_provider_specific_params(): """Test getting provider-specific parameters""" from litellm.proxy.guardrails.guardrail_endpoints import _get_fields_from_model diff --git a/tests/unit/proxy/guardrails/test_guardrail_registry.py b/tests/unit/proxy/guardrails/test_guardrail_registry.py index c7b13f0bd00..3bc6d163e5d 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_registry.py +++ b/tests/unit/proxy/guardrails/test_guardrail_registry.py @@ -1,7 +1,7 @@ import json from collections.abc import Iterable, Iterator from typing import ClassVar, Final -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from pydantic import ValidationError @@ -160,6 +160,43 @@ def test_duplicate_config_guardrail_names_get_distinct_stable_ids(): registry_module.guardrail_initializer_registry.pop("dup_name_test", None) +def test_initialize_guardrail_treats_invalid_stored_scope_as_both(): + from litellm.proxy.guardrails import guardrail_registry as registry_module + + guardrail_type: Final = "invalid_stored_scope_test" + + def _initializer(litellm_params: LitellmParams, guardrail: Guardrail) -> CustomGuardrail: + return CustomGuardrail( + guardrail_name=guardrail["guardrail_name"], + event_hook=GuardrailEventHooks(litellm_params.mode), + default_on=True, + ) + + registry_module.guardrail_initializer_registry[guardrail_type] = _initializer + try: + handler: Final = InMemoryGuardrailHandler() + guardrail: Final = Guardrail( + guardrail_id="invalid-stored-scope", + guardrail_name="invalid-stored-scope", + litellm_params={ + "guardrail": guardrail_type, + "mode": "pre_call", + "default_on": True, + "stream_scope": "sometimes", + }, + ) + + parsed_guardrail: Final = handler.initialize_guardrail(guardrail=guardrail, source="db") + callback: Final = handler.guardrail_id_to_custom_guardrail["invalid-stored-scope"] + + assert parsed_guardrail["litellm_params"].stream_scope is None + assert callback is not None + assert callback.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is True + assert callback.should_run_guardrail(data={"stream": True}, event_type=GuardrailEventHooks.pre_call) is True + finally: + registry_module.guardrail_initializer_registry.pop(guardrail_type, None) + + def _register_mode_following_initializer(guardrail_type: str): """Registers like the shipped initializers do: construct, then add the instance to litellm's callbacks.""" import litellm @@ -1605,6 +1642,35 @@ def test_sync_guardrail_from_db_applies_db_dict_params_to_live_instance(): cb_list[:] = snapshot +def test_configure_callback_scoping_copies_stream_scope_when_constructor_omits_it(): + from litellm.proxy.guardrails.guardrail_registry import _configure_callback_scoping + + class _CtorWithoutStreamScope(CustomGuardrail): + def __init__(self) -> None: + super().__init__( + guardrail_name="scoped", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + + instance = _CtorWithoutStreamScope() + params = LitellmParams(guardrail="bedrock", mode="post_call", stream_scope="streaming") + _configure_callback_scoping(instance, "scoped", params) + + assert instance.stream_scope_default == "streaming" + assert instance.should_run_guardrail({"stream": True}, GuardrailEventHooks.post_call) is True + assert instance.should_run_guardrail({"stream": False}, GuardrailEventHooks.post_call) is False + + +def test_configure_callback_scoping_tolerates_a_custom_logger_callback(): + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.guardrails.guardrail_registry import _configure_callback_scoping + + callback: Final = CustomLogger() + _configure_callback_scoping(callback, "logger-backed", LitellmParams(guardrail="custom", mode="pre_call")) # pyright: ignore[reportArgumentType] # module-path guardrails may be plain CustomLogger + assert "stream_scope_by_hook" not in vars(callback) + + _ENCRYPTED_PREFIX = "litellm_enc::" diff --git a/tests/unit/proxy/lens/test_endpoints.py b/tests/unit/proxy/lens/test_endpoints.py index fee53320802..e9ad63a701e 100644 --- a/tests/unit/proxy/lens/test_endpoints.py +++ b/tests/unit/proxy/lens/test_endpoints.py @@ -1,5 +1,6 @@ import asyncio -from collections.abc import Callable, Mapping +from collections.abc import AsyncGenerator, Callable, Mapping +from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final @@ -49,7 +50,7 @@ from litellm.proxy.lens.models import ( TraceIdentity, Worker, ) -from litellm.proxy.lens.repository import DueLens, Row +from litellm.proxy.lens.repository import Database, DueLens, Row from litellm.proxy.lens.signals import SignalConfig, StoredTraceSignal from litellm.proxy.lens.state import claim_job, queue_job, replace_job from litellm.rust_bridge.trace.generated.models import ExecutionRow, LensSampleParams @@ -69,6 +70,10 @@ class ResultDatabase: self.stored = stored self.completed: tuple[ReviewVersion, ...] = () + @asynccontextmanager + async def transaction(self) -> AsyncGenerator[Database]: + yield self + async def query_raw(self, query: str, *args: object) -> tuple[Row, ...]: if query.startswith("SELECT data FROM"): return (Row(data=self.stored.model_dump(mode="json")),) diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 1887021a53a..d329db75fb4 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -23,6 +23,7 @@ from starlette.datastructures import FormData import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail from tests._master_key import MASTER_KEY as SHARED_MASTER_KEY from litellm.caching.caching import DualCache from litellm.types.utils import CallTypesLiteral @@ -47,6 +48,7 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( get_vertex_base_url, is_azure_ai_search_service_level_index_create, gigachat_proxy_route, + handle_bedrock_passthrough_router_model, llm_passthrough_factory_proxy_route, milvus_proxy_route, mistral_proxy_route, @@ -61,9 +63,32 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( from litellm.proxy._types import LitellmUserRoles, SpecialHeaders, UserAPIKeyAuth from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials +def _assert_bedrock_processing_data_classification(data: dict[str, object], is_streaming: bool) -> None: + streaming_guardrail: Final = CustomGuardrail( + guardrail_name="streaming-only", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + stream_scope="streaming", + ) + non_streaming_guardrail: Final = CustomGuardrail( + guardrail_name="non-streaming-only", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + stream_scope="non_streaming", + ) + provider_body: Final = data["data"] + + assert streaming_guardrail.should_run_guardrail(data, GuardrailEventHooks.pre_call) is is_streaming + assert non_streaming_guardrail.should_run_guardrail(data, GuardrailEventHooks.pre_call) is not is_streaming + assert isinstance(provider_body, dict) + assert "is_streaming_request" not in provider_body + assert "litellm_server_streaming_classification" not in provider_body + + class TestVertexPassthroughGetVertexBaseUrl: """Module-local get_vertex_base_url (trailing slash); rules match common_utils.""" @@ -413,9 +438,7 @@ class TestVertexAIPassThroughHandler: # Mock the vertex handler for global location mock_handler = Mock() - mock_handler.get_default_base_target_url.return_value = ( - "https://aiplatform.googleapis.com/" - ) + mock_handler.get_default_base_target_url.return_value = "https://aiplatform.googleapis.com/" mock_get_handler.return_value = mock_handler # Mock create_pass_through_route to return a function that returns a mock response @@ -1237,9 +1260,7 @@ class TestVertexAIDiscoveryPassThroughHandler: # Mock the discovery handler mock_handler = Mock() - mock_handler.get_default_base_target_url.return_value = ( - "https://discoveryengine.googleapis.com" - ) + mock_handler.get_default_base_target_url.return_value = "https://discoveryengine.googleapis.com" mock_get_handler.return_value = mock_handler # Mock create_pass_through_route to return a function that returns a mock response @@ -1459,6 +1480,170 @@ class TestBedrockLLMProxyRoute: assert call_kwargs["model"] == "anthropic.claude-3-sonnet-20240229-v1:0" assert result == "success" + @pytest.mark.asyncio + @pytest.mark.parametrize( + "action, is_streaming", + [ + ("converse-stream", True), + ("invoke-with-response-stream", True), + ("converse", False), + ("invoke", False), + ], + ) + async def test_bedrock_direct_actions_classify_guardrail_stream_scope( + self, action: str, is_streaming: bool + ) -> None: + mock_request: Final = Mock() + mock_request.method = "POST" + mock_processor: Final = Mock() + mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value="success") + request_body: Final = {"messages": [{"role": "user", "content": "test"}]} + + with ( + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._read_request_body", + return_value=request_body, + ), + patch( + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing", + return_value=mock_processor, + ) as processor_constructor, + ): + result: Final = await bedrock_llm_proxy_route( + endpoint=f"model/test-model/{action}", + request=mock_request, + fastapi_response=Mock(), + user_api_key_dict=Mock(), + ) + + assert result == "success" + processing_data: Final = processor_constructor.call_args.kwargs["data"] + _assert_bedrock_processing_data_classification(processing_data, is_streaming) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "action, is_streaming", + [ + ("converse-stream", True), + ("invoke-with-response-stream", True), + ("converse", False), + ("invoke", False), + ], + ) + async def test_bedrock_router_actions_classify_guardrail_stream_scope( + self, action: str, is_streaming: bool + ) -> None: + mock_request: Final = Mock() + mock_request.method = "POST" + mock_processor: Final = Mock() + mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value="success") + request_body: Final = {"messages": [{"role": "user", "content": "test"}]} + + with patch( + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing", + return_value=mock_processor, + ) as processor_constructor: + result: Final = await handle_bedrock_passthrough_router_model( + model="test-model", + endpoint=f"model/test-model/{action}", + request=mock_request, + request_body=request_body, + llm_router=Mock(), + user_api_key_dict=Mock(), + proxy_logging_obj=Mock(), + general_settings={}, + proxy_config=None, + select_data_generator=None, + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + version=None, + ) + + assert result == "success" + processing_data: Final = processor_constructor.call_args.kwargs["data"] + _assert_bedrock_processing_data_classification(processing_data, is_streaming) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("model_id", "action"), + [ + ("my-converse-stream-model", "converse"), + ("my-invoke-with-response-stream-model", "invoke"), + ], + ) + async def test_bedrock_direct_model_id_does_not_imply_streaming(self, model_id: str, action: str) -> None: + mock_request: Final = Mock() + mock_request.method = "POST" + mock_processor: Final = Mock() + mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value="success") + request_body: Final = {"messages": [{"role": "user", "content": "test"}]} + + with ( + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._read_request_body", + return_value=request_body, + ), + patch( + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing", + return_value=mock_processor, + ) as processor_constructor, + ): + result: Final = await bedrock_llm_proxy_route( + endpoint=f"/model/{model_id}/{action}", + request=mock_request, + fastapi_response=Mock(), + user_api_key_dict=Mock(), + ) + + assert result == "success" + processing_data: Final = processor_constructor.call_args.kwargs["data"] + _assert_bedrock_processing_data_classification(processing_data, is_streaming=False) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("model_id", "action"), + [ + ("my-converse-stream-model", "converse"), + ("my-invoke-with-response-stream-model", "invoke"), + ], + ) + async def test_bedrock_router_model_id_does_not_imply_streaming(self, model_id: str, action: str) -> None: + mock_request: Final = Mock() + mock_request.method = "POST" + mock_processor: Final = Mock() + mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value="success") + request_body: Final = {"messages": [{"role": "user", "content": "test"}]} + + with patch( + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing", + return_value=mock_processor, + ) as processor_constructor: + result: Final = await handle_bedrock_passthrough_router_model( + model=model_id, + endpoint=f"/model/{model_id}/{action}", + request=mock_request, + request_body=request_body, + llm_router=Mock(), + user_api_key_dict=Mock(), + proxy_logging_obj=Mock(), + general_settings={}, + proxy_config=None, + select_data_generator=None, + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + version=None, + ) + + assert result == "success" + processing_data: Final = processor_constructor.call_args.kwargs["data"] + _assert_bedrock_processing_data_classification(processing_data, is_streaming=False) + @pytest.mark.asyncio async def test_bedrock_error_handling_returns_actual_error(self): """ @@ -1879,7 +2064,6 @@ class TestBedrockAgentRuntimePassthroughToggle: class TestBedrockAgentRuntimePassthroughVirtualKeyLeak: - VKEY: Final = "sk-litellm-victim-key" MASTER_KEY: Final = SHARED_MASTER_KEY ENDPOINT: Final = "knowledgebases/KB1234567/retrieve" @@ -2075,10 +2259,10 @@ class TestVLLMProxyRoute: "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model", return_value=True, ) - @patch("litellm.proxy.proxy_server.llm_router") # test-quality-ok: patching litellm internal for unit test isolation - async def test_vllm_proxy_route_with_router_model( - self, mock_llm_router, mock_is_router, mock_get_body - ): + @patch( + "litellm.proxy.proxy_server.llm_router" + ) # test-quality-ok: patching litellm internal for unit test isolation + async def test_vllm_proxy_route_with_router_model(self, mock_llm_router, mock_is_router, mock_get_body): mock_request = MagicMock(spec=Request) mock_request.method = "POST" mock_request.headers = {"content-type": "application/json"} @@ -2111,9 +2295,7 @@ class TestVLLMProxyRoute: @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.llm_passthrough_factory_proxy_route" ) - async def test_vllm_proxy_route_fallback_to_factory( - self, mock_factory_route, mock_is_router, mock_get_body - ): + async def test_vllm_proxy_route_fallback_to_factory(self, mock_factory_route, mock_is_router, mock_get_body): mock_request = MagicMock(spec=Request) mock_fastapi_response = MagicMock(spec=Response) mock_user_api_key_dict = MagicMock() @@ -2140,10 +2322,10 @@ class TestGigachatProxyRoute: "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model", return_value=True, ) - @patch("litellm.proxy.proxy_server.llm_router") # test-quality-ok: patching litellm internal for unit test isolation - async def test_gigachat_proxy_route_with_router_model( - self, mock_llm_router, mock_is_router, mock_get_body - ): + @patch( + "litellm.proxy.proxy_server.llm_router" + ) # test-quality-ok: patching litellm internal for unit test isolation + async def test_gigachat_proxy_route_with_router_model(self, mock_llm_router, mock_is_router, mock_get_body): mock_request = MagicMock(spec=Request) mock_request.method = "POST" mock_request.headers = {"content-type": "application/json"} @@ -2396,21 +2578,25 @@ class TestGigachatProxyRoute: return _inner() - with patch.object( - processor, - "common_processing_pre_call_logic", - new=AsyncMock( - return_value=( - processor.data, - processor.data["litellm_logging_obj"], - ) + with ( + patch.object( + processor, + "common_processing_pre_call_logic", + new=AsyncMock( + return_value=( + processor.data, + processor.data["litellm_logging_obj"], + ) + ), + ), + patch( # test-quality-ok: patching litellm internal for unit test isolation + "litellm.proxy.common_request_processing.route_request", + new=_fake_route_request, + ), + patch( # test-quality-ok: patching litellm internal for unit test isolation + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.get_custom_headers", + return_value={"x-litellm-call-id": "call-123"}, ), - ), patch( # test-quality-ok: patching litellm internal for unit test isolation - "litellm.proxy.common_request_processing.route_request", - new=_fake_route_request, - ), patch( # test-quality-ok: patching litellm internal for unit test isolation - "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.get_custom_headers", - return_value={"x-litellm-call-id": "call-123"}, ): result = await processor.base_passthrough_process_llm_request( request=mock_request, @@ -4223,7 +4409,9 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak: @pytest.mark.parametrize( ("credential", "authenticated"), [ - pytest.param("modified_key", UserAPIKeyAuth(api_key="modified_key"), id="custom-auth-echoing-opaque-credential"), + pytest.param( + "modified_key", UserAPIKeyAuth(api_key="modified_key"), id="custom-auth-echoing-opaque-credential" + ), pytest.param( LITELLM_JWT, UserAPIKeyAuth(api_key=LITELLM_JWT, user_id="jwt-subject"), @@ -4298,7 +4486,9 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak: (b"x-goog-api-key", b"AIza-real-google-api-key"), (b"content-type", b"application/json"), ], - authenticated=UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, user_role=LitellmUserRoles.PROXY_ADMIN), + authenticated=UserAPIKeyAuth( + api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, user_role=LitellmUserRoles.PROXY_ADMIN + ), ) assert raised is None assert forwarded is not None @@ -4440,10 +4630,14 @@ class TestAnthropicPassthroughVirtualKeyLeak: raised, forwarded = await self._run( monkeypatch, [(header, value), (b"anthropic-version", b"2023-06-01"), (b"content-type", b"application/json")], - authenticated=UserAPIKeyAuth(api_key="sk-ant-api03-callers-own-key", user_role=LitellmUserRoles.INTERNAL_USER), + authenticated=UserAPIKeyAuth( + api_key="sk-ant-api03-callers-own-key", user_role=LitellmUserRoles.INTERNAL_USER + ), master_key=None, ) - assert raised is None, "with no master key the proxy authenticated nothing, so nothing of the caller's is a LiteLLM secret" + assert raised is None, ( + "with no master key the proxy authenticated nothing, so nothing of the caller's is a LiteLLM secret" + ) assert forwarded is not None assert forwarded.get(header.decode()) == value.decode() @@ -5231,7 +5425,9 @@ class TestTranscribeProxyRoute: ) -> None: with respx.mock(assert_all_called=False) as upstream: route = upstream.post(TRANSCRIBE_UPSTREAM) - response = transcribe_client.post("/transcribe/StartTranscriptionJob", json={**dict(self.START_JOB_BODY), **body}) + response = transcribe_client.post( + "/transcribe/StartTranscriptionJob", json={**dict(self.START_JOB_BODY), **body} + ) assert response.status_code == 403 assert member in response.json()["detail"] @@ -6873,9 +7069,7 @@ class TestAzureBodyModelGroupRelay: AZURE_SPEECH_SHORT_AUDIO_ENDPOINT: Final = "/speech/recognition/conversation/cognitiveservices/v1" AZURE_SPEECH_BATCH_ENDPOINT: Final = "/speechtotext/v3.2/transcriptions" AZURE_SPEECH_FAST_ENDPOINT: Final = "/speechtotext/transcriptions:transcribe" -AZURE_SPEECH_PCM16_HEADER: Final = ( - b"RIFF\x24\x0c\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x01\x00\x80\x3e\x00\x00\x00\x7d\x00\x00\x02\x00\x10\x00data\x00\x0c\x00\x00" -) +AZURE_SPEECH_PCM16_HEADER: Final = b"RIFF\x24\x0c\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x01\x00\x80\x3e\x00\x00\x00\x7d\x00\x00\x02\x00\x10\x00data\x00\x0c\x00\x00" AZURE_SPEECH_WAV_BYTES: Final = AZURE_SPEECH_PCM16_HEADER + b"\x00" * 3072 AZURE_SPEECH_WAV_SECONDS: Final = 3072 / (16000 * 2) AZURE_SPEECH_NON_UTF8_WAV_BYTES: Final = AZURE_SPEECH_PCM16_HEADER + bytes(range(256)) * 12 @@ -6912,9 +7106,9 @@ class TestAzureSpeechProxyRoute: def test_short_audio_forwards_raw_wav_bytes_with_server_key(self, azure_speech_client: TestClient) -> None: with respx.mock(assert_all_called=True) as upstream: - route = upstream.post( - f"https://eastus.stt.speech.microsoft.com{AZURE_SPEECH_SHORT_AUDIO_ENDPOINT}" - ).mock(return_value=httpx.Response(200, json=AZURE_SPEECH_TRANSCRIPT)) + route = upstream.post(f"https://eastus.stt.speech.microsoft.com{AZURE_SPEECH_SHORT_AUDIO_ENDPOINT}").mock( + return_value=httpx.Response(200, json=AZURE_SPEECH_TRANSCRIPT) + ) response = azure_speech_client.post( f"/azure_speech{AZURE_SPEECH_SHORT_AUDIO_ENDPOINT}", @@ -7390,9 +7584,9 @@ class TestAzureSpeechRawBodyThroughRealAuth: self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, body: bytes ) -> None: with respx.mock(assert_all_called=True) as upstream, caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): - route = upstream.post( - f"https://eastus.stt.speech.microsoft.com{AZURE_SPEECH_SHORT_AUDIO_ENDPOINT}" - ).mock(return_value=httpx.Response(200, json=AZURE_SPEECH_TRANSCRIPT)) + route = upstream.post(f"https://eastus.stt.speech.microsoft.com{AZURE_SPEECH_SHORT_AUDIO_ENDPOINT}").mock( + return_value=httpx.Response(200, json=AZURE_SPEECH_TRANSCRIPT) + ) response = self._post_wav( monkeypatch, f"/azure_speech{AZURE_SPEECH_SHORT_AUDIO_ENDPOINT}", "sk-master-key", body=body @@ -7416,9 +7610,9 @@ class TestAzureSpeechRawBodyThroughRealAuth: ) -> None: boundary: Final = "lit7939boundary" multipart_body: Final = ( - f"--{boundary}\r\nContent-Disposition: form-data; name=\"definition\"\r\n\r\n".encode() + f'--{boundary}\r\nContent-Disposition: form-data; name="definition"\r\n\r\n'.encode() + json.dumps({"locales": ["en-US"]}).encode() - + f"\r\n--{boundary}\r\nContent-Disposition: form-data; name=\"audio\"; filename=\"eagle.wav\"\r\n" + + f'\r\n--{boundary}\r\nContent-Disposition: form-data; name="audio"; filename="eagle.wav"\r\n' "Content-Type: audio/wav\r\n\r\n".encode() + AZURE_SPEECH_NON_UTF8_WAV_BYTES + f"\r\n--{boundary}--\r\n".encode() @@ -8170,9 +8364,7 @@ class TestTinyFishProxyRoute: assert response.status_code == 200 assert json.loads(route.calls.last.request.content)["use_vault"] is True - def test_returns_401_on_missing_api_key( - self, tinyfish_client: TestClient, monkeypatch: pytest.MonkeyPatch - ) -> None: + def test_returns_401_on_missing_api_key(self, tinyfish_client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("TINYFISH_API_KEY") with respx.mock: diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index cd279a6ed4b..a6ce8ee6fbc 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -5,7 +5,7 @@ import logging import os import sys import zlib -from collections.abc import Callable, Mapping +from collections.abc import AsyncIterator, Callable, Mapping from contextlib import ExitStack, contextmanager from dataclasses import dataclass from io import BytesIO @@ -27,6 +27,7 @@ from starlette.datastructures import UploadFile as StarletteUploadFile import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features @@ -55,6 +56,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.types import utils as types_utils +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, @@ -1638,6 +1640,223 @@ async def test_pass_through_request_streaming_marks_logging_obj_as_stream(): assert logging_obj.model_call_details["stream"] is True +@pytest.mark.asyncio +@pytest.mark.parametrize("body_stream", [None, True], ids=["stream-absent", "stream-true"]) +async def test_pass_through_request_preserves_caller_streaming_request_field(body_stream): + captured_hook_data: dict[str, object] = {} + + async def capture_pre_call(user_api_key_dict, data, call_type, endpoint_type: EndpointType): + captured_hook_data.update(data) + return data + + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor" + ) as mock_chunk_processor: + mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=capture_pre_call) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) + + upstream_response = MagicMock() + upstream_response.status_code = 200 + upstream_response.headers = {} + upstream_response.raise_for_status = MagicMock() + + async_client = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + async def _empty_chunks(*args, **kwargs): + return + yield # pragma: no cover + + mock_chunk_processor.return_value = _empty_chunks() + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = httpx.URL( + "http://test-proxy.com/gemini/v1beta/models/gemini-pro:streamGenerateContent" + ) + mock_request.scope = {"path": "/gemini/v1beta/models/gemini-pro:streamGenerateContent"} + request_body: Final = { + "contents": [{"parts": [{"text": "hi"}]}], + "is_streaming_request": "caller-value", + **({"stream": True} if body_stream is True else {}), + } + mock_request.body = AsyncMock(return_value=json.dumps(request_body).encode()) + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + + await pass_through_request( + request=mock_request, + target="https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:streamGenerateContent", + custom_headers={}, + user_api_key_dict=MagicMock(), + stream=True, + ) + + assert captured_hook_data.get("is_streaming_request") == "caller-value" + assert captured_hook_data.get("stream") is body_stream + + upstream_json = async_client.build_request.call_args.kwargs["json"] + assert upstream_json["is_streaming_request"] == "caller-value" + assert "litellm_server_streaming_classification" not in upstream_json + assert upstream_json["contents"] == request_body["contents"] + + +@pytest.mark.asyncio +async def test_streaming_pass_through_drops_marker_after_hook_rebuilds_body_from_json(): + async def json_rebuilding_pre_call(user_api_key_dict, data, call_type, endpoint_type: EndpointType): + rebuilt = json.loads(json.dumps({k: v for k, v in data.items() if k != "litellm_logging_obj"})) + return {**rebuilt, "litellm_logging_obj": data["litellm_logging_obj"]} + + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor" + ) as mock_chunk_processor: + mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=json_rebuilding_pre_call) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) + + upstream_response = MagicMock() + upstream_response.status_code = 200 + upstream_response.headers = {} + upstream_response.raise_for_status = MagicMock() + + async_client = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + async def _empty_chunks(*args, **kwargs): + return + yield # pragma: no cover + + mock_chunk_processor.return_value = _empty_chunks() + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = httpx.URL("http://test-proxy.com/openai/v1/chat/completions") + mock_request.scope = {"path": "/openai/v1/chat/completions"} + request_body: Final = { + "model": "gpt-5-mini", + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + } + mock_request.body = AsyncMock(return_value=json.dumps(request_body).encode()) + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + + await pass_through_request( + request=mock_request, + target="https://api.openai.com/v1/chat/completions", + custom_headers={}, + user_api_key_dict=MagicMock(), + ) + + upstream_json = async_client.build_request.call_args.kwargs["json"] + assert upstream_json == request_body, upstream_json + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route_stream", "body_stream", "expected_streaming"), + [ + (True, False, False), + (False, True, True), + (True, None, True), + (None, None, False), + ], + ids=["body-disables-route-stream", "body-enables-streaming", "route-enables-absent-body", "both-absent"], +) +async def test_passthrough_guardrails_follow_effective_relay_stream_decision( + route_stream: bool | None, + body_stream: bool | None, + expected_streaming: bool, +): + async def return_pre_call_data( + user_api_key_dict: UserAPIKeyAuth, + data: dict[str, object], + call_type: str, + endpoint_type: EndpointType, + ) -> dict[str, object]: + return dict(data) + + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor" + ) as mock_chunk_processor: + mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=return_pre_call_data) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) + + upstream_response = MagicMock() + upstream_response.status_code = 200 + upstream_response.headers = {} + upstream_response.aread = AsyncMock(return_value=b"{}") + upstream_response.text = "{}" + upstream_response.raise_for_status = MagicMock() + + async_client = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + async_client.request = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + async def empty_chunks() -> AsyncIterator[bytes]: + yield b"" + + mock_chunk_processor.return_value = empty_chunks() + + request_body: Final = { + "message": "hello", + **({"stream": body_stream} if body_stream is not None else {}), + } + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = httpx.URL("http://test-proxy.com/guardrail-stream-scope") + mock_request.scope = {"path": "/guardrail-stream-scope"} + mock_request.body = AsyncMock(return_value=json.dumps(request_body).encode()) + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + + await pass_through_request( + request=mock_request, + target="http://upstream.test/guardrail-stream-scope", + custom_headers={}, + user_api_key_dict=MagicMock(), + stream=route_stream, + ) + hook_data: Final[dict[str, object]] = mock_proxy_logging.pre_call_hook.call_args.kwargs["data"] + + streaming_guardrail: Final = CustomGuardrail( + guardrail_name="streaming-only", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + stream_scope="streaming", + ) + non_streaming_guardrail: Final = CustomGuardrail( + guardrail_name="non-streaming-only", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + stream_scope="non_streaming", + ) + assert streaming_guardrail.should_run_guardrail(hook_data, GuardrailEventHooks.pre_call) is expected_streaming + assert ( + non_streaming_guardrail.should_run_guardrail(hook_data, GuardrailEventHooks.pre_call) is not expected_streaming + ) + + @pytest.mark.asyncio async def test_pass_through_request_sse_response_marks_logging_obj_as_stream(): """ diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index 6aa7eca0f15..d13abbd379b 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -96,6 +96,20 @@ class AlwaysPassGuardrail(CustomGuardrail): return None +class StreamScopedPassGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str, stream_scope: object): + super().__init__( + guardrail_name=guardrail_name, + event_hook="pre_call", + default_on=True, + stream_scope=stream_scope, + ) + self.calls = 0 + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.calls += 1 + + class PassthroughBlockGuardrail(CustomGuardrail): """Mock guardrail that blocks using the legacy passthrough contract.""" @@ -378,6 +392,46 @@ async def test_block_carries_original_guardrail_exception(monkeypatch): assert result.original_exception.detail == "Content policy violation" +@pytest.mark.asyncio +async def test_pipeline_step_honors_stream_scope(monkeypatch): + stream_only = StreamScopedPassGuardrail(guardrail_name="stream-only", stream_scope="streaming") + later = AlwaysFailGuardrail(guardrail_name="later-block") + monkeypatch.setattr(litellm, "callbacks", [stream_only, later]) + steps = [ + PipelineStep(guardrail="stream-only", on_fail="block", on_pass="allow"), + PipelineStep(guardrail="later-block", on_fail="block", on_pass="allow"), + ] + + skipped = await PipelineExecutor.execute_steps( + steps=steps, + mode="pre_call", + data={"messages": [{"role": "user", "content": "hi"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="stream-scope", + ) + assert stream_only.calls == 0 + assert later.calls == 1 + assert skipped.step_results[0].outcome == "skip" + assert skipped.step_results[0].action_taken == "next" + assert skipped.terminal_action == "block" + + stream_only.calls = 0 + later.calls = 0 + ran = await PipelineExecutor.execute_steps( + steps=steps, + mode="pre_call", + data={"messages": [{"role": "user", "content": "hi"}], "stream": True}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="stream-scope", + ) + assert stream_only.calls == 1 + assert later.calls == 0 + assert ran.step_results[0].outcome == "pass" + assert ran.terminal_action == "allow" + + @pytest.mark.asyncio async def test_unsupported_mode_yields_error_outcome_without_exception(monkeypatch): """An unexpected hook mode must surface as an error outcome (carrying no diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index c5bb1ed5b61..9aa9110f7a4 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -18,7 +18,7 @@ from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATIO from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth +from litellm.proxy._types import SpendLogsMetadataFields, SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_tracking_utils import ( @@ -37,9 +37,11 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_guardrail_information_for_spend_logs, _sanitize_request_body_for_spend_logs_payload, _scrub_raw_model_from_error_information, + configured_spend_logs_metadata_fields, get_logging_payload, get_spend_logs_id, should_store_prompts_and_responses_in_spend_logs, + spend_log_row_with_retained_metadata, ) from litellm.proxy.utils import hash_token from litellm.types.router import GenericLiteLLMParams @@ -79,10 +81,16 @@ def test_classifier_audit_spend_storage_obeys_privacy_and_truncation(monkeypatch "classifier_input": {"system": "rubric" * 1000, "messages": [{"role": "user", "content": "ask"}]}, "originating_request_masked": {"input": "source-only", "api_key": "REDACTED"}, } - stored: Final = json.loads(_get_proxy_server_request_for_spend_logs_payload( - metadata={}, litellm_params={"proxy_server_request": {"body": {"model": "classifier"}}}, - kwargs={"standard_logging_object": audit, "standard_callback_dynamic_params": {"turn_off_message_logging": redact}}, - )) + stored: Final = json.loads( + _get_proxy_server_request_for_spend_logs_payload( + metadata={}, + litellm_params={"proxy_server_request": {"body": {"model": "classifier"}}}, + kwargs={ + "standard_logging_object": audit, + "standard_callback_dynamic_params": {"turn_off_message_logging": redact}, + }, + ) + ) if not store_prompts or redact: assert "classifier_input" not in stored assert "originating_request_masked" not in stored @@ -208,9 +216,7 @@ def test_batch_lifecycle_rows_derive_the_same_session_from_the_batch_id(): from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id create_session: Final = _get_batch_trace_session_id(call_type="acreate_batch", request_id="batch-uid-1") - cost_session: Final = _get_batch_trace_session_id( - call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost" - ) + cost_session: Final = _get_batch_trace_session_id(call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost") assert create_session == cost_session == "batch-uid-1" @@ -2911,7 +2917,11 @@ def test_proxy_server_request_payload_redacts_provider_credentials(mock_should_s "extra_headers": {"Authorization": "Bearer canary-extra-header"}, "tools": [ {"type": "function", "function": {"name": "f", "parameters": tool_parameters}}, - {"type": "mcp", "server_url": "https://mcp.example.com", "headers": {"Authorization": "canary-mcp"}}, + { + "type": "mcp", + "server_url": "https://mcp.example.com", + "headers": {"Authorization": "canary-mcp"}, + }, ], "fallbacks": [{"model": "azure-b", **credentials}], "metadata": metadata, @@ -5465,7 +5475,7 @@ ANTHROPIC_MESSAGES_SSE_CHUNKS: Final = ( 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n', 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' '"usage":{"output_tokens":4}}\n\n', - "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + 'event: message_stop\ndata: {"type":"message_stop"}\n\n', ) @@ -5503,9 +5513,7 @@ def test_spend_log_request_id_is_the_message_id_a_non_streaming_messages_caller_ """ logging_obj = _anthropic_messages_logging_obj(stream=False) - logged_response = logging_obj._handle_anthropic_messages_response_logging( - result=ANTHROPIC_MESSAGES_RESPONSE - ) + logged_response = logging_obj._handle_anthropic_messages_response_logging(result=ANTHROPIC_MESSAGES_RESPONSE) assert logged_response.id == "msg_01Lit6806NonStreaming" assert ( @@ -5581,9 +5589,7 @@ def test_spend_log_request_id_still_falls_back_to_litellm_call_id_without_a_prov end_time=datetime.datetime.now(timezone.utc), logging_obj=logging_obj, ) - assert logging_obj.model_call_details["complete_streaming_response"].id == ( - "6806cafe-0000-4000-8000-000000000001" - ) + assert logging_obj.model_call_details["complete_streaming_response"].id == ("6806cafe-0000-4000-8000-000000000001") def test_spend_log_request_id_for_chat_completions_is_untouched(): @@ -5665,6 +5671,7 @@ def test_failed_agent_request_keeps_registered_display_name(): assert payload["status"] == "failure" assert payload["model_id"] == "registered-agent" + _CLI_SESSION_ALIAS: Final = "cli-session-alice" _CLI_SESSION_TOKEN: Final = "cli-session-Qm7xJ2kP9sLw4vT1nR8yAa" @@ -5796,20 +5803,109 @@ def test_baseline_estimate_metadata_comes_from_the_logging_stamp() -> None: def test_untrusted_agent_label_cannot_replace_verified_billing_identity(billing_agent: str | None) -> None: kwargs = { "model": "gpt-4", - "litellm_params": {"metadata": { - "user_api_key": "test-key", - "agent_id": "header-selected-agent", - "billing_agent_id": billing_agent, - }}, + "litellm_params": { + "metadata": { + "user_api_key": "test-key", + "agent_id": "header-selected-agent", + "billing_agent_id": billing_agent, + } + }, } payload = get_logging_payload( - kwargs=kwargs, response_obj={"id": "request"}, - start_time=datetime.datetime.now(timezone.utc), end_time=datetime.datetime.now(timezone.utc), + kwargs=kwargs, + response_obj={"id": "request"}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), ) assert payload["agent_id"] == "header-selected-agent" assert payload["billing_agent_id"] == billing_agent +def _spend_log_row(metadata: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({"request_id": "req-1", "spend": 0.5, "metadata": json.dumps(dict(metadata))}) + + +_STORED_METADATA: Final = MappingProxyType( + { + "status": "success", + "cold_storage_object_key": "logs/req-1.json", + "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"max_tokens": 10}}, + "usage_object": {"prompt_tokens": 3}, + "user_api_key_alias": "alias", + } +) + + +def test_spend_log_row_keeps_every_metadata_field_when_unconfigured() -> None: + row: Final = _spend_log_row(_STORED_METADATA) + + assert spend_log_row_with_retained_metadata(row, None) is row + + +def test_spend_log_row_drops_excluded_metadata_fields_only() -> None: + row: Final = _spend_log_row(_STORED_METADATA) + + stored: Final = spend_log_row_with_retained_metadata( + row, SpendLogsMetadataFields(exclude=("model_map_information", "user_api_key_alias")) + ) + + assert json.loads(cast(str, stored["metadata"])) == { + "status": "success", + "cold_storage_object_key": "logs/req-1.json", + "usage_object": {"prompt_tokens": 3}, + } + assert {name: value for name, value in stored.items() if name != "metadata"} == { + "request_id": "req-1", + "spend": 0.5, + } + + +def test_spend_log_row_include_keeps_listed_and_always_kept_fields() -> None: + stored: Final = spend_log_row_with_retained_metadata( + _spend_log_row(_STORED_METADATA), SpendLogsMetadataFields(include=("usage_object",)) + ) + + assert json.loads(cast(str, stored["metadata"])) == { + "status": "success", + "cold_storage_object_key": "logs/req-1.json", + "usage_object": {"prompt_tokens": 3}, + } + + +@pytest.mark.parametrize( + "configured", + [ + {"include": ["usage_object"], "exclude": ["model_map_information"]}, + {}, + {"exclude": ["model_map_informaton"]}, + {"include": ["usage_object", "not_a_field"]}, + {"exclude": ["status"]}, + {"exclude": ["cold_storage_object_key"]}, + {"exclude": ["model_map_information"], "drop": ["usage_object"]}, + ], +) +def test_spend_logs_metadata_fields_rejects_ambiguous_or_lossy_config(configured: dict[str, list[str]]) -> None: + from pydantic import ValidationError + + from litellm.proxy._types import ConfigGeneralSettings + + with pytest.raises(ValidationError): + ConfigGeneralSettings.model_validate({"spend_logs_metadata_fields": configured}) + + +def test_configured_spend_logs_metadata_fields_ignores_invalid_runtime_value() -> None: + with patch( + "litellm.proxy.proxy_server.general_settings", + {"spend_logs_metadata_fields": {"include": ["usage_object"], "exclude": ["status"]}}, + ): + assert configured_spend_logs_metadata_fields() is None + with patch( + "litellm.proxy.proxy_server.general_settings", + {"spend_logs_metadata_fields": {"exclude": ["model_map_information"]}}, + ): + assert configured_spend_logs_metadata_fields() == SpendLogsMetadataFields(exclude=("model_map_information",)) + + class _OcrUsageInfoDict(TypedDict, total=False): pages_processed: ReadOnly[int] doc_size_bytes: ReadOnly[int] diff --git a/tests/unit/proxy/test_fallback_management_endpoints.py b/tests/unit/proxy/test_fallback_management_endpoints.py index 054dafbf5a7..3ee1b64580d 100644 --- a/tests/unit/proxy/test_fallback_management_endpoints.py +++ b/tests/unit/proxy/test_fallback_management_endpoints.py @@ -8,17 +8,27 @@ Tests: 4. Validation tests (invalid models, duplicate fallbacks, etc.) """ +import copy +import json +from collections.abc import AsyncIterator +from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException +from litellm import Router +from litellm.proxy.utils import evict_config_param, get_config_param from litellm.proxy.management_endpoints.fallback_management_endpoints import ( FallbackCreateRequest, + FallbackDeleteResponse, + FallbackResponse, create_fallback, delete_fallback, get_fallback, ) +from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict class TestFallbackCreateRequest: @@ -104,6 +114,268 @@ class TestFallbackCreateRequest: assert request.fallback_type == "content_policy" +TEAM_ID: Final = "team-a" +PRIMARY_INTERNAL_NAME: Final = f"model_name_{TEAM_ID}_1a6437cb-4cab-432c-8099-1d7411731a8b" +FALLBACK_INTERNAL_NAME: Final = f"model_name_{TEAM_ID}_83151607-5556-4bbf-ac65-c474dbc64eba" +PRIMARY_DEPLOYMENT_ID: Final = "team-primary-id" +FALLBACK_DEPLOYMENT_ID: Final = "team-fallback-id" + +FallbackRules = list[dict[str, list[str]] | dict[str, object] | str] +RouterSettings = dict[str, FallbackRules | None] + + +def _litellm_params(mock_response: str | None) -> LiteLLMParamsTypedDict: + if mock_response is None: + return {"model": "openai/gpt-5.4-mini", "api_key": "fake"} + return {"model": "openai/gpt-5.4-mini", "api_key": "fake", "mock_response": mock_response} + + +def _team_scoped_deployment( + internal_name: str, public_name: str, deployment_id: str, mock_response: str | None = None +) -> DeploymentTypedDict: + return { + "model_name": internal_name, + "litellm_params": _litellm_params(mock_response), + "model_info": {"id": deployment_id, "team_id": TEAM_ID, "team_public_model_name": public_name}, + } + + +def _team_router(primary_mock_response: str | None = None, fallback_mock_response: str | None = None) -> Router: + return Router( + model_list=[ + {"model_name": "gpt-5.4-mini", "litellm_params": _litellm_params(None)}, + _team_scoped_deployment( + PRIMARY_INTERNAL_NAME, "team-primary", PRIMARY_DEPLOYMENT_ID, primary_mock_response + ), + _team_scoped_deployment( + FALLBACK_INTERNAL_NAME, "team-fallback", FALLBACK_DEPLOYMENT_ID, fallback_mock_response + ), + ], + num_retries=0, + ) + + +class TestCreateFallbackForTeamScopedModels: + """POST /fallback takes the public name a caller invokes a team-scoped model by, not only the stored internal one""" + + @pytest.fixture + def router(self) -> Router: + return _team_router() + + @pytest.fixture + def prisma_client(self) -> MagicMock: + client: Final = MagicMock() + client.db.litellm_config.upsert = AsyncMock() + return client + + @pytest.fixture + def proxy_config(self) -> MagicMock: + config: Final = MagicMock() + config.get_config = AsyncMock(return_value={"router_settings": {}}) + return config + + async def _create( + self, request: FallbackCreateRequest, router: Router, prisma_client: MagicMock, proxy_config: MagicMock + ) -> FallbackResponse: + with ( + patch("litellm.proxy.proxy_server.llm_router", router), + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + return await create_fallback(request, MagicMock()) + + async def test_public_names_create_a_rule_keyed_on_the_public_name( + self, router: Router, prisma_client: MagicMock, proxy_config: MagicMock + ) -> None: + request: Final = FallbackCreateRequest(model="team-primary", fallback_models=["team-fallback"]) + + response: Final = await self._create(request, router, prisma_client, proxy_config) + + assert response.model == "team-primary" + assert response.fallback_models == ["team-fallback"] + assert router.fallbacks == [{"team-primary": ["team-fallback"]}] + persisted: Final = json.loads( + prisma_client.db.litellm_config.upsert.call_args.kwargs["data"]["create"]["param_value"] + ) + assert persisted["fallbacks"] == [{"team-primary": ["team-fallback"]}] + + async def test_internal_names_keep_working( + self, router: Router, prisma_client: MagicMock, proxy_config: MagicMock + ) -> None: + request: Final = FallbackCreateRequest(model=PRIMARY_INTERNAL_NAME, fallback_models=[FALLBACK_INTERNAL_NAME]) + + response: Final = await self._create(request, router, prisma_client, proxy_config) + + assert router.fallbacks == [{PRIMARY_INTERNAL_NAME: [FALLBACK_INTERNAL_NAME]}] + assert response.model == PRIMARY_INTERNAL_NAME + + async def test_public_fallback_target_behind_a_gateway_primary( + self, router: Router, prisma_client: MagicMock, proxy_config: MagicMock + ) -> None: + request: Final = FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]) + + await self._create(request, router, prisma_client, proxy_config) + + assert router.fallbacks == [{"gpt-5.4-mini": ["team-fallback"]}] + + async def test_unknown_name_is_rejected_and_the_error_names_the_public_names( + self, router: Router, prisma_client: MagicMock, proxy_config: MagicMock + ) -> None: + request: Final = FallbackCreateRequest(model="team-missing", fallback_models=["team-fallback"]) + + with pytest.raises(HTTPException) as exc_info: + await self._create(request, router, prisma_client, proxy_config) + + assert exc_info.value.status_code == 404 + assert {"team-primary", "team-fallback", "gpt-5.4-mini"} <= set(exc_info.value.detail["available_models"]) + + async def test_a_team_request_fails_over_to_the_public_name_fallback( + self, prisma_client: MagicMock, proxy_config: MagicMock + ) -> None: + router: Final = _team_router(primary_mock_response="litellm.RateLimitError", fallback_mock_response="pong") + request: Final = FallbackCreateRequest(model="team-primary", fallback_models=["team-fallback"]) + await self._create(request, router, prisma_client, proxy_config) + + response: Final = await router.acompletion( + model="team-primary", + messages=[{"role": "user", "content": "ping"}], + metadata={"user_api_key_team_id": TEAM_ID}, + ) + + assert response.choices[0].message.content == "pong" + assert response._hidden_params["model_id"] == FALLBACK_DEPLOYMENT_ID + + +def _config_row(router_settings: RouterSettings) -> SimpleNamespace: + return SimpleNamespace(param_name="router_settings", param_value=router_settings) + + +class _StoredRouterSettings: + """The LiteLLM_Config router_settings row, with the proxy objects that read it through the config cache""" + + def __init__(self, router_settings: RouterSettings | None) -> None: + self.row: SimpleNamespace | None = None if router_settings is None else _config_row(router_settings) + self.prisma_client: Final = MagicMock() + self.prisma_client.get_generic_data = AsyncMock(side_effect=lambda **_: self.row) + self.prisma_client.db.litellm_config.upsert = AsyncMock(side_effect=self._upsert) + self.proxy_config: Final = MagicMock() + self.proxy_config.get_config = AsyncMock(side_effect=self._get_config) + + async def _upsert(self, where: dict[str, str], data: dict[str, dict[str, str]]) -> None: + self.row = _config_row(json.loads(data["update"]["param_value"])) + + async def _get_config(self) -> dict[str, RouterSettings]: + row: Final = await get_config_param(self.prisma_client, "router_settings") + return {"router_settings": copy.deepcopy(row.param_value) if row is not None else {}} + + def written_by_another_instance(self, router_settings: RouterSettings) -> None: + self.row = _config_row(router_settings) + + def stored_fallbacks(self) -> FallbackRules: + assert self.row is not None + return self.row.param_value["fallbacks"] + + async def cached_fallbacks(self) -> FallbackRules: + row: Final = await get_config_param(self.prisma_client, "router_settings") + return row.param_value["fallbacks"] + + +TEAM_RULE: Final = {"team-primary": ["team-fallback"]} +GATEWAY_RULE: Final = {"gpt-5.4-mini": ["team-fallback"]} +MALFORMED_GATEWAY_RULE: Final = {"gpt-5.4-mini": "team-fallback"} +NON_STANDARD_RULES: Final = [ + "claude-3-haiku", + {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "retry"}]}, +] + + +class TestFallbackWritesSeeTheLatestStoredRules: + """A write reads the rules the database holds now and leaves no stale copy in the config cache behind""" + + @pytest.fixture(autouse=True) + async def clean_config_cache(self) -> AsyncIterator[None]: + await evict_config_param("router_settings") + yield + await evict_config_param("router_settings") + + async def _create(self, request: FallbackCreateRequest, stored: _StoredRouterSettings) -> FallbackResponse: + with ( + patch("litellm.proxy.proxy_server.llm_router", _team_router()), + patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), + patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + return await create_fallback(request, MagicMock()) + + async def _delete(self, model: str, stored: _StoredRouterSettings) -> FallbackDeleteResponse: + with ( + patch("litellm.proxy.proxy_server.llm_router", _team_router()), + patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), + patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + return await delete_fallback(model, "general", MagicMock()) + + async def test_a_second_create_keeps_the_first_rule(self) -> None: + stored: Final = _StoredRouterSettings(None) + + await self._create(FallbackCreateRequest(model="team-primary", fallback_models=["team-fallback"]), stored) + await self._create(FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), stored) + + assert stored.stored_fallbacks() == [TEAM_RULE, GATEWAY_RULE] + assert await stored.cached_fallbacks() == [TEAM_RULE, GATEWAY_RULE] + + async def test_a_create_keeps_a_rule_another_instance_stored_since_this_one_last_read(self) -> None: + stored: Final = _StoredRouterSettings({}) + await get_config_param(stored.prisma_client, "router_settings") + stored.written_by_another_instance({"fallbacks": [TEAM_RULE]}) + + await self._create(FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), stored) + + assert stored.stored_fallbacks() == [TEAM_RULE, GATEWAY_RULE] + + async def test_a_delete_keeps_a_rule_another_instance_stored_since_this_one_last_read(self) -> None: + stored: Final = _StoredRouterSettings({"fallbacks": [TEAM_RULE]}) + await get_config_param(stored.prisma_client, "router_settings") + stored.written_by_another_instance({"fallbacks": [TEAM_RULE, GATEWAY_RULE]}) + + await self._delete("team-primary", stored) + + assert stored.stored_fallbacks() == [GATEWAY_RULE] + assert await stored.cached_fallbacks() == [GATEWAY_RULE] + + async def test_writes_keep_the_non_standard_rules_the_router_accepts(self) -> None: + stored: Final = _StoredRouterSettings({"fallbacks": [*NON_STANDARD_RULES, TEAM_RULE]}) + + await self._create(FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), stored) + assert stored.stored_fallbacks() == [*NON_STANDARD_RULES, TEAM_RULE, GATEWAY_RULE] + + await self._delete("team-primary", stored) + assert stored.stored_fallbacks() == [*NON_STANDARD_RULES, GATEWAY_RULE] + + async def test_a_create_replaces_a_same_key_rule_whatever_its_targets_shape(self) -> None: + stored: Final = _StoredRouterSettings({"fallbacks": [MALFORMED_GATEWAY_RULE, TEAM_RULE]}) + + await self._create(FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), stored) + + assert stored.stored_fallbacks() == [GATEWAY_RULE, TEAM_RULE] + + async def test_a_delete_removes_a_same_key_rule_whatever_its_targets_shape(self) -> None: + stored: Final = _StoredRouterSettings({"fallbacks": [MALFORMED_GATEWAY_RULE, TEAM_RULE]}) + + await self._delete("gpt-5.4-mini", stored) + + assert stored.stored_fallbacks() == [TEAM_RULE] + + async def test_a_create_treats_a_null_rule_list_as_empty(self) -> None: + stored: Final = _StoredRouterSettings({"fallbacks": None}) + + await self._create(FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), stored) + + assert stored.stored_fallbacks() == [GATEWAY_RULE] + + @pytest.mark.asyncio class TestCreateFallback: """Test the create_fallback endpoint""" @@ -113,6 +385,7 @@ class TestCreateFallback: """Create a mock router""" router = MagicMock() router.model_names = {"gpt-3.5-turbo", "gpt-4", "claude-3-haiku"} + router.team_public_model_names = frozenset() router.fallbacks = [] router.context_window_fallbacks = [] router.content_policy_fallbacks = [] diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index c6c473ae71e..818ca50fed9 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -16,6 +16,7 @@ from pydantic import ValidationError as PydanticValidationError from starlette.datastructures import Headers import litellm +from litellm.constants import SERVER_STREAMING_CLASSIFICATION_KEY, SERVER_STREAMING_CLASSIFICATION_MARKER from litellm.proxy._types import AddTeamCallback, ProxyException, TeamCallbackMetadata, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import ( KeyAndTeamLoggingSettings, @@ -850,8 +851,13 @@ def test_initial_snapshot_refresh_clears_a_previous_guardrail_checkpoint() -> No from litellm.proxy.litellm_pre_call_utils import refresh_proxy_server_request_body_snapshot logging_obj: Final = Logging( - model="test-model", messages=[], stream=False, call_type="acompletion", - start_time=datetime.now(), litellm_call_id="new-request", function_id="new-request", + model="test-model", + messages=[], + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="new-request", + function_id="new-request", ) logging_obj.shadow_eval_request_snapshot = GuardrailRequestSnapshot.capture( {"messages": [{"role": "user", "content": "previous request"}]}, @@ -913,8 +919,14 @@ async def test_post_guardrail_snapshot_preserves_logging_only_masking_in_spend_l } data: Final = {"messages": messages, "metadata": metadata, "proxy_server_request": {}} logging_obj: Final = Logging( - model="test-model", messages=messages, stream=False, call_type="acompletion", - start_time=datetime.now(), litellm_call_id="mask-spend", function_id="mask-spend", kwargs=data, + model="test-model", + messages=messages, + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="mask-spend", + function_id="mask-spend", + kwargs=data, ) data["litellm_logging_obj"] = logging_obj refresh_proxy_server_request_body_snapshot(data, guardrails_applied=True) @@ -928,9 +940,13 @@ async def test_post_guardrail_snapshot_preserves_logging_only_masking_in_spend_l kwargs, _ = await guardrail.async_logging_hook( kwargs=logging_obj.model_call_details, result=None, call_type="acompletion" ) - stored: Final = json.loads(_get_proxy_server_request_for_spend_logs_payload( - metadata={}, litellm_params=kwargs["litellm_params"], kwargs=kwargs, - )) + stored: Final = json.loads( + _get_proxy_server_request_for_spend_logs_payload( + metadata={}, + litellm_params=kwargs["litellm_params"], + kwargs=kwargs, + ) + ) assert kwargs["messages"] == [{"role": "user", "content": "email [EMAIL]"}] assert stored["messages"] == kwargs["messages"] @@ -1089,6 +1105,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "messages": [{"role": "user", "content": "hello"}], "mock_response": "free response", "mock_tool_calls": [{"id": "call_1"}], + "is_streaming_request": "caller-value", "disable_global_guardrails": True, "enable_prompt_caching": True, "routing_decision": {"cause": "forged", "routed_model": "spoofed"}, @@ -1114,6 +1131,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): assert "enable_prompt_caching" not in updated assert "routing_decision" not in updated assert "litellm_gateway_injected_cache" not in updated + assert updated["is_streaming_request"] == "caller-value" assert "weights" not in updated assert "_router_weights" not in updated assert "weights" not in updated["proxy_server_request"]["body"] @@ -8803,20 +8821,27 @@ async def test_mcp_credentials_only_removed_from_logging_copies(path: str, custo request.headers = Headers(request.headers) settings: Final = {"mcp_client_side_auth_header_name": custom_auth, "user_header_name": "x-user-id"} server: Final = MCPServer( - server_id="header-test", name="header-test", transport="http", url="https://example.com/mcp", + server_id="header-test", + name="header-test", + transport="http", + url="https://example.com/mcp", extra_headers=["x-service-token", "x-user-id"], ) with ( patch("litellm.proxy.proxy_server.general_settings", settings), patch.dict( "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager.config_mcp_servers", - {"header-test": server}, clear=True, + {"header-test": server}, + clear=True, ), ): updated: Final = await add_litellm_data_to_request( data={"model": "test-model", "messages": [{"role": "user", "content": "hello"}]}, - request=request, user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), - proxy_config=MagicMock(), general_settings=settings, version="test", + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings=settings, + version="test", ) for header_dict in _all_header_dicts(updated, metadata_name): assert not any(value in json.dumps(header_dict) for value in secrets.values()) @@ -8836,7 +8861,10 @@ def test_signoz_callback_vars_are_scoped_to_the_signoz_callback(): data=AddTeamCallback( callback_name="signoz", callback_type="success", - callback_vars={"signoz_ingestion_key": "team-key", "signoz_ingestion_endpoint": "https://ingest.eu.signoz.cloud:443"}, + callback_vars={ + "signoz_ingestion_key": "team-key", + "signoz_ingestion_endpoint": "https://ingest.eu.signoz.cloud:443", + }, ), team_callback_settings_obj=None, ) @@ -8856,6 +8884,54 @@ def test_signoz_callback_vars_are_scoped_to_the_signoz_callback(): assert under_other.callback_vars == {"langfuse_host": "https://cloud.langfuse.com"} +def test_body_snapshot_excludes_the_server_streaming_marker() -> None: + from litellm.constants import SERVER_STREAMING_CLASSIFICATION_KEY, SERVER_STREAMING_CLASSIFICATION_MARKER + from litellm.proxy.litellm_pre_call_utils import refresh_proxy_server_request_body_snapshot + + proxy_request: Final = {"body": {}} + data: Final = { + "messages": [{"role": "user", "content": "hi"}], + SERVER_STREAMING_CLASSIFICATION_KEY: SERVER_STREAMING_CLASSIFICATION_MARKER, + "proxy_server_request": proxy_request, + } + + refresh_proxy_server_request_body_snapshot(data) + + assert proxy_request == {"body": {"messages": [{"role": "user", "content": "hi"}]}} + + +@pytest.mark.parametrize( + "marker", + [ + SERVER_STREAMING_CLASSIFICATION_MARKER, + json.loads(json.dumps(SERVER_STREAMING_CLASSIFICATION_MARKER)), + ], + ids=["enum", "json-string"], +) +def test_body_snapshot_drops_only_the_marker_and_keeps_caller_value(marker: str) -> None: + from litellm.proxy.litellm_pre_call_utils import refresh_proxy_server_request_body_snapshot + + marker_request: Final = {"body": {}} + marker_data: Final = { + "messages": [{"role": "user", "content": "hi"}], + SERVER_STREAMING_CLASSIFICATION_KEY: marker, + "proxy_server_request": marker_request, + } + + refresh_proxy_server_request_body_snapshot(marker_data) + + assert SERVER_STREAMING_CLASSIFICATION_KEY not in marker_request["body"], marker_request + + caller_request: Final = {"body": {}} + caller_data: Final = { + "messages": [{"role": "user", "content": "hi"}], + SERVER_STREAMING_CLASSIFICATION_KEY: "caller-value", + "proxy_server_request": caller_request, + } + + refresh_proxy_server_request_body_snapshot(caller_data) + + assert caller_request["body"][SERVER_STREAMING_CLASSIFICATION_KEY] == "caller-value", caller_request def test_arize_otlp_protocol_on_a_key_logging_entry_reaches_the_destination(monkeypatch): from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.proxy.litellm_pre_call_utils import resolve_tenant_otel_destinations diff --git a/tests/unit/proxy/test_pricing_field_strip.py b/tests/unit/proxy/test_pricing_field_strip.py index a0e25e91f37..395fdc3a16e 100644 --- a/tests/unit/proxy/test_pricing_field_strip.py +++ b/tests/unit/proxy/test_pricing_field_strip.py @@ -26,7 +26,6 @@ from litellm.proxy.litellm_pre_call_utils import ( from litellm.types.utils import CustomPricingLiteLLMParams - def _make_request_mock() -> Request: request_mock = MagicMock(spec=Request) request_mock.url.path = "/v1/chat/completions" @@ -58,9 +57,7 @@ class TestStripClientPricingOverrides: # The strip set is built from the model so additions are picked up # automatically — this test guards against the model and the strip # set drifting apart if someone replaces the auto-derivation later. - assert _CLIENT_PRICING_CONTROL_FIELDS == frozenset( - CustomPricingLiteLLMParams.model_fields.keys() - ) + assert _CLIENT_PRICING_CONTROL_FIELDS == frozenset(CustomPricingLiteLLMParams.model_fields.keys()) # Sanity: the obvious top-level pricing fields are in the set. for field in ( "input_cost_per_token", @@ -184,9 +181,7 @@ class TestStripClientPricingOverrides: verbose_proxy_logger.setLevel(logging.DEBUG) with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): _strip_client_pricing_overrides({"model": "gpt-4", "temperature": 0.7}) - assert not any( - "pricing" in record.getMessage().lower() for record in caplog.records - ) + assert not any("pricing" in record.getMessage().lower() for record in caplog.records) @pytest.mark.asyncio @@ -211,6 +206,26 @@ async def test_add_litellm_data_to_request_strips_root_pricing_fields(): assert "output_cost_per_token" not in updated +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_preserves_caller_streaming_request(): + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "is_streaming_request": True, + } + + updated = await add_litellm_data_to_request( + data=data, + request=_make_request_mock(), + user_api_key_dict=_user_api_key_auth(), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated["is_streaming_request"] is True + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_strips_client_disconnect_metadata(): data = { @@ -318,9 +333,7 @@ async def test_add_litellm_data_to_request_skips_strip_with_team_opt_in(): "input_cost_per_token": 0.0001, } - user_auth = _user_api_key_auth( - team_metadata={"allow_client_pricing_override": True} - ) + user_auth = _user_api_key_auth(team_metadata={"allow_client_pricing_override": True}) updated = await add_litellm_data_to_request( data=data, request=_make_request_mock(), diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index e7a32816464..e4c4ddce64f 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from dotenv import load_dotenv +from pydantic import JsonValue load_dotenv() @@ -20,6 +21,7 @@ from litellm import Router from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.anthropic_endpoints.endpoints import ( count_tokens as anthropic_count_tokens, @@ -753,41 +755,45 @@ def test_vertex_ai_partner_models_token_counting_endpoint(vertex_location): ) +class _FailingRuntimeHandler(BedrockCountTokensHandler): + def __init__(self, error: Exception) -> None: + super().__init__() + self._error = error + + async def handle_count_tokens_request( + self, + request_data: dict[str, object], + litellm_params: dict[str, object], + resolved_model: str, + client: AsyncHTTPHandler | None = None, + ) -> dict[str, JsonValue]: + raise self._error + + @pytest.mark.asyncio async def test_bedrock_token_counter_error_propagation_bedrock_error(): """ Test that BedrockTokenCounter properly returns error response when BedrockError is raised. Verifies that the status code and error message are preserved. """ - counter = BedrockTokenCounter() + counter = BedrockTokenCounter( + runtime_handler=_FailingRuntimeHandler(BedrockError(status_code=429, message="Rate limit exceeded")) + ) - # Mock the handler to raise BedrockError with specific status code - with patch.object( - counter, "count_tokens", wraps=counter.count_tokens - ) as mock_count: - # We need to patch at the handler level - with patch( - "litellm.llms.bedrock.count_tokens.bedrock_token_counter.BedrockCountTokensHandler" - ) as MockHandler: - mock_handler_instance = MockHandler.return_value - mock_handler_instance.handle_count_tokens_request = AsyncMock( - side_effect=BedrockError(status_code=429, message="Rate limit exceeded") - ) + result = await counter.count_tokens( + model_to_use="anthropic.claude-3-sonnet", + messages=[{"role": "user", "content": "hello"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="bedrock/anthropic.claude-3-sonnet", + ) - result = await counter.count_tokens( - model_to_use="anthropic.claude-3-sonnet", - messages=[{"role": "user", "content": "hello"}], - contents=None, - deployment={"litellm_params": {}}, - request_model="bedrock/anthropic.claude-3-sonnet", - ) - - assert result is not None - assert result.error is True - assert result.status_code == 429 - assert "Rate limit exceeded" in result.error_message - assert result.tokenizer_type == "bedrock_api" - assert result.total_tokens == 0 + assert result is not None + assert result.error is True + assert result.status_code == 429 + assert "Rate limit exceeded" in result.error_message + assert result.tokenizer_type == "bedrock_api" + assert result.total_tokens == 0 @pytest.mark.asyncio @@ -795,28 +801,20 @@ async def test_bedrock_token_counter_error_propagation_generic_exception(): """ Test that BedrockTokenCounter returns error response with 500 status for generic exceptions. """ - counter = BedrockTokenCounter() + counter = BedrockTokenCounter(runtime_handler=_FailingRuntimeHandler(Exception("Unexpected error"))) - with patch( - "litellm.llms.bedrock.count_tokens.bedrock_token_counter.BedrockCountTokensHandler" - ) as MockHandler: - mock_handler_instance = MockHandler.return_value - mock_handler_instance.handle_count_tokens_request = AsyncMock( - side_effect=Exception("Unexpected error") - ) + result = await counter.count_tokens( + model_to_use="anthropic.claude-3-sonnet", + messages=[{"role": "user", "content": "hello"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="bedrock/anthropic.claude-3-sonnet", + ) - result = await counter.count_tokens( - model_to_use="anthropic.claude-3-sonnet", - messages=[{"role": "user", "content": "hello"}], - contents=None, - deployment={"litellm_params": {}}, - request_model="bedrock/anthropic.claude-3-sonnet", - ) - - assert result is not None - assert result.error is True - assert result.status_code == 500 - assert "Unexpected error" in result.error_message + assert result is not None + assert result.error is True + assert result.status_code == 500 + assert "Unexpected error" in result.error_message @pytest.mark.asyncio diff --git a/tests/unit/proxy/test_proxy_utils.py b/tests/unit/proxy/test_proxy_utils.py index ba61e393281..92bb163d6ec 100644 --- a/tests/unit/proxy/test_proxy_utils.py +++ b/tests/unit/proxy/test_proxy_utils.py @@ -2,7 +2,7 @@ import asyncio import json import os from datetime import datetime -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, Final, List, Optional, Union from unittest.mock import Mock import pytest @@ -3329,3 +3329,22 @@ def test_handle_exception_on_proxy_preserves_auth_error_status_code(): result = handle_exception_on_proxy(auth_error) assert int(result.code) == 401, f"Expected 401, got {result.code}" + + +class _RedisDown: + async def async_delete_cache(self, key: str) -> None: + raise ConnectionError("redis is down") + + +async def test_evict_config_param_clears_the_local_layer_and_survives_a_redis_outage() -> None: + from litellm.caching.caching import DualCache + from litellm.proxy.utils import _config_cache_key, evict_config_param + + cache: Final = DualCache(redis_cache=_RedisDown()) + await cache.in_memory_cache.async_set_cache( + _config_cache_key("router_settings"), {"param_name": "router_settings", "param_value": {"fallbacks": []}} + ) + + await evict_config_param("router_settings", cache=cache) + + assert await cache.in_memory_cache.async_get_cache(_config_cache_key("router_settings")) is None diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_config_param_cache.py b/tests/unit/proxy/utils/prisma_and_spend/test_config_param_cache.py index 0d408de9ec6..cdaf5c1ee3b 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_config_param_cache.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_config_param_cache.py @@ -12,8 +12,9 @@ Symbols pinned here: from __future__ import annotations +import logging from types import SimpleNamespace -from typing import Any, List +from typing import Any, Final, List from unittest.mock import AsyncMock, MagicMock import pytest @@ -214,14 +215,23 @@ async def test_invalidate_config_param_evicts_from_cache( @pytest.mark.asyncio -async def test_invalidate_config_param_propagates_cache_error( - _swap_config_cache: Any, +async def test_invalidate_config_param_survives_a_cache_error( + _swap_config_cache: Any, caplog: pytest.LogCaptureFixture ) -> None: _swap_config_cache.async_delete_cache = AsyncMock( side_effect=ConnectionError("redis down") ) - with pytest.raises(ConnectionError): + caplog.set_level(logging.WARNING, logger="LiteLLM Proxy") + utils_mod.verbose_proxy_logger.addHandler(caplog.handler) + try: await invalidate_config_param("p5") + finally: + utils_mod.verbose_proxy_logger.removeHandler(caplog.handler) + actual: Final = { + "delete_calls": _swap_config_cache.async_delete_cache.await_count, + "warned": "config cache eviction of p5 failed: redis down" in caplog.text, + } + assert actual == {"delete_calls": 1, "warned": True} @pytest.mark.asyncio diff --git a/tests/unit/rust_bridge/chat_completions/test_route_host.py b/tests/unit/rust_bridge/chat_completions/test_route_host.py index 92eed17de0f..402c032f340 100644 --- a/tests/unit/rust_bridge/chat_completions/test_route_host.py +++ b/tests/unit/rust_bridge/chat_completions/test_route_host.py @@ -4,7 +4,7 @@ from typing import Final import pytest import litellm -from litellm.rust_bridge.chat_completions.route_host import arguments, connection_defaults, response +from litellm.rust_bridge.chat_completions.route_host import connection_defaults, response from litellm.types.utils import ModelResponse @@ -35,12 +35,6 @@ def test_response_builds_the_public_model_response() -> None: assert built.usage.total_tokens == 5 -def test_arguments_are_the_public_kwargs_view() -> None: - kwargs: Final = MappingProxyType({"metadata": {"user_id": "u"}}) - - assert arguments(kwargs) is kwargs - - @pytest.mark.parametrize( ("provider", "global_key", "provider_key", "expected_key", "expected_base"), ( diff --git a/tests/unit/rust_bridge/messages/test_route_host.py b/tests/unit/rust_bridge/messages/test_route_host.py index dde76ee5e82..16b9634beb9 100644 --- a/tests/unit/rust_bridge/messages/test_route_host.py +++ b/tests/unit/rust_bridge/messages/test_route_host.py @@ -1,8 +1,7 @@ from types import MappingProxyType from typing import Final -from litellm.rust_bridge.messages.route_host import arguments, response -from litellm.rust_bridge.public_call import NativeCall +from litellm.rust_bridge.messages.route_host import response import pytest import litellm from litellm.rust_bridge.messages import route_host @@ -29,26 +28,6 @@ def test_response_is_a_detached_public_messages_dict() -> None: assert "_hidden_params" not in native -def test_arguments_preserve_the_bound_view() -> None: - kwargs: Final = MappingProxyType({"litellm_metadata": {"user_id": "u"}}) - request: Final = NativeCall( - args=(), - kwargs=kwargs, - bound={ - "model": "claude-sonnet-4-5", - "messages": [{"role": "user", "content": "hi"}], - "max_tokens": 16, - "stream": None, - "api_key": None, - "api_base": None, - "custom_llm_provider": "anthropic", - **kwargs, - }, - ) - - assert arguments(request.bound) is request.bound - - def test_settings_project_caller_configuration_without_resolving_a_model(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "drop_params", False) monkeypatch.setattr(litellm, "reasoning_auto_summary", True) diff --git a/tests/unit/rust_bridge/native_route_wheel_test.py b/tests/unit/rust_bridge/native_route_wheel_test.py index 9e7523aa29e..268eca0f663 100644 --- a/tests/unit/rust_bridge/native_route_wheel_test.py +++ b/tests/unit/rust_bridge/native_route_wheel_test.py @@ -20,13 +20,6 @@ from typing import Final REQUEST_STARTED: Final = threading.Event() REQUEST_CANCELLED: Final = threading.Event() -ANTHROPIC_RESPONSE: Final = ( - b'{"id":"msg_native","type":"message","role":"assistant",' - b'"model":"claude-sonnet-4-5","content":[{"type":"text","text":"native-message"}],' - b'"stop_reason":"end_turn","stop_sequence":null,' - b'"usage":{"input_tokens":2,"output_tokens":3}}' -) - class NativeRouteServer(ThreadingHTTPServer): request_queue_size = 64 @@ -49,7 +42,7 @@ class NativeRouteHandler(BaseHTTPRequestHandler): return status: Final = 429 if outcome == "429" else 200 - response_body: Final = native_response(status, route) + response_body: Final = native_response(status) self.send_response(status) self.send_header("content-type", "application/json") @@ -78,32 +71,23 @@ def assert_native_request( headers: HTTPMessage, body: object, ) -> None: - if route not in {"transcription", "chat_completions"}: + if route != "transcription": raise AssertionError(f"unexpected route marker: {route!r}") if outcome not in {"success", "429", "hang"}: raise AssertionError(f"unexpected outcome marker: {outcome!r}") if not isinstance(body, dict): raise TypeError(f"{route} sent {type(body).__name__}, expected a JSON object") - if route == "transcription": - assert path == "/model/mistral.voxtral-mini-3b-2507/converse" - assert headers.get("authorization", "").startswith("AWS4-HMAC-SHA256 ") - assert headers.get("x-amz-date") - assert body["messages"][0]["content"][0]["audio"]["source"]["bytes"] == "AQI=" - assert "The audio language is en" in body["messages"][0]["content"][1]["text"] - return - assert path == "/v1/messages" - assert headers.get("x-api-key") == "sk-native" - assert body["model"] == "claude-sonnet-4-5" - assert body["max_tokens"] == 17 - assert body["messages"][0]["content"] == [{"type": "text", "text": "hello-from-chat"}] + assert path == "/model/mistral.voxtral-mini-3b-2507/converse" + assert headers.get("authorization", "").startswith("AWS4-HMAC-SHA256 ") + assert headers.get("x-amz-date") + assert body["messages"][0]["content"][0]["audio"]["source"]["bytes"] == "AQI=" + assert "The audio language is en" in body["messages"][0]["content"][1]["text"] -def native_response(status: int, route: str | None) -> bytes: +def native_response(status: int) -> bytes: if status == 429: return b'{"error":"native-rate-limit"}' - if route == "transcription": - return b'{"output":{"message":{"content":[{"text":"native-transcription"}]}}}' - return ANTHROPIC_RESPONSE + return b'{"output":{"message":{"content":[{"text":"native-transcription"}]}}}' def load_native(native_path: Path) -> object: @@ -138,38 +122,23 @@ def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]: "language": "en", }, } - if route == "chat_completions": - return common | { - "model": "anthropic/claude-sonnet-4-5", - "messages": [{"role": "user", "content": "hello-from-chat"}], - "optional_params": {"max_tokens": 17}, - "api_key": "sk-native", - } raise AssertionError(f"unknown route: {route}") def assert_success(route: str, response: object) -> None: if not isinstance(response, dict): raise TypeError(f"{route} returned {type(response).__name__}, expected dict") - actual: Final = success_value(route, response) - expected: Final = "native-transcription" if route == "transcription" else "native-message" - if actual != expected: - raise AssertionError(f"{route} returned {actual!r}, expected {expected!r}") - - -def success_value(route: str, response: dict[object, object]) -> object: - if route == "transcription": - return response["text"] - return response["choices"][0]["message"]["content"] + if response["text"] != "native-transcription": + raise AssertionError(f"{route} returned {response['text']!r}, expected 'native-transcription'") def assert_rate_limit(route: str, error: BaseException) -> None: - if error.args != (429, native_response(429, route).decode()): + if error.args != (429, native_response(429).decode()): raise AssertionError(f"{route} returned the wrong 429 error: {error!r}") def exercise_sync(native: object, api_base: str) -> None: - for route in ("transcription", "chat_completions"): + for route in ("transcription",): function: Final = getattr(native, route) assert_success(route, function(route_call(route, api_base, "success"))) try: @@ -181,7 +150,7 @@ def exercise_sync(native: object, api_base: str) -> None: async def exercise_async(native: object, api_base: str) -> None: - for route in ("transcription", "chat_completions"): + for route in ("transcription",): function: Final = getattr(native, f"a{route}") assert_success(route, await function(route_call(route, api_base, "success"))) try: @@ -192,13 +161,11 @@ async def exercise_async(native: object, api_base: str) -> None: raise AssertionError(f"a{route} accepted a 429 response") responses: Final = await asyncio.wait_for( - asyncio.gather( - *(native.achat_completions(route_call("chat_completions", api_base, "success")) for _ in range(32)) - ), + asyncio.gather(*(native.atranscription(route_call("transcription", api_base, "success")) for _ in range(32))), timeout=15, ) for response in responses: - assert_success("chat_completions", response) + assert_success("transcription", response) def exercise_routes(native_path: Path, api_base: str) -> object: @@ -210,8 +177,8 @@ def exercise_routes(native_path: Path, api_base: str) -> object: def exercise_signal(native: object, api_base: str) -> int: try: - native.chat_completions( - route_call("chat_completions", api_base, "hang"), + native.transcription( + route_call("transcription", api_base, "hang"), ) except KeyboardInterrupt: sys.stdout.write("KeyboardInterrupt\n") diff --git a/tests/unit/rust_bridge/responses/test_route_host.py b/tests/unit/rust_bridge/responses/test_route_host.py index 1d67f2c368c..e110f16350b 100644 --- a/tests/unit/rust_bridge/responses/test_route_host.py +++ b/tests/unit/rust_bridge/responses/test_route_host.py @@ -5,8 +5,7 @@ import pytest from pydantic import ValidationError import litellm -from litellm.rust_bridge.responses.route_host import arguments, connection_defaults, map_failure, response -from litellm.rust_bridge.public_call import NativeCall +from litellm.rust_bridge.responses.route_host import connection_defaults, map_failure, response from litellm.types.llms.openai import ResponsesAPIResponse @@ -42,26 +41,6 @@ def test_response_rejects_a_payload_missing_required_fields() -> None: response(MappingProxyType({"object": "response"})) -def test_arguments_preserve_the_bound_view() -> None: - kwargs: Final = MappingProxyType({"litellm_metadata": {"user_id": "u"}}) - request: Final = NativeCall( - args=(), - kwargs=kwargs, - bound={ - "model": "gpt-4o", - "input": "hi", - "stream": None, - "api_key": None, - "api_base": None, - "custom_llm_provider": "openai", - "extra_headers": None, - **kwargs, - }, - ) - - assert arguments(request.bound) is request.bound - - @pytest.mark.parametrize( ("global_key", "provider_key", "expected"), ( diff --git a/tests/unit/test_router/test_router_callback_hook_sequence.py b/tests/unit/test_router/test_router_callback_hook_sequence.py new file mode 100644 index 00000000000..a7c41964950 --- /dev/null +++ b/tests/unit/test_router/test_router_callback_hook_sequence.py @@ -0,0 +1,509 @@ +import asyncio +import inspect +import json +from collections import Counter +from collections.abc import Callable, Mapping, Sequence +from datetime import datetime +from typing import Final, Literal, NamedTuple, TypedDict + +from typing_extensions import ReadOnly + +import httpx +import pytest +import respx +from pydantic import TypeAdapter + +import litellm +from litellm.caching.caching import Cache +from litellm.integrations.custom_logger import CustomLogger +from litellm.router import Router + +_PRIMARY: Final = "https://hooks-primary.openai.azure.com" +_FALLBACK: Final = "https://hooks-fallback.openai.azure.com" +_API_VERSION: Final = "2024-10-21" +_MESSAGES: Final = [{"role": "user", "content": "Hi - i'm openai"}] +_USAGE: Final = {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7} +_COMPLETION: Final = { + "id": "chatcmpl-hooks", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4.1-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop"}], + "usage": _USAGE, +} +_EMBEDDING_VECTOR: Final = [0.1, 0.2, 0.3] +_EMBEDDING: Final = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": _EMBEDDING_VECTOR}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 3, "total_tokens": 3}, +} +_STREAM: Final = ( + "".join( + f"data: {json.dumps(chunk)}\n\n" + for chunk in ( + { + "id": "chatcmpl-hooks", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-4.1-mini", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hel"}, "finish_reason": None}], + }, + { + "id": "chatcmpl-hooks", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-4.1-mini", + "choices": [{"index": 0, "delta": {"content": "lo"}, "finish_reason": "stop"}], + }, + ) + ) + + "data: [DONE]\n\n" +) +_AUTH_ERROR: Final = httpx.Response( + 401, + json={ + "error": {"message": "Incorrect API key provided", "type": "invalid_request_error", "code": "invalid_api_key"} + }, +) + +_OUR_MODEL_GROUPS: Final = frozenset({"hooks-group", "primary-group", "fallback-group"}) + +_State = Literal[ + "sync_pre_api_call", + "post_api_call", + "async_stream", + "sync_success", + "async_success", + "sync_failure", + "async_failure", +] + + +class _HookEvent(NamedTuple): + state: _State + model: object + kwargs: Mapping[str, object] + response: object + + +def _router_context_problems(kwargs: Mapping[str, object]) -> tuple[str, ...]: + litellm_params: Final = kwargs.get("litellm_params") + if not isinstance(litellm_params, dict): + return ("litellm_params",) + metadata: Final = litellm_params.get("metadata") + model_info: Final = litellm_params.get("model_info") + checks: Final = { + "metadata": isinstance(metadata, dict), + "model_group": isinstance(metadata, dict) and isinstance(metadata.get("model_group"), str), + "deployment": isinstance(metadata, dict) and isinstance(metadata.get("deployment"), str), + "model_info": isinstance(model_info, dict), + "model_info id": isinstance(model_info, dict) and isinstance(model_info.get("id"), str), + "proxy_server_request": isinstance(litellm_params.get("proxy_server_request"), (str, type(None))), + "preset_cache_key": isinstance(litellm_params.get("preset_cache_key"), (str, type(None))), + "stream_response": isinstance(litellm_params.get("stream_response"), dict), + } + return tuple(name for name, ok in checks.items() if not ok) + + +def _request_problems(kwargs: Mapping[str, object]) -> tuple[str, ...]: + checks: Final = { + "model": isinstance(kwargs.get("model"), str), + "messages": isinstance(kwargs.get("messages"), list), + "optional_params": isinstance(kwargs.get("optional_params"), dict), + "start_time": isinstance(kwargs.get("start_time"), (datetime, type(None))), + "stream": isinstance(kwargs.get("stream"), bool), + "user": isinstance(kwargs.get("user"), (str, type(None))), + } + return (*(name for name, ok in checks.items() if not ok), *_router_context_problems(kwargs)) + + +def _call_detail_problems(kwargs: Mapping[str, object]) -> tuple[str, ...]: + original_response: Final = kwargs.get("original_response") + checks: Final = { + "input": isinstance(kwargs.get("input"), (list, dict, str)), + "api_key": isinstance(kwargs.get("api_key"), (str, type(None))), + "original_response": isinstance(original_response, (str, litellm.CustomStreamWrapper, type(None))) + or inspect.iscoroutine(original_response) + or inspect.isasyncgen(original_response), + "additional_args": isinstance(kwargs.get("additional_args"), (dict, type(None))), + "log_event_type": isinstance(kwargs.get("log_event_type"), str), + } + return tuple(name for name, ok in checks.items() if not ok) + + +def _is_from_this_test(kwargs: Mapping[str, object]) -> bool: + litellm_params: Final = kwargs.get("litellm_params") + metadata: Final = litellm_params.get("metadata") if isinstance(litellm_params, dict) else None + return isinstance(metadata, dict) and metadata.get("model_group") in _OUR_MODEL_GROUPS + + +class _HookRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: Final[list[_HookEvent]] = [] + self.errors: Final[list[str]] = [] + self.loop: asyncio.AbstractEventLoop | None = None + self.waiters: Final[list[tuple[Callable[[Sequence[_State]], bool], asyncio.Event]]] = [] + + @property + def states(self) -> list[_State]: + return [event.state for event in self.events] + + def _record( + self, state: _State, model: object, kwargs: Mapping[str, object], response: object, problems: Sequence[str] + ) -> None: + if not _is_from_this_test(kwargs): + return + self.errors.extend(f"{state}: {problem}" for problem in problems) + self.events.append(_HookEvent(state, model, kwargs, response)) + if self.loop is not None: + self.loop.call_soon_threadsafe(self._notify) + + def _notify(self) -> None: + for predicate, event in self.waiters: + if predicate(tuple(self.states)): + event.set() + + async def until(self, predicate: Callable[[Sequence[_State]], bool]) -> tuple[_State, ...]: + self.loop = asyncio.get_running_loop() + if not predicate(tuple(self.states)): + event: Final = asyncio.Event() + self.waiters.append((predicate, event)) + await asyncio.wait_for(event.wait(), timeout=10) + return tuple(self.states) + + def log_pre_api_call(self, model: object, messages: object, kwargs: Mapping[str, object]) -> None: + problems: Final = ( + *(("model",) if not isinstance(model, str) else ()), + *(("messages",) if not isinstance(messages, list) else ()), + *_request_problems(kwargs), + ) + self._record("sync_pre_api_call", model, kwargs, messages, problems) + + def log_post_api_call( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + problems: Final = ( + *(("start_time",) if not isinstance(start_time, datetime) else ()), + *(("end_time",) if end_time is not None else ()), + *(("response_obj",) if response_obj is not None else ()), + *_request_problems(kwargs), + *_call_detail_problems(kwargs), + ) + self._record("post_api_call", kwargs.get("model"), kwargs, response_obj, problems) + + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record("sync_success", kwargs.get("model"), kwargs, response_obj, ()) + + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record("sync_failure", kwargs.get("model"), kwargs, response_obj, ()) + + async def async_log_stream_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record("async_stream", kwargs.get("model"), kwargs, response_obj, ()) + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + problems: Final = ( + *(("times",) if not (isinstance(start_time, datetime) and isinstance(end_time, datetime)) else ()), + *( + ("response_obj",) + if not isinstance(response_obj, (litellm.ModelResponse, litellm.EmbeddingResponse)) + else () + ), + *(("cache_hit",) if not isinstance(kwargs.get("cache_hit"), (bool, type(None))) else ()), + *_request_problems(kwargs), + *_call_detail_problems(kwargs), + ) + self._record("async_success", kwargs.get("model"), kwargs, response_obj, problems) + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + problems: Final = ( + *(("times",) if not (isinstance(start_time, datetime) and isinstance(end_time, datetime)) else ()), + *(("response_obj",) if response_obj is not None else ()), + *(("exception",) if not isinstance(kwargs.get("exception"), Exception) else ()), + *_request_problems(kwargs), + *_call_detail_problems(kwargs), + ) + self._record("async_failure", kwargs.get("model"), kwargs, response_obj, problems) + + +def _settled(terminal: int, posts: int) -> Callable[[Sequence[_State]], bool]: + def satisfied(states: Sequence[_State]) -> bool: + terminals: Final = sum(state in ("async_success", "async_failure") for state in states) + return terminals >= terminal and states.count("post_api_call") >= posts + + return satisfied + + +@pytest.fixture +def recorder(monkeypatch: pytest.MonkeyPatch) -> _HookRecorder: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + hook_recorder: Final = _HookRecorder() + monkeypatch.setattr(litellm, "callbacks", [hook_recorder]) + return hook_recorder + + +def _router(model: str, api_base: str = _PRIMARY) -> Router: + return Router( + model_list=[ + { + "model_name": "hooks-group", + "litellm_params": { + "model": model, + "api_key": "sk-unit-test", + "api_base": api_base, + "api_version": _API_VERSION, + }, + "model_info": {"base_model": model}, + } + ], + num_retries=0, + ) + + +class _RouterMetadata(TypedDict): + model_group: ReadOnly[str] + deployment: ReadOnly[str] + + +class _RouterModelInfo(TypedDict): + id: ReadOnly[str] + + +class _RouterParams(TypedDict): + metadata: ReadOnly[_RouterMetadata] + model_info: ReadOnly[_RouterModelInfo] + + +def _router_params(event: _HookEvent) -> _RouterParams: + return TypeAdapter(_RouterParams).validate_python(event.kwargs["litellm_params"]) + + +def _model_group(event: _HookEvent) -> str: + return _router_params(event)["metadata"]["model_group"] + + +def _model_id(event: _HookEvent) -> str: + return _router_params(event)["model_info"]["id"] + + +def _of_state(recorder: _HookRecorder, state: _State) -> list[_HookEvent]: + return [event for event in recorder.events if event.state == state] + + +def _completion(event: _HookEvent) -> litellm.ModelResponse: + assert isinstance(event.response, litellm.ModelResponse) + return event.response + + +@pytest.mark.asyncio +async def test_router_chat_success_streaming_and_failure_fire_the_hooks_in_order( + recorder: _HookRecorder, respx_mock: respx.MockRouter +) -> None: + route: Final = respx_mock.post(url__startswith=f"{_PRIMARY}/openai/deployments/gpt-4.1-mini/chat/completions") + router: Final = _router("azure/gpt-4.1-mini") + + route.mock(return_value=httpx.Response(200, json=_COMPLETION)) + await router.acompletion(model="hooks-group", messages=_MESSAGES) + await recorder.until(_settled(1, posts=1)) + assert recorder.states == ["sync_pre_api_call", "post_api_call", "async_success"] + pre, post, success = recorder.events + assert pre.model == "gpt-4.1-mini" + assert pre.response == _MESSAGES + assert {_model_group(event) for event in recorder.events} == {"hooks-group"} + assert len({_model_id(event) for event in recorder.events}) == 1 + assert post.kwargs["messages"] == _MESSAGES + assert success.kwargs["stream"] is False + assert _completion(success).choices[0].message.content == "hello" + assert _completion(success).usage.total_tokens == _USAGE["total_tokens"] + + route.mock(return_value=httpx.Response(200, text=_STREAM, headers={"content-type": "text/event-stream"})) + stream: Final = await router.acompletion(model="hooks-group", messages=_MESSAGES, stream=True) + assert "".join([chunk.choices[0].delta.content or "" async for chunk in stream]) == "hello" + await recorder.until(_settled(2, posts=2)) + streamed: Final = recorder.events[3:] + assert sorted(event.state for event in streamed[:2]) == ["post_api_call", "sync_pre_api_call"] + assert [event.state for event in streamed[2:]] == ["async_success"] + assert all(event.kwargs["stream"] is True for event in streamed) + assert _completion(streamed[2]).choices[0].message.content == "hello" + + route.mock(return_value=_AUTH_ERROR) + with pytest.raises(litellm.AuthenticationError): + await router.acompletion(model="hooks-group", messages=_MESSAGES) + await recorder.until(_settled(3, posts=3)) + failed: Final = recorder.events[6:] + assert [event.state for event in failed] == ["sync_pre_api_call", "post_api_call", "async_failure"] + assert failed[2].response is None + assert isinstance(failed[2].kwargs["exception"], litellm.AuthenticationError) + assert recorder.errors == [] + + +@pytest.mark.asyncio +async def test_router_embedding_success_and_failure_fire_the_hooks_in_order( + recorder: _HookRecorder, respx_mock: respx.MockRouter +) -> None: + route: Final = respx_mock.post(url__startswith=f"{_PRIMARY}/openai/deployments/text-embedding-3-small/embeddings") + router: Final = _router("azure/text-embedding-3-small") + + route.mock(return_value=httpx.Response(200, json=_EMBEDDING)) + await router.aembedding(model="hooks-group", input=["hello"]) + await recorder.until(_settled(1, posts=1)) + assert recorder.states == ["sync_pre_api_call", "post_api_call", "async_success"] + assert {event.model for event in recorder.events} == {"text-embedding-3-small"} + assert {_model_group(event) for event in recorder.events} == {"hooks-group"} + embedding: Final = recorder.events[2].response + assert isinstance(embedding, litellm.EmbeddingResponse) + assert embedding.model_dump()["data"][0]["embedding"] == _EMBEDDING_VECTOR + assert embedding.usage.prompt_tokens == 3 + + route.mock(return_value=_AUTH_ERROR) + with pytest.raises(litellm.AuthenticationError): + await router.aembedding(model="hooks-group", input=["hello"]) + await recorder.until(_settled(2, posts=2)) + assert recorder.states[3:] == ["sync_pre_api_call", "post_api_call", "async_failure"] + assert recorder.events[5].response is None + assert isinstance(recorder.events[5].kwargs["exception"], litellm.AuthenticationError) + assert recorder.errors == [] + + +@pytest.mark.asyncio +async def test_router_fallback_fires_failure_then_success_hooks( + recorder: _HookRecorder, respx_mock: respx.MockRouter +) -> None: + respx_mock.post(url__startswith=f"{_PRIMARY}/openai/deployments/gpt-4.1-mini/chat/completions").mock( + return_value=_AUTH_ERROR + ) + fallback: Final = respx_mock.post( + url__startswith=f"{_FALLBACK}/openai/deployments/gpt-4.1-mini/chat/completions" + ).mock(return_value=httpx.Response(200, json=_COMPLETION)) + router: Final = Router( + model_list=[ + { + "model_name": "primary-group", + "litellm_params": { + "model": "azure/gpt-4.1-mini", + "api_key": "my-bad-key", + "api_base": _PRIMARY, + "api_version": _API_VERSION, + }, + }, + { + "model_name": "fallback-group", + "litellm_params": { + "model": "azure/gpt-4.1-mini", + "api_key": "sk-unit-test", + "api_base": _FALLBACK, + "api_version": _API_VERSION, + }, + }, + ], + fallbacks=[{"primary-group": ["fallback-group"]}], + num_retries=0, + ) + + await router.acompletion(model="primary-group", messages=_MESSAGES) + await recorder.until(_settled(2, posts=2)) + + assert fallback.call_count == 1 + assert recorder.states == [ + "sync_pre_api_call", + "post_api_call", + "async_failure", + "sync_pre_api_call", + "post_api_call", + "async_success", + ] + assert [_model_group(event) for event in recorder.events] == ["primary-group"] * 3 + ["fallback-group"] * 3 + assert _model_id(recorder.events[0]) != _model_id(recorder.events[3]) + assert isinstance(recorder.events[2].kwargs["exception"], litellm.AuthenticationError) + assert _completion(recorder.events[5]).choices[0].message.content == "hello" + assert recorder.errors == [] + + +@pytest.mark.asyncio +async def test_router_completion_cache_hit_fires_a_second_success_hook( + recorder: _HookRecorder, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "cache", Cache()) + route: Final = respx_mock.post(url__startswith=f"{_PRIMARY}/openai/deployments/gpt-4.1-mini/chat/completions").mock( + return_value=httpx.Response(200, json=_COMPLETION) + ) + router: Final = _router("azure/gpt-4.1-mini") + + await router.acompletion(model="hooks-group", messages=_MESSAGES, caching=True) + await recorder.until(_settled(1, posts=1)) + await router.acompletion(model="hooks-group", messages=_MESSAGES, caching=True) + await recorder.until(_settled(2, posts=1)) + + assert route.call_count == 1 + assert recorder.states == ["sync_pre_api_call", "post_api_call", "async_success", "async_success"] + first, second = _of_state(recorder, "async_success") + assert first.kwargs.get("cache_hit") is not True + assert second.kwargs.get("cache_hit") is True + assert _completion(second).choices[0].message.content == "hello" + assert recorder.errors == [] + + +@pytest.mark.asyncio +async def test_router_streaming_cache_hit_still_fires_the_success_hook( + recorder: _HookRecorder, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "cache", Cache()) + route: Final = respx_mock.post(url__startswith=f"{_PRIMARY}/openai/deployments/gpt-4.1-mini/chat/completions").mock( + return_value=httpx.Response(200, text=_STREAM, headers={"content-type": "text/event-stream"}) + ) + router: Final = _router("azure/gpt-4.1-mini") + + first: Final = await router.acompletion(model="hooks-group", messages=_MESSAGES, stream=True, caching=True) + first_text: Final = "".join([chunk.choices[0].delta.content or "" async for chunk in first]) + await recorder.until(_settled(1, posts=1)) + states_after_first: Final = len(recorder.states) + second: Final = await router.acompletion(model="hooks-group", messages=_MESSAGES, stream=True, caching=True) + second_text: Final = "".join([chunk.choices[0].delta.content or "" async for chunk in second]) + await recorder.until(_settled(2, posts=1)) + + assert route.call_count == 1 + assert first_text == second_text == "hello" + assert sorted(recorder.states[:2]) == ["post_api_call", "sync_pre_api_call"] + assert recorder.states[2:states_after_first] == ["async_success"] + assert recorder.states[states_after_first:] == ["async_success"] + first_success, second_success = _of_state(recorder, "async_success") + assert first_success.kwargs.get("cache_hit") is not True + assert second_success.kwargs.get("cache_hit") is True + assert _completion(second_success).choices[0].message.content == "hello" + assert recorder.errors == [] + + +@pytest.mark.asyncio +async def test_router_embedding_cache_hit_fires_a_second_success_hook( + recorder: _HookRecorder, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "cache", Cache()) + route: Final = respx_mock.post( + url__startswith=f"{_PRIMARY}/openai/deployments/text-embedding-3-small/embeddings" + ).mock(return_value=httpx.Response(200, json=_EMBEDDING)) + router: Final = _router("azure/text-embedding-3-small") + + await router.aembedding(model="hooks-group", input=["hello"], caching=True) + await recorder.until(_settled(1, posts=1)) + await router.aembedding(model="hooks-group", input=["hello"], caching=True) + await recorder.until(_settled(2, posts=1)) + + assert route.call_count == 1 + assert recorder.states == ["sync_pre_api_call", "post_api_call", "async_success", "async_success"] + first, second = _of_state(recorder, "async_success") + assert first.kwargs.get("cache_hit") is not True + assert second.kwargs.get("cache_hit") is True + assert isinstance(second.response, litellm.EmbeddingResponse) + assert second.response.model_dump()["data"][0]["embedding"] == _EMBEDDING_VECTOR + assert recorder.errors == [] diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 8a7c9825f0a..58fd2250d2a 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1020,6 +1020,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_web_search": {"type": "boolean"}, "supports_bedrock_runtime_chat_completions_tools_with_reasoning": {"type": "boolean"}, "supports_bedrock_runtime_chat_completions_response_format": {"type": "boolean"}, + "supports_bedrock_runtime_chat_completions_inline_reasoning": {"type": "boolean"}, "supports_url_context": {"type": "boolean"}, "supports_multimodal": {"type": "boolean"}, "uses_embed_content": {"type": "boolean"}, diff --git a/tests/unit/types/test_guardrails_case_normalization.py b/tests/unit/types/test_guardrails_case_normalization.py index 26c1d395320..8c8192be3be 100644 --- a/tests/unit/types/test_guardrails_case_normalization.py +++ b/tests/unit/types/test_guardrails_case_normalization.py @@ -2,12 +2,19 @@ Test case normalization in LitellmParams for all guardrail types """ -from typing import Literal +import logging +from typing import Final, Literal import pytest from pydantic import ValidationError -from litellm.types.guardrails import BaseLitellmParams, LitellmParams +from litellm.types.guardrails import ( + BaseLitellmParams, + LitellmParams, + runtime_stream_scope, + stored_stream_scope, + with_tolerated_stream_scope, +) class TestLitellmParamsCaseNormalization: @@ -184,3 +191,80 @@ class TestSensitiveDataRoutingValidation: on_sensitive_data="BLOCK", ) assert params.on_sensitive_data == "block" + + +class TestStreamScopeValidation: + def test_scalar_is_case_normalized(self): + params = LitellmParams(guardrail="bedrock", mode="post_call", stream_scope="Streaming") + assert params.stream_scope == "streaming" + + def test_map_keys_and_values_are_normalized(self): + params = LitellmParams( + guardrail="bedrock", + mode=["pre_call", "post_call"], + stream_scope={"Pre_Call": "Both", "POST_CALL": "Non_Streaming"}, + ) + assert params.stream_scope == {"pre_call": "both", "post_call": "non_streaming"} + + def test_invalid_scalar_is_rejected(self): + with pytest.raises(ValidationError, match="stream_scope must be one of"): + LitellmParams(guardrail="bedrock", mode="post_call", stream_scope="chunks") + + def test_invalid_map_key_is_rejected(self): + with pytest.raises(ValidationError, match="stream_scope keys must be guardrail modes"): + LitellmParams(guardrail="bedrock", mode="post_call", stream_scope={"not_a_mode": "both"}) + + def test_invalid_map_value_is_rejected(self): + with pytest.raises(ValidationError, match="stream_scope must be one of"): + LitellmParams(guardrail="bedrock", mode="post_call", stream_scope={"post_call": "sometimes"}) + + def test_runtime_stream_scope_normalizes_direct_constructor_maps(self): + default, by_hook = runtime_stream_scope({"Pre_Call": "streaming"}) + assert default == "both" + assert dict(by_hook) == {"pre_call": "streaming"} + + def test_runtime_stream_scope_rejects_invalid_direct_input(self): + with pytest.raises(ValueError, match="stream_scope must be one of"): + runtime_stream_scope("chunks") + + @pytest.mark.parametrize( + "value, expected", + [ + ("streaming", "streaming"), + ("non_streaming", "non_streaming"), + ("both", "both"), + ({"pre_call": "streaming"}, {"pre_call": "streaming"}), + ("sometimes", None), + ({"pre_call": "sometimes"}, None), + ], + ) + def test_stored_stream_scope_tolerates_invalid_values( + self, + value: object, + expected: object, + caplog: pytest.LogCaptureFixture, + ) -> None: + caplog.set_level(logging.WARNING) + + result: Final = stored_stream_scope(value) + + assert result == expected + if expected is None: + assert f"Ignoring invalid stored stream_scope value of type {type(value).__name__}" in caplog.text + assert "sometimes" not in caplog.text + + def test_tolerated_stream_scope_rewrites_only_the_scope_field(self) -> None: + params: Final = { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "stream_scope": "sometimes", + } + + tolerated: Final = with_tolerated_stream_scope(params) + + assert tolerated == { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "stream_scope": None, + } + assert params["stream_scope"] == "sometimes" diff --git a/tests/unit/types/test_openai_decisions.py b/tests/unit/types/test_openai_decisions.py deleted file mode 100644 index d9387d644d5..00000000000 --- a/tests/unit/types/test_openai_decisions.py +++ /dev/null @@ -1,150 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping -from typing import Final - -import pytest -from pydantic import TypeAdapter, ValidationError - -from litellm.types.openai_decisions import ( - ChoiceAnswer, - DecisionsRequestBody, - DecisionsResponse, - RefusalAnswer, - ScoreQuestion, -) - -_REQUEST: Final[Mapping[str, object]] = { - "input": [ - { - "role": "user", - "content": [ - {"type": "input_text", "text": "Is this receipt a valid business expense?"}, - {"type": "input_image", "image_url": "https://example.com/receipt.png", "detail": "high"}, - ], - } - ], - "questions": [ - {"type": "predicate", "name": "is_expense", "instructions": "Is this a business expense?"}, - { - "type": "choice", - "name": "approve", - "instructions": "Should this be approved?", - "choices": [{"value": True, "description": "approve"}, {"value": False, "description": "reject"}], - }, - { - "type": "score", - "name": "risk", - "instructions": "How risky is this expense?", - "levels": [{"label": "low"}, {"label": "high", "description": "needs a manager"}], - }, - ], - "safety_identifier": "user-123", -} -_RESPONSE: Final[Mapping[str, object]] = { - "model": "gpt-6-luna", - "answers": [ - {"type": "predicate", "name": "is_expense", "probability": 0.92}, - { - "type": "choice", - "name": "approve", - "choice": True, - "probabilities": [{"value": True, "probability": 0.7}, {"value": False, "probability": 0.3}], - "confidence": 0.7, - }, - {"type": "refusal", "name": "risk"}, - ], - "usage": { - "input_tokens": 120, - "input_tokens_details": {"cached_tokens": 100, "cache_write_tokens": 0}, - "output_tokens": 12, - "output_tokens_details": {"reasoning_tokens": 4}, - "total_tokens": 132, - }, -} -_REQUEST_ADAPTER: Final[TypeAdapter[DecisionsRequestBody]] = TypeAdapter(DecisionsRequestBody) -_RESPONSE_ADAPTER: Final[TypeAdapter[DecisionsResponse]] = TypeAdapter(DecisionsResponse) - - -def test_the_documented_request_round_trips_with_its_boolean_choices_and_image_part() -> None: - request: Final = _REQUEST_ADAPTER.validate_python(_REQUEST) - - assert request.model_dump(mode="json", exclude_none=True) == _REQUEST - assert isinstance(request.questions[2], ScoreQuestion) - - -def test_the_documented_response_keeps_answer_order_refusals_and_token_details() -> None: - response: Final = _RESPONSE_ADAPTER.validate_python(_RESPONSE) - - assert response.model_dump(mode="json") == _RESPONSE - assert isinstance(response.answers[1], ChoiceAnswer) - assert isinstance(response.answers[2], RefusalAnswer) - - -_OFF_SPEC: Final[tuple[tuple[str, object], ...]] = ( - ("questions", [{"type": "noul", "name": "q", "instructions": "x"}]), - ("questions", [{"type": "choice", "name": "q", "instructions": "x", "choices": [{"value": 1}]}]), - ("questions", [{"type": "score", "name": "q", "levels": [{"label": "low"}]}]), - ("input", {"state": "not an OpenAI input"}), -) - - -@pytest.mark.parametrize(("field", "value"), _OFF_SPEC) -def test_requests_off_the_spec_are_rejected(field: str, value: object) -> None: - with pytest.raises(ValidationError): - _REQUEST_ADAPTER.validate_python({**_REQUEST, field: value}) - - -_EMPTY_COLLECTIONS: Final[tuple[tuple[str, list[object]], ...]] = ( - ("questions", []), - ("questions", [{"type": "choice", "name": "q", "instructions": "x", "choices": []}]), - ("questions", [{"type": "score", "name": "q", "instructions": "x", "levels": []}]), -) - - -@pytest.mark.parametrize(("field", "value"), _EMPTY_COLLECTIONS) -def test_empty_collections_are_left_for_the_provider_to_judge(field: str, value: list[object]) -> None: - request: Final = _REQUEST_ADAPTER.validate_python({**_REQUEST, field: value}) - - assert request.model_dump(mode="json", exclude_none=True)[field] == value - - -def test_choice_values_keep_their_type_so_a_string_true_and_a_boolean_true_stay_distinct() -> None: - answer: Final = { - "type": "choice", - "name": "approve", - "choice": "true", - "probabilities": [{"value": "true", "probability": 0.6}, {"value": True, "probability": 0.4}], - "confidence": 0.6, - } - - response: Final = _RESPONSE_ADAPTER.validate_python({**_RESPONSE, "answers": [answer]}) - - assert response.model_dump(mode="json")["answers"] == [answer] - with pytest.raises(ValidationError): - _RESPONSE_ADAPTER.validate_python({**_RESPONSE, "answers": [{**answer, "choice": 1}]}) - - -_OFF_SPEC_RESPONSE: Final[tuple[str, ...]] = ("model", "usage") - - -@pytest.mark.parametrize("field", _OFF_SPEC_RESPONSE) -def test_responses_missing_a_required_field_are_rejected(field: str) -> None: - with pytest.raises(ValidationError): - _RESPONSE_ADAPTER.validate_python({k: v for k, v in _RESPONSE.items() if k != field}) - - -def test_usage_without_token_details_is_rejected() -> None: - usage: Final = {"input_tokens": 120, "output_tokens": 12, "total_tokens": 132} - - with pytest.raises(ValidationError): - _RESPONSE_ADAPTER.validate_python({**_RESPONSE, "usage": usage}) - - -def test_hidden_params_live_outside_the_wire_body() -> None: - response: Final = _RESPONSE_ADAPTER.validate_python(_RESPONSE) - - response.set_hidden_params({"custom_llm_provider": "openai"}) - - assert response.hidden_params == {"custom_llm_provider": "openai"} - assert "_hidden_params" not in response.model_dump(mode="json") diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailReadOnlyDetails.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailReadOnlyDetails.tsx new file mode 100644 index 00000000000..5fe88b4870c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailReadOnlyDetails.tsx @@ -0,0 +1,68 @@ +import { Badge } from "@/components/ui/badge"; +import { GuardrailModeRows } from "./GuardrailModeDisplay"; +import { GuardrailStreamScopeDetail } from "./StreamScopeFields"; +import ToolPermissionRulesEditor, { type ToolPermissionConfig } from "./tool_permission/ToolPermissionRulesEditor"; + +export const GuardrailReadOnlyDetails = ({ + guardrailId, + guardrailName, + displayName, + litellmParams, + streamScope, + defaultOn, + piiEntityCount, + createdAt, + updatedAt, + showToolPermission, + toolPermissionConfig, +}: { + guardrailId: string; + guardrailName: string; + displayName: string; + litellmParams: { mode?: unknown; logging_only_scope?: string | null }; + streamScope: unknown; + defaultOn: boolean | undefined; + piiEntityCount: number; + createdAt: string; + updatedAt: string; + showToolPermission: boolean; + toolPermissionConfig: ToolPermissionConfig; +}) => ( +
+
+

Guardrail ID

+
{guardrailId}
+
+
+

Guardrail Name

+
{guardrailName || "Unnamed Guardrail"}
+
+
+

Provider

+
{displayName}
+
+ + +
+

Default On

+ {defaultOn ? "Yes" : "No"} +
+ {piiEntityCount > 0 && ( +
+

PII Protection

+
+ {piiEntityCount} PII entities configured +
+
+ )} +
+

Created At

+
{createdAt}
+
+
+

Last Updated

+
{updatedAt}
+
+ {showToolPermission && } +
+); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/StreamScopeFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/StreamScopeFields.tsx new file mode 100644 index 00000000000..1a8f55ce477 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/StreamScopeFields.tsx @@ -0,0 +1,95 @@ +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { GuardrailField, labelWithHint, type GuardrailFormControl } from "./GuardrailFormField"; +import { + STREAM_SCOPE_OPTIONS, + formatGuardrailStreamScope, + type GuardrailStreamScope, + isGuardrailStreamScope, +} from "./guardrail_info_helpers"; + +const STREAM_SCOPE_ITEMS = STREAM_SCOPE_OPTIONS.map((option) => ({ + label: option.label, + value: option.value, +})); + +const REQUEST_SHAPE_HINT = + "Run this guardrail on streaming requests, non-streaming requests, or both, for each selected mode."; + +export const StreamScopeFields = ({ + modes, + value, + onChange, +}: { + modes: string[]; + value: Record; + onChange: (next: Record) => void; +}) => { + if (modes.length === 0) return null; + + return ( +
+ {modes.map((mode) => { + const selected = value[mode] ?? "both"; + return ( +
+ + +
+ ); + })} +
+ ); +}; + +export const StreamScopeFormField = ({ control, modes }: { control: GuardrailFormControl; modes: string[] }) => ( + + {({ value, onChange }) => ( + | undefined) ?? {}} + onChange={onChange} + /> + )} + +); + +export const GuardrailStreamScopeCaption = ({ raw }: { raw: unknown }) => { + const label = formatGuardrailStreamScope(raw); + if (!label) return null; + return

{label}

; +}; + +export const GuardrailStreamScopeDetail = ({ raw }: { raw: unknown }) => { + const label = formatGuardrailStreamScope(raw); + if (!label) return null; + return ( +
+

Request shape

+
{label}
+
+ ); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.characterization.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.characterization.test.tsx index 230be13ef33..9628c284bd2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.characterization.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.characterization.test.tsx @@ -89,6 +89,23 @@ describe("AddGuardrailForm create payload characterization", () => { }); }); + it("sends stream_scope when a mode is restricted to streaming requests", async () => { + const user = userEvent.setup({ delay: null }); + renderForm(); + + await user.type(await screen.findByLabelText("Guardrail Name"), "my-bedrock"); + await pickProvider(user, "Bedrock Guardrail"); + await chooseSelectOption(user, screen.getByLabelText("pre_call applies to"), "Streaming only"); + await user.type(await screen.findByPlaceholderText("The guardrail id on Bedrock"), "gr-123"); + await user.click(screen.getByRole("button", { name: "Next" })); + await user.click(await screen.findByRole("button", { name: "Create Guardrail" })); + + await waitFor(() => expect(networking.createGuardrailCall).toHaveBeenCalledTimes(1)); + expect(payload()).toMatchObject({ + litellm_params: { stream_scope: "streaming" }, + }); + }); + it("switches mode from the seeded string to an array once the user touches the multi select", async () => { const user = userEvent.setup({ delay: null }); renderForm(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx index 43a542289f6..098245fa34b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx @@ -23,11 +23,14 @@ import { shouldRenderContentFilterConfigSettings, shouldRenderLLMJudgeFields, shouldRenderPIIConfigSettings, - supportsDirectionalLoggingOnlyScope, + streamScopePayload, toModeArray, + type GuardrailStreamScope, + supportsDirectionalLoggingOnlyScope, type LoggingOnlyScope, type LoggingOnlyScopeChoice, } from "./guardrail_info_helpers"; +import { StreamScopeFormField } from "./StreamScopeFields"; import { Logo } from "@/components/molecules/logo/Logo"; import { MultiSelect } from "@/components/shared/MultiSelect"; import { FieldGroup } from "@/components/ui/field"; @@ -172,6 +175,7 @@ const INITIAL_VALUES: GuardrailFormValues = { logging_only_scope_choice: "default", skip_system_message_choice: "inherit", skip_tool_message_choice: "inherit", + stream_scope_by_mode: {}, }; const ALWAYS_ON_ITEMS = [ @@ -466,6 +470,14 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a guardrail_info: {}, }; + const streamScope = streamScopePayload( + toModeArray(values.mode), + (values.stream_scope_by_mode as Record | undefined) ?? {}, + ); + if (streamScope !== undefined) { + guardrailData.litellm_params.stream_scope = streamScope; + } + const skipForCreate = choiceToSkipSystemForCreate(asSkipChoice(values.skip_system_message_choice)); if (skipForCreate !== undefined) { guardrailData.litellm_params.skip_system_message_in_guardrail = skipForCreate; @@ -773,6 +785,8 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a )} + + { expect(screen.getByRole("button", { name: /update guardrail/i })).toBeInTheDocument(); }); + it("should open without crashing for a tag-scoped mode and show it read-only", async () => { + renderModal({ + editData: { + guardrail_id: "g-tag", + guardrail_name: "tag-mode-guardrail", + litellm_params: { + mode: { tags: { "team-a": "pre_call" }, default: "post_call" }, + default_on: true, + custom_code: "def apply_guardrail(): pass", + }, + }, + }); + + expect(await screen.findByText("Edit Custom Guardrail")).toBeInTheDocument(); + const modeInput = screen.getByLabelText("Mode (tag-scoped)"); + expect(modeInput).toBeDisabled(); + expect(modeInput).toHaveValue("post_call, pre_call (tag-based)"); + expect(screen.getByText("Mode (tag-scoped, read-only)")).toBeInTheDocument(); + expect(screen.queryByText(/applies to/)).not.toBeInTheDocument(); + }); + + it("should omit mode and stream_scope from the update payload for a tag-scoped guardrail", async () => { + const user = userEvent.setup(); + renderModal({ + editData: { + guardrail_id: "g-tag", + guardrail_name: "tag-mode-guardrail", + litellm_params: { + mode: { tags: { "team-a": "pre_call" }, default: "post_call" }, + default_on: true, + custom_code: "def apply_guardrail(): pass", + }, + }, + }); + + const nameInput = await screen.findByDisplayValue("tag-mode-guardrail"); + await user.clear(nameInput); + await user.type(nameInput, "renamed-guardrail"); + await user.click(screen.getByRole("button", { name: /update guardrail/i })); + + await waitFor(() => { + expect(mockUpdate).toHaveBeenCalledTimes(1); + }); + const [token, guardrailId, payload] = mockUpdate.mock.calls[0] as [ + string, + string, + Record>, + ]; + expect(token).toBe("test-token"); + expect(guardrailId).toBe("g-tag"); + expect(payload.guardrail_name).toBe("renamed-guardrail"); + expect(payload.litellm_params.custom_code).toBe("def apply_guardrail(): pass"); + expect(payload.litellm_params).not.toHaveProperty("mode"); + expect(payload.litellm_params).not.toHaveProperty("stream_scope"); + }); + it("should keep save disabled until a guardrail name is entered", async () => { const user = userEvent.setup(); renderModal(); @@ -111,8 +167,7 @@ describe("CustomCodeModal", () => { expect(await screen.findByDisplayValue(/async def apply_guardrail/)).toBeInTheDocument(); - const comboboxes = screen.getAllByRole("combobox"); - await user.click(comboboxes[comboboxes.length - 1]); + await user.click(screen.getByRole("combobox", { name: "Template" })); const options = await screen.findAllByText("Block SSN"); await user.click(options[options.length - 1]); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx index 9deef4aeee3..c89540596bd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx @@ -37,152 +37,30 @@ import { import { Switch } from "@/components/ui/switch"; import { Textarea } from "@/components/ui/textarea"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; - -// Code templates -const CODE_TEMPLATES = { - empty: { - name: "Empty Template", - code: `async def apply_guardrail(inputs, request_data, input_type): - # inputs: {texts, images, tools, tool_calls, structured_messages, model} - # request_data: {model, user_id, team_id, end_user_id, metadata} - # input_type: "request" or "response" - return allow()`, - }, - blockSSN: { - name: "Block SSN", - code: `def apply_guardrail(inputs, request_data, input_type): - for text in inputs["texts"]: - if regex_match(text, r"\\d{3}-\\d{2}-\\d{4}"): - return block("SSN detected") - return allow()`, - }, - redactEmail: { - name: "Redact Emails", - code: `def apply_guardrail(inputs, request_data, input_type): - pattern = r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}" - modified = [] - for text in inputs["texts"]: - modified.append(regex_replace(text, pattern, "[EMAIL REDACTED]")) - return modify(texts=modified)`, - }, - blockSQL: { - name: "Block SQL Injection", - code: `def apply_guardrail(inputs, request_data, input_type): - if input_type != "request": - return allow() - for text in inputs["texts"]: - if contains_code_language(text, ["sql"]): - return block("SQL code not allowed") - return allow()`, - }, - validateJSON: { - name: "Validate JSON", - code: `def apply_guardrail(inputs, request_data, input_type): - if input_type != "response": - return allow() - - schema = {"type": "object", "required": ["name", "value"]} - - for text in inputs["texts"]: - obj = json_parse(text) - if obj is None: - return block("Invalid JSON response") - if not json_schema_valid(obj, schema): - return block("Response missing required fields") - return allow()`, - }, - externalAPI: { - name: "External API Check (async)", - code: `async def apply_guardrail(inputs, request_data, input_type): - # Call an external moderation API (async for non-blocking) - for text in inputs["texts"]: - response = await http_post( - "https://api.example.com/moderate", - body={"text": text, "user_id": request_data["user_id"]}, - headers={"Authorization": "Bearer YOUR_API_KEY"}, - timeout=10 - ) - - if not response["success"]: - # API call failed, allow by default or block - return allow() - - if response["body"].get("flagged"): - return block(response["body"].get("reason", "Content flagged")) - - return allow()`, - }, -}; - -// Available primitives organized by category -const PRIMITIVES = { - "Return Values": [ - { name: "allow()", desc: "Let request/response through" }, - { name: "block(reason)", desc: "Reject with message" }, - { name: "flag(reason, metadata={})", desc: "Let through, record a non-blocking violation" }, - { name: "modify(texts=[], images=[], tool_calls=[])", desc: "Transform content" }, - ], - "HTTP Requests (async)": [ - { name: "await http_request(url, method, headers, body)", desc: "Make async HTTP request" }, - { name: "await http_get(url, headers)", desc: "Async GET request" }, - { name: "await http_post(url, body, headers)", desc: "Async POST request" }, - ], - "Regex Functions": [ - { name: "regex_match(text, pattern)", desc: "Returns True if pattern found" }, - { name: "regex_replace(text, pattern, replacement)", desc: "Replace all matches" }, - { name: "regex_find_all(text, pattern)", desc: "Return list of matches" }, - ], - "JSON Functions": [ - { name: "json_parse(text)", desc: "Parse JSON string, returns None on error" }, - { name: "json_stringify(obj)", desc: "Convert to JSON string" }, - { name: "json_schema_valid(obj, schema)", desc: "Validate against JSON schema" }, - ], - "URL Functions": [ - { name: "extract_urls(text)", desc: "Extract all URLs from text" }, - { name: "is_valid_url(url)", desc: "Check if URL is valid" }, - { name: "all_urls_valid(text)", desc: "Check all URLs in text are valid" }, - ], - "Code Detection": [ - { name: "detect_code(text)", desc: "Returns True if code detected" }, - { name: "detect_code_languages(text)", desc: "Returns list of detected languages" }, - { name: 'contains_code_language(text, ["sql"])', desc: "Check for specific languages" }, - ], - "Text Utilities": [ - { name: "contains(text, substring)", desc: "Check if substring exists" }, - { name: "contains_any(text, [substr1, substr2])", desc: "Check if any substring exists" }, - { name: "word_count(text)", desc: "Count words" }, - { name: "char_count(text)", desc: "Count characters" }, - { name: "lower(text) / upper(text) / trim(text)", desc: "String transforms" }, - ], -}; - -const MODE_OPTIONS = [ - { value: "pre_call", label: "pre_call (Request)" }, - { value: "post_call", label: "post_call (Response)" }, - { value: "during_call", label: "during_call (Parallel)" }, - { value: "logging_only", label: "logging_only" }, - { value: "pre_mcp_call", label: "pre_mcp_call (Before MCP Tool Call)" }, - { value: "post_mcp_call", label: "post_mcp_call (After MCP Tool Call)" }, - { value: "during_mcp_call", label: "during_mcp_call (During MCP Tool Call)" }, -]; - -const TEMPLATE_ITEMS = Object.entries(CODE_TEMPLATES).map(([key, template]) => ({ - value: key, - label: template.name, -})); - -type ModeOption = (typeof MODE_OPTIONS)[number]; - -const MODE_OPTION_BY_VALUE: Record = Object.fromEntries( - MODE_OPTIONS.map((option) => [option.value, option]), -); +import { StreamScopeFields } from "../StreamScopeFields"; +import { + formatGuardrailMode, + streamScopeByModeFromConfig, + streamScopeForUpdate, + streamScopePayload, + type GuardrailStreamScope, +} from "../guardrail_info_helpers"; +import { + CODE_TEMPLATES, + MODE_OPTION_BY_VALUE, + MODE_OPTIONS, + PRIMITIVES, + TEMPLATE_ITEMS, + type ModeOption, +} from "./custom_code_catalog"; // Data for editing an existing guardrail + export interface EditGuardrailData { guardrail_id: string; guardrail_name: string; litellm_params: { - mode?: string | string[]; + mode?: string | string[] | Record; default_on?: boolean; custom_code?: string; logging_only_scope?: LoggingOnlyScope | null; @@ -204,6 +82,7 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS const isEditMode = !!editData; const [guardrailName, setGuardrailName] = useState(""); const [mode, setMode] = useState(["pre_call"]); + const [streamScopeByMode, setStreamScopeByMode] = useState>({}); const [loggingOnlyScopeChoice, setLoggingOnlyScopeChoice] = useState("default"); const [defaultOn, setDefaultOn] = useState(false); const [selectedTemplate, setSelectedTemplate] = useState("empty"); @@ -315,11 +194,14 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS setCode(CODE_TEMPLATES[templateKey as keyof typeof CODE_TEMPLATES].code); }; - // Normalize mode from API (string or string[]) to string[] - const normalizeMode = (m: string | string[] | undefined): string[] => { + // Normalize mode from API (string or string[]) to string[]. + // A tag-scoped mode dict ({ tags, default }) is managed outside this editor, so it + // contributes no editable modes and is displayed read-only instead. + const normalizeMode = (m: string | string[] | Record | undefined): string[] => { if (m === undefined || m === null) return ["pre_call"]; if (Array.isArray(m)) return m.length ? m : ["pre_call"]; - return [m]; + if (typeof m === "string") return [m]; + return []; }; // Reset form when modal opens or editData changes @@ -329,6 +211,12 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS // Edit mode: populate with existing data setGuardrailName(editData.guardrail_name || ""); setMode(normalizeMode(editData.litellm_params?.mode)); + setStreamScopeByMode( + streamScopeByModeFromConfig( + editData.litellm_params?.stream_scope, + normalizeMode(editData.litellm_params?.mode), + ), + ); setLoggingOnlyScopeChoice(loggingOnlyScopeToChoice(editData.litellm_params?.logging_only_scope)); setDefaultOn(editData.litellm_params?.default_on || false); setCode(editData.litellm_params?.custom_code || CODE_TEMPLATES.empty.code); @@ -337,6 +225,7 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS // Create mode: reset to defaults setGuardrailName(""); setMode(["pre_call"]); + setStreamScopeByMode({}); setLoggingOnlyScopeChoice("default"); setDefaultOn(false); setSelectedTemplate("empty"); @@ -411,11 +300,21 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS if (defaultOn !== editData.litellm_params?.default_on) { updateData.litellm_params.default_on = defaultOn; } + const nextStreamScope = streamScopeForUpdate( + mode, + streamScopeByMode, + editData.litellm_params?.stream_scope, + existingMode, + ); + if (nextStreamScope !== undefined) { + updateData.litellm_params.stream_scope = nextStreamScope; + } await updateGuardrailCall(accessToken, editData.guardrail_id, updateData); toast.success("Custom code guardrail updated successfully"); } else { // Create new guardrail + const streamScope = streamScopePayload(mode, streamScopeByMode); const guardrailData = { guardrail_name: guardrailName, litellm_params: { @@ -423,6 +322,7 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS mode: mode, default_on: defaultOn, custom_code: code, + ...(streamScope !== undefined ? { stream_scope: streamScope } : {}), ...getCustomCodeLoggingOnlyScopeCreate(mode, loggingOnlyScopeChoice), }, guardrail_info: {}, @@ -511,6 +411,11 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS const lineCount = code.split("\n").length; const selectedModeOptions = mode.map((value) => MODE_OPTION_BY_VALUE[value]).filter(Boolean); + const rawEditMode = editData?.litellm_params?.mode; + const tagScopedModeLabel = + rawEditMode !== null && typeof rawEditMode === "object" && !Array.isArray(rawEditMode) + ? formatGuardrailMode(rawEditMode) || "-" + : null; return ( !open && onClose()}> @@ -533,32 +438,38 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS />
- - setMode(options.map((option) => option.value))} - multiple - > - } className="w-full"> - {selectedModeOptions.map((option) => ( - - {option.label} - - ))} - - - - No matching modes - - {(option: ModeOption) => ( - + + {tagScopedModeLabel ? ( + + ) : ( + setMode(options.map((option) => option.value))} + multiple + > + } className="w-full"> + {selectedModeOptions.map((option) => ( + {option.label} - - )} - - - + + ))} + + + + No matching modes + + {(option: ModeOption) => ( + + {option.label} + + )} + + + + )}
{mode.includes("logging_only") && ( @@ -600,6 +511,11 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS + {mode.length > 0 && ( +
+ +
+ )} {/* Main Content */}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/custom_code_catalog.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/custom_code_catalog.ts new file mode 100644 index 00000000000..aad3df181a7 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/custom_code_catalog.ts @@ -0,0 +1,136 @@ +export const CODE_TEMPLATES = { + empty: { + name: "Empty Template", + code: `async def apply_guardrail(inputs, request_data, input_type): + # inputs: {texts, images, tools, tool_calls, structured_messages, model} + # request_data: {model, user_id, team_id, end_user_id, metadata} + # input_type: "request" or "response" + return allow()`, + }, + blockSSN: { + name: "Block SSN", + code: `def apply_guardrail(inputs, request_data, input_type): + for text in inputs["texts"]: + if regex_match(text, r"\\d{3}-\\d{2}-\\d{4}"): + return block("SSN detected") + return allow()`, + }, + redactEmail: { + name: "Redact Emails", + code: `def apply_guardrail(inputs, request_data, input_type): + pattern = r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}" + modified = [] + for text in inputs["texts"]: + modified.append(regex_replace(text, pattern, "[EMAIL REDACTED]")) + return modify(texts=modified)`, + }, + blockSQL: { + name: "Block SQL Injection", + code: `def apply_guardrail(inputs, request_data, input_type): + if input_type != "request": + return allow() + for text in inputs["texts"]: + if contains_code_language(text, ["sql"]): + return block("SQL code not allowed") + return allow()`, + }, + validateJSON: { + name: "Validate JSON", + code: `def apply_guardrail(inputs, request_data, input_type): + if input_type != "response": + return allow() + + schema = {"type": "object", "required": ["name", "value"]} + + for text in inputs["texts"]: + obj = json_parse(text) + if obj is None: + return block("Invalid JSON response") + if not json_schema_valid(obj, schema): + return block("Response missing required fields") + return allow()`, + }, + externalAPI: { + name: "External API Check (async)", + code: `async def apply_guardrail(inputs, request_data, input_type): + # Call an external moderation API (async for non-blocking) + for text in inputs["texts"]: + response = await http_post( + "https://api.example.com/moderate", + body={"text": text, "user_id": request_data["user_id"]}, + headers={"Authorization": "Bearer YOUR_API_KEY"}, + timeout=10 + ) + + if not response["success"]: + # API call failed, allow by default or block + return allow() + + if response["body"].get("flagged"): + return block(response["body"].get("reason", "Content flagged")) + + return allow()`, + }, +}; + +export const PRIMITIVES = { + "Return Values": [ + { name: "allow()", desc: "Let request/response through" }, + { name: "block(reason)", desc: "Reject with message" }, + { name: "flag(reason, metadata={})", desc: "Let through, record a non-blocking violation" }, + { name: "modify(texts=[], images=[], tool_calls=[])", desc: "Transform content" }, + ], + "HTTP Requests (async)": [ + { name: "await http_request(url, method, headers, body)", desc: "Make async HTTP request" }, + { name: "await http_get(url, headers)", desc: "Async GET request" }, + { name: "await http_post(url, body, headers)", desc: "Async POST request" }, + ], + "Regex Functions": [ + { name: "regex_match(text, pattern)", desc: "Returns True if pattern found" }, + { name: "regex_replace(text, pattern, replacement)", desc: "Replace all matches" }, + { name: "regex_find_all(text, pattern)", desc: "Return list of matches" }, + ], + "JSON Functions": [ + { name: "json_parse(text)", desc: "Parse JSON string, returns None on error" }, + { name: "json_stringify(obj)", desc: "Convert to JSON string" }, + { name: "json_schema_valid(obj, schema)", desc: "Validate against JSON schema" }, + ], + "URL Functions": [ + { name: "extract_urls(text)", desc: "Extract all URLs from text" }, + { name: "is_valid_url(url)", desc: "Check if URL is valid" }, + { name: "all_urls_valid(text)", desc: "Check all URLs in text are valid" }, + ], + "Code Detection": [ + { name: "detect_code(text)", desc: "Returns True if code detected" }, + { name: "detect_code_languages(text)", desc: "Returns list of detected languages" }, + { name: 'contains_code_language(text, ["sql"])', desc: "Check for specific languages" }, + ], + "Text Utilities": [ + { name: "contains(text, substring)", desc: "Check if substring exists" }, + { name: "contains_any(text, [substr1, substr2])", desc: "Check if any substring exists" }, + { name: "word_count(text)", desc: "Count words" }, + { name: "char_count(text)", desc: "Count characters" }, + { name: "lower(text) / upper(text) / trim(text)", desc: "String transforms" }, + ], +}; + +export const MODE_OPTIONS = [ + { value: "pre_call", label: "pre_call (Request)" }, + { value: "post_call", label: "post_call (Response)" }, + { value: "during_call", label: "during_call (Parallel)" }, + { value: "logging_only", label: "logging_only" }, + { value: "pre_mcp_call", label: "pre_mcp_call (Before MCP Tool Call)" }, + { value: "post_mcp_call", label: "post_mcp_call (After MCP Tool Call)" }, + { value: "during_mcp_call", label: "during_mcp_call (During MCP Tool Call)" }, +]; + +export const TEMPLATE_ITEMS = Object.entries(CODE_TEMPLATES).map(([key, template]) => ({ + value: key, + label: template.name, +})); + +export type ModeOption = (typeof MODE_OPTIONS)[number]; + +export const MODE_OPTION_BY_VALUE: Record = Object.fromEntries( + MODE_OPTIONS.map((option) => [option.value, option]), +); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx index 73d9fb8123e..c67758c0d98 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx @@ -33,7 +33,8 @@ import { } from "./GuardrailFormField"; import ContentFilterManager, { formatContentFilterDataForAPI } from "./content_filter/ContentFilterManager"; import CustomCodeModal, { EditGuardrailData } from "./custom_code/CustomCodeModal"; -import { GuardrailModeCard, GuardrailModeRows } from "./GuardrailModeDisplay"; +import { GuardrailModeCard } from "./GuardrailModeDisplay"; +import { GuardrailReadOnlyDetails } from "./GuardrailReadOnlyDetails"; import { getLoggingOnlyScopeUpdate, getGuardrailLogoAndName, @@ -41,10 +42,15 @@ import { loggingOnlyScopeToChoice, skipSystemMessageToChoice, skipToolMessageToChoice, + streamScopeByModeFromConfig, + streamScopeForUpdate, supportsDirectionalLoggingOnlyScope, + toModeArray, type SkipSystemMessageChoice, type SkipToolMessageChoice, + type GuardrailStreamScope, } from "./guardrail_info_helpers"; +import { GuardrailStreamScopeCaption, StreamScopeFormField } from "./StreamScopeFields"; import GuardrailOptionalParams from "./guardrail_optional_params"; import GuardrailProviderFields from "./guardrail_provider_fields"; import PiiConfiguration from "./pii_configuration"; @@ -236,6 +242,13 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, "skip_tool_message_choice", skipToolMessageToChoice(guardrailData.litellm_params?.skip_tool_message_in_guardrail), ); + form.setValue( + "stream_scope_by_mode", + streamScopeByModeFromConfig( + guardrailData.litellm_params?.stream_scope, + toModeArray(guardrailData.litellm_params?.mode), + ), + ); form.setValue( "guardrail_info", guardrailData.guardrail_info ? JSON.stringify(guardrailData.guardrail_info, null, 2) : "", @@ -304,6 +317,16 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, updateData.litellm_params.default_on = values.default_on; } + const modes = toModeArray(guardrailData.litellm_params?.mode); + const nextStreamScope = streamScopeForUpdate( + modes, + (values.stream_scope_by_mode as Record | undefined) ?? {}, + guardrailData.litellm_params?.stream_scope, + ); + if (nextStreamScope !== undefined) { + updateData.litellm_params.stream_scope = nextStreamScope; + } + const prevSkipChoice = skipSystemMessageToChoice(guardrailData.litellm_params?.skip_system_message_in_guardrail); const nextSkipChoice = values.skip_system_message_choice as SkipSystemMessageChoice | undefined; if (nextSkipChoice !== undefined && nextSkipChoice !== prevSkipChoice) { @@ -566,6 +589,7 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, +

Created At

@@ -723,6 +747,11 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, )} + + = ({ guardrailId, onClose, ) : ( -
-
-

Guardrail ID

-
{guardrailData.guardrail_id}
-
-
-

Guardrail Name

-
{guardrailData.guardrail_name || "Unnamed Guardrail"}
-
-
-

Provider

-
{displayName}
-
- -
-

Default On

- - {guardrailData.litellm_params?.default_on ? "Yes" : "No"} - -
- - {guardrailData.litellm_params?.pii_entities_config && - Object.keys(guardrailData.litellm_params.pii_entities_config).length > 0 && ( -
-

PII Protection

-
- - {Object.keys(guardrailData.litellm_params.pii_entities_config).length} PII entities - configured - -
-
- )} - -
-

Created At

-
{formatDate(guardrailData.created_at)}
-
-
-

Last Updated

-
{formatDate(guardrailData.updated_at)}
-
- - {guardrailData.litellm_params?.guardrail === "tool_permission" && ( - - )} -
+ )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx index 96ade4b8c1b..d149ef98328 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx @@ -15,6 +15,11 @@ import { skipToolMessageToChoice, choiceToSkipToolForCreate, formatGuardrailMode, + formatGuardrailStreamScope, + streamScopeByModeFromConfig, + streamScopeForMode, + streamScopeForUpdate, + streamScopePayload, loggingOnlyScopeToChoice, choiceToLoggingOnlyScope, getLoggingOnlyScopeUpdate, @@ -369,4 +374,56 @@ describe("guardrail_info_helpers", () => { expect(choiceToSkipToolForCreate("no")).toBe(false); }); }); + + describe("stream_scope helpers", () => { + it("treats omitted config as both for every mode", () => { + expect(streamScopeForMode(undefined, "pre_call")).toBe("both"); + expect(streamScopeByModeFromConfig(undefined, ["pre_call", "post_call"])).toEqual({ + pre_call: "both", + post_call: "both", + }); + }); + + it("applies a scalar to every selected mode and omits both-only payloads", () => { + expect(streamScopeForMode("streaming", "post_call")).toBe("streaming"); + expect(streamScopePayload(["pre_call", "post_call"], { pre_call: "both", post_call: "both" })).toBeUndefined(); + expect(streamScopePayload(["pre_call", "post_call"], { pre_call: "streaming", post_call: "streaming" })).toBe( + "streaming", + ); + }); + + it("keeps a mixed map instead of collapsing it to a scalar", () => { + expect(streamScopePayload(["pre_call", "post_call"], { pre_call: "both", post_call: "streaming" })).toEqual({ + post_call: "streaming", + }); + }); + + it("formats scalar and per-mode stream scopes for display", () => { + expect(formatGuardrailStreamScope(undefined)).toBe(""); + expect(formatGuardrailStreamScope("both")).toBe("Streaming and non-streaming"); + expect(formatGuardrailStreamScope("non_streaming")).toBe("Non-streaming only"); + expect(formatGuardrailStreamScope({ post_call: "streaming", pre_call: "both" })).toBe( + "post_call: Streaming only, pre_call: Streaming and non-streaming", + ); + }); + + it("emits both on update only when a prior restriction is cleared", () => { + expect(streamScopeForUpdate(["post_call"], { post_call: "streaming" }, undefined)).toBe("streaming"); + expect(streamScopeForUpdate(["post_call"], { post_call: "both" }, "streaming")).toBe("both"); + expect(streamScopeForUpdate(["post_call"], { post_call: "both" }, undefined)).toBeUndefined(); + }); + + it("keeps stored restrictions for modes outside the current selection", () => { + expect( + streamScopeForUpdate( + ["pre_call"], + { pre_call: "streaming" }, + { pre_call: "streaming", post_call: "non_streaming" }, + ), + ).toBeUndefined(); + expect( + streamScopeForUpdate(["pre_call"], { pre_call: "both" }, { pre_call: "streaming", post_call: "non_streaming" }), + ).toEqual({ post_call: "non_streaming" }); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index f55148b7315..a898f3b9abd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -343,3 +343,86 @@ export function choiceToSkipToolForCreate(choice: SkipToolMessageChoice | undefi if (choice === "no") return false; return undefined; } + +export const GUARDRAIL_STREAM_SCOPES = ["both", "streaming", "non_streaming"] as const; +export type GuardrailStreamScope = (typeof GUARDRAIL_STREAM_SCOPES)[number]; + +export const STREAM_SCOPE_OPTIONS: { value: GuardrailStreamScope; label: string }[] = [ + { value: "both", label: "Streaming and non-streaming" }, + { value: "streaming", label: "Streaming only" }, + { value: "non_streaming", label: "Non-streaming only" }, +]; + +export const isGuardrailStreamScope = (value: unknown): value is GuardrailStreamScope => + value === "both" || value === "streaming" || value === "non_streaming"; + +export const streamScopeForMode = (raw: unknown, mode: string): GuardrailStreamScope => { + if (isGuardrailStreamScope(raw)) return raw; + if (raw !== null && typeof raw === "object" && !Array.isArray(raw)) { + const value = (raw as Record)[mode]; + if (isGuardrailStreamScope(value)) return value; + } + return "both"; +}; + +export const streamScopeByModeFromConfig = (raw: unknown, modes: string[]): Record => + Object.fromEntries(modes.map((mode) => [mode, streamScopeForMode(raw, mode)])); + +export const streamScopePayload = ( + modes: string[], + scopes: Record, +): GuardrailStreamScope | Record | undefined => { + const perMode: Record = Object.fromEntries( + modes.map((mode) => [mode, scopes[mode] ?? "both"]), + ); + const values = Object.values(perMode); + if (values.length === 0 || values.every((scope) => scope === "both")) return undefined; + const unique = new Set(values); + if (unique.size === 1) return values[0]; + return Object.fromEntries(Object.entries(perMode).filter((entry) => entry[1] !== "both")); +}; + +export const formatGuardrailStreamScope = (raw: unknown): string => { + if (isGuardrailStreamScope(raw)) { + return STREAM_SCOPE_OPTIONS.find((option) => option.value === raw)?.label ?? raw; + } + if (raw !== null && typeof raw === "object" && !Array.isArray(raw)) { + const entries = Object.entries(raw as Record).filter( + (entry): entry is [string, GuardrailStreamScope] => isGuardrailStreamScope(entry[1]), + ); + if (entries.length === 0) return ""; + return entries.map(([mode, scope]) => `${mode}: ${formatGuardrailStreamScope(scope)}`).join(", "); + } + return ""; +}; + +export const streamScopeForUpdate = ( + modes: string[], + nextByMode: Record, + previousRaw: unknown, + previousModes: string[] = modes, +): GuardrailStreamScope | Record | undefined => { + const previousMap: Record = + previousRaw !== null && typeof previousRaw === "object" && !Array.isArray(previousRaw) + ? (previousRaw as Record) + : {}; + const preserved: Record = Object.fromEntries( + Object.entries(previousMap).filter( + (entry): entry is [string, GuardrailStreamScope] => !modes.includes(entry[0]) && isGuardrailStreamScope(entry[1]), + ), + ); + const nextModes = [...modes, ...Object.keys(preserved)]; + const nextStreamScope = streamScopePayload(nextModes, { + ...preserved, + ...Object.fromEntries(modes.map((mode) => [mode, nextByMode[mode] ?? "both"])), + }); + const previousCompareModes = Object.keys(previousMap).length > 0 ? Object.keys(previousMap) : previousModes; + const previousStreamScope = streamScopePayload( + previousCompareModes, + streamScopeByModeFromConfig(previousRaw, previousCompareModes), + ); + if (JSON.stringify(nextStreamScope ?? "both") === JSON.stringify(previousStreamScope ?? "both")) { + return undefined; + } + return nextStreamScope ?? "both"; +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts index 7e7089e685f..32dac9c2f65 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts @@ -110,13 +110,14 @@ export const useKeys = ( page: number, pageSize: number, options: KeyListCallOptions = {}, + queryOptions: { enabled?: boolean } = {}, ): UseQueryResult => { const { accessToken } = useAuthorized(); return useQuery({ queryKey: keyKeys.list({ page, limit: pageSize, ...options }), queryFn: async () => await keyListCall(accessToken!, page, pageSize, options), - enabled: Boolean(accessToken), + enabled: Boolean(accessToken) && (queryOptions.enabled ?? true), staleTime: 30000, // 30 seconds placeholderData: keepPreviousData, }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts index 6489bc2171d..09a0f214da3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts @@ -1,5 +1,5 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { renderHook, waitFor } from "@testing-library/react"; +import { act, renderHook, waitFor } from "@testing-library/react"; import React, { ReactNode } from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { @@ -11,6 +11,7 @@ import { useAutoRouters, useInfiniteModelInfo, useModelHub, + useModelAccessGroupNames, useModelsInfo, usePlainChatModelGroups, useSelectedTeamModels, @@ -565,6 +566,127 @@ describe("useUserModels", () => { }); }); +describe("useModelAccessGroupNames", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + vi.clearAllMocks(); + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("fetches and returns the caller's model access group names", async () => { + vi.mocked(modelAvailableCall).mockResolvedValue({ + data: [ + { id: "repro-access-group", object: "model", created: 0, owned_by: "litellm" }, + { id: "another-access-group", object: "model", created: 0, owned_by: "litellm" }, + ], + }); + + const { result } = renderHook(() => useModelAccessGroupNames(), { wrapper }); + + await waitFor(() => { + expect(result.current).toEqual(new Set(["repro-access-group", "another-access-group"])); + }); + + expect(modelAvailableCall).toHaveBeenCalledWith( + "test-access-token", + "test-user-id", + "Admin", + false, + null, + true, + true, + ); + }); + + it("returns undefined while the access-group lookup is pending", () => { + vi.mocked(modelAvailableCall).mockReturnValue(new Promise(() => undefined)); + + const { result } = renderHook(() => useModelAccessGroupNames(), { wrapper }); + + expect(result.current).toBeUndefined(); + }); + + it("returns undefined until authorization is ready", () => { + const unauthorizedContext = { + accessToken: null, + userId: null, + userRole: null, + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }; + mockUseAuthorized.mockReturnValue(unauthorizedContext); + + const { result } = renderHook(() => useModelAccessGroupNames(), { wrapper }); + + expect(result.current).toBeUndefined(); + expect(modelAvailableCall).not.toHaveBeenCalled(); + }); + + it("returns an empty set when the access-group lookup fails", async () => { + vi.mocked(modelAvailableCall).mockRejectedValue(new Error("lookup failed")); + + const { result } = renderHook(() => useModelAccessGroupNames(), { wrapper }); + + await waitFor(() => expect(result.current).toBeDefined()); + expect(result.current?.size).toBe(0); + }); + + it("keeps cached access-group names after a failed refetch", async () => { + vi.mocked(modelAvailableCall).mockResolvedValueOnce({ + data: [{ id: "repro-access-group", object: "model", created: 0, owned_by: "litellm" }], + }); + + const { result } = renderHook(() => useModelAccessGroupNames(), { wrapper }); + + await waitFor(() => { + expect(result.current?.has("repro-access-group")).toBe(true); + }); + + const queryKey = queryClient + .getQueryCache() + .getAll() + .find((query) => { + return query.queryKey[0] === "modelAccessGroupNames"; + })?.queryKey; + expect(queryKey).toBeDefined(); + if (!queryKey) throw new Error("The access-group query was not created"); + + vi.mocked(modelAvailableCall).mockRejectedValueOnce(new Error("refetch failed")); + await act(async () => { + await queryClient.refetchQueries({ queryKey }); + }); + + await waitFor(() => { + expect(queryClient.getQueryCache().find({ queryKey })?.state.status).toBe("error"); + }); + expect(result.current).toEqual(new Set(["repro-access-group"])); + }); +}); + describe("useSelectedTeamModels", () => { let queryClient: QueryClient; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index b5cc329d4c9..8bc2a50b4d8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -1,4 +1,5 @@ import { useQuery, useInfiniteQuery, useQueryClient, UseQueryResult } from "@tanstack/react-query"; +import { useMemo } from "react"; import { createQueryKeys } from "../common/queryKeysFactory"; import { modelInfoCall, modelHubCall, modelAvailableCall } from "@/components/networking"; import useAuthorized from "../useAuthorized"; @@ -29,6 +30,7 @@ const allProxyModelsKeys = createQueryKeys("allProxyModels"); const selectedTeamModelsKeys = createQueryKeys("selectedTeamModels"); const infiniteModelKeys = createQueryKeys("infiniteModels"); const userModelsKeys = createQueryKeys("userModels"); +const modelAccessGroupNameKeys = createQueryKeys("modelAccessGroupNames"); export const useModelsInfo = ( page: number = 1, @@ -265,6 +267,31 @@ export const useUserModels = (): UseQueryResult => { }); }; +export const useModelAccessGroupNames = (): ReadonlySet | undefined => { + const { accessToken, userId, userRole } = useAuthorized(); + const { data, isError } = useQuery({ + queryKey: modelAccessGroupNameKeys.list({}), + queryFn: async () => { + const response: AllProxyModelsResponse = await modelAvailableCall( + accessToken!, + userId!, + userRole!, + false, + null, + true, + true, + ); + return response.data.map((model) => model.id); + }, + enabled: Boolean(accessToken && userId && userRole), + }); + return useMemo(() => { + if (data !== undefined) return new Set(data); + if (isError) return new Set(); + return undefined; + }, [data, isError]); +}; + export const useSelectedTeamModels = (teamID: string | null) => { const { accessToken, userId, userRole } = useAuthorized(); return useQuery({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts index 7d1d035b4d4..2f46e548a7b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts @@ -118,13 +118,14 @@ export const useTeamsTable = ( }; export const teamKeys = createQueryKeys("teams"); -export const useTeams = (): UseQueryResult => { +export const useTeams = (queryOptions: { enabled?: boolean } = {}): UseQueryResult => { const { accessToken, userId, userRole } = useAuthorized(); - return useQuery({ + const teamsQueryOptions = { queryKey: teamKeys.list({}), queryFn: async () => await fetchTeams(accessToken!, userId, userRole, null), - enabled: Boolean(accessToken), - }); + enabled: Boolean(accessToken) && queryOptions.enabled !== false, + }; + return useQuery(teamsQueryOptions); }; const ALL_TEAMS_PAGE_SIZE = 100; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 4762dd2fed1..81d290cdfe4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -23,6 +23,7 @@ import { Sheet, SheetContent, SheetTitle, SheetTrigger } from "@/components/ui/s import { Button } from "@/components/ui/button"; import { Menu } from "lucide-react"; import { useMediaQuery } from "usehooks-ts"; +import { CommandPaletteProvider } from "@/components/CommandPalette/CommandPaletteProvider"; const pluginApiClient = createApiClient({ getBaseUrl: () => getProxyBaseUrl() ?? "" }); @@ -168,26 +169,28 @@ function DashboardShell({ children }: { children: React.ReactNode }) { setMobileNavigationKey(null)} /> -
- - } - > - - - } - /> - - - - - - -
{children}
-
+ +
+ + } + > + + + } + /> + + + + + + +
{children}
+
+
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx index 865cc808627..86b5659d570 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx @@ -84,6 +84,7 @@ vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useAllProxyModels: vi.fn(() => ({ data: { data: [] }, isLoading: false })), + useModelAccessGroupNames: vi.fn(() => new Set()), })); vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx index 410758491fe..844ac9f44e2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx @@ -1,4 +1,5 @@ -import { renderWithProviders, screen, waitFor } from "../../../../../tests/test-utils"; +import { useSyncExternalStore } from "react"; +import { act, renderWithProviders, screen, waitFor } from "../../../../../tests/test-utils"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import AddModelPanel from "./AddModelPanel"; @@ -64,6 +65,13 @@ vi.mock("@/app/(dashboard)/hooks/providers/useProviderFields", () => ({ { key: "api_base", label: "API Base", field_type: "text", required: false }, ], }, + { + provider: "Anthropic", + provider_display_name: "Anthropic", + litellm_provider: "anthropic", + default_model_placeholder: "claude-3-opus", + credential_fields: [{ key: "api_key", label: "API Key", field_type: "password", required: false }], + }, ], isLoading: false, error: null, @@ -74,6 +82,26 @@ vi.mock("@/components/vector_store_management/VectorStoreSelector", () => ({ default: () =>
, })); +type Catalog = Record; + +const createCatalogFeed = () => { + let current: Catalog | undefined; + const listeners = new Set<() => void>(); + const subscribe = (listener: () => void) => { + listeners.add(listener); + return () => { + listeners.delete(listener); + }; + }; + return { + useData: () => ({ data: useSyncExternalStore(subscribe, () => current) }), + publish: (next: Catalog) => { + current = next; + listeners.forEach((listener) => listener()); + }, + }; +}; + const lastCreatedModel = () => modelCreateCall.mock.calls.at(-1)?.[1]; const PROXY_ADMIN = { @@ -165,6 +193,70 @@ describe("AddModelPanel submit payload contract", () => { expect(screen.queryByText("openai-gpt-4o-deployment-id")).not.toBeInTheDocument(); }); + it("offers the catalog models once the catalog arrives after the provider was picked", async () => { + const catalog = createCatalogFeed(); + mockUseModelCostMap.mockImplementation(() => catalog.useData()); + const { user } = await setup(); + await user.click(screen.getByRole("combobox", { name: /provider/i })); + await user.click(await screen.findByText("OpenAI")); + expect(await screen.findByPlaceholderText("gpt-3.5-turbo")).toBeInTheDocument(); + + act(() => catalog.publish({ "gpt-4o-2024-08-06": { litellm_provider: "openai" } })); + await user.click(await screen.findByPlaceholderText("Select models")); + + expect(await screen.findByText("gpt-4o-2024-08-06")).toBeInTheDocument(); + }); + + it("keeps a model name typed before the catalog arrives instead of swapping the field under the user", async () => { + const catalog = createCatalogFeed(); + mockUseModelCostMap.mockImplementation(() => catalog.useData()); + const { user } = await setup(); + await user.click(screen.getByRole("combobox", { name: /provider/i })); + await user.click(await screen.findByText("OpenAI")); + await user.type(await screen.findByPlaceholderText("gpt-3.5-turbo"), "my-fine-tune"); + + act(() => catalog.publish({ "gpt-4o-2024-08-06": { litellm_provider: "openai" } })); + + await waitFor(() => expect(screen.getByPlaceholderText("gpt-3.5-turbo")).toHaveValue("my-fine-tune")); + expect(screen.queryByPlaceholderText("Select models")).not.toBeInTheDocument(); + }); + + it("offers the catalog models when a name typed before the catalog arrived was cleared again", async () => { + const catalog = createCatalogFeed(); + mockUseModelCostMap.mockImplementation(() => catalog.useData()); + const { user } = await setup(); + await user.click(screen.getByRole("combobox", { name: /provider/i })); + await user.click(await screen.findByText("OpenAI")); + const typed = await screen.findByPlaceholderText("gpt-3.5-turbo"); + await user.type(typed, "my-fine-tune"); + await user.clear(typed); + + act(() => catalog.publish({ "gpt-4o-2024-08-06": { litellm_provider: "openai" } })); + await user.click(await screen.findByPlaceholderText("Select models")); + + expect(await screen.findByText("gpt-4o-2024-08-06")).toBeInTheDocument(); + }); + + it("swaps the offered models when the provider changes while the catalog stays the same", async () => { + const catalog = createCatalogFeed(); + catalog.publish({ + "gpt-4o-2024-08-06": { litellm_provider: "openai" }, + "claude-sonnet-4-5": { litellm_provider: "anthropic" }, + }); + mockUseModelCostMap.mockImplementation(() => catalog.useData()); + const { user } = await setup(); + const provider = screen.getByRole("combobox", { name: /provider/i }); + await user.click(provider); + await user.click(await screen.findByText("OpenAI")); + await user.clear(provider); + await user.type(provider, "Anthropic"); + await user.click(await screen.findByText("Anthropic")); + await user.click(await screen.findByPlaceholderText("Select models")); + + expect(await screen.findByText("claude-sonnet-4-5")).toBeInTheDocument(); + expect(screen.queryByText("gpt-4o-2024-08-06")).not.toBeInTheDocument(); + }); + it("sends only the always-mounted fields while Advanced Settings stays closed", async () => { const { fillRequired, submit } = await setup(); await fillRequired(); 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 9dacabe506b..9fbd138bda4 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 @@ -1,7 +1,7 @@ "use client"; -import { useState } from "react"; -import { useForm } from "react-hook-form"; +import { useMemo, useState } from "react"; +import { useForm, useWatch } from "react-hook-form"; import { useQueryClient } from "@tanstack/react-query"; import AddModelForm from "@/components/add_model/AddModelForm"; import { handleAddModelSubmit } from "@/components/add_model/handle_add_model_submit"; @@ -27,7 +27,11 @@ export default function AddModelPanel() { const { data: credentialsResponse } = useCredentials(); const { data: teams } = useTeams(); const [selectedProvider, setSelectedProvider] = useState(Providers.Anthropic); - const [providerModels, setProviderModels] = useState([]); + const pickedProvider = useWatch({ control: form.control, name: "custom_llm_provider" }); + const providerModels = useMemo( + () => (typeof pickedProvider === "string" ? getProviderModels(pickedProvider, modelCostMapData) : []), + [pickedProvider, modelCostMapData], + ); const [showAdvancedSettings, setShowAdvancedSettings] = useState(false); const refresh = () => queryClient.invalidateQueries({ queryKey: ["models", "list"] }); @@ -57,9 +61,6 @@ export default function AddModelPanel() { selectedProvider={selectedProvider} setSelectedProvider={setSelectedProvider} providerModels={providerModels} - setProviderModelsFn={(provider) => - setProviderModels(provider === null ? [] : getProviderModels(provider, modelCostMapData)) - } getPlaceholder={getPlaceholder} showAdvancedSettings={showAdvancedSettings} setShowAdvancedSettings={setShowAdvancedSettings} diff --git a/ui/litellm-dashboard/src/components/CommandPalette/CommandPalette.integration.test.tsx b/ui/litellm-dashboard/src/components/CommandPalette/CommandPalette.integration.test.tsx new file mode 100644 index 00000000000..c722194d496 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CommandPalette/CommandPalette.integration.test.tsx @@ -0,0 +1,442 @@ +import { act, fireEvent, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import type { KeyResponse } from "@/components/key_team_helpers/key_list"; +import { writeStorage } from "@/lib/storage"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; +import { keyDetailHref } from "@/utils/entityLinks"; +import { uiHref } from "@/utils/uiHref"; +import { renderWithProviders, testQueryClient } from "../../../tests/test-utils"; +import { COMMAND_PALETTE_HINT_KEY } from "./CommandPaletteHint"; +import { CommandPaletteProvider } from "./CommandPaletteProvider"; +import { CommandPaletteTrigger } from "./CommandPaletteTrigger"; + +const mocks = vi.hoisted(() => ({ + pathname: "/api-keys", + push: vi.fn<(href: string) => void>(), + useKeys: vi.fn(), + teamAlias: "platform-team", + isPlaceholderData: false, +})); + +vi.mock("next/navigation", () => ({ + usePathname: () => mocks.pathname, + useRouter: () => ({ push: mocks.push }), +})); + +vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ + useKeys: mocks.useKeys, +})); +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useTeams: () => ({ + data: [{ team_id: "team-1", team_alias: mocks.teamAlias, members_with_roles: [] }], + }), +})); +vi.mock("@/app/(dashboard)/hooks/useIsOrgAdmin", () => ({ + default: () => false, +})); +vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ + useUISettings: () => ({ data: { values: {} } }), +})); + +const HIGH_VOLUME_KEY = keyFixture("key-prod-0", "high-volume-prod", "sk-...4zoA", "team-1"); +const PROD_KEY = keyFixture("key-prod-1", "prod-backend", "sk-...AHeA", "team-1"); +const STAGING_KEY = keyFixture("key-staging-1", "staging-agent", "sk-...stgB", "unknown-team"); +const KEY_FIXTURES = [HIGH_VOLUME_KEY, PROD_KEY, STAGING_KEY]; +const UI_SETTINGS_RESPONSE = { + server_root_path: "", + proxy_base_url: null, + admin_ui_disabled: false, + auto_redirect_to_sso: false, + sso_configured: false, + is_control_plane: false, + workers: [], +}; +const PROD_KEY_LIST_OPTIONS = { search: "prod", sortBy: "created_at", sortOrder: "desc", expand: "user" }; +const cachedKeyResults = new Map(); + +function keyFixture(token: string, alias: string, keyName: string, teamId: string | null = null): KeyResponse { + return { + token, + token_id: token, + key_name: keyName, + key_alias: alias, + spend: 2.5, + team_id: teamId, + team_alias: "", + } as KeyResponse; +} + +function sessionCookie() { + const encode = (value: object) => + btoa(JSON.stringify(value)).replaceAll("=", "").replaceAll("+", "-").replaceAll("/", "_"); + const claims = { + key: "sk-session-test", + user_id: "test-admin", + user_role: "proxy_admin", + premium_user: true, + auth_header_name: "X-Gateway-Session", + exp: Date.now() / 1000 + 3600, + }; + document.cookie = `token=${encode({ alg: "none" })}.${encode(claims)}.test; Path=/`; +} + +function renderPalette() { + return renderWithProviders( + + + , + ); +} + +beforeEach(() => { + testQueryClient.clear(); + window.localStorage.clear(); + cachedKeyResults.clear(); + vi.clearAllMocks(); + mocks.pathname = "/api-keys"; + mocks.teamAlias = "platform-team"; + mocks.isPlaceholderData = false; + sessionCookie(); + vi.stubGlobal( + "fetch", + vi.fn(async (input: RequestInfo | URL) => + String(input).includes("/v2/team/list") + ? Response.json({ + teams: [{ team_id: "team-1", team_alias: "platform-team" }], + total_pages: 1, + }) + : Response.json(UI_SETTINGS_RESPONSE), + ), + ); + vi.mocked(useKeys).mockImplementation((_page, _pageSize, options) => { + const query = options.search?.toLowerCase() ?? ""; + const keys = + cachedKeyResults.get(query) ?? + KEY_FIXTURES.filter( + (key) => !query || key.key_alias.toLowerCase().includes(query) || key.key_name.toLowerCase().includes(query), + ); + cachedKeyResults.set(query, keys); + return { + data: { keys, total_count: keys.length, current_page: 1, total_pages: 1 }, + isFetching: false, + isError: false, + isPlaceholderData: mocks.isPlaceholderData, + } as ReturnType; + }); +}); + +afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); + testQueryClient.clear(); + vi.unstubAllGlobals(); + document.cookie = "token=; Max-Age=0; Path=/"; +}); + +describe("CommandPalette integration", () => { + it("shows the discovery hint on the keys route", async () => { + renderPalette(); + + expect(await screen.findByText("to search keys")).toBeVisible(); + }); + + it("opens the palette from the hint and persists that it was seen", async () => { + const user = userEvent.setup(); + renderPalette(); + + await user.click(screen.getByRole("button", { name: /to search keys/ })); + + expect(await screen.findByRole("combobox", { name: "Search" })).toBeVisible(); + expect(screen.queryByText("to search keys")).not.toBeInTheDocument(); + expect(window.localStorage.getItem(COMMAND_PALETTE_HINT_KEY.name)).toBe("true"); + }); + + it("hides and persists the discovery hint when Ctrl+K opens the palette", async () => { + renderPalette(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + + expect(await screen.findByRole("combobox", { name: "Search" })).toBeVisible(); + expect(screen.queryByText("to search keys")).not.toBeInTheDocument(); + expect(window.localStorage.getItem(COMMAND_PALETTE_HINT_KEY.name)).toBe("true"); + }); + + it("dismisses the discovery hint without opening the palette", async () => { + const user = userEvent.setup(); + renderPalette(); + + await user.click(screen.getByRole("button", { name: "Dismiss search hint" })); + + expect(screen.queryByRole("dialog", { name: "Command palette" })).not.toBeInTheDocument(); + expect(screen.queryByText("to search keys")).not.toBeInTheDocument(); + expect(window.localStorage.getItem(COMMAND_PALETTE_HINT_KEY.name)).toBe("true"); + }); + + it("hides the discovery hint in memory when storage cannot persist its dismissal", async () => { + const user = userEvent.setup(); + vi.spyOn(Storage.prototype, "setItem").mockImplementation(() => { + throw new Error("Storage unavailable"); + }); + renderPalette(); + + await user.click(await screen.findByRole("button", { name: "Dismiss search hint" })); + + expect(screen.queryByText("to search keys")).not.toBeInTheDocument(); + expect(window.localStorage.getItem(COMMAND_PALETTE_HINT_KEY.name)).toBeNull(); + }); + + it("keeps the discovery hint hidden after closing the palette when storage is unavailable", async () => { + const user = userEvent.setup(); + vi.spyOn(Storage.prototype, "setItem").mockImplementation(() => { + throw new Error("Storage unavailable"); + }); + renderPalette(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const input = await screen.findByRole("combobox", { name: "Search" }); + expect(screen.queryByText("to search keys")).not.toBeInTheDocument(); + + await user.keyboard("{Escape}"); + + expect(screen.queryByText("to search keys")).not.toBeInTheDocument(); + expect(window.localStorage.getItem(COMMAND_PALETTE_HINT_KEY.name)).toBeNull(); + }); + + it("does not show the discovery hint when it has already been seen", () => { + writeStorage(COMMAND_PALETTE_HINT_KEY, true); + renderPalette(); + + expect(screen.queryByText("to search keys")).not.toBeInTheDocument(); + }); + + it("shows global-search copy outside the keys routes", async () => { + mocks.pathname = "/logs"; + renderPalette(); + + expect(await screen.findByText("to search or jump to a page")).toBeVisible(); + }); + + it("opens and focuses with Ctrl+K, rejects extra modifiers, and toggles closed", async () => { + renderPalette(); + fireEvent.keyDown(document, { key: "k", ctrlKey: true, shiftKey: true }); + fireEvent.keyDown(document, { key: "k", ctrlKey: true, altKey: true }); + expect(screen.queryByRole("dialog", { name: "Command palette" })).not.toBeInTheDocument(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const input = await screen.findByRole("combobox", { name: "Search" }); + expect(input).toHaveFocus(); + fireEvent.change(input, { target: { value: "stale" } }); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + expect(screen.queryByRole("dialog", { name: "Command palette" })).not.toBeInTheDocument(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const reopenedInput = await screen.findByRole("combobox", { name: "Search" }); + expect(reopenedInput).toHaveValue(""); + }); + + it("searches keys, shows key details, and opens the selected key", async () => { + const user = userEvent.setup(); + renderPalette(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const input = await screen.findByRole("combobox", { name: "Search" }); + expect(input).toHaveAttribute("placeholder", "Search keys by alias or ID…"); + expect(screen.getAllByText("Virtual Keys")).toHaveLength(2); + + await user.type(input, "prod"); + await waitFor(() => { + expect(vi.mocked(useKeys)).toHaveBeenLastCalledWith(1, 8, expect.objectContaining(PROD_KEY_LIST_OPTIONS), { + enabled: true, + }); + }); + const firstKeyOption = await screen.findByRole("option", { name: /high-volume-prod/ }); + const keyOption = await screen.findByRole("option", { name: /prod-backend/ }); + expect(firstKeyOption).toHaveAttribute("aria-selected", "true"); + expect(keyOption).toHaveTextContent("platform-team"); + expect(keyOption).toHaveTextContent("$2.50"); + expect(screen.queryByRole("option", { name: /staging-agent/ })).not.toBeInTheDocument(); + + fireEvent.keyDown(input, { key: "ArrowDown" }); + expect(keyOption).toHaveAttribute("aria-selected", "true"); + fireEvent.keyDown(input, { key: "Enter" }); + expect(mocks.push).toHaveBeenCalledWith(keyDetailHref(PROD_KEY.token)); + expect(screen.queryByRole("dialog", { name: "Command palette" })).not.toBeInTheDocument(); + }); + + it("opens the first matching key when Enter is pressed without navigating", async () => { + const user = userEvent.setup(); + renderPalette(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const input = await screen.findByRole("combobox", { name: "Search" }); + await user.type(input, "prod"); + const firstKeyOption = await screen.findByRole("option", { name: /high-volume-prod/ }); + expect(firstKeyOption).toHaveAttribute("aria-selected", "true"); + + fireEvent.keyDown(input, { key: "Enter" }); + expect(mocks.push).toHaveBeenCalledWith(keyDetailHref(HIGH_VOLUME_KEY.token)); + }); + + it("does not activate stale key results while the new query is debouncing", () => { + vi.useFakeTimers(); + renderPalette(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const input = screen.getByRole("combobox", { name: "Search" }); + fireEvent.change(input, { target: { value: "prod" } }); + vi.advanceTimersByTime(DEBOUNCE_WAIT_MS / 2); + + expect(screen.getByText("Searching…")).toBeVisible(); + expect(screen.queryByRole("option", { name: /high-volume-prod/ })).not.toBeInTheDocument(); + expect(screen.getByRole("option", { name: /Filter the keys table/ })).toBeVisible(); + + fireEvent.keyDown(input, { key: "Enter" }); + + expect(mocks.push).toHaveBeenCalledWith(`${uiHref("api-keys")}?key_search=prod`); + expect(mocks.push).not.toHaveBeenCalledWith(keyDetailHref(HIGH_VOLUME_KEY.token)); + expect(mocks.push).not.toHaveBeenCalledWith(keyDetailHref(PROD_KEY.token)); + }); + + it("does not show placeholder key rows after the query debounce completes", () => { + vi.useFakeTimers(); + mocks.isPlaceholderData = true; + renderPalette(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const input = screen.getByRole("combobox", { name: "Search" }); + fireEvent.change(input, { target: { value: "prod" } }); + act(() => { + vi.advanceTimersByTime(DEBOUNCE_WAIT_MS); + }); + + expect(screen.getByText("Searching…")).toBeVisible(); + expect(screen.queryByRole("option", { name: /high-volume-prod/ })).not.toBeInTheDocument(); + }); + + it("keeps the same key selected when its subtitle changes", async () => { + const { rerender } = renderPalette(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const input = await screen.findByRole("combobox", { name: "Search" }); + const secondKeyOption = await screen.findByRole("option", { name: /prod-backend/ }); + + fireEvent.keyDown(input, { key: "ArrowDown" }); + expect(secondKeyOption).toHaveAttribute("aria-selected", "true"); + + mocks.teamAlias = "renamed-platform-team"; + rerender( + + + , + ); + + const updatedSecondKeyOption = await screen.findByRole("option", { name: /prod-backend/ }); + expect(updatedSecondKeyOption).toHaveTextContent("renamed-platform-team"); + expect(updatedSecondKeyOption).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("option", { name: /high-volume-prod/ })).toHaveAttribute("aria-selected", "false"); + }); + + it("omits an unknown team ID from the key subtitle", async () => { + const user = userEvent.setup(); + renderPalette(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const input = await screen.findByRole("combobox", { name: "Search" }); + await user.type(input, "staging"); + const stagingOption = await screen.findByRole("option", { name: /staging-agent/ }); + expect(stagingOption).not.toHaveTextContent("unknown-team"); + }); + + it("keeps only non-empty groups and the filter action when no keys match", async () => { + const user = userEvent.setup(); + renderPalette(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const input = await screen.findByRole("combobox", { name: "Search" }); + await user.type(input, "no-match"); + + expect(await screen.findByText("No keys match “no-match”")).toBeVisible(); + expect(screen.getByRole("group", { name: "Keys" })).toHaveTextContent("No keys match “no-match”"); + expect(screen.getByRole("option", { name: /Filter the keys table/ })).toBeVisible(); + expect(screen.queryByRole("group", { name: "Pages" })).not.toBeInTheDocument(); + }); + + it("shows one global no-results message when the query has no matches", async () => { + const user = userEvent.setup(); + mocks.pathname = "/teams"; + renderPalette(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const input = await screen.findByRole("combobox", { name: "Search" }); + await user.type(input, "no-match"); + + expect(await screen.findByText("No results for “no-match”")).toBeVisible(); + expect(screen.queryByRole("group")).not.toBeInTheDocument(); + expect(screen.queryByRole("option", { name: /Search virtual keys/ })).not.toBeInTheDocument(); + }); + + it("opens the virtual-key table with the current search applied", async () => { + const user = userEvent.setup(); + renderPalette(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const input = await screen.findByRole("combobox", { name: "Search" }); + await user.type(input, "prod"); + await waitFor(() => { + expect(vi.mocked(useKeys)).toHaveBeenLastCalledWith(1, 8, expect.objectContaining({ search: "prod" }), { + enabled: true, + }); + }); + + fireEvent.keyDown(input, { key: "ArrowDown" }); + fireEvent.keyDown(input, { key: "ArrowDown" }); + fireEvent.keyDown(input, { key: "Enter" }); + expect(mocks.push).toHaveBeenCalledWith(`${uiHref("api-keys")}?key_search=prod`); + }); + + it("opens a matching global page from the teams route", async () => { + const user = userEvent.setup(); + mocks.pathname = "/teams"; + renderPalette(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const input = await screen.findByRole("combobox", { name: "Search" }); + expect(input).toHaveAttribute("placeholder", "Search pages and actions…"); + await user.type(input, "logs"); + + const logsOption = await screen.findByRole("option", { name: /^Logs/ }); + expect(logsOption).toBeVisible(); + fireEvent.keyDown(input, { key: "Enter" }); + expect(mocks.push).toHaveBeenCalledWith(uiHref("logs")); + }); + + it("clears the global query when switching to key search", async () => { + const user = userEvent.setup(); + mocks.pathname = "/teams"; + renderPalette(); + + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const input = await screen.findByRole("combobox", { name: "Search" }); + await user.type(input, "keys"); + + expect(await screen.findByRole("option", { name: /Search virtual keys/ })).toBeVisible(); + fireEvent.keyDown(input, { key: "Enter" }); + + const keysInput = await screen.findByRole("combobox", { name: "Search" }); + expect(keysInput).toHaveValue(""); + expect(await screen.findByText("Recent keys")).toBeVisible(); + }); + + it("switches from an empty keys scope to global search on Backspace", async () => { + renderPalette(); + fireEvent.keyDown(document, { key: "k", ctrlKey: true }); + const input = await screen.findByRole("combobox", { name: "Search" }); + expect(input).toHaveAttribute("placeholder", "Search keys by alias or ID…"); + + fireEvent.keyDown(input, { key: "Backspace" }); + expect(input).toHaveAttribute("placeholder", "Search pages and actions…"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/CommandPalette/CommandPalette.tsx b/ui/litellm-dashboard/src/components/CommandPalette/CommandPalette.tsx new file mode 100644 index 00000000000..c8a1d8d0f6c --- /dev/null +++ b/ui/litellm-dashboard/src/components/CommandPalette/CommandPalette.tsx @@ -0,0 +1,319 @@ +"use client"; + +import { useEffect, useMemo, useState, type KeyboardEvent as ReactKeyboardEvent, type ReactNode } from "react"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; +import { FileText, KeyRound, LoaderCircle, Search } from "lucide-react"; +import { usePathname, useRouter } from "next/navigation"; +import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { Dialog, DialogContent, DialogTitle } from "@/components/ui/dialog"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; +import { keyDetailHref } from "@/utils/entityLinks"; +import { routeSegmentForPathname, uiHref } from "@/utils/uiHref"; +import type { KeyResponse } from "@/components/key_team_helpers/key_list"; +import { CommandPaletteResultRow, type CommandPaletteRowData } from "./CommandPaletteResultRow"; +import { useVisibleMenuGroups } from "./useVisibleMenuGroups"; +import { flattenNavItems, matchNavItems, paletteScopeForRoute, type PaletteNavItem, type PaletteScope } from "./utils"; + +type PaletteEntry = CommandPaletteRowData & + ( + | { kind: "key"; key: KeyResponse } + | { kind: "filter" } + | { kind: "search-keys" } + | { kind: "page"; page: PaletteNavItem } + ); + +interface PaletteGroup { + label: string; + entries: PaletteEntry[]; + emptyMessage?: string; + isLoading?: boolean; +} + +const shortcutChipClass = + "rounded-sm border border-border border-b-2 bg-muted px-[3px] font-mono text-muted-foreground"; + +export function CommandPalette({ open, setOpen }: { open: boolean; setOpen: (open: boolean) => void }) { + if (!open) return null; + return ; +} + +function CommandPaletteDialog({ setOpen }: { setOpen: (open: boolean) => void }) { + const pathname = usePathname(); + const router = useRouter(); + const [query, setQuery] = useState(""); + const [activeSelection, setActiveSelection] = useState<{ + id: string; + query: string; + scope: PaletteScope; + } | null>(null); + const [scope, setScope] = useState(() => paletteScopeForRoute(routeSegmentForPathname(pathname))); + const [debouncedQuery] = useDebouncedValue(query, { wait: DEBOUNCE_WAIT_MS }); + const trimmedQuery = query.trim(); + const debouncedSearch = debouncedQuery.trim(); + const { data: teams } = useTeams({ enabled: false }); + const teamAliases = useMemo( + () => new Map((teams ?? []).map((team) => [team.team_id, team.team_alias] as const)), + [teams], + ); + const keyListOptions = { + search: debouncedSearch || undefined, + sortBy: "created_at", + sortOrder: "desc", + expand: "user", + }; + const keyResults = useKeys(1, 8, keyListOptions, { enabled: scope === "keys" }); + const isPendingQuery = trimmedQuery !== debouncedSearch; + const hasPlaceholderKeyResults = keyResults.isPlaceholderData; + const isInitialKeyLoad = keyResults.isFetching && !keyResults.data; + const isLoadingKeys = scope === "keys" && (isPendingQuery || hasPlaceholderKeyResults || isInitialKeyLoad); + const visibleGroups = useVisibleMenuGroups(); + const navItems = useMemo(() => flattenNavItems(visibleGroups), [visibleGroups]); + const navMatches = useMemo( + () => (scope === "global" || trimmedQuery ? matchNavItems(navItems, query) : []), + [navItems, query, scope, trimmedQuery], + ); + const groups = useMemo(() => { + const pages: PaletteEntry[] = navMatches.map((page) => ({ + id: `page:${page.route}`, + kind: "page", + title: page.label, + subtitle: page.section, + icon: page.icon ?? , + typeLabel: "Page", + page, + })); + + if (scope === "global") { + const searchKeysAction = !trimmedQuery || "search virtual keys".includes(trimmedQuery.toLowerCase()); + return [ + ...(searchKeysAction + ? [ + { + label: "Actions", + entries: [ + { + id: "action:search-keys", + kind: "search-keys" as const, + title: "Search virtual keys", + subtitle: "Find a key by alias or ID", + icon: , + typeLabel: "Action" as const, + }, + ], + }, + ] + : []), + ...(pages.length > 0 ? [{ label: "Pages", entries: pages }] : []), + ]; + } + + const keyEntries: PaletteEntry[] = (isLoadingKeys ? [] : keyResults.data?.keys ?? []).map((key) => { + const teamAlias = key.team_alias || (key.team_id ? teamAliases.get(key.team_id) : undefined); + const spend = typeof key.spend === "number" && Number.isFinite(key.spend) ? `$${key.spend.toFixed(2)}` : null; + const details = [teamAlias, spend].filter((detail): detail is string => Boolean(detail)); + const subtitle: ReactNode = ( + <> + {key.key_name} + {details.length > 0 && · {details.join(" · ")}} + + ); + return { + id: `key:${key.token}`, + kind: "key", + title: key.key_alias || "Unnamed key", + subtitle, + icon: , + typeLabel: "Key", + key, + }; + }); + const keysGroup: PaletteGroup = { + label: trimmedQuery ? "Keys" : "Recent keys", + entries: keyEntries, + emptyMessage: trimmedQuery ? `No keys match “${trimmedQuery}”` : "No recent keys", + isLoading: isLoadingKeys, + }; + + if (!trimmedQuery) return [keysGroup]; + + const filterEntry: PaletteEntry = { + id: "action:filter-keys", + kind: "filter", + title: `Filter the keys table for “${trimmedQuery}”`, + subtitle: "Open Virtual Keys with this search applied", + icon: , + typeLabel: "Action", + }; + + return [ + keysGroup, + { label: "Actions", entries: [filterEntry] }, + ...(pages.length > 0 ? [{ label: "Pages", entries: pages }] : []), + ]; + }, [isLoadingKeys, keyResults.data?.keys, navMatches, scope, teamAliases, trimmedQuery]); + const selectableEntries = useMemo(() => groups.flatMap((group) => group.entries), [groups]); + const activeEntryId = activeSelection?.query === query && activeSelection.scope === scope ? activeSelection.id : null; + const selectedIndex = selectableEntries.findIndex((entry) => entry.id === activeEntryId); + const activeIndex = selectableEntries.length === 0 ? -1 : Math.max(0, selectedIndex); + const activeOptionId = activeIndex >= 0 ? `command-palette-option-${activeIndex}` : undefined; + const hasKeyError = scope === "keys" && keyResults.isError && keyResults.data === undefined; + + useEffect(() => { + if (activeOptionId) document.getElementById(activeOptionId)?.scrollIntoView?.({ block: "nearest" }); + }, [activeOptionId]); + + const activate = (entry: PaletteEntry) => { + switch (entry.kind) { + case "key": + router.push(keyDetailHref(entry.key.token)); + setOpen(false); + return; + case "filter": + router.push(`${uiHref("api-keys")}?key_search=${encodeURIComponent(trimmedQuery)}`); + setOpen(false); + return; + case "search-keys": + setQuery(""); + setScope("keys"); + return; + case "page": + router.push(uiHref(entry.page.route)); + setOpen(false); + return; + } + }; + + const handleInputKeyDown = (event: ReactKeyboardEvent) => { + if (event.key === "Backspace" && scope === "keys" && query === "") { + event.preventDefault(); + setScope("global"); + return; + } + + if (event.key === "Enter") { + event.preventDefault(); + const activeEntry = selectableEntries[activeIndex]; + if (activeEntry) activate(activeEntry); + return; + } + + if (event.key !== "ArrowDown" && event.key !== "ArrowUp") return; + event.preventDefault(); + if (selectableEntries.length === 0) return; + const currentIndex = activeIndex < 0 ? 0 : activeIndex; + const nextIndex = + event.key === "ArrowDown" + ? (currentIndex + 1) % selectableEntries.length + : (currentIndex - 1 + selectableEntries.length) % selectableEntries.length; + setActiveSelection({ id: selectableEntries[nextIndex].id, query, scope }); + }; + + const renderGroupContents = (group: PaletteGroup, firstOptionIndex: number): ReactNode => { + if (group.isLoading) { + return ( +
+ + Searching… +
+ ); + } + if (group.entries.length === 0) + return
{group.emptyMessage}
; + return group.entries.map((entry, entryIndex) => { + const currentIndex = firstOptionIndex + entryIndex; + return ( + activate(entry)} + onHover={() => setActiveSelection({ id: entry.id, query, scope })} + /> + ); + }); + }; + + return ( + + + Command palette +
+ + {scope === "keys" && ( + + + Virtual Keys + + )} + { + setQuery(event.currentTarget.value); + }} + onKeyDown={handleInputKeyDown} + /> +
+
+ {hasKeyError && ( +
+ Unable to search keys +
+ )} + {!hasKeyError && scope === "global" && trimmedQuery && groups.length === 0 && ( +
No results for “{trimmedQuery}”
+ )} + {!hasKeyError && + groups.map((group, groupIndex) => { + const firstOptionIndex = groups + .slice(0, groupIndex) + .reduce((count, previousGroup) => count + previousGroup.entries.length, 0); + return ( +
+
+ {group.label} +
+ {renderGroupContents(group, firstOptionIndex)} +
+ ); + })} +
+
+
+ LiteLLM + {scope === "keys" ? "Virtual Keys" : "All pages"} +
+
+ + ↵ Open + + + ↑↓ Navigate + + + esc Close + +
+
+ + + ); +} diff --git a/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteHint.tsx b/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteHint.tsx new file mode 100644 index 00000000000..0a4c3fcda18 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteHint.tsx @@ -0,0 +1,44 @@ +"use client"; + +import { Search, X } from "lucide-react"; +import { z } from "zod"; +import { storageKey } from "@/lib/storage"; +import type { PaletteScope } from "./utils"; +import { useCommandPalette } from "./CommandPaletteProvider"; +import { useClientMounted, useShortcutLabel } from "./useShortcutLabel"; + +export const COMMAND_PALETTE_HINT_KEY = storageKey("local", "litellmCommandPaletteHintSeen", z.boolean(), false); + +interface CommandPaletteHintProps { + scope: PaletteScope; + onDismiss: () => void; +} + +export function CommandPaletteHint({ scope, onDismiss }: CommandPaletteHintProps) { + const { setOpen } = useCommandPalette(); + const shortcut = useShortcutLabel(); + const mounted = useClientMounted(); + + if (!mounted) return null; + + return ( +
+ + +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteProvider.tsx b/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteProvider.tsx new file mode 100644 index 00000000000..b2ec2625c45 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteProvider.tsx @@ -0,0 +1,65 @@ +"use client"; + +import { createContext, useCallback, useContext, useEffect, useMemo, useState, type ReactNode } from "react"; +import { usePathname } from "next/navigation"; +import { useStoredValue } from "@/lib/storage"; +import { routeSegmentForPathname } from "@/utils/uiHref"; +import { CommandPalette } from "./CommandPalette"; +import { CommandPaletteHint, COMMAND_PALETTE_HINT_KEY } from "./CommandPaletteHint"; +import { paletteScopeForRoute } from "./utils"; + +interface CommandPaletteContextValue { + open: boolean; + setOpen: (open: boolean) => void; + toggle: () => void; +} + +const CommandPaletteContext = createContext(null); + +export function useCommandPalette(): CommandPaletteContextValue { + const context = useContext(CommandPaletteContext); + if (!context) throw new Error("useCommandPalette must be used within CommandPaletteProvider"); + return context; +} + +export function CommandPaletteProvider({ children }: { children: ReactNode }) { + const [open, setOpenState] = useState(false); + const [hintSeen, setHintSeen] = useStoredValue(COMMAND_PALETTE_HINT_KEY); + const [hintDismissedInMemory, setHintDismissedInMemory] = useState(false); + const pathname = usePathname(); + const scope = paletteScopeForRoute(routeSegmentForPathname(pathname)); + const markHintSeen = useCallback(() => { + setHintDismissedInMemory(true); + setHintSeen(true); + }, [setHintSeen]); + const setOpen = useCallback( + (nextOpen: boolean) => { + if (nextOpen) markHintSeen(); + setOpenState(nextOpen); + }, + [markHintSeen], + ); + const toggle = useCallback(() => setOpen(!open), [open, setOpen]); + const value = useMemo(() => ({ open, setOpen, toggle }), [open, setOpen, toggle]); + const shouldShowHint = !hintSeen && !hintDismissedInMemory && !open; + + useEffect(() => { + const handleKeyDown = (event: KeyboardEvent) => { + if ((event.metaKey || event.ctrlKey) && event.key.toLowerCase() === "k" && !event.shiftKey && !event.altKey) { + event.preventDefault(); + toggle(); + } + }; + + document.addEventListener("keydown", handleKeyDown); + return () => document.removeEventListener("keydown", handleKeyDown); + }, [toggle]); + + return ( + + {children} + + {shouldShowHint && } + + ); +} diff --git a/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteResultRow.tsx b/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteResultRow.tsx new file mode 100644 index 00000000000..b4d8e2b3e74 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteResultRow.tsx @@ -0,0 +1,50 @@ +"use client"; + +import type { ReactNode } from "react"; +import { cn } from "@/lib/cva.config"; + +export interface CommandPaletteRowData { + id: string; + title: string; + subtitle?: ReactNode; + icon: ReactNode; + typeLabel: "Key" | "Page" | "Action"; +} + +export function CommandPaletteResultRow({ + item, + optionId, + active, + onActivate, + onHover, +}: { + item: CommandPaletteRowData; + optionId: string; + active: boolean; + onActivate: () => void; + onHover: () => void; +}) { + return ( + + ); +} diff --git a/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteTrigger.tsx b/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteTrigger.tsx new file mode 100644 index 00000000000..bd594800736 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteTrigger.tsx @@ -0,0 +1,50 @@ +"use client"; + +import { Search } from "lucide-react"; +import { usePathname } from "next/navigation"; +import { Button } from "@/components/ui/button"; +import { routeSegmentForPathname } from "@/utils/uiHref"; +import { paletteScopeForRoute } from "./utils"; +import { useCommandPalette } from "./CommandPaletteProvider"; +import { useShortcutLabel } from "./useShortcutLabel"; + +export function CommandPaletteTrigger({ mobile = false }: { mobile?: boolean }) { + const pathname = usePathname(); + const { open, setOpen } = useCommandPalette(); + const shortcut = useShortcutLabel(); + const isKeysRoute = paletteScopeForRoute(routeSegmentForPathname(pathname)) === "keys"; + const label = isKeysRoute ? "Search keys…" : "Search or jump to…"; + + if (mobile) { + return ( + + ); + } + + return ( + + ); +} diff --git a/ui/litellm-dashboard/src/components/CommandPalette/useShortcutLabel.ts b/ui/litellm-dashboard/src/components/CommandPalette/useShortcutLabel.ts new file mode 100644 index 00000000000..78e1865c90e --- /dev/null +++ b/ui/litellm-dashboard/src/components/CommandPalette/useShortcutLabel.ts @@ -0,0 +1,20 @@ +"use client"; + +import { useSyncExternalStore } from "react"; +import { shortcutLabel } from "./utils"; + +const subscribeToShortcut = () => () => {}; +const getClientShortcut = () => + typeof navigator === "undefined" ? "Ctrl K" : shortcutLabel(navigator.platform, navigator.userAgent); +const getServerShortcut = () => "Ctrl K"; +const subscribeToMount = () => () => {}; +const getClientMounted = () => true; +const getServerMounted = () => false; + +export function useShortcutLabel(): string { + return useSyncExternalStore(subscribeToShortcut, getClientShortcut, getServerShortcut); +} + +export function useClientMounted(): boolean { + return useSyncExternalStore(subscribeToMount, getClientMounted, getServerMounted); +} diff --git a/ui/litellm-dashboard/src/components/CommandPalette/useVisibleMenuGroups.ts b/ui/litellm-dashboard/src/components/CommandPalette/useVisibleMenuGroups.ts new file mode 100644 index 00000000000..8c3b9382934 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CommandPalette/useVisibleMenuGroups.ts @@ -0,0 +1,34 @@ +"use client"; + +import { useMemo } from "react"; +import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; +import useIsOrgAdmin from "@/app/(dashboard)/hooks/useIsOrgAdmin"; +import { isUserTeamAdminForAnyTeam } from "@/utils/roles"; +import { menuGroups, visibleMenuGroups, type MenuVisibilityContext } from "@/components/leftnav"; + +export const useVisibleMenuGroups = () => { + const { userId, userRole, isViewOnly } = useAuthorized(); + const isOrgAdmin = useIsOrgAdmin(); + const { data: teams } = useTeams({ enabled: false }); + const { data: settings } = useUISettings(); + const values = settings?.values; + const isTeamAdmin = useMemo(() => isUserTeamAdminForAnyTeam(teams ?? null, userId ?? ""), [teams, userId]); + + return useMemo(() => { + const context: MenuVisibilityContext = { + userRole, + isViewOnly, + isOrgAdmin, + isTeamAdmin, + enabledPagesInternalUsers: values?.enabled_ui_pages_internal_users ?? null, + enableProjectsUI: Boolean(values?.enable_projects_ui), + disableAgentsForInternalUsers: Boolean(values?.disable_agents_for_internal_users), + allowAgentsForTeamAdmins: Boolean(values?.allow_agents_for_team_admins), + disableVectorStoresForInternalUsers: Boolean(values?.disable_vector_stores_for_internal_users), + allowVectorStoresForTeamAdmins: Boolean(values?.allow_vector_stores_for_team_admins), + }; + return visibleMenuGroups(menuGroups, context); + }, [isOrgAdmin, isTeamAdmin, isViewOnly, userRole, values]); +}; diff --git a/ui/litellm-dashboard/src/components/CommandPalette/utils.test.ts b/ui/litellm-dashboard/src/components/CommandPalette/utils.test.ts new file mode 100644 index 00000000000..5bb61756f80 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CommandPalette/utils.test.ts @@ -0,0 +1,91 @@ +import { describe, expect, it } from "vitest"; +import { visibleMenuGroups, type MenuGroup, menuGroups } from "@/components/leftnav"; +import { flattenNavItems, matchNavItems, paletteScopeForRoute, shortcutLabel, type PaletteNavItem } from "./utils"; + +describe("paletteScopeForRoute", () => { + it.each([ + ["api-keys", "keys"], + ["", "keys"], + ["teams", "global"], + ] as const)("uses the expected palette scope for %s", (routeSegment, expected) => { + expect(paletteScopeForRoute(routeSegment)).toBe(expected); + }); +}); + +describe("flattenNavItems", () => { + it("flattens visible leaves and skips external links", () => { + const groups: MenuGroup[] = [ + { + groupLabel: "Tools", + items: [ + { key: "search-tools", page: "search-tools", label: "Search Tools" }, + { key: "docs", page: "docs", label: "Docs", external_url: "https://example.com" }, + { key: "tools", page: "tools", label: "Tools", children: [{ key: "logs", page: "logs", label: "Logs" }] }, + ], + }, + ]; + + expect(flattenNavItems(groups).map(({ label }) => label)).toEqual(["Search Tools", "Logs"]); + }); +}); + +describe("visibleMenuGroups", () => { + const context = { + userRole: "Internal User", + isViewOnly: false, + isOrgAdmin: false, + isTeamAdmin: false, + }; + + it("applies the internal-user page allowlist", () => { + const visible = visibleMenuGroups(menuGroups, { + ...context, + enabledPagesInternalUsers: ["api-keys"], + }); + + expect(flattenNavItems(visible).map(({ route }) => route)).toEqual(["api-keys"]); + }); + + it("hides Projects when its UI setting is disabled", () => { + const projectAdminContext = { ...context, userRole: "Admin", isTeamAdmin: true }; + const visibleWithFlag = visibleMenuGroups(menuGroups, { ...projectAdminContext, enableProjectsUI: true }); + const visibleWithoutFlag = visibleMenuGroups(menuGroups, { ...projectAdminContext, enableProjectsUI: false }); + + expect(flattenNavItems(visibleWithFlag).some(({ route }) => route === "projects")).toBe(true); + expect(flattenNavItems(visibleWithoutFlag).some(({ route }) => route === "projects")).toBe(false); + }); + + it("keeps admin pages visible despite the internal-user page allowlist", () => { + const visible = visibleMenuGroups(menuGroups, { + ...context, + userRole: "Admin", + enabledPagesInternalUsers: ["api-keys"], + }); + + expect(flattenNavItems(visible).some(({ route }) => route === "users")).toBe(true); + expect(flattenNavItems(visible).some(({ route }) => route === "organizations")).toBe(true); + }); +}); + +describe("matchNavItems", () => { + const items: PaletteNavItem[] = [ + { key: "section", label: "Teams", section: "Logs", route: "teams" }, + { key: "substring", label: "Audit Logs", section: "Access Control", route: "audit-logs" }, + { key: "prefix", label: "Logging", section: "Observability", route: "logging" }, + ]; + + it("ranks label prefixes above label substrings and section matches", () => { + expect(matchNavItems(items, "log").map(({ key }) => key)).toEqual(["prefix", "substring", "section"]); + }); + + it("returns every item in its original order for an empty query", () => { + expect(matchNavItems(items, " ")).toEqual(items); + }); +}); + +describe("shortcutLabel", () => { + it("uses the command key on Mac and Ctrl elsewhere", () => { + expect(shortcutLabel("MacIntel", "Mozilla/5.0")).toBe("⌘K"); + expect(shortcutLabel("Linux x86_64", "Mozilla/5.0 (Windows NT 10.0)")).toBe("Ctrl K"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/CommandPalette/utils.ts b/ui/litellm-dashboard/src/components/CommandPalette/utils.ts new file mode 100644 index 00000000000..ad02e155b0e --- /dev/null +++ b/ui/litellm-dashboard/src/components/CommandPalette/utils.ts @@ -0,0 +1,59 @@ +import type { ReactNode } from "react"; +import { labelText, routeOf, sectionText, type MenuGroup } from "@/components/leftnav"; + +export type PaletteScope = "keys" | "global"; + +export interface PaletteNavItem { + key: string; + label: string; + section: string; + route: string; + icon?: ReactNode; +} + +export const paletteScopeForRoute = (routeSegment: string): PaletteScope => + routeSegment === "" || routeSegment === "api-keys" ? "keys" : "global"; + +const matchRank = (item: PaletteNavItem, query: string): number => { + const label = item.label.toLowerCase(); + if (label.startsWith(query)) return 0; + if (label.includes(query)) return 1; + if (item.section.toLowerCase().includes(query)) return 2; + return -1; +}; + +export const flattenNavItems = (groups: readonly MenuGroup[]): PaletteNavItem[] => + groups.flatMap((group) => { + const flattenItems = (items: MenuGroup["items"]): PaletteNavItem[] => + items.flatMap((item) => { + if (item.external_url) return []; + if (item.children) return flattenItems(item.children); + return [ + { + key: item.key, + label: labelText(item), + section: sectionText(group.groupLabel), + route: routeOf(item), + icon: item.icon, + }, + ]; + }); + + return flattenItems(group.items); + }); + +export const matchNavItems = (items: readonly PaletteNavItem[], query: string): PaletteNavItem[] => { + const normalizedQuery = query.trim().toLowerCase(); + if (!normalizedQuery) return [...items]; + + return items + .map((item, index) => { + return { item, index, rank: matchRank(item, normalizedQuery) }; + }) + .filter(({ rank }) => rank >= 0) + .sort((left, right) => left.rank - right.rank || left.index - right.index) + .map(({ item }) => item); +}; + +export const shortcutLabel = (platform: string, userAgent: string): string => + [platform, userAgent].some((value) => value.toLowerCase().includes("mac")) ? "⌘K" : "Ctrl K"; diff --git a/ui/litellm-dashboard/src/components/DashboardHeader.integration.test.tsx b/ui/litellm-dashboard/src/components/DashboardHeader.integration.test.tsx index 756997b7ede..cb58670c20a 100644 --- a/ui/litellm-dashboard/src/components/DashboardHeader.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/DashboardHeader.integration.test.tsx @@ -3,6 +3,7 @@ import { render, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { PluginModeProvider } from "@/contexts/PluginModeContext"; +import { CommandPaletteProvider } from "@/components/CommandPalette/CommandPaletteProvider"; import { DashboardHeader } from "./DashboardHeader"; vi.mock("next/navigation", () => ({ usePathname: () => "/ui/logs" })); @@ -57,7 +58,9 @@ async function openTools() { render( - + + + , ); diff --git a/ui/litellm-dashboard/src/components/DashboardHeader.test.tsx b/ui/litellm-dashboard/src/components/DashboardHeader.test.tsx index 63f7b7aa136..822cec9deb3 100644 --- a/ui/litellm-dashboard/src/components/DashboardHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/DashboardHeader.test.tsx @@ -1,7 +1,8 @@ -import { afterEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { act, fireEvent, render, screen, within } from "@testing-library/react"; import { DashboardHeader } from "./DashboardHeader"; import { NAV_PRODUCT_LINK_CLASS } from "@/components/Navbar/navProductLinkClass"; +import { CommandPaletteProvider } from "@/components/CommandPalette/CommandPaletteProvider"; const { mockUsePluginMode, mockUseUISettings, state } = vi.hoisted(() => { const state = { @@ -31,31 +32,43 @@ vi.mock("@/components/Navbar/NotificationsBell/NotificationsBell", () => ({ Noti vi.mock("@/components/Navbar/WorkerDropdown/WorkerDropdown", () => ({ default: () => null })); vi.mock("@/components/liteadmin/LiteAdmin", () => ({ default: () => })); +const renderDashboardHeader = () => + render( + + + , + ); + describe("DashboardHeader breadcrumb", () => { + beforeEach(() => { + localStorage.clear(); + }); + afterEach(() => { state.plugins = []; state.enableChatUI = false; state.pathname = "/ui/logs"; state.isDesktop = false; + localStorage.clear(); }); it("titles the breadcrumb from the current route, not from a sidebar page id", () => { state.pathname = "/ui/models-and-endpoints"; - render(); + renderDashboardHeader(); expect(screen.getByText("Models + Endpoints")).toBeInTheDocument(); }); it("titles the dashboard root as Virtual Keys", () => { state.pathname = "/ui/"; - render(); + renderDashboardHeader(); expect(screen.getByText("Virtual Keys")).toBeInTheDocument(); }); it("roots the breadcrumb in the AI Gateway selector (with a Chat option) and drops the static section crumb when the selector is available", async () => { state.enableChatUI = true; - render(); + renderDashboardHeader(); expect(screen.getByText("Logs")).toBeInTheDocument(); expect(screen.queryByText("Observability")).not.toBeInTheDocument(); @@ -68,7 +81,7 @@ describe("DashboardHeader breadcrumb", () => { }); it("keeps the AI Gateway selector at the root even when there is nothing to switch to (discovery)", () => { - render(); + renderDashboardHeader(); expect(screen.getByRole("button", { name: /AI Gateway/i })).toBeInTheDocument(); expect(screen.getByText("Logs")).toBeInTheDocument(); @@ -76,7 +89,7 @@ describe("DashboardHeader breadcrumb", () => { }); it("styles Docs with the shared product-link class instead of a muted toolbar button", () => { - render(); + renderDashboardHeader(); const docs = screen.getByRole("link", { name: "Docs" }); for (const cls of NAV_PRODUCT_LINK_CLASS.trim().split(/\s+/)) { @@ -86,7 +99,7 @@ describe("DashboardHeader breadcrumb", () => { }); it("renders the tools divider centered rather than stretched to the top of the row", () => { - const { container } = render(); + const { container } = renderDashboardHeader(); const separators = container.querySelectorAll('[data-slot="separator"][data-orientation="vertical"]'); expect(separators).toHaveLength(1); @@ -95,7 +108,7 @@ describe("DashboardHeader breadcrumb", () => { }); it("places LiteAdmin in the header tools ahead of Docs", () => { - render(); + renderDashboardHeader(); const liteAdmin = within(screen.getByRole("banner")).getByRole("button", { name: "LiteAdmin" }); expect(liteAdmin.compareDocumentPosition(screen.getByRole("link", { name: "Docs" }))).toBe( @@ -104,7 +117,7 @@ describe("DashboardHeader breadcrumb", () => { }); it("keeps the gateway selector and tools available from the compact header menu", async () => { - render(); + renderDashboardHeader(); fireEvent.click(screen.getByRole("button", { name: "More options" })); const tools = await screen.findByRole("dialog", { name: "Gateway tools" }); expect(within(tools).getByRole("button", { name: "AI Gateway" })).toBeInTheDocument(); @@ -113,16 +126,24 @@ describe("DashboardHeader breadcrumb", () => { }); it("closes mobile tools when switching to desktop and keeps them closed when returning", async () => { - const { rerender } = render(); + const { rerender } = renderDashboardHeader(); fireEvent.click(screen.getByRole("button", { name: "More options" })); expect(await screen.findByRole("dialog", { name: "Gateway tools" })).toBeInTheDocument(); state.isDesktop = true; - rerender(); + rerender( + + + , + ); expect(screen.queryByRole("dialog", { name: "Gateway tools" })).not.toBeInTheDocument(); state.isDesktop = false; - rerender(); + rerender( + + + , + ); expect(screen.queryByRole("dialog", { name: "Gateway tools" })).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/DashboardHeader.tsx b/ui/litellm-dashboard/src/components/DashboardHeader.tsx index 43440b0c92d..7580f8b69c0 100644 --- a/ui/litellm-dashboard/src/components/DashboardHeader.tsx +++ b/ui/litellm-dashboard/src/components/DashboardHeader.tsx @@ -27,6 +27,7 @@ import { useMediaQuery } from "usehooks-ts"; import { Ellipsis } from "lucide-react"; import { Button } from "@/components/ui/button"; import { Popover, PopoverContent, PopoverTitle, PopoverTrigger } from "@/components/ui/popover"; +import { CommandPaletteTrigger } from "@/components/CommandPalette/CommandPaletteTrigger"; // Top bar for the dashboard shell. Sits only over the content column (the brand // lives in the sidebar header); mirrors the design's breadcrumb-left / tools-right layout. @@ -62,6 +63,8 @@ export function DashboardHeader({ navigationTrigger }: { navigationTrigger?: Rea + +
{showWorkerSwitch && ( <> @@ -78,6 +81,7 @@ export function DashboardHeader({ navigationTrigger }: { navigationTrigger?: Rea
+ }> 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 12efc5beb7c..220d46647f6 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 @@ -164,7 +164,6 @@ const createTestProps = (userRole = "proxy_admin", userId = "user-1", isTeamAdmi mountedValues: () => projectMountedValues(registry, form.getValues), handleOk: vi.fn().mockResolvedValue(true), setSelectedProvider: vi.fn(), - setProviderModelsFn: vi.fn(), getPlaceholder: vi.fn((provider: string) => `Enter ${provider} model name`), setShowAdvancedSettings: vi.fn(), selectedProvider: Providers.OpenAI, diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index a416bf87602..b5b6c439dc2 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -52,7 +52,6 @@ interface AddModelFormProps { selectedProvider: string | null; setSelectedProvider: (provider: string | null) => void; providerModels: string[]; - setProviderModelsFn: (provider: string | null) => void; getPlaceholder: (provider: string) => string; showAdvancedSettings: boolean; setShowAdvancedSettings: (show: boolean) => void; @@ -76,7 +75,6 @@ const AddModelForm: React.FC = ({ selectedProvider, setSelectedProvider, providerModels, - setProviderModelsFn, getPlaceholder, showAdvancedSettings, setShowAdvancedSettings, @@ -166,7 +164,6 @@ const AddModelForm: React.FC = ({ const applyProviderSelection = (provider: string | null) => { setSelectedProvider(provider); - setProviderModelsFn(provider); form.setValue("model", []); form.setValue("model_name", undefined); }; diff --git a/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx b/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx index 8cf4b077f39..529c0e2a63e 100644 --- a/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx +++ b/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx @@ -131,12 +131,12 @@ const LiteLLMModelNameField: React.FC = ({ } }} /> - ) : providerModels.length > 0 ? ( + ) : providerModels.length > 0 && !(typeof control.value === "string" && control.value !== "") ? ( { control.onChange(value); handleModelChange(value); diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 4cba12c9f7e..256fe578912 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -103,7 +103,7 @@ interface SidebarProps { allowVectorStoresForTeamAdmins?: boolean; } -interface MenuItem { +export interface MenuItem { key: string; page: string; route?: string; @@ -114,12 +114,92 @@ interface MenuItem { external_url?: string; } -interface MenuGroup { +export interface MenuGroup { groupLabel: string; items: MenuItem[]; roles?: string[]; } +export interface MenuVisibilityContext { + userRole: string; + isViewOnly: boolean; + isOrgAdmin: boolean; + isTeamAdmin: boolean; + enabledPagesInternalUsers?: string[] | null; + enableProjectsUI?: boolean; + disableAgentsForInternalUsers?: boolean; + allowAgentsForTeamAdmins?: boolean; + disableVectorStoresForInternalUsers?: boolean; + allowVectorStoresForTeamAdmins?: boolean; +} + +const adminPageVisibility = (item: MenuItem, context: MenuVisibilityContext, isAdmin: boolean): boolean | null => { + if (item.key !== "organizations" && item.key !== "users") return null; + const hasRoleAccess = !item.roles || item.roles.includes(context.userRole) || context.isOrgAdmin; + if (!hasRoleAccess) return false; + if (!isAdmin && context.enabledPagesInternalUsers != null) { + return context.enabledPagesInternalUsers.includes(item.page); + } + return true; +}; + +const projectPageIsVisible = (item: MenuItem, context: MenuVisibilityContext): boolean => { + if (item.key !== "projects") return true; + if (!context.enableProjectsUI) return false; + return canViewProjectsPage({ + userRole: context.userRole, + isOrgAdmin: context.isOrgAdmin, + isTeamAdmin: context.isTeamAdmin, + }); +}; + +const pageIsDisabledForInternalUsers = ( + item: MenuItem, + pageKey: "agents" | "vector-stores", + isAdmin: boolean, + context: MenuVisibilityContext, +): boolean => { + const isDisabled = + pageKey === "agents" ? context.disableAgentsForInternalUsers : context.disableVectorStoresForInternalUsers; + const allowedForTeamAdmins = + pageKey === "agents" ? context.allowAgentsForTeamAdmins : context.allowVectorStoresForTeamAdmins; + if (item.key !== pageKey || isAdmin || !isDisabled) return false; + return !(allowedForTeamAdmins && context.isTeamAdmin); +}; + +const pageIsEnabledForInternalUser = (item: MenuItem, context: MenuVisibilityContext): boolean => { + const enabledPages = context.enabledPagesInternalUsers; + if (enabledPages == null) return true; + if (item.children?.some((child) => enabledPages.includes(child.page))) return true; + return enabledPages.includes(item.page); +}; + +const menuItemIsVisible = (item: MenuItem, context: MenuVisibilityContext, isAdmin: boolean): boolean => { + if (item.children && item.children.length === 0) return false; + if (item.key === "llm-playground" && context.isViewOnly) return false; + const adminVisibility = adminPageVisibility(item, context, isAdmin); + if (adminVisibility !== null) return adminVisibility; + if (!projectPageIsVisible(item, context)) return false; + if (pageIsDisabledForInternalUsers(item, "agents", isAdmin, context)) return false; + if (pageIsDisabledForInternalUsers(item, "vector-stores", isAdmin, context)) return false; + if (item.roles && !item.roles.includes(context.userRole)) return false; + if (!isAdmin && context.enabledPagesInternalUsers != null) return pageIsEnabledForInternalUser(item, context); + return true; +}; + +export const visibleMenuGroups = (groups: readonly MenuGroup[], context: MenuVisibilityContext): MenuGroup[] => { + const isAdmin = isAdminRole(context.userRole); + const filterItems = (items: MenuItem[]): MenuItem[] => + items + .map((item) => ({ ...item, children: item.children ? filterItems(item.children) : undefined })) + .filter((item) => menuItemIsVisible(item, context, isAdmin)); + + return groups + .filter((group) => !group.roles || group.roles.includes(context.userRole)) + .map((group) => ({ groupLabel: group.groupLabel, items: filterItems(group.items) })) + .filter((group) => group.items.length > 0); +}; + // Menu groups organized by category - defined outside component for export. // Shape (key/page/label/roles/children) is consumed by page_utils.ts; only the // icons changed to lucide as part of the sidebar redesign. @@ -390,7 +470,7 @@ const menuGroups: MenuGroup[] = [ const HOME_ROUTE = "api-keys"; -const routeOf = (item: MenuItem): string => item.route ?? item.page; +export const routeOf = (item: MenuItem): string => item.route ?? item.page; const routeForPathname = (pathname: string): string => routeSegmentForPathname(pathname) || HOME_ROUTE; @@ -422,6 +502,8 @@ const SECTION_DISPLAY: Record = { SETTINGS: "Settings", }; +export const sectionText = (groupLabel: string): string => SECTION_DISPLAY[groupLabel] ?? groupLabel; + const prettify = (key: string): string => key .split(/[-_]/) @@ -435,7 +517,7 @@ export const getBreadcrumb = (pathname: string): { section: string | null; title const route = routeForPathname(pathname); for (const group of menuGroups) { for (const item of group.items) { - const section = SECTION_DISPLAY[group.groupLabel] ?? group.groupLabel; + const section = sectionText(group.groupLabel); if (routeOf(item) === route) return { section, title: labelText(item) }; const child = item.children?.find((c) => routeOf(c) === route); if (child) return { section, title: labelText(child) }; @@ -486,56 +568,32 @@ const Sidebar_: React.FC = ({ const isTeamAdmin = useMemo(() => isUserTeamAdminForAnyTeam(teams ?? null, userId ?? ""), [teams, userId]); - const filterItemsByRole = (items: MenuItem[]): MenuItem[] => { - const isAdmin = isAdminRole(userRole); - return items - .map((item) => ({ ...item, children: item.children ? filterItemsByRole(item.children) : undefined })) - .filter((item) => { - // A parent whose children were all filtered out renders as a leaf link - // to its own page id, which is not a real route. Drop it instead. - if (item.children && item.children.length === 0) return false; - if (item.key === "llm-playground" && isViewOnly) return false; - if (item.key === "organizations" || item.key === "users") { - const hasRoleAccess = !item.roles || item.roles.includes(userRole) || isOrgAdmin; - if (!hasRoleAccess) return false; - if (!isAdmin && enabledPagesInternalUsers != null) return enabledPagesInternalUsers.includes(item.page); - return true; - } - if ( - item.key === "projects" && - !(enableProjectsUI && canViewProjectsPage({ userRole, isOrgAdmin, isTeamAdmin })) - ) - return false; - if ( - !isAdmin && - item.key === "agents" && - disableAgentsForInternalUsers && - !(allowAgentsForTeamAdmins && isTeamAdmin) - ) - return false; - if ( - !isAdmin && - item.key === "vector-stores" && - disableVectorStoresForInternalUsers && - !(allowVectorStoresForTeamAdmins && isTeamAdmin) - ) - return false; - if (item.roles && !item.roles.includes(userRole)) return false; - if (!isAdmin && enabledPagesInternalUsers != null) { - if (item.children && item.children.length > 0) { - const hasVisibleChildren = item.children.some((child) => enabledPagesInternalUsers.includes(child.page)); - if (hasVisibleChildren) return true; - } - return enabledPagesInternalUsers.includes(item.page); - } - return true; - }); - }; - - const visibleGroups = menuGroups - .filter((group) => !group.roles || group.roles.includes(userRole)) - .map((group) => ({ groupLabel: group.groupLabel, items: filterItemsByRole(group.items) })) - .filter((group) => group.items.length > 0); + const visibleGroups = useMemo(() => { + const context: MenuVisibilityContext = { + userRole, + isViewOnly, + isOrgAdmin, + isTeamAdmin, + enabledPagesInternalUsers, + enableProjectsUI, + disableAgentsForInternalUsers, + allowAgentsForTeamAdmins, + disableVectorStoresForInternalUsers, + allowVectorStoresForTeamAdmins, + }; + return visibleMenuGroups(menuGroups, context); + }, [ + allowAgentsForTeamAdmins, + allowVectorStoresForTeamAdmins, + disableAgentsForInternalUsers, + disableVectorStoresForInternalUsers, + enableProjectsUI, + enabledPagesInternalUsers, + isOrgAdmin, + isTeamAdmin, + isViewOnly, + userRole, + ]); const toggleGroup = (key: string) => { if (collapsed) { diff --git a/ui/litellm-dashboard/src/components/policies/types.ts b/ui/litellm-dashboard/src/components/policies/types.ts index 430864f93df..2830a22950c 100644 --- a/ui/litellm-dashboard/src/components/policies/types.ts +++ b/ui/litellm-dashboard/src/components/policies/types.ts @@ -102,7 +102,7 @@ export interface PolicyAttachmentListResponse { export interface PipelineStepResult { guardrail_name: string; - outcome: "pass" | "fail" | "error"; + outcome: "pass" | "fail" | "error" | "skip"; action_taken: string; modified_data: Record | null; error_detail: string | null; diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx index 456dc91b13f..3335ebf965a 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx @@ -75,6 +75,7 @@ vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useAllProxyModels: vi.fn(), + useModelAccessGroupNames: vi.fn(() => new Set()), })); vi.mock("@/app/(dashboard)/hooks/teams/useTeams", async (importOriginal) => ({ diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index d721bdab2a5..b9df537884a 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -81,6 +81,7 @@ vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useAllProxyModels: vi.fn(), + useModelAccessGroupNames: vi.fn(() => new Set()), })); vi.mock("@/app/(dashboard)/hooks/teams/useTeams", async (importOriginal) => ({ @@ -231,7 +232,7 @@ vi.mock("../key_team_helpers/filter_helpers", () => ({ fetchAllOrganizations: vi.fn().mockResolvedValue([]), })); -import { useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels"; +import { useAllProxyModels, useModelAccessGroupNames } from "@/app/(dashboard)/hooks/models/useModels"; import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { teamKeys, teamsTableKeys, useTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; @@ -242,6 +243,7 @@ import { useAccessGroups } from "@/app/(dashboard)/hooks/accessGroups/useAccessG import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; const mockUseAllProxyModels = vi.mocked(useAllProxyModels); +const mockUseModelAccessGroupNames = vi.mocked(useModelAccessGroupNames); const mockUseKeys = vi.mocked(useKeys); const mockUseTeam = vi.mocked(useTeam); const mockUseOrganization = vi.mocked(useOrganization); @@ -292,6 +294,7 @@ const createMockTeamData = (overrides = {}) => ({ }); const seedDefaultMocks = () => { + mockUseModelAccessGroupNames.mockReturnValue(new Set()); mockUseAllProxyModels.mockReturnValue({ data: { data: [] }, isLoading: false, @@ -361,6 +364,35 @@ describe("TeamInfoView", () => { }); describe("display and rendering", () => { + it("links direct model chips to their matching access-group or model filter", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ models: ["repro-access-group", "gpt-4.1"] }), + ); + mockUseModelAccessGroupNames.mockReturnValue(new Set(["repro-access-group"])); + + renderWithProviders(); + + expect(await screen.findByRole("link", { name: "repro-access-group" })).toHaveAttribute( + "href", + expect.stringMatching(/\?access_group=repro-access-group$/), + ); + expect(screen.getByRole("link", { name: "gpt-4.1" })).toHaveAttribute( + "href", + expect.stringMatching(/\?model_group=gpt-4\.1$/), + ); + + await userEvent.setup({ delay: null }).click(screen.getByRole("tab", { name: "Settings" })); + const settings = await screen.findByRole("tabpanel", { name: "Settings" }); + expect(within(settings).getByRole("link", { name: "repro-access-group" })).toHaveAttribute( + "href", + expect.stringMatching(/\?access_group=repro-access-group$/), + ); + expect(within(settings).getByRole("link", { name: "gpt-4.1" })).toHaveAttribute( + "href", + expect.stringMatching(/\?model_group=gpt-4\.1$/), + ); + }); + it("should render", async () => { vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 4c62423a474..badfbee0cf1 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -26,7 +26,7 @@ import { ArrowLeftIcon } from "@heroicons/react/outline"; import { StatusBadge, type StatusTone } from "@/components/shared/table_cells/status_badge"; import { BadgeLink } from "@/components/shared/BadgeLink"; import { Badge } from "@/components/ui/badge"; -import { modelGroupHref } from "@/utils/entityLinks"; +import { modelGroupHref, modelOrAccessGroupHref } from "@/utils/entityLinks"; import { Card } from "@/components/ui/card"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; import { Input as UIInput } from "@/components/ui/input"; @@ -119,6 +119,7 @@ import { TEAM_INFO_TAB_LABELS, } from "./tabVisibilityUtils"; import TeamMembersComponent from "./TeamMemberTab"; +import { useModelAccessGroupNames } from "@/app/(dashboard)/hooks/models/useModels"; import { isValidThreshold, TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY, @@ -154,8 +155,14 @@ const TEAM_MODEL_BADGE_TONES: Record = { "access-group": "success", }; -const teamModelBadgeHref = (badge: TeamModelBadge): string | undefined => - badge.kind === "direct" || badge.kind === "access-group" ? modelGroupHref(badge.label) : undefined; +const teamModelBadgeHref = ( + badge: TeamModelBadge, + accessGroupNames: ReadonlySet | undefined, +): string | undefined => { + if (badge.kind === "direct") return modelOrAccessGroupHref(badge.label, accessGroupNames); + if (badge.kind === "access-group") return modelGroupHref(badge.label); + return undefined; +}; export type McpGrantResolution = | { readonly kind: "resolved"; readonly serverIds: ReadonlySet } @@ -622,6 +629,7 @@ const TeamInfoView: React.FC = ({ const routerSettingsRef = React.useRef(null); const [organization, setOrganization] = useState(null); const { userRole } = useAuthorized(); + const accessGroupNames = useModelAccessGroupNames(); const { data: allMcpServers = [], isError: mcpServersFailed, isLoading: mcpServersLoading } = useMCPServers(); const { data: allMcpToolsets = [], isError: mcpToolsetsFailed, isLoading: mcpToolsetsLoading } = useMCPToolsets(); const { data: allAccessGroups = [], isError: accessGroupsFailed, isLoading: accessGroupsLoading } = useAccessGroups(); @@ -1355,7 +1363,7 @@ const TeamInfoView: React.FC = ({ @@ -2191,7 +2199,7 @@ const TeamInfoView: React.FC = ({

Models

{info.models.map((model, index) => ( - + {model} ))} @@ -2202,7 +2210,7 @@ const TeamInfoView: React.FC = ({

Default Member Models

{info.default_team_member_models.map((model, index) => ( - + {model} ))} diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index 09042dea930..88f7fbb9f8c 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -23,6 +23,10 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: mockUseAuthorized, })); +vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ + useModelAccessGroupNames: vi.fn(() => new Set()), +})); + vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ useOrganizations: () => ({ data: [] }), })); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index 48f2f47ceda..f4e1fe288ec 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -1,7 +1,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import useTeams from "@/app/(dashboard)/hooks/useTeams"; import { renderWithProviders } from "../../../tests/test-utils"; -import { screen, waitFor } from "@testing-library/react"; +import { screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; @@ -9,6 +9,12 @@ import { keyDeleteCall, keyUpdateCall } from "../networking"; import { QueryClient } from "@tanstack/react-query"; import KeyInfoView, { needsLifetimeSpendBackfill } from "./key_info_view"; +vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ + useModelAccessGroupNames: vi.fn(() => new Set()), +})); + +import { useModelAccessGroupNames } from "@/app/(dashboard)/hooks/models/useModels"; + const editViewMocks = vi.hoisted(() => ({ onSubmit: undefined as ((v: Record) => Promise) | undefined, })); @@ -639,6 +645,7 @@ describe("KeyInfoView", () => { beforeEach(() => { vi.mocked(useTeams).mockReturnValue({ teams: [mockTeam], setTeams: vi.fn() }); vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock); + vi.mocked(useModelAccessGroupNames).mockReturnValue(new Set()); }); it("links the key's team by alias, resolved from the teams list, to the team page", async () => { @@ -703,6 +710,37 @@ describe("KeyInfoView", () => { ); }); + it("links access-group model chips to the access-group filter", async () => { + vi.mocked(useModelAccessGroupNames).mockReturnValue(new Set(["repro-access-group"])); + const keyData = { ...MOCK_KEY_DATA, models: ["repro-access-group"] }; + renderWithProviders( + {}} keyId="test-key-id" onKeyDataUpdate={() => {}} teams={[]} />, + ); + + expect(await screen.findByRole("link", { name: "repro-access-group" })).toHaveAttribute( + "href", + expect.stringMatching(/\?access_group=repro-access-group$/), + ); + + await userEvent.setup({ delay: null }).click(screen.getByRole("tab", { name: "Settings" })); + const settings = await screen.findByRole("tabpanel", { name: "Settings" }); + expect(within(settings).getByRole("link", { name: "repro-access-group" })).toHaveAttribute( + "href", + expect.stringMatching(/\?access_group=repro-access-group$/), + ); + }); + + it("keeps access-group model chips unlinked while access-group names are loading", async () => { + vi.mocked(useModelAccessGroupNames).mockReturnValue(undefined); + const keyData = { ...MOCK_KEY_DATA, models: ["repro-access-group"] }; + renderWithProviders( + {}} keyId="test-key-id" onKeyDataUpdate={() => {}} teams={[]} />, + ); + + expect(await screen.findAllByText("repro-access-group")).not.toHaveLength(0); + expect(screen.queryByRole("link", { name: "repro-access-group" })).not.toBeInTheDocument(); + }); + it("keeps the all-proxy-models grant chip non-clickable", async () => { const keyData = { ...MOCK_KEY_DATA, models: ["all-proxy-models"] }; renderWithProviders( diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index cfb5e9fa1f8..04a792bb5ca 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -1,4 +1,5 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { useModelAccessGroupNames } from "@/app/(dashboard)/hooks/models/useModels"; import { useProjects } from "@/app/(dashboard)/hooks/projects/useProjects"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import { useApplyUserBudgetToTeamKeys } from "@/app/(dashboard)/hooks/uiSettings/useApplyUserBudgetToTeamKeys"; @@ -14,7 +15,7 @@ import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from " import { HoverCard, HoverCardContent, HoverCardTrigger } from "@/components/ui/hover-card"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { EntityLink } from "@/components/shared/EntityLink"; -import { modelGroupHref, teamDetailHref } from "@/utils/entityLinks"; +import { modelOrAccessGroupHref, teamDetailHref } from "@/utils/entityLinks"; import { BadgeLink } from "@/components/shared/BadgeLink"; import { KeyInfoHeader } from "./KeyInfoHeader"; import KeySavingsTab from "./KeySavingsTab"; @@ -95,6 +96,7 @@ export default function KeyInfoView({ backButtonText = "Back to Keys", }: KeyInfoViewProps) { const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized(); + const accessGroupNames = useModelAccessGroupNames(); const activityDateRange = useActivityDateRange(); const queryClient = useQueryClient(); const canEditGuardrails = premiumUser || (userRole != null && rolesWithWriteAccess.includes(userRole)); @@ -752,7 +754,11 @@ export default function KeyInfoView({
{currentKeyData.models && currentKeyData.models.length > 0 ? ( currentKeyData.models.map((model, index) => ( - + {model} )) @@ -1104,7 +1110,11 @@ export default function KeyInfoView({
{currentKeyData.models && currentKeyData.models.length > 0 ? ( currentKeyData.models.map((model, index) => ( - + {model} )) diff --git a/ui/litellm-dashboard/src/components/ui/dialog.tsx b/ui/litellm-dashboard/src/components/ui/dialog.tsx index a4c17736595..e5af118a701 100644 --- a/ui/litellm-dashboard/src/components/ui/dialog.tsx +++ b/ui/litellm-dashboard/src/components/ui/dialog.tsx @@ -28,7 +28,10 @@ function DialogOverlay({ className, ...props }: DialogPrimitive.Backdrop.Props) - + ({ serverRootPath: "" })); -import { modelGroupHref, teamDetailHref, userDetailHref } from "./entityLinks"; +import { accessGroupHref, modelGroupHref, modelOrAccessGroupHref, teamDetailHref, userDetailHref } from "./entityLinks"; describe("userDetailHref", () => { it("targets the users page filtered to the encoded user id", () => { @@ -39,3 +39,30 @@ describe("modelGroupHref", () => { }, ); }); + +describe("accessGroupHref", () => { + it("targets the models page filtered to the encoded access group", () => { + expect(accessGroupHref("a b/c")).toMatch(/\/models-and-endpoints\?access_group=a%20b%2Fc$/); + }); +}); + +describe("modelOrAccessGroupHref", () => { + it("uses the access-group filter for a known access group", () => { + expect(modelOrAccessGroupHref("repro-access-group", new Set(["repro-access-group"]))).toMatch( + /\?access_group=repro-access-group$/, + ); + }); + + it("uses the model-group filter for a name outside the access-group set", () => { + expect(modelOrAccessGroupHref("gpt-4.1", new Set(["repro-access-group"]))).toMatch(/\?model_group=gpt-4\.1$/); + }); + + it("keeps grant sentinels without a link unless they are access groups", () => { + expect(modelOrAccessGroupHref("all-team-models", new Set())).toBeUndefined(); + }); + + it("keeps model names unlinked until access group names are available", () => { + expect(modelOrAccessGroupHref("repro-access-group", undefined)).toBeUndefined(); + expect(modelOrAccessGroupHref("gpt-4.1", undefined)).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/entityLinks.ts b/ui/litellm-dashboard/src/utils/entityLinks.ts index 33f2aa34976..69d47a8f497 100644 --- a/ui/litellm-dashboard/src/utils/entityLinks.ts +++ b/ui/litellm-dashboard/src/utils/entityLinks.ts @@ -29,3 +29,15 @@ export function modelGroupHref(modelGroup: string): string | undefined { if (MODEL_GRANT_SENTINELS.has(modelGroup)) return undefined; return `${uiHref("models-and-endpoints")}?model_group=${encodeURIComponent(modelGroup)}`; } + +export function accessGroupHref(accessGroup: string): string { + return `${uiHref("models-and-endpoints")}?access_group=${encodeURIComponent(accessGroup)}`; +} + +export function modelOrAccessGroupHref( + name: string, + accessGroupNames: ReadonlySet | undefined, +): string | undefined { + if (accessGroupNames === undefined) return undefined; + return accessGroupNames.has(name) ? accessGroupHref(name) : modelGroupHref(name); +}