diff --git a/.circleci/scripts/run_integration.sh b/.circleci/scripts/run_integration.sh index 0d6cdcabd57..08b0281b30f 100644 --- a/.circleci/scripts/run_integration.sh +++ b/.circleci/scripts/run_integration.sh @@ -121,6 +121,10 @@ start_proxy() { "LITELLM_MODEL_COST_MAP_URL=$INTEGRATION_UPSTREAM_URL/_cost_map" "MODEL_COST_MAP_MIN_MODEL_COUNT=1" "MODEL_COST_MAP_MAX_SHRINK_RATIO=0" + "GEMINI_API_BASE=$INTEGRATION_UPSTREAM_URL" + "ANTHROPIC_API_BASE=$INTEGRATION_UPSTREAM_URL" + "GEMINI_API_KEY=sk-scripted-provider" + "ANTHROPIC_API_KEY=sk-scripted-provider" ) else cost_map_env=("LITELLM_LOCAL_MODEL_COST_MAP=True") diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index 4fb8f068eb0..0adbc015ad0 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -205,7 +205,7 @@ def main( native_module: Final = load_native_module(native_path) native_module_loads: Final = native_module is not None panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") - native_size_limit: Final = 25_000_000 + native_size_limit: Final = 30_000_000 native_size_within_limit: Final = native_member.file_size <= native_size_limit validations: Final = ( (f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG), @@ -222,7 +222,7 @@ def main( ("Python extension entry point is present", extension_entry_point_present), ("Native module loads", native_module_loads), ("Production module omits the panic test hook", panic_test_hook_absent), - ("Native extension does not exceed 25 MB", native_size_within_limit), + ("Native extension does not exceed 30 MB", native_size_within_limit), ("Wheel contents are valid", not unexpected_members), ) @@ -267,7 +267,7 @@ def main( ), ( not native_size_within_limit, - f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB", + f"native extension exceeds 30 MB: {native_member.file_size / 1_000_000:.2f} MB", ), (bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"), ) diff --git a/README.md b/README.md index 1624d408419..e927c80b8b4 100644 --- a/README.md +++ b/README.md @@ -307,6 +307,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse | [Deepgram (`deepgram`)](https://docs.litellm.ai/docs/providers/deepgram) | ✅ | ✅ | ✅ | | | ✅ | | | | | | [DeepInfra (`deepinfra`)](https://docs.litellm.ai/docs/providers/deepinfra) | ✅ | ✅ | ✅ | | | | | | | | | [Deepseek (`deepseek`)](https://docs.litellm.ai/docs/providers/deepseek) | ✅ | ✅ | ✅ | | | | | | | | +| [Eden AI (`edenai`)](https://docs.litellm.ai/docs/providers/edenai) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | [ElevenLabs (`elevenlabs`)](https://docs.litellm.ai/docs/providers/elevenlabs) | ✅ | ✅ | ✅ | | | ✅ | ✅ | | | | | [Empower (`empower`)](https://docs.litellm.ai/docs/providers/empower) | ✅ | ✅ | ✅ | | | | | | | | | [Fal AI (`fal_ai`)](https://docs.litellm.ai/docs/providers/fal_ai) | ✅ | ✅ | ✅ | | ✅ | | | | | | @@ -356,7 +357,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse | [Petals (`petals`)](https://docs.litellm.ai/docs/providers/petals) | ✅ | ✅ | ✅ | | | | | | | | | [Pinstripes (`pinstripes`)](https://docs.litellm.ai/docs/providers/pinstripes) | ✅ | ✅ | ✅ | | | | | | | | | [Predibase (`predibase`)](https://docs.litellm.ai/docs/providers/predibase) | ✅ | ✅ | ✅ | | | | | | | | -| [Qwen AI Platform (`qwen_ai_platform`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ | +| [Qianwen AI Platform (`qwen_ai_platform`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ | | [QwenCloud (`qwencloud`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ | | [Recraft (`recraft`)](https://docs.litellm.ai/docs/providers/recraft) | | | | | ✅ | | | | | | | [Replicate (`replicate`)](https://docs.litellm.ai/docs/providers/replicate) | ✅ | ✅ | ✅ | | | | | | | | diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 0a1208787eb..89dde59c018 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -927,6 +927,12 @@ dependencies = [ "libc", ] +[[package]] +name = "crc16" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "338089f42c427b86394a5ee60ff321da23a5c89c9d89514c829687b26359fcff" + [[package]] name = "crc32fast" version = "1.5.1" @@ -3781,9 +3787,11 @@ checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed" dependencies = [ "arcstr", "combine", + "crc16", "itoa", "num-bigint 0.5.1", "percent-encoding", + "rand 0.10.2", "rustls 0.23.42", "rustls-native-certs", "ryu", diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index 5818f75ff3d..ea937098698 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -7,7 +7,7 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -redis = { version = "1.7.0", features = ["tls-rustls"] } +redis = { version = "1.7.0", features = ["cluster", "tls-rustls"] } r2d2 = "0.8.10" tokio.workspace = true diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index a960c383bf4..e2e2656fcbb 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -9,8 +9,14 @@ use litellm_cache::{ }; use redis::Commands; +use crate::topology::RedisTopology; + +mod connection; mod operations; +pub(crate) use connection::ConnectionRef; +use connection::{ClusterConnectionManager, ConnectionManager}; + pub use operations::{ RedisArg, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript, }; @@ -19,40 +25,6 @@ const DEFAULT_TTL: Duration = Duration::from_secs(600); const REDIS_TIMEOUT: Duration = Duration::from_secs(5); const REDIS_POOL_SIZE: u32 = 16; -struct PooledConnection { - connection: redis::Connection, - failed: bool, -} - -/// Pools connections without a checkout PING, which would double every operation's round trips. -/// A timed-out command leaves its reply on the socket while redis still reports the connection -/// open, so any connection whose operation failed is discarded instead of being reused. -struct ConnectionManager(redis::Client); - -impl r2d2::ManageConnection for ConnectionManager { - type Connection = PooledConnection; - type Error = redis::RedisError; - - fn connect(&self) -> Result { - let connection = self.0.get_connection()?; - connection.set_read_timeout(Some(REDIS_TIMEOUT))?; - connection.set_write_timeout(Some(REDIS_TIMEOUT))?; - Ok(PooledConnection { - connection, - failed: false, - }) - } - - fn is_valid(&self, connection: &mut PooledConnection) -> Result<(), redis::RedisError> { - redis::cmd("PING").query::(&mut connection.connection)?; - Ok(()) - } - - fn has_broken(&self, connection: &mut PooledConnection) -> bool { - connection.failed || !redis::ConnectionLike::is_open(&connection.connection) - } -} - const INCREMENT_SCRIPT: &str = concat!( "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ", "if redis.call('TTL', KEYS[1]) == -1 then ", @@ -70,42 +42,10 @@ const CLAIM_ATTEMPTS: usize = 8; enum Connections { Pool(r2d2::Pool), + Cluster(r2d2::Pool), Fixed(Mutex), } -struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike); - -impl redis::ConnectionLike for ConnectionRef<'_> { - fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult { - self.0.req_packed_command(cmd) - } - - fn req_packed_commands( - &mut self, - cmd: &[u8], - offset: usize, - count: usize, - ) -> redis::RedisResult> { - self.0.req_packed_commands(cmd, offset, count) - } - - fn get_db(&self) -> i64 { - self.0.get_db() - } - - fn supports_pipelining(&self) -> bool { - self.0.supports_pipelining() - } - - fn check_connection(&mut self) -> bool { - self.0.check_connection() - } - - fn is_open(&self) -> bool { - self.0.is_open() - } -} - impl Connections where C: redis::ConnectionLike + Send + 'static, @@ -117,13 +57,19 @@ where match self { Self::Pool(pool) => { let mut pooled = pool.get().map_err(|_| Error::Unavailable)?; - let result = operation(&mut ConnectionRef(&mut pooled.connection)); + let result = operation(&mut ConnectionRef::Node(&mut pooled.connection)); + pooled.failed = matches!(result, Err(Error::Unavailable)); + result + } + Self::Cluster(pool) => { + let mut pooled = pool.get().map_err(|_| Error::Unavailable)?; + let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection)); pooled.failed = matches!(result, Err(Error::Unavailable)); result } Self::Fixed(connection) => { let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; - operation(&mut ConnectionRef(&mut *connection)) + operation(&mut ConnectionRef::Node(&mut *connection)) } } } @@ -134,27 +80,46 @@ pub struct RedisCache { default_ttl: Duration, codec: S, namespace: Option, + topology: RedisTopology, } impl RedisCache { pub fn new(url: &str, default_ttl: Option, codec: S) -> Result { - let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; - let pool = r2d2::Pool::builder() - .max_size(REDIS_POOL_SIZE) - .min_idle(Some(0)) - .connection_timeout(REDIS_TIMEOUT) - .test_on_check_out(false) - .build(ConnectionManager(client)) - .map_err(|_| Error::Unavailable)?; + Self::connect(url, &RedisTopology::Standalone, default_ttl, codec) + } + + pub fn connect( + url: &str, + topology: &RedisTopology, + default_ttl: Option, + codec: S, + ) -> Result { + let connections = match topology { + RedisTopology::Standalone => Connections::Pool(pool(ConnectionManager::open(url)?)?), + RedisTopology::Cluster { startup_nodes } => { + Connections::Cluster(pool(ClusterConnectionManager::open(url, startup_nodes)?)?) + } + }; Ok(Self { - connections: Arc::new(Connections::Pool(pool)), + connections: Arc::new(connections), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), codec, namespace: None, + topology: topology.clone(), }) } } +fn pool(manager: M) -> Result, Error> { + r2d2::Pool::builder() + .max_size(REDIS_POOL_SIZE) + .min_idle(Some(0)) + .connection_timeout(REDIS_TIMEOUT) + .test_on_check_out(false) + .build(manager) + .map_err(|_| Error::Unavailable) +} + impl RedisCache where S: CacheCodec, @@ -166,6 +131,7 @@ where default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), codec, namespace: None, + topology: RedisTopology::Standalone, } } @@ -180,6 +146,10 @@ where self.namespace.as_deref() } + pub fn topology(&self) -> &RedisTopology { + &self.topology + } + fn namespaced_key(&self, key: &str) -> String { namespaced_key(self.namespace.as_deref(), key) } @@ -200,26 +170,14 @@ where } fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> { - let mut cursor = 0u64; - loop { - let (next_cursor, keys): (u64, Vec) = redis::cmd("SCAN") - .cursor_arg(cursor) - .arg("MATCH") - .arg(pattern) - .arg("COUNT") - .arg(1000) - .query(connection) - .map_err(|_| Error::Unavailable)?; + connection.scan(pattern, 1000, |connection, keys| { if !keys.is_empty() { connection .del::<_, usize>(keys) .map_err(|_| Error::Unavailable)?; } - if next_cursor == 0 { - return Ok(()); - } - cursor = next_cursor; - } + Ok(true) + }) } fn decode_response(&self, value: redis::Value) -> Result, Error> { @@ -350,19 +308,19 @@ where }) .collect::, _>>()?; let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); + if entries.is_empty() { + return Ok(()); + } Self::run_blocking(Arc::clone(&self.connections), move |connection| { - let mut pipeline = redis::pipe(); - for (key, payload) in entries { - pipeline - .cmd("SETEX") - .arg(key) - .arg(ttl) - .arg(payload) - .ignore(); - } - pipeline - .query::<()>(connection) - .map_err(|_| Error::Unavailable) + let commands = entries + .into_iter() + .map(|(key, payload)| { + let mut command = redis::cmd("SETEX"); + command.arg(key).arg(ttl).arg(payload); + command + }) + .collect(); + connection.pipeline(commands).map(drop) }) .await } @@ -373,7 +331,7 @@ where async fn test_connection(&self) -> Result { match Self::run_blocking(Arc::clone(&self.connections), |connection| { - Ok(match redis::cmd("PING").query::(connection) { + Ok(match connection.ping() { Ok(_) => CacheConnectionResult { status: CacheConnectionStatus::Success, message: "Redis cache connection test successful".into(), diff --git a/litellm-rust/crates/cache-redis/src/cache/connection.rs b/litellm-rust/crates/cache-redis/src/cache/connection.rs new file mode 100644 index 00000000000..1834f1d94e5 --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/cache/connection.rs @@ -0,0 +1,392 @@ +use std::collections::HashMap; + +use litellm_cache::Error; +use redis::{ + ConnectionAddr, ConnectionInfo, ConnectionLike, IntoConnectionInfo, + cluster::{ClusterClient, ClusterClientBuilder, ClusterConnection, NodeAddress}, + cluster_routing::{ + MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo, SingleNodeRoutingInfo, Slot, + }, +}; + +use super::REDIS_TIMEOUT; +use crate::topology::RedisNode; + +pub(super) struct PooledConnection { + pub(super) connection: C, + pub(super) failed: bool, +} + +/// Pools connections without a checkout PING, which would double every operation's round trips. +/// A timed-out command leaves its reply on the socket while redis still reports the connection +/// open, so any connection whose operation failed is discarded instead of being reused. +pub(super) struct ConnectionManager(redis::Client); + +impl ConnectionManager { + pub(super) fn open(url: &str) -> Result { + redis::Client::open(url) + .map(Self) + .map_err(|_| Error::Unavailable) + } +} + +impl r2d2::ManageConnection for ConnectionManager { + type Connection = PooledConnection; + type Error = redis::RedisError; + + fn connect(&self) -> Result { + let connection = self.0.get_connection()?; + connection.set_read_timeout(Some(REDIS_TIMEOUT))?; + connection.set_write_timeout(Some(REDIS_TIMEOUT))?; + Ok(PooledConnection { + connection, + failed: false, + }) + } + + fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> { + redis::cmd("PING").query::(&mut connection.connection)?; + Ok(()) + } + + fn has_broken(&self, connection: &mut Self::Connection) -> bool { + connection.failed || !redis::ConnectionLike::is_open(&connection.connection) + } +} + +pub(super) struct ClusterConnectionManager(ClusterClient); + +impl ClusterConnectionManager { + pub(super) fn open(url: &str, startup_nodes: &[RedisNode]) -> Result { + if startup_nodes.is_empty() { + return Err(Error::Unavailable); + } + let info = url.into_connection_info().map_err(|_| Error::Unavailable)?; + let nodes = startup_nodes + .iter() + .map(|node| node_info(&info, node)) + .collect::, _>>()?; + ClusterClientBuilder::new(nodes) + .connection_timeout(REDIS_TIMEOUT) + .response_timeout(REDIS_TIMEOUT) + .build() + .map(Self) + .map_err(|_| Error::Unavailable) + } +} + +fn node_info(info: &ConnectionInfo, node: &RedisNode) -> Result { + let addr = match info.addr() { + ConnectionAddr::Tcp(..) => ConnectionAddr::Tcp(node.host.clone(), node.port), + ConnectionAddr::TcpTls { + insecure, + tls_params, + .. + } => ConnectionAddr::TcpTls { + host: node.host.clone(), + port: node.port, + insecure: *insecure, + tls_params: tls_params.clone(), + }, + _ => return Err(Error::Unavailable), + }; + Ok(info.clone().set_addr(addr)) +} + +impl r2d2::ManageConnection for ClusterConnectionManager { + type Connection = PooledConnection; + type Error = redis::RedisError; + + fn connect(&self) -> Result { + let connection = self.0.get_connection()?; + connection.set_read_timeout(Some(REDIS_TIMEOUT))?; + connection.set_write_timeout(Some(REDIS_TIMEOUT))?; + Ok(PooledConnection { + connection, + failed: false, + }) + } + + fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> { + redis::cmd("PING").query::(&mut connection.connection)?; + Ok(()) + } + + fn has_broken(&self, connection: &mut Self::Connection) -> bool { + connection.failed || !redis::ConnectionLike::is_open(&connection.connection) + } +} + +pub(crate) enum ConnectionRef<'a> { + Node(&'a mut dyn redis::ConnectionLike), + Cluster(&'a mut ClusterConnection), +} + +impl redis::ConnectionLike for ConnectionRef<'_> { + fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult { + match self { + Self::Node(connection) => connection.req_packed_command(cmd), + Self::Cluster(connection) => connection.req_packed_command(cmd), + } + } + + fn req_packed_commands( + &mut self, + cmd: &[u8], + offset: usize, + count: usize, + ) -> redis::RedisResult> { + match self { + Self::Node(connection) => connection.req_packed_commands(cmd, offset, count), + Self::Cluster(connection) => connection.req_packed_commands(cmd, offset, count), + } + } + + fn get_db(&self) -> i64 { + match self { + Self::Node(connection) => connection.get_db(), + Self::Cluster(connection) => redis::ConnectionLike::get_db(*connection), + } + } + + fn supports_pipelining(&self) -> bool { + match self { + Self::Node(connection) => connection.supports_pipelining(), + Self::Cluster(connection) => redis::ConnectionLike::supports_pipelining(*connection), + } + } + + fn check_connection(&mut self) -> bool { + match self { + Self::Node(connection) => connection.check_connection(), + Self::Cluster(connection) => connection.check_connection(), + } + } + + fn is_open(&self) -> bool { + match self { + Self::Node(connection) => connection.is_open(), + Self::Cluster(connection) => redis::ConnectionLike::is_open(*connection), + } + } +} + +impl ConnectionRef<'_> { + pub(crate) fn pipeline( + &mut self, + commands: Vec, + ) -> Result, Error> { + match self { + Self::Node(connection) => { + let mut pipeline = redis::pipe(); + for command in &commands { + pipeline.add_command(command.clone()); + } + pipeline + .query::>(*connection) + .map_err(|_| Error::Unavailable) + } + Self::Cluster(connection) => { + let mut replies: Vec> = vec![None; commands.len()]; + for indices in slot_groups(&commands).into_values() { + let mut pipeline = redis::pipe(); + for index in &indices { + pipeline.add_command(commands[*index].clone()); + } + let values = connection + .req_packed_commands(&pipeline.get_packed_pipeline(), 0, indices.len()) + .map_err(|_| Error::Unavailable)?; + if values.len() != indices.len() { + return Err(Error::Unavailable); + } + for (index, value) in indices.into_iter().zip(values) { + replies[index] = Some(value); + } + } + replies + .into_iter() + .collect::>>() + .ok_or(Error::Unavailable) + } + } + } + + pub(crate) fn scan( + &mut self, + pattern: &str, + count: usize, + mut visit: impl FnMut(&mut Self, Vec) -> Result, + ) -> Result<(), Error> { + let pages = match self { + Self::Node(connection) => { + let page = scan_command(0, pattern, count) + .query::(*connection) + .map_err(|_| Error::Unavailable)?; + vec![(None, page)] + } + Self::Cluster(connection) => connection + .route_command( + &scan_command(0, pattern, count), + RoutingInfo::MultiNode(( + MultipleNodeRoutingInfo::AllMasters, + Some(ResponsePolicy::Special), + )), + ) + .map_err(|_| Error::Unavailable) + .and_then(primary_pages)? + .into_iter() + .map(|(node, page)| (Some(node), page)) + .collect(), + }; + for (node, (mut cursor, mut keys)) in pages { + loop { + if !visit(self, keys)? { + return Ok(()); + } + if cursor == 0 { + break; + } + (cursor, keys) = self.scan_page(node.as_ref(), cursor, pattern, count)?; + } + } + Ok(()) + } + + pub(crate) fn ping(&mut self) -> Result { + let command = redis::cmd("PING"); + match self { + Self::Node(connection) => command + .query::(*connection) + .map(|response| response == "PONG"), + Self::Cluster(connection) => connection + .route_command( + &command, + RoutingInfo::MultiNode(( + MultipleNodeRoutingInfo::AllNodes, + Some(ResponsePolicy::AllSucceeded), + )), + ) + .map(|_| true), + } + } + + pub(crate) fn node_text(&mut self, command: &redis::Cmd) -> Result { + match self { + Self::Node(connection) => command.query(*connection).map_err(|_| Error::Unavailable), + Self::Cluster(connection) => { + let value = connection + .route_command( + command, + RoutingInfo::MultiNode(( + MultipleNodeRoutingInfo::AllNodes, + Some(ResponsePolicy::Special), + )), + ) + .map_err(|_| Error::Unavailable)?; + let redis::Value::Map(entries) = value else { + return Err(Error::Unavailable); + }; + let mut replies = entries + .into_iter() + .map(|(node, reply)| { + Ok(( + redis::from_redis_value::(node) + .map_err(|_| Error::Unavailable)?, + redis::from_redis_value::(reply) + .map_err(|_| Error::Unavailable)?, + )) + }) + .collect::, Error>>()?; + replies.sort(); + Ok(replies + .into_iter() + .map(|(_, reply)| reply) + .collect::>() + .join("\n")) + } + } + } + + pub(crate) fn flushall(&mut self) -> Result<(), Error> { + let command = redis::cmd("FLUSHALL"); + match self { + Self::Node(connection) => command.query(*connection).map_err(|_| Error::Unavailable), + Self::Cluster(connection) => connection + .route_command( + &command, + RoutingInfo::MultiNode(( + MultipleNodeRoutingInfo::AllMasters, + Some(ResponsePolicy::AllSucceeded), + )), + ) + .map(|_| ()) + .map_err(|_| Error::Unavailable), + } + } + + fn scan_page( + &mut self, + node: Option<&NodeAddress>, + cursor: u64, + pattern: &str, + count: usize, + ) -> Result { + let command = scan_command(cursor, pattern, count); + match (self, node) { + (Self::Node(connection), None) => { + command.query(*connection).map_err(|_| Error::Unavailable) + } + (Self::Cluster(connection), Some(node)) => connection + .route_command( + &command, + RoutingInfo::SingleNode(SingleNodeRoutingInfo::ByAddress { + host: node.host().to_string(), + port: node.port(), + }), + ) + .map_err(|_| Error::Unavailable) + .and_then(|value| redis::from_redis_value(value).map_err(|_| Error::Unavailable)), + _ => Err(Error::Unavailable), + } + } +} + +type ScanPage = (u64, Vec); + +fn primary_pages(value: redis::Value) -> Result, Error> { + let redis::Value::Map(entries) = value else { + return Err(Error::Unavailable); + }; + entries + .into_iter() + .map(|(node, page)| { + let node = redis::from_redis_value::(node).map_err(|_| Error::Unavailable)?; + let node = NodeAddress::try_from(node.as_str()).map_err(|_| Error::Unavailable)?; + let page = redis::from_redis_value::(page).map_err(|_| Error::Unavailable)?; + Ok((node, page)) + }) + .collect() +} + +fn scan_command(cursor: u64, pattern: &str, count: usize) -> redis::Cmd { + let mut command = redis::cmd("SCAN"); + command + .cursor_arg(cursor) + .arg("MATCH") + .arg(pattern) + .arg("COUNT") + .arg(count); + command +} + +fn slot_groups(commands: &[redis::Cmd]) -> HashMap> { + let mut groups: HashMap> = HashMap::new(); + for (index, command) in commands.iter().enumerate() { + let key = match command.args_iter().nth(1) { + Some(redis::Arg::Simple(key)) => key, + _ => b"", + }; + groups.entry(Slot::for_key(key)).or_default().push(index); + } + groups +} diff --git a/litellm-rust/crates/cache-redis/src/cache/operations.rs b/litellm-rust/crates/cache-redis/src/cache/operations.rs index d8d9ae24c4c..9a7023338bf 100644 --- a/litellm-rust/crates/cache-redis/src/cache/operations.rs +++ b/litellm-rust/crates/cache-redis/src/cache/operations.rs @@ -183,20 +183,13 @@ where } pub fn sync_ping(&self) -> Result { - self.connections.execute(|connection| { - redis::cmd("PING") - .query::(connection) - .map(|response| response == "PONG") - .map_err(|_| Error::Unavailable) - }) + self.connections + .execute(|connection| connection.ping().map_err(|_| Error::Unavailable)) } pub async fn ping(&self) -> Result { Self::run_blocking(Arc::clone(&self.connections), |connection| { - redis::cmd("PING") - .query::(connection) - .map(|response| response == "PONG") - .map_err(|_| Error::Unavailable) + connection.ping().map_err(|_| Error::Unavailable) }) .await } @@ -216,24 +209,13 @@ where pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result, Error> { let pattern = format!("{}*", self.namespaced_key(pattern)); Self::run_blocking(Arc::clone(&self.connections), move |connection| { - let mut cursor = 0u64; let mut matches = Vec::new(); - loop { - let (next_cursor, keys): (u64, Vec) = redis::cmd("SCAN") - .cursor_arg(cursor) - .arg("MATCH") - .arg(&pattern) - .arg("COUNT") - .arg(count) - .query(connection) - .map_err(|_| Error::Unavailable)?; + connection.scan(&pattern, count, |_, keys| { matches.extend(keys); - if matches.len() >= count || next_cursor == 0 { - matches.truncate(count); - return Ok(matches); - } - cursor = next_cursor; - } + Ok(matches.len() < count) + })?; + matches.truncate(count); + Ok(matches) }) .await } @@ -250,13 +232,18 @@ where let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl)); Self::run_blocking(Arc::clone(&self.connections), move |connection| { - let mut pipeline = redis::pipe(); - pipeline.cmd("SADD").arg(&key).arg(values); - pipeline.cmd("EXPIRE").arg(&key).arg(ttl).ignore(); - pipeline - .query::<(usize,)>(connection) - .map(|(added,)| added) - .map_err(|_| Error::Unavailable) + let mut sadd = redis::cmd("SADD"); + sadd.arg(&key).arg(values); + let mut expire = redis::cmd("EXPIRE"); + expire.arg(&key).arg(ttl); + let replies = connection.pipeline(vec![sadd, expire])?; + replies + .into_iter() + .next() + .map(redis::from_redis_value::) + .transpose() + .map_err(|_| Error::Unavailable)? + .ok_or(Error::Unavailable) }) .await } @@ -293,11 +280,19 @@ where return Ok(Vec::new()); } Self::run_blocking(Arc::clone(&self.connections), move |connection| { - let mut pipeline = redis::pipe(); - for (key, values) in operations { - pipeline.cmd("RPUSH").arg(key).arg(values); - } - pipeline.query(connection).map_err(|_| Error::Unavailable) + let commands = operations + .into_iter() + .map(|(key, values)| { + let mut command = redis::cmd("RPUSH"); + command.arg(key).arg(values); + command + }) + .collect(); + connection + .pipeline(commands)? + .into_iter() + .map(|value| redis::from_redis_value(value).map_err(|_| Error::Unavailable)) + .collect() }) .await } @@ -339,16 +334,18 @@ where .map(|(_, count)| count.is_some()) .collect::>(); let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { - let mut pipeline = redis::pipe(); - for (key, count) in operations { - let command = pipeline.cmd("LPOP").arg(key); - if let Some(count) = count { - command.arg(count); - } - } - pipeline - .query::>(connection) - .map_err(|_| Error::Unavailable) + let commands = operations + .into_iter() + .map(|(key, count)| { + let mut command = redis::cmd("LPOP"); + command.arg(key); + if let Some(count) = count { + command.arg(count); + } + command + }) + .collect(); + connection.pipeline(commands) }) .await?; values @@ -381,28 +378,17 @@ where } pub fn client_list(&self) -> Result { - self.connections.execute(|connection| { - redis::cmd("CLIENT") - .arg("LIST") - .query(connection) - .map_err(|_| Error::Unavailable) - }) + self.connections + .execute(|connection| connection.node_text(redis::cmd("CLIENT").arg("LIST"))) } pub fn info(&self) -> Result { - self.connections.execute(|connection| { - redis::cmd("INFO") - .query(connection) - .map_err(|_| Error::Unavailable) - }) + self.connections + .execute(|connection| connection.node_text(&redis::cmd("INFO"))) } pub fn flushall(&self) -> Result<(), Error> { - self.connections.execute(|connection| { - redis::cmd("FLUSHALL") - .query(connection) - .map_err(|_| Error::Unavailable) - }) + self.connections.execute(|connection| connection.flushall()) } } @@ -441,14 +427,27 @@ where return Ok(Vec::new()); } Self::run_blocking(Arc::clone(&self.connections), move |connection| { - let mut pipeline = redis::pipe(); + let mut commands = Vec::with_capacity(operations.len() * 2); + let mut increments = Vec::with_capacity(operations.len()); for (key, amount, ttl) in operations { - pipeline.cmd("INCRBYFLOAT").arg(&key).arg(amount); + let mut increment = redis::cmd("INCRBYFLOAT"); + increment.arg(&key).arg(amount); + increments.push(commands.len()); + commands.push(increment); if let Some(ttl) = ttl { - pipeline.cmd("EXPIRE").arg(key).arg(ttl).ignore(); + let mut expire = redis::cmd("EXPIRE"); + expire.arg(key).arg(ttl); + commands.push(expire); } } - pipeline.query(connection).map_err(|_| Error::Unavailable) + let mut replies = connection.pipeline(commands)?; + increments + .into_iter() + .map(|index| { + redis::from_redis_value(std::mem::take(&mut replies[index])) + .map_err(|_| Error::Unavailable) + }) + .collect() }) .await } diff --git a/litellm-rust/crates/cache-redis/tests/cluster.rs b/litellm-rust/crates/cache-redis/tests/cluster.rs new file mode 100644 index 00000000000..2c3fc818b66 --- /dev/null +++ b/litellm-rust/crates/cache-redis/tests/cluster.rs @@ -0,0 +1,492 @@ +//! Contract tests against a real Redis Cluster. Set `LITELLM_TEST_REDIS_CLUSTER_NODES` to a +//! comma separated `host:port` list (for example `127.0.0.1:7000,127.0.0.1:7001`) to run them. + +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, CacheConnectionStatus, CacheScript, ClaimCache, + CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec, + ScriptCache, +}; +use litellm_cache_redis::{ + RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisNode, RedisRpushOperation, + RedisTopology, +}; +use redis::cluster_routing::Slot; + +type Cache = RedisCache>; + +fn topology() -> Option { + let nodes = std::env::var("LITELLM_TEST_REDIS_CLUSTER_NODES").ok()?; + let startup_nodes = nodes + .split(',') + .map(|node| { + let (host, port) = node.trim().rsplit_once(':').expect("host:port"); + RedisNode { + host: host.to_string(), + port: port.parse().expect("port"), + } + }) + .collect(); + Some(RedisTopology::Cluster { startup_nodes }) +} + +fn namespace(label: &str) -> String { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + format!("cluster-test:{label}:{nanos}") +} + +fn cluster_url() -> String { + std::env::var("LITELLM_TEST_REDIS_CLUSTER_URL") + .unwrap_or_else(|_| "redis://127.0.0.1:7000".into()) +} + +fn cluster_cache(label: &str) -> Option { + let topology = topology()?; + Some( + Cache::connect( + &cluster_url(), + &topology, + Some(Duration::from_secs(120)), + JsonCodec::new(), + ) + .expect("cluster connection") + .with_namespace(Some(namespace(label))), + ) +} + +fn counter_cache(label: &str) -> Option>> { + let topology = topology()?; + Some( + RedisCache::connect( + &cluster_url(), + &topology, + Some(Duration::from_secs(60)), + JsonCodec::new(), + ) + .expect("cluster connection") + .with_namespace(Some(namespace(label))), + ) +} + +fn multi_slot_keys(count: usize) -> Vec { + let keys: Vec = (0..count).map(|index| format!("key-{index}")).collect(); + let slots: std::collections::HashSet = keys.iter().map(Slot::for_key).collect(); + assert!(slots.len() > 1, "keys must span multiple slots"); + keys +} + +macro_rules! cluster_or_skip { + ($label:expr) => { + match cluster_cache($label) { + Some(cache) => cache, + None => return, + } + }; +} + +#[test] +fn constructor_rejects_clusters_without_startup_nodes() { + let error = Cache::connect( + "redis://127.0.0.1:7000", + &RedisTopology::Cluster { + startup_nodes: Vec::new(), + }, + None, + JsonCodec::new(), + ) + .err(); + assert!(matches!(error, Some(Error::Unavailable))); +} + +#[test] +fn constructor_rejects_unix_socket_urls_for_clusters() { + let error = Cache::connect( + "redis+unix:///tmp/redis.sock", + &RedisTopology::Cluster { + startup_nodes: vec![RedisNode { + host: "127.0.0.1".into(), + port: 7000, + }], + }, + None, + JsonCodec::new(), + ) + .err(); + assert!(matches!(error, Some(Error::Unavailable))); +} + +#[test] +fn single_key_operations_round_trip_with_ttl_rounding() { + let cache = cluster_or_skip!("single"); + let context = ExactCacheContext { + ttl: Some(Duration::from_millis(1500)), + }; + let keys = multi_slot_keys(12); + for (index, key) in keys.iter().enumerate() { + cache + .set_cache(key, serde_json::json!({ "index": index }), &context) + .unwrap(); + } + for (index, key) in keys.iter().enumerate() { + assert_eq!( + cache.get_cache(key, &context).unwrap(), + Some(serde_json::json!({ "index": index })) + ); + } + let runtime = tokio::runtime::Runtime::new().unwrap(); + let ttl = runtime.block_on(cache.async_get_ttl(&keys[0])).unwrap(); + assert_eq!(ttl, Some(2)); + cache.delete_cache(&keys[0]).unwrap(); + assert_eq!(cache.get_cache(&keys[0], &context).unwrap(), None); + assert!(cache.sync_ping().unwrap()); +} + +#[tokio::test] +async fn batch_reads_span_slots_and_preserve_order_with_malformed_entries() { + let cache = cluster_or_skip!("batch"); + let context = ExactCacheContext::default(); + let keys = multi_slot_keys(40); + for (index, key) in keys.iter().enumerate() { + if index % 5 == 0 { + continue; + } + cache + .async_set_cache(key, serde_json::json!(index), context.clone()) + .await + .unwrap(); + } + let mut raw = redis::cluster::ClusterClient::new(vec![cluster_url()]) + .unwrap() + .get_connection() + .unwrap(); + let malformed = format!("{}:{}", cache.namespace().unwrap(), keys[1]); + redis::cmd("SET") + .arg(&malformed) + .arg("not json") + .exec(&mut raw) + .unwrap(); + + let entries = cache + .async_batch_get_cache(keys.clone(), context.clone()) + .await + .unwrap(); + assert_eq!(entries.len(), keys.len()); + for (index, entry) in entries.iter().enumerate() { + let expected = if index == 1 { + BatchEntry::Invalid + } else if index % 5 == 0 { + BatchEntry::Miss + } else { + BatchEntry::Hit(serde_json::json!(index)) + }; + assert_eq!(*entry, expected, "entry {index}"); + } + let sync_entries = cache.batch_get_cache(&keys, &context).unwrap(); + assert_eq!(sync_entries, entries); + + cache.delete_cache_keys(keys.clone()).await.unwrap(); + let entries = cache.async_batch_get_cache(keys, context).await.unwrap(); + assert!(entries.iter().all(|entry| *entry == BatchEntry::Miss)); +} + +#[tokio::test] +async fn pipelines_group_by_slot_and_return_results_in_submission_order() { + let cache = cluster_or_skip!("pipeline"); + let keys = multi_slot_keys(30); + let entries = keys + .iter() + .enumerate() + .map(|(index, key)| (key.clone(), serde_json::json!(index))) + .collect(); + cache + .async_set_cache_pipeline(entries, ExactCacheContext::default()) + .await + .unwrap(); + let hits = cache + .async_batch_get_cache(keys.clone(), ExactCacheContext::default()) + .await + .unwrap(); + assert!( + hits.iter() + .enumerate() + .all(|(index, entry)| *entry == BatchEntry::Hit(serde_json::json!(index))) + ); + + let queues: Vec = keys.iter().map(|key| format!("queue:{key}")).collect(); + let pushed = cache + .async_rpush_pipeline( + queues + .iter() + .enumerate() + .map(|(index, key)| RedisRpushOperation { + key: key.clone(), + values: (0..=index) + .map(|value| RedisArg::Integer(value as i64)) + .collect(), + }) + .collect(), + ) + .await + .unwrap(); + assert_eq!(pushed, (1..=keys.len()).collect::>()); + let popped = cache + .async_lpop_pipeline( + queues + .iter() + .enumerate() + .map(|(index, key)| RedisLpopOperation { + key: key.clone(), + count: (index % 2 == 0).then_some(2), + }) + .collect(), + ) + .await + .unwrap(); + for (index, result) in popped.into_iter().enumerate() { + match result { + RedisLpopResult::Value(value) => { + assert_eq!(index % 2, 1, "queue {index}"); + assert_eq!(value, b"0"); + } + RedisLpopResult::Values(values) => { + assert_eq!(index % 2, 0, "queue {index}"); + let expected: Vec> = (0..=index) + .take(2) + .map(|value| value.to_string().into_bytes()) + .collect(); + assert_eq!(values, expected); + } + other => panic!("queue {index}: {other:?}"), + } + } + + let counters: Vec = keys.iter().map(|key| format!("counter:{key}")).collect(); + let Some(counter) = counter_cache("counter") else { + return; + }; + let totals = counter + .async_increment_pipeline( + counters + .iter() + .enumerate() + .map(|(index, key)| IncrementOperation { + key: key.clone(), + amount: index as f64 + 0.5, + ttl: (index % 3 == 0).then_some(Duration::from_secs(30)), + }) + .collect(), + ) + .await + .unwrap(); + let expected: Vec = (0..keys.len()).map(|index| index as f64 + 0.5).collect(); + assert_eq!(totals, expected); + assert_eq!(counter.async_get_ttl(&counters[0]).await.unwrap(), Some(30)); + assert_eq!(counter.async_get_ttl(&counters[1]).await.unwrap(), None); + counter.async_flush_cache().await.unwrap(); + cache.async_flush_cache().await.unwrap(); +} + +#[tokio::test] +async fn scan_and_scoped_flush_cover_every_primary() { + let cache = cluster_or_skip!("flush"); + let other = cluster_or_skip!("other"); + let context = ExactCacheContext::default(); + let keys = multi_slot_keys(60); + for key in &keys { + cache + .async_set_cache(key, serde_json::json!(true), context.clone()) + .await + .unwrap(); + other + .async_set_cache(key, serde_json::json!(true), context.clone()) + .await + .unwrap(); + } + let mut scanned = cache.async_scan_iter("key-", 1000).await.unwrap(); + scanned.sort(); + let mut expected: Vec = keys + .iter() + .map(|key| format!("{}:{key}", cache.namespace().unwrap())) + .collect(); + expected.sort(); + assert_eq!(scanned, expected); + assert_eq!(cache.async_scan_iter("key-", 7).await.unwrap().len(), 7); + + cache.flush_cache().unwrap(); + let flushed = cache + .async_batch_get_cache(keys.clone(), context.clone()) + .await + .unwrap(); + assert!(flushed.iter().all(|entry| *entry == BatchEntry::Miss)); + let kept = other.async_batch_get_cache(keys, context).await.unwrap(); + assert!( + kept.iter() + .all(|entry| *entry == BatchEntry::Hit(serde_json::json!(true))) + ); + other.async_flush_cache().await.unwrap(); +} + +fn ping_calls_per_node(startup: &redis::Client) -> Vec<(String, u64)> { + let mut connection = startup.get_connection().unwrap(); + let nodes: String = redis::cmd("CLUSTER") + .arg("NODES") + .query(&mut connection) + .unwrap(); + let mut counts: Vec<(String, u64)> = nodes + .lines() + .map(|line| { + let address = line.split_whitespace().nth(1).unwrap(); + let address = address.split('@').next().unwrap(); + let mut node = redis::Client::open(format!("redis://{address}")) + .unwrap() + .get_connection() + .unwrap(); + let stats: String = redis::cmd("INFO") + .arg("commandstats") + .query(&mut node) + .unwrap(); + let calls = stats + .lines() + .find_map(|stat| stat.strip_prefix("cmdstat_ping:calls=")) + .and_then(|rest| rest.split(',').next()) + .map_or(0, |calls| calls.parse().unwrap()); + (address.to_string(), calls) + }) + .collect(); + counts.sort(); + counts +} + +#[tokio::test] +async fn ping_reaches_every_node() { + let cache = cluster_or_skip!("ping"); + let startup = redis::Client::open(cluster_url()).unwrap(); + let before = ping_calls_per_node(&startup); + assert!(before.len() >= 2, "{before:?}"); + assert!(cache.ping().await.unwrap()); + let after = ping_calls_per_node(&startup); + for ((node, calls_before), (_, calls_after)) in before.iter().zip(&after) { + assert!(calls_after > calls_before, "{node} was not pinged"); + } + assert!(cache.sync_ping().unwrap()); + let result = cache.test_connection().await.unwrap(); + assert_eq!(result.status, CacheConnectionStatus::Success); +} + +#[tokio::test] +async fn counters_claims_scripts_and_sets_work_on_the_cluster() { + let Some(counter) = counter_cache("counter") else { + return; + }; + let context = ExactCacheContext::default(); + assert_eq!( + counter + .increment_cache("spend", 1.5, context.clone()) + .unwrap(), + 1.5 + ); + assert_eq!( + counter + .async_increment("spend", 2.0, context.clone()) + .await + .unwrap(), + 3.5 + ); + assert_eq!( + counter + .increment_with_floor("budget", -3, Duration::from_secs(30)) + .unwrap(), + 0 + ); + assert_eq!( + counter + .async_increment_with_floor("budget", 7, Duration::from_secs(30)) + .await + .unwrap(), + 7 + ); + assert_eq!(counter.async_set_max("peak", 4.0, None).await.unwrap(), 4.0); + assert_eq!(counter.async_set_max("peak", 2.0, None).await.unwrap(), 4.0); + counter.flush_cache().unwrap(); + + let cache = cluster_or_skip!("claim"); + let owner = serde_json::json!("owner-a"); + let rival = serde_json::json!("owner-b"); + assert_eq!( + cache + .claim_cache("lock", owner.clone(), &[], context.clone()) + .unwrap(), + owner + ); + assert_eq!( + cache + .async_claim_cache("lock", rival.clone(), vec![owner.clone()], context.clone()) + .await + .unwrap(), + owner + ); + assert_eq!( + cache + .claim_cache("lock", rival.clone(), &[], context.clone()) + .unwrap(), + owner + ); + assert_eq!( + cache + .async_claim_cache("lock", rival.clone(), vec![rival.clone()], context.clone()) + .await + .unwrap(), + rival + ); + + let script = cache + .async_register_script("return redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])".into()); + let reply = script + .invoke( + vec!["scripted".into()], + vec![RedisArg::Bytes(b"payload".to_vec()), RedisArg::Integer(5)], + ) + .await + .unwrap(); + assert_eq!(reply, redis::Value::Okay); + assert_eq!(cache.async_get_ttl("scripted").await.unwrap(), Some(5)); + let evaluated: redis::Value = cache + .async_eval( + "return redis.call('GET', KEYS[1])".into(), + vec!["scripted".into()], + Vec::new(), + ) + .await + .unwrap(); + assert_eq!(evaluated, redis::Value::BulkString(b"payload".to_vec())); + + assert_eq!( + cache + .async_set_cache_sadd( + "members", + vec![ + RedisArg::Bytes(b"a".to_vec()), + RedisArg::Bytes(b"b".to_vec()) + ], + Some(Duration::from_secs(9)), + ) + .await + .unwrap(), + 2 + ); + assert_eq!(cache.async_get_ttl("members").await.unwrap(), Some(9)); + + let result = cache.test_connection().await.unwrap(); + assert_eq!(result.status, CacheConnectionStatus::Success); + assert!(cache.ping().await.unwrap()); + let info = cache.info().unwrap(); + assert!(info.matches("redis_version").count() > 1, "{info}"); + assert!(cache.client_list().unwrap().contains("id=")); + cache.async_flush_cache().await.unwrap(); + assert_eq!(cache.async_get_ttl("members").await.unwrap(), None); + assert_eq!(cache.get_cache("lock", &context).unwrap(), None); +} diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 17c4cb40903..745edc978d3 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -1,10 +1,11 @@ use std::{path::PathBuf, time::Duration}; use litellm_cache::CacheType; +use litellm_cache_redis::{RedisNode, RedisTopology}; use pyo3::{ exceptions::{PyTypeError, PyValueError}, prelude::*, - types::{PyAny, PyDict, PyString}, + types::{PyAny, PyDict, PyList, PyString}, }; use super::{native::NativeResponseCache, request::duration}; @@ -74,9 +75,21 @@ pub(super) struct RedisCacheConfig { pub(super) default_ttl: Duration, pub(super) namespace: Option, pub(super) flush_size: usize, + pub(super) topology: RedisTopology, pub(super) connection: RedisConnectionConfig, } +struct RedisClientProjection<'py> { + topology: RedisTopology, + host: String, + port: u16, + pool_size: usize, + resolved: Bound<'py, PyDict>, + tls: Option, +} + +const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31; + pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), @@ -195,6 +208,9 @@ impl NativeCacheConfig { CacheBackendConfig::Redis(_) if service.kind() != "redis" => { Some("facade and native backend types must match") } + CacheBackendConfig::Redis(config) if service.topology() != Some(&config.topology) => { + Some("facade and native backend topologies must match") + } CacheBackendConfig::Redis(config) => (service.namespace() != config.namespace.as_deref()) .then_some("facade and native backend namespaces must match"), @@ -245,9 +261,6 @@ fn project_redis( backend: &Bound<'_, PyAny>, ) -> PyResult> { let source = backend.getattr("redis_kwargs")?.cast_into::()?; - if has_value(&source, "startup_nodes")? { - return Ok(Err(UnsupportedCacheConfig::RedisTopology)); - } if has_value(&source, "sentinel_nodes")? { return Ok(Err(UnsupportedCacheConfig::RedisTopology)); } @@ -287,26 +300,25 @@ fn project_redis( } let client = backend.getattr("redis_client")?; - let pool = client.getattr("connection_pool")?; - if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? { - return Ok(Err(UnsupportedCacheConfig::RedisConnection)); - } - let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; - for key in ["credential_provider", "redis_connect_func"] { - if has_value(&resolved, key)? { - return Ok(Err(UnsupportedCacheConfig::RedisCredentials)); - } - } - let connection_class = resolved - .get_item("connection_class")? - .unwrap_or(pool.getattr("connection_class")?); - let tls = if class_is(&connection_class, "redis.connection", "Connection")? { - None - } else if class_is(&connection_class, "redis.connection", "SSLConnection")? { - Some(project_tls(&resolved)?) + let projection = if has_value(&source, "startup_nodes")? { + project_cluster_client(&source, &client)? } else { - return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + project_standalone_client(&client)? }; + let RedisClientProjection { + topology, + host, + port, + pool_size, + resolved, + tls, + } = match projection { + Ok(projection) => projection, + Err(reason) => return Ok(Err(reason)), + }; + if has_value(&resolved, "credential_provider")? { + return Ok(Err(UnsupportedCacheConfig::RedisCredentials)); + } let protocol = match optional_i64(&resolved, "protocol")?.unwrap_or(2) { 2 => RedisProtocol::Resp2, @@ -319,15 +331,15 @@ fn project_redis( default_ttl: duration(backend.getattr("default_ttl")?.extract::()?)?, namespace: optional_attribute_string(backend, "namespace")?, flush_size: backend.getattr("redis_flush_size")?.extract::()?, + topology, connection: RedisConnectionConfig { - host: required_string(&resolved, "host")?, - port: u16::try_from(required_i64(&resolved, "port")?) - .map_err(|_| PyValueError::new_err("invalid Redis port"))?, + host, + port, database: optional_i64(&resolved, "db")?.unwrap_or(0), username: optional_dict_string(&resolved, "username")?, password: optional_dict_string(&resolved, "password")?, protocol, - pool_size: pool.getattr("max_connections")?.extract::()?, + pool_size, read_timeout: optional_dict_duration(&resolved, "socket_timeout")?, connect_timeout: optional_dict_duration(&resolved, "socket_connect_timeout")?, socket_keepalive: optional_bool(&resolved, "socket_keepalive")?, @@ -338,6 +350,128 @@ fn project_redis( })) } +#[inline(never)] +fn project_standalone_client<'py>( + client: &Bound<'py, PyAny>, +) -> PyResult, UnsupportedCacheConfig>> { + let pool = client.getattr("connection_pool")?; + if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } + let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; + if has_value(&resolved, "redis_connect_func")? { + return Ok(Err(UnsupportedCacheConfig::RedisCredentials)); + } + let connection_class = resolved + .get_item("connection_class")? + .unwrap_or(pool.getattr("connection_class")?); + let tls = if class_is(&connection_class, "redis.connection", "Connection")? { + None + } else if class_is(&connection_class, "redis.connection", "SSLConnection")? { + Some(project_tls(&resolved)?) + } else { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + }; + Ok(Ok(RedisClientProjection { + topology: RedisTopology::Standalone, + host: required_string(&resolved, "host")?, + port: port(required_i64(&resolved, "port")?)?, + pool_size: pool.getattr("max_connections")?.extract::()?, + resolved, + tls, + })) +} + +#[inline(never)] +fn project_cluster_client<'py>( + source: &Bound<'py, PyDict>, + client: &Bound<'py, PyAny>, +) -> PyResult, UnsupportedCacheConfig>> { + let Some(startup_nodes) = startup_nodes(source)? else { + return Ok(Err(UnsupportedCacheConfig::RedisTopology)); + }; + if !instance_class_is(client, "redis.cluster", "RedisCluster")? { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } + let nodes = client.getattr("nodes_manager")?; + if !class_is( + &nodes.getattr("connection_pool_class")?, + "redis.connection", + "ConnectionPool", + )? { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } + let resolved = nodes.getattr("connection_kwargs")?.cast_into::()?; + if let Some(connect) = resolved.get_item("redis_connect_func")? + && !connect.is_none() + { + let own_hook = connect + .getattr("__self__") + .is_ok_and(|owner| owner.is(client)) + && connect + .getattr("__func__") + .and_then(|function| Ok(function.is(&client.get_type().getattr("on_connect")?))) + .unwrap_or(false); + if !own_hook { + return Ok(Err(UnsupportedCacheConfig::RedisCredentials)); + } + } + let tls = if optional_bool(&resolved, "ssl")?.unwrap_or(false) { + Some(project_tls(&resolved)?) + } else { + None + }; + let first = &startup_nodes[0]; + Ok(Ok(RedisClientProjection { + host: first.host.clone(), + port: first.port, + pool_size: optional_i64(&resolved, "max_connections")? + .map(|value| { + usize::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis pool size")) + }) + .transpose()? + .unwrap_or(REDIS_PY_DEFAULT_MAX_CONNECTIONS), + topology: RedisTopology::Cluster { startup_nodes }, + resolved, + tls, + })) +} + +#[inline(never)] +fn startup_nodes(source: &Bound<'_, PyDict>) -> PyResult>> { + let Some(nodes) = source.get_item("startup_nodes")? else { + return Ok(None); + }; + let Ok(nodes) = nodes.cast_into::() else { + return Ok(None); + }; + if nodes.is_empty() { + return Ok(None); + } + let mut parsed = Vec::with_capacity(nodes.len()); + for node in nodes.iter() { + let Ok(node) = node.cast_into::() else { + return Ok(None); + }; + if node.len() != 2 || !has_value(&node, "host")? || !has_value(&node, "port")? { + return Ok(None); + } + let (Ok(host), Ok(port)) = ( + required_string(&node, "host"), + required_i64(&node, "port").and_then(port), + ) else { + return Ok(None); + }; + parsed.push(RedisNode { host, port }); + } + Ok(Some(parsed)) +} + +#[inline(never)] +fn port(value: i64) -> PyResult { + u16::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis port")) +} + #[inline(never)] fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { Ok(RedisTlsConfig { @@ -506,12 +640,27 @@ mod tests { use pyo3::{prelude::*, types::PyDict}; + use litellm_cache_redis::{RedisNode, RedisTopology}; + use super::{ CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement, DiskCacheConfig, NativeCacheConfig, RedisProtocol, }; use crate::cache::native::NativeResponseCache; + fn cluster_facade<'py>(py: Python<'py>, startup_nodes: &str, hook: &str) -> Bound<'py, PyAny> { + facade( + py, + &format!( + "RedisCluster = type('RedisCluster', (), {{'__module__': 'redis.cluster', 'on_connect': lambda self, connection: None}})\n\ + client = RedisCluster()\n\ + client.nodes_manager = SimpleNamespace(connection_pool_class=ConnectionPool, connection_kwargs={{'password': 'secret', 'redis_connect_func': {hook}, 'protocol': 3, 'ssl': True, 'ssl_cert_reqs': 'none'}})\n\ + backend = SimpleNamespace(default_ttl=120, namespace='team', redis_flush_size=100, redis_kwargs={{'startup_nodes': {startup_nodes}, 'password': 'secret'}}, redis_client=client)\n\ + facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=100, semantic_cache_scope='key', cache=backend)" + ), + ) + } + fn facade<'py>(py: Python<'py>, body: &str) -> Bound<'py, PyAny> { let locals = PyDict::new(py); py.run( @@ -630,7 +779,6 @@ mod tests { assert_eq!(reason.message(), "native Redis credentials require Python"); }); } - #[test] fn projects_builtin_disk_configuration_and_rejects_custom_stores() { Python::initialize(); @@ -710,4 +858,87 @@ mod tests { ); }); } + + #[test] + fn projects_cluster_startup_nodes_as_redis_topology() { + Python::initialize(); + Python::attach(|py| { + let facade = cluster_facade( + py, + "[{'host': 'node-a', 'port': 7000}, {'host': 'node-b', 'port': 7001}]", + "client.on_connect", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("cluster startup nodes should project natively"); + }; + let CacheBackendConfig::Redis(redis) = &config.backend else { + panic!("expected Redis configuration"); + }; + let expected = RedisTopology::Cluster { + startup_nodes: vec![ + RedisNode { + host: "node-a".into(), + port: 7000, + }, + RedisNode { + host: "node-b".into(), + port: 7001, + }, + ], + }; + assert_eq!(redis.topology, expected); + assert_eq!(redis.connection.host, "node-a"); + assert_eq!(redis.connection.port, 7000); + assert_eq!(redis.connection.password.as_deref(), Some("secret")); + assert_eq!(redis.connection.protocol, RedisProtocol::Resp3); + assert_eq!( + redis + .connection + .tls + .as_ref() + .unwrap() + .certificate_requirement, + CertificateRequirement::None + ); + }); + } + + #[test] + fn malformed_startup_nodes_and_foreign_connect_hooks_stay_on_python() { + Python::initialize(); + Python::attach(|py| { + for (startup_nodes, hook, message) in [ + ( + "[{'host': 'node-a', 'port': 7000, 'server_type': 'primary'}]", + "client.on_connect", + "native Redis topology is not implemented", + ), + ( + "[{'host': 'node-a', 'port': 'seven'}]", + "client.on_connect", + "native Redis topology is not implemented", + ), + ( + "[]", + "client.on_connect", + "native Redis topology is not implemented", + ), + ( + "[{'host': 'node-a', 'port': 7000}]", + "lambda connection: None", + "native Redis credentials require Python", + ), + ] { + let facade = cluster_facade(py, startup_nodes, hook); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("{startup_nodes} with {hook} must stay on Python"); + }; + assert_eq!(reason.message(), message, "{startup_nodes} with {hook}"); + } + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index a22a1754169..0bb5c7571a4 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -1,3 +1,4 @@ +use litellm_cache_redis::RedisTopology; use litellm_host_python::from_py; use pyo3::{ PyTraverseError, PyVisit, @@ -29,7 +30,8 @@ struct RedisPoolGuard { reference: Py, connection_class: Py, connection_kwargs: Py, - max_connections: usize, + max_connections: Option, + attributes: RedisPoolAttributes, } struct DiskStoreGuard { @@ -37,6 +39,24 @@ struct DiskStoreGuard { directory: String, } +struct RedisPoolAttributes { + pool: &'static str, + connection_class: &'static str, + max_connections: Option<&'static str>, +} + +const STANDALONE_POOL: RedisPoolAttributes = RedisPoolAttributes { + pool: "connection_pool", + connection_class: "connection_class", + max_connections: Some("max_connections"), +}; + +const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes { + pool: "nodes_manager", + connection_class: "connection_pool_class", + max_connections: None, +}; + pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, @@ -144,31 +164,40 @@ impl ObjectGuard { } impl RedisPoolGuard { - fn capture(backend: &Bound<'_, PyAny>) -> PyResult { - let pool = backend - .getattr("redis_client")? - .getattr("connection_pool")?; + fn capture(backend: &Bound<'_, PyAny>, attributes: RedisPoolAttributes) -> PyResult { + let pool = backend.getattr("redis_client")?.getattr(attributes.pool)?; Ok(Self { reference: pool.clone().unbind(), - connection_class: pool.getattr("connection_class")?.unbind(), + connection_class: pool.getattr(attributes.connection_class)?.unbind(), connection_kwargs: pool .getattr("connection_kwargs")? .call_method0("copy")? .unbind(), - max_connections: pool.getattr("max_connections")?.extract::()?, + max_connections: Self::max_connections(&pool, &attributes)?, + attributes, }) } + fn max_connections( + pool: &Bound<'_, PyAny>, + attributes: &RedisPoolAttributes, + ) -> PyResult> { + attributes + .max_connections + .map(|name| pool.getattr(name)?.extract::()) + .transpose() + } + fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { let pool = backend .getattr("redis_client")? - .getattr("connection_pool")?; + .getattr(self.attributes.pool)?; Ok(self.reference.bind(py).is(&pool) && self .connection_class .bind(py) - .is(&pool.getattr("connection_class")?) - && self.max_connections == pool.getattr("max_connections")?.extract::()? + .is(&pool.getattr(self.attributes.connection_class)?) + && self.max_connections == Self::max_connections(&pool, &self.attributes)? && self .connection_kwargs .bind(py) @@ -215,10 +244,16 @@ impl FacadeGuard { "only exact built-in Cache facades can be registered", )); } - let (module, name, cache_kind) = match kind { - "memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), - "redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"), - "disk" => ("litellm.caching.disk_cache", "DiskCache", "disk"), + let cluster = matches!(service.topology(), Some(RedisTopology::Cluster { .. })); + let (module, name, cache_kind) = match (kind, cluster) { + ("memory", _) => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), + ("redis", false) => ("litellm.caching.redis_cache", "RedisCache", "redis"), + ("redis", true) => ( + "litellm.caching.redis_cluster_cache", + "RedisClusterCache", + "redis", + ), + ("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"), _ => unreachable!(), }; let backend = facade.getattr("cache")?; @@ -264,9 +299,11 @@ impl FacadeGuard { "redis_flush_size", ], )?, - redis_pool: (kind == "redis") - .then(|| RedisPoolGuard::capture(&backend)) - .transpose()?, + redis_pool: match (kind, cluster) { + ("redis", false) => Some(RedisPoolGuard::capture(&backend, STANDALONE_POOL)?), + ("redis", true) => Some(RedisPoolGuard::capture(&backend, CLUSTER_POOL)?), + _ => None, + }, disk_store: (kind == "disk") .then(|| DiskStoreGuard::capture(&backend)) .transpose()?, diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 2cb2045577c..ba058ab0d77 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,3 +1,4 @@ +use litellm_cache_redis::{RedisNode, RedisTopology}; use litellm_host_python::release_gil; use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; @@ -34,16 +35,28 @@ impl CacheTestHandle { } #[staticmethod] - #[pyo3(signature = (url, *, ttl_seconds=60.0, namespace=None))] + #[pyo3(signature = (url, *, ttl_seconds=60.0, namespace=None, startup_nodes=None))] fn redis( py: Python<'_>, url: String, ttl_seconds: f64, namespace: Option, + startup_nodes: Option>, ) -> PyResult { let ttl = Some(duration(ttl_seconds)?); - let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace)) - .map_err(cache_error)?; + let topology = match startup_nodes { + None => RedisTopology::Standalone, + Some(nodes) => RedisTopology::Cluster { + startup_nodes: nodes + .into_iter() + .map(|(host, port)| RedisNode { host, port }) + .collect(), + }, + }; + let service = release_gil(py, move || { + NativeResponseCache::redis(&url, &topology, ttl, namespace) + }) + .map_err(cache_error)?; Ok(Self { service, guard: None, diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 14993a6e8e3..e764fde8b39 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -3,7 +3,7 @@ use std::{path::Path, sync::Arc, time::Duration}; use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; use litellm_cache_disk::DiskCache; use litellm_cache_memory::InMemoryCache; -use litellm_cache_redis::RedisCache; +use litellm_cache_redis::{RedisCache, RedisTopology}; use litellm_cache_response::{ CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer, }; @@ -36,10 +36,12 @@ impl NativeResponseCache { pub fn redis( url: &str, + topology: &RedisTopology, ttl: Option, namespace: Option, ) -> Result { - let backend = RedisCache::new(url, ttl, ResponseCacheCodec)?.with_namespace(namespace); + let backend = + RedisCache::connect(url, topology, ttl, ResponseCacheCodec)?.with_namespace(namespace); Ok(Self::Redis { cache: Arc::new(ResponseCache::new(Arc::new(backend))), buffer: None, @@ -77,6 +79,13 @@ impl NativeResponseCache { } } + pub fn topology(&self) -> Option<&RedisTopology> { + match self { + Self::Memory(_) => None, + Self::Redis { cache, .. } => Some(cache.backend().topology()), + } + } + pub fn capacity(&self) -> Option { match self { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), diff --git a/litellm/__init__.py b/litellm/__init__.py index d202bd41cfe..44515472648 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -689,6 +689,7 @@ recraft_models: Set = set() cometapi_models: Set = set() oci_models: Set = set() vercel_ai_gateway_models: Set = set() +edenai_models: Set = set() # mutable-ok: filled from the price map at import, like the sibling provider sets volcengine_models: Set = set() wandb_models: Set = set(WANDB_MODELS) ovhcloud_models: Set = set() @@ -763,6 +764,8 @@ def _populate_provider_model_sets(model_cost_map: Dict) -> None: openrouter_models.add(key) elif value.get("litellm_provider") == "vercel_ai_gateway": vercel_ai_gateway_models.add(key) + elif value.get("litellm_provider") == "edenai": + edenai_models.add(key) elif value.get("litellm_provider") == "datarobot": datarobot_models.add(key) elif value.get("litellm_provider") == "vertex_ai-text-models": @@ -1111,6 +1114,7 @@ model_list = list( | oci_models | heroku_models | vercel_ai_gateway_models + | edenai_models | volcengine_models | wandb_models | ovhcloud_models @@ -1139,6 +1143,7 @@ def _build_models_by_provider() -> dict: "baseten": baseten_models, "openrouter": openrouter_models, "vercel_ai_gateway": vercel_ai_gateway_models, + "edenai": edenai_models, "datarobot": datarobot_models, "vertex_ai": vertex_chat_models | vertex_text_models @@ -2117,6 +2122,30 @@ if TYPE_CHECKING: from .llms.vercel_ai_gateway.chat.transformation import ( VercelAIGatewayConfig as VercelAIGatewayConfig, ) + from .llms.edenai.chat.transformation import ( + EdenAIChatConfig as EdenAIChatConfig, + ) + from .llms.edenai.responses.transformation import ( + EdenAIResponsesAPIConfig as EdenAIResponsesAPIConfig, + ) + from .llms.edenai.messages.transformation import ( + EdenAIAnthropicMessagesConfig as EdenAIAnthropicMessagesConfig, + ) + from .llms.edenai.embedding.transformation import ( + EdenAIEmbeddingConfig as EdenAIEmbeddingConfig, + ) + from .llms.edenai.audio_transcription.transformation import ( + EdenAIAudioTranscriptionConfig as EdenAIAudioTranscriptionConfig, + ) + from .llms.edenai.text_to_speech.transformation import ( + EdenAITextToSpeechConfig as EdenAITextToSpeechConfig, + ) + from .llms.edenai.image_generation.transformation import ( + EdenAIImageGenerationConfig as EdenAIImageGenerationConfig, + ) + from .llms.edenai.videos.transformation import ( + EdenAIVideoConfig as EdenAIVideoConfig, + ) from .llms.ovhcloud.chat.transformation import ( OVHCloudChatConfig as OVHCloudChatConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index bca04a17250..db4eb8bdb33 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -327,6 +327,14 @@ LLM_CONFIG_NAMES: Final = ( "InceptionChatConfig", "HyperbolicChatConfig", "VercelAIGatewayConfig", + "EdenAIChatConfig", + "EdenAIResponsesAPIConfig", + "EdenAIAnthropicMessagesConfig", + "EdenAIEmbeddingConfig", + "EdenAIAudioTranscriptionConfig", + "EdenAITextToSpeechConfig", + "EdenAIImageGenerationConfig", + "EdenAIVideoConfig", "OVHCloudChatConfig", "OVHCloudEmbeddingConfig", "CometAPIEmbeddingConfig", @@ -1232,6 +1240,17 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.vercel_ai_gateway.chat.transformation", "VercelAIGatewayConfig", ), + "EdenAIChatConfig": (".llms.edenai.chat.transformation", "EdenAIChatConfig"), + "EdenAIResponsesAPIConfig": (".llms.edenai.responses.transformation", "EdenAIResponsesAPIConfig"), + "EdenAIAnthropicMessagesConfig": (".llms.edenai.messages.transformation", "EdenAIAnthropicMessagesConfig"), + "EdenAIEmbeddingConfig": (".llms.edenai.embedding.transformation", "EdenAIEmbeddingConfig"), + "EdenAIAudioTranscriptionConfig": ( + ".llms.edenai.audio_transcription.transformation", + "EdenAIAudioTranscriptionConfig", + ), + "EdenAITextToSpeechConfig": (".llms.edenai.text_to_speech.transformation", "EdenAITextToSpeechConfig"), + "EdenAIImageGenerationConfig": (".llms.edenai.image_generation.transformation", "EdenAIImageGenerationConfig"), + "EdenAIVideoConfig": (".llms.edenai.videos.transformation", "EdenAIVideoConfig"), "OVHCloudChatConfig": (".llms.ovhcloud.chat.transformation", "OVHCloudChatConfig"), "OVHCloudEmbeddingConfig": ( ".llms.ovhcloud.embedding.transformation", diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index 1331de4c266..917bfbd5ae9 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -11,6 +11,7 @@ "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", + "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": "fast-mode-2026-02-01", "files-api-2025-04-14": "files-api-2025-04-14", @@ -44,6 +45,7 @@ "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", + "dangerous-tool-use-2026-09-03": null, "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": null, "files-api-2025-04-14": "files-api-2025-04-14", @@ -76,6 +78,7 @@ "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": null, + "dangerous-tool-use-2026-09-03": null, "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": null, "files-api-2025-04-14": null, @@ -109,6 +112,7 @@ "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", + "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": null, "files-api-2025-04-14": null, @@ -143,6 +147,7 @@ "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", + "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": null, "files-api-2025-04-14": null, @@ -177,6 +182,7 @@ "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", + "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", "effort-2025-11-24": null, "fast-mode-2026-02-01": null, "files-api-2025-04-14": null, @@ -210,6 +216,7 @@ "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", + "dangerous-tool-use-2026-09-03": null, "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": "fast-mode-2026-02-01", "files-api-2025-04-14": "files-api-2025-04-14", diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index c810278f566..a09aacb8a23 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -80,6 +80,8 @@ class _AsyncRedisCommands(Protocol): def ttl(self, name: str) -> Awaitable[int]: ... + def expire(self, name: str, time: int) -> Awaitable[bool]: ... + def rpush(self, name: str, *values: str | bytes | float) -> Awaitable[int]: ... def lpop(self, name: str, count: int | None = None) -> Awaitable[object]: ... @@ -1948,6 +1950,14 @@ class RedisCache(BaseCache): _record_swallowed_redis_failure(self._circuit_breaker, e) return None + @_redis_circuit_breaker_guard + async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: + """EXPIRE an existing key without touching its value. False when the key is absent.""" + _used_ttl: Final = self.get_ttl(ttl=ttl) + if _used_ttl is None: + return False + return await self._async_commands().expire(self.check_and_fix_namespace(key=key), _used_ttl) + @_redis_circuit_breaker_guard async def async_rpush( self, diff --git a/litellm/constants.py b/litellm/constants.py index 72495b389d7..88cc5b04743 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -750,6 +750,7 @@ LITELLM_CHAT_PROVIDERS: Final = [ "inception", "vercel_ai_gateway", "wandb", + "edenai", "ovhcloud", "lemonade", "docker_model_runner", @@ -925,6 +926,7 @@ openai_compatible_endpoints: Final[list] = [ "https://api.hyperbolic.xyz/v1", "https://ai-gateway.helicone.ai/", "https://ai-gateway.vercel.sh/v1", + "https://api.edenai.run/v3", "https://api.inference.wandb.ai/v1", "https://api.clarifai.com/v2/ext/openai/v1", "https://api.libertai.io/v1", @@ -994,6 +996,7 @@ openai_compatible_providers: Final[list] = [ "hyperbolic", "vercel_ai_gateway", "aiml", + "edenai", "wandb", "cometapi", "clarifai", diff --git a/litellm/images/main.py b/litellm/images/main.py index 81547a153c3..1f722eb752a 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -388,6 +388,7 @@ def image_generation( litellm.LlmProviders.DASHSCOPE, litellm.LlmProviders.QWENCLOUD, litellm.LlmProviders.QWEN_AI_PLATFORM, + litellm.LlmProviders.EDENAI, ): if image_generation_config is None: raise ValueError(f"image generation config is not supported for {custom_llm_provider}") diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index c23b3291365..d3ad7234d93 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -427,7 +427,7 @@ class LLMCallSpanData: # plain ``.get`` — no repeated ``isinstance`` guards. raw_response: Final = payload.get("response") response: Final = cast(Mapping[str, object], raw_response if isinstance(raw_response, dict) else {}) - choices_out: Final = _dicts(response.get("choices")) or _responses_choices(response) + choices_out: Final = _dicts(response.get("choices")) or _responses_choices(response) or _ocr_choices(response) # ``finish_reasons`` is metadata, not content, so derive it from # ``choices_out`` before gating. The raw message/choice bodies are only # retained when content capture is enabled (see ``capture_span_content``); @@ -752,6 +752,22 @@ def _responses_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: return (choice,) +def _ocr_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: + markdowns: Final = tuple( + text for page in _dicts(response.get("pages")) if (text := as_str(page.get("markdown"))) is not None + ) + if not markdowns: + return () + message: Final[_AssistantMessage] = { + "role": "assistant", + "content": "\n\n".join(markdowns), + "refusal": None, + "tool_calls": None, + } + choice: Final[_Choice] = {"message": message, "finish_reason": None} + return (choice,) + + def _responses_parts_text(parts: tuple[Mapping[str, object], ...], part_type: str, field: str) -> str | None: texts: Final = tuple( text for part in parts if part.get("type") == part_type if (text := as_str(part.get(field))) is not None diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index d29b1fc74ef..ce6f77f78a0 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -4,7 +4,8 @@ import copy import logging import re from collections.abc import Iterable, Mapping -from typing import TYPE_CHECKING, Any, Final, Literal +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol import httpx from pydantic import TypeAdapter, ValidationError @@ -703,3 +704,24 @@ def redact_nested_match_and_regex_keys( except Exception: return payload return redacted + + +RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost" +_NO_HEADERS: Final[Mapping[str, object]] = MappingProxyType({}) + + +class _CarriesHiddenParams(Protocol): + _hidden_params: dict[str, object] # mutable-ok: the responses billed here keep hidden params in a plain dict + + +def set_response_cost_in_hidden_params(response: _CarriesHiddenParams, cost: float | None) -> None: + """Record a provider-reported cost where the cost calculator looks before the price map.""" + if cost is None: + return + hidden_params: Final = response._hidden_params # pyright: ignore[reportPrivateUsage] # no public accessor + additional_headers: Final[object] = hidden_params.get("additional_headers") + merged: Final[dict[str, object]] = { # mutable-ok: assigned into the plain-dict hidden params + **(additional_headers if isinstance(additional_headers, Mapping) else _NO_HEADERS), + RESPONSE_COST_HEADER: cost, + } + hidden_params["additional_headers"] = merged # rebind-ok: the caller's record is the point diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index b7067a45117..5868e79323a 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -362,6 +362,9 @@ def get_llm_provider( elif endpoint == "https://ai-gateway.vercel.sh/v1": custom_llm_provider = "vercel_ai_gateway" dynamic_api_key = get_secret_str("VERCEL_AI_GATEWAY_API_KEY") + elif endpoint == "https://api.edenai.run/v3": + custom_llm_provider = "edenai" # rebind-ok: api_base detection resolves the provider in place + dynamic_api_key = get_secret_str("EDENAI_API_KEY") elif endpoint == "https://api.inference.wandb.ai/v1": custom_llm_provider = "wandb" dynamic_api_key = get_secret_str("WANDB_API_KEY") @@ -853,6 +856,9 @@ def _get_openai_compatible_provider_info( api_base, dynamic_api_key, ) = litellm.VercelAIGatewayConfig()._get_openai_compatible_provider_info(api_base, api_key) + elif custom_llm_provider == "edenai": + api_base = litellm.EdenAIChatConfig.get_api_base(api_base) # rebind-ok: chain resolves in place + dynamic_api_key = litellm.EdenAIChatConfig.get_api_key(api_key) # rebind-ok: chain resolves in place elif custom_llm_provider == "aiml": ( api_base, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ba8addbaaa0..b34f1b3aafd 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -69,7 +69,11 @@ from litellm.litellm_core_utils.classifier_logging import ( classifier_input_snapshot, is_classifier_call, ) -from litellm.litellm_core_utils.core_helpers import is_expected_client_error, reconstruct_model_name +from litellm.litellm_core_utils.core_helpers import ( + is_expected_client_error, + reconstruct_model_name, + set_response_cost_in_hidden_params, +) from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.internal_call_metadata import ( MODEL_ACCESS_GROUP_METADATA_KEY, @@ -3918,6 +3922,7 @@ class Logging(LiteLLMLoggingBaseClass): ): ## return unified Usage object if isinstance(result.response.usage, ResponseAPIUsage): + set_response_cost_in_hidden_params(result.response, result.response.usage.cost) transformed_usage: Final = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( result.response.usage ) diff --git a/litellm/llms/base_llm/videos/transformation.py b/litellm/llms/base_llm/videos/transformation.py index f725b295d0f..a765d493347 100644 --- a/litellm/llms/base_llm/videos/transformation.py +++ b/litellm/llms/base_llm/videos/transformation.py @@ -272,6 +272,19 @@ class BaseVideoConfig(ABC): ) -> VideoObject: pass + async def async_transform_video_status_retrieve_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str | None = None, + ) -> VideoObject: + """Async transform video status retrieve response.""" + return self.transform_video_status_retrieve_response( + raw_response=raw_response, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + ) + def transform_video_create_character_request( self, name: str, diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 4b52a3bafe6..1f37fafde01 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -533,6 +533,9 @@ class AmazonAnthropicClaudeMessagesConfig( if anthropic_model_info.is_eager_input_streaming_used(tools): beta_set.add(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER) + if anthropic_messages_optional_request_params.get("safeguards") is not None: + beta_set.add(ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value) + self._filter_context_management_for_bedrock_invoke( anthropic_messages_request=anthropic_messages_request, beta_set=beta_set, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 49a332e62bb..7f1ff5298ba 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -8881,7 +8881,7 @@ class BaseLLMHTTPHandler: url=url, headers=headers, ) - return video_status_provider_config.transform_video_status_retrieve_response( + return await video_status_provider_config.async_transform_video_status_retrieve_response( raw_response=response, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, diff --git a/litellm/llms/dashscope/common_utils.py b/litellm/llms/dashscope/common_utils.py index 9ed9c276e43..952960e8207 100644 --- a/litellm/llms/dashscope/common_utils.py +++ b/litellm/llms/dashscope/common_utils.py @@ -103,7 +103,7 @@ def missing_dashscope_family_key_message(custom_llm_provider: str) -> str: ) if custom_llm_provider == "qwen_ai_platform": return ( - "Missing API key for Qwen AI Platform. Set QWEN_AI_PLATFORM_API_KEY or " + "Missing API key for Qianwen AI Platform. Set QWEN_AI_PLATFORM_API_KEY or " "DASHSCOPE_API_KEY environment variable or pass api_key parameter." ) return "Missing API key for DashScope. Set DASHSCOPE_API_KEY environment variable or pass api_key parameter." diff --git a/litellm/llms/dashscope/qwen_ai_platform.py b/litellm/llms/dashscope/qwen_ai_platform.py index 6998a57e2b7..864f0f459e4 100644 --- a/litellm/llms/dashscope/qwen_ai_platform.py +++ b/litellm/llms/dashscope/qwen_ai_platform.py @@ -23,7 +23,7 @@ def _require_qwen_ai_platform_api_key(api_key: str | None) -> str: resolved: Final = _resolve_qwen_ai_platform_api_key(api_key) if resolved is None: raise ValueError( - "Qwen AI Platform API key is required. Set 'QWEN_AI_PLATFORM_API_KEY' or 'DASHSCOPE_API_KEY' env var " + "Qianwen AI Platform API key is required. Set 'QWEN_AI_PLATFORM_API_KEY' or 'DASHSCOPE_API_KEY' env var " "or pass api_key explicitly." ) return resolved diff --git a/litellm/llms/edenai/audio_transcription/transformation.py b/litellm/llms/edenai/audio_transcription/transformation.py new file mode 100644 index 00000000000..fc8a13d5ccd --- /dev/null +++ b/litellm/llms/edenai/audio_transcription/transformation.py @@ -0,0 +1,91 @@ +""" +Support for OpenAI's `/v1/audio/transcriptions` endpoint on Eden AI, served at `/v3/audio/transcriptions` +with the real per-request cost at the top level of the JSON body. + +Docs: https://www.edenai.co/docs/api-reference/audio/audio-transcriptions +""" + +from collections.abc import Mapping +from typing import Final + +import httpx + +from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params +from litellm.llms.base_llm.audio_transcription.transformation import AudioTranscriptionRequestData +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.openai.transcriptions.whisper_transformation import OpenAIWhisperAudioTranscriptionConfig +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import FileTypes, TranscriptionResponse +from litellm.utils import convert_to_model_response_object + +from ..common_utils import EdenAIException, authorized_headers, endpoint_url, reported_cost + + +def _form_fields(model: str, optional_params: Mapping[str, object]) -> dict[str, object]: # mutable-ok: httpx form data + """LiteLLM parks non-OpenAI params, `model` included, under `extra_body` for the OpenAI SDK; a + multipart body carries them as top-level text fields instead.""" + extras: Final = optional_params.get("extra_body") + nested: Final = extras.items() if isinstance(extras, Mapping) else () + fields: Final = (*optional_params.items(), *nested, ("model", model)) + return {key: value for key, value in fields if key != "extra_body"} # mutable-ok: httpx form data + + +class EdenAIAudioTranscriptionConfig(OpenAIWhisperAudioTranscriptionConfig): + @property + def has_native_transcription_endpoint(self) -> bool: + return True + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + stream: bool | None = None, + ) -> str: + return endpoint_url(api_base, "audio/transcriptions") + + def validate_environment( + self, + headers: dict[str, object], # mutable-ok: inherited contract + model: str, + messages: list[AllMessageValues], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, object]: # mutable-ok: inherited contract + return authorized_headers(headers, api_key, model) + + def transform_audio_transcription_request( + self, + model: str, + audio_file: FileTypes, + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + ) -> AudioTranscriptionRequestData: + """Eden reports `duration` and `cost` on every body, so the Whisper default of `verbose_json`, + which the gpt-4o-transcribe models reject, is not needed for cost tracking.""" + audio: Final = process_audio_file(audio_file) + files: Final = {"file": (audio.filename, audio.file_content, audio.content_type)} # mutable-ok: httpx contract + return AudioTranscriptionRequestData(data=_form_fields(model, optional_params), files=files) + + def transform_audio_transcription_response(self, raw_response: httpx.Response) -> TranscriptionResponse: + if "application/json" not in raw_response.headers.get("content-type", ""): + return TranscriptionResponse(text=raw_response.text) + body: Final = raw_response.json() + response: Final[TranscriptionResponse] = convert_to_model_response_object( + response_object=body, model_response_object=TranscriptionResponse(), response_type="audio_transcription" + ) + set_response_cost_in_hidden_params(response, reported_cost(body)) + return response + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract + ) -> BaseLLMException: + return EdenAIException(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/edenai/chat/transformation.py b/litellm/llms/edenai/chat/transformation.py new file mode 100644 index 00000000000..4fd5a9d550b --- /dev/null +++ b/litellm/llms/edenai/chat/transformation.py @@ -0,0 +1,145 @@ +""" +Support for OpenAI's `/v1/chat/completions` endpoint on Eden AI. + +Eden AI is an OpenAI-compatible gateway (one key across 1000+ models), so requests go through the +shared HTTP handler untouched. Every Eden response reports the real per-request cost at the top +level of the body; the only translation here lifts that number into LiteLLM's cost tracking. + +Docs: https://www.edenai.co/docs +""" + +from collections.abc import AsyncIterator, Iterator, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +import httpx +from pydantic import BaseModel, TypeAdapter + +import litellm +from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ModelResponse, ModelResponseStream, Usage + +from ..common_utils import EdenAIException, reported_cost, resolve_api_base, resolve_api_key + +if TYPE_CHECKING: + import tiktoken + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_OPTIONAL_MAPPING: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) + + +class _EdenAIModel(BaseModel): + id: str + + +class _EdenAIModelCatalog(BaseModel): + data: tuple[_EdenAIModel, ...] + + +def _stream_options_with_usage(request: Mapping[str, object]) -> Mapping[str, object]: + current: Final = _OPTIONAL_MAPPING.validate_python(request.get("stream_options")) or MappingProxyType({}) + return MappingProxyType({**current, "include_usage": True}) + + +class EdenAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): + def chunk_parser(self, chunk: dict[str, object]) -> ModelResponseStream: # mutable-ok: inherited contract + parsed: Final = super().chunk_parser(chunk) + cost: Final = reported_cost(chunk) + usage: Final[object] = getattr(parsed, "usage", None) + if cost is not None and isinstance(usage, Usage): + usage.cost = cost + return parsed + + +class EdenAIChatConfig(OpenAIGPTConfig): + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: inherited contract + reasoning: Final[tuple[str, ...]] = ( + ("reasoning_effort",) + if litellm.supports_reasoning(model=model, custom_llm_provider=litellm.LlmProviders.EDENAI.value) + else () + ) + return [*super().get_supported_openai_params(model), *reasoning] # mutable-ok: inherited contract + + @staticmethod + def get_api_key(api_key: str | None = None) -> str | None: + return resolve_api_key(api_key) + + @staticmethod + def get_api_base(api_base: str | None = None) -> str: + return resolve_api_base(api_base) + + def transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + headers: dict[str, object], # mutable-ok: inherited contract + ) -> dict[str, object]: # mutable-ok: inherited contract + request: Final[dict[str, object]] = super().transform_request( # mutable-ok: inherited contract + model, messages, optional_params, litellm_params, headers + ) + if not request.get("stream"): + return request + return {**request, "stream_options": dict(_stream_options_with_usage(request))} # mutable-ok: JSON body + + def transform_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: "LiteLLMLoggingObj", + request_data: dict[str, object], # mutable-ok: inherited contract + messages: list[AllMessageValues], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + encoding: "tiktoken.Encoding | None", + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ModelResponse: + response: Final = super().transform_response( + model=model, + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data=request_data, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=encoding, + api_key=api_key, + json_mode=json_mode, + ) + set_response_cost_in_hidden_params(response, reported_cost(raw_response.content)) + return response + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract + ) -> BaseLLMException: + return EdenAIException(message=error_message, status_code=status_code, headers=headers) + + def get_model_response_iterator( + self, + streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, + sync_stream: bool, + json_mode: bool | None = False, + ) -> EdenAIChatCompletionStreamingHandler: + return EdenAIChatCompletionStreamingHandler( + streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode + ) + + def get_models( + self, api_key: str | None = None, api_base: str | None = None + ) -> list[str]: # mutable-ok: inherited contract + response: Final = litellm.module_level_client.get(url=f"{self.get_api_base(api_base)}/models") + if not response.is_success: + raise EdenAIException(status_code=response.status_code, message=response.text, headers=response.headers) + catalog: Final = _EdenAIModelCatalog.model_validate(response.json()) + return [f"edenai/{model.id}" for model in catalog.data] # mutable-ok: inherited contract diff --git a/litellm/llms/edenai/common_utils.py b/litellm/llms/edenai/common_utils.py new file mode 100644 index 00000000000..a97354cc30b --- /dev/null +++ b/litellm/llms/edenai/common_utils.py @@ -0,0 +1,80 @@ +""" +Pieces shared by every Eden AI endpoint: credentials, the exception class, and the per-request +`cost` Eden reports at the top level of each response body, or in a header when the body is binary. +""" + +from collections.abc import Container, Mapping +from types import MappingProxyType +from typing import Final + +from pydantic import AliasChoices, BaseModel, Field, ValidationError + +import litellm +from litellm.exceptions import AuthenticationError +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str +from litellm.types.utils import LlmProviders + +EDENAI_API_BASE: Final = "https://api.edenai.run/v3" +EDENAI_COST_HEADER: Final = "x-edenai-cost" + + +class EdenAIException(BaseLLMException): + pass + + +class _EdenAIExtras(BaseModel): + cost: float | None = Field(default=None, validation_alias=AliasChoices("cost", EDENAI_COST_HEADER)) + + +def resolve_api_base(api_base: str | None) -> str: + return api_base or get_secret_str("EDENAI_API_BASE") or EDENAI_API_BASE + + +def resolve_api_key(api_key: str | None) -> str | None: + return api_key or get_secret_str("EDENAI_API_KEY") + + +def require_api_key(api_key: str | None, model: str) -> str: + resolved: Final = resolve_api_key(api_key or litellm.api_key) + if resolved is None: + raise AuthenticationError( + message="Missing Eden AI API key: set EDENAI_API_KEY or pass api_key", + llm_provider=LlmProviders.EDENAI.value, + model=model, + ) + return resolved + + +def reported_cost(payload: object) -> float | None: + try: + extras: Final = ( + _EdenAIExtras.model_validate_json(payload) + if isinstance(payload, bytes) + else _EdenAIExtras.model_validate(payload) + ) + except ValidationError: + return None + return extras.cost + + +def authorized_headers( + headers: Mapping[str, object], api_key: str | None, model: str +) -> dict[str, object]: # mutable-ok: header contract + return {**headers, "Authorization": f"Bearer {require_api_key(api_key, model)}"} # mutable-ok: header contract + + +def json_headers( + headers: Mapping[str, object], api_key: str | None, model: str +) -> dict[str, object]: # mutable-ok: header contract + """The shared HTTP handler sends some JSON bodies as raw content, so the type must be set here.""" + authorized: Final = authorized_headers(headers, api_key, model) + return {**authorized, "Content-Type": "application/json"} # mutable-ok: header contract + + +def endpoint_url(api_base: str | None, path: str) -> str: + return f"{resolve_api_base(api_base).rstrip('/')}/{path}" + + +def pick(params: Mapping[str, object], keys: Container[str]) -> Mapping[str, object]: + return MappingProxyType({key: value for key, value in params.items() if key in keys}) diff --git a/litellm/llms/edenai/embedding/transformation.py b/litellm/llms/edenai/embedding/transformation.py new file mode 100644 index 00000000000..1c2cc937875 --- /dev/null +++ b/litellm/llms/edenai/embedding/transformation.py @@ -0,0 +1,97 @@ +""" +Support for OpenAI's `/v1/embeddings` endpoint on Eden AI, served at `/v3/embeddings` with the real +per-request cost at the top level of the body. + +Docs: https://www.edenai.co/docs/v3/llms/embeddings +""" + +from typing import TYPE_CHECKING, Final + +import httpx + +from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig +from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues +from litellm.types.utils import EmbeddingResponse +from litellm.utils import convert_to_model_response_object + +from ..common_utils import EdenAIException, endpoint_url, json_headers, pick, reported_cost + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_SUPPORTED_PARAMS: Final = ("dimensions", "encoding_format", "user") + + +class EdenAIEmbeddingConfig(BaseEmbeddingConfig): + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: inherited contract + return list(_SUPPORTED_PARAMS) # mutable-ok: inherited contract + + def map_openai_params( + self, + non_default_params: dict[str, object], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: inherited contract + return {**optional_params, **pick(non_default_params, _SUPPORTED_PARAMS)} # mutable-ok: inherited contract + + def validate_environment( + self, + headers: dict[str, object], # mutable-ok: inherited contract + model: str, + messages: list[AllMessageValues], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, object]: # mutable-ok: inherited contract + return json_headers(headers, api_key, model) + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + stream: bool | None = None, + ) -> str: + return endpoint_url(api_base, "embeddings") + + def transform_embedding_request( + self, + model: str, + input: AllEmbeddingInputValues, + optional_params: dict[str, object], # mutable-ok: inherited contract + headers: dict[str, object], # mutable-ok: inherited contract + ) -> dict[str, object]: # mutable-ok: inherited contract + return {"model": model, "input": input, **optional_params} # mutable-ok: inherited contract + + def transform_embedding_response( + self, + model: str, + raw_response: httpx.Response, + model_response: EmbeddingResponse, + logging_obj: "LiteLLMLoggingObj", + api_key: str | None, + request_data: dict[str, object], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + ) -> EmbeddingResponse: + body: Final = raw_response.json() + logging_obj.post_call(original_response=body) + response: Final[EmbeddingResponse] = convert_to_model_response_object( + response_object=body, model_response_object=model_response, response_type="embedding" + ) + set_response_cost_in_hidden_params(response, reported_cost(body)) + return response + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract + ) -> BaseLLMException: + return EdenAIException(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/edenai/image_generation/transformation.py b/litellm/llms/edenai/image_generation/transformation.py new file mode 100644 index 00000000000..2965c4041de --- /dev/null +++ b/litellm/llms/edenai/image_generation/transformation.py @@ -0,0 +1,115 @@ +""" +Support for OpenAI's `/v1/images/generations` endpoint on Eden AI, served at `/v3/images/generations` +for every image model in the catalog with the real per-request cost at the top level of the body. + +Docs: https://www.edenai.co/docs/v3/llms/image-generation +""" + +from typing import TYPE_CHECKING, Final + +import httpx + +from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.image_generation.transformation import BaseImageGenerationConfig +from litellm.types.llms.openai import AllMessageValues, OpenAIImageGenerationOptionalParams +from litellm.types.utils import ImageResponse +from litellm.utils import convert_to_model_response_object + +from ..common_utils import EdenAIException, endpoint_url, json_headers, pick, reported_cost + +if TYPE_CHECKING: + import tiktoken + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_SUPPORTED_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = ( + "background", + "moderation", + "n", + "output_compression", + "output_format", + "quality", + "response_format", + "size", + "style", + "user", +) + + +class EdenAIImageGenerationConfig(BaseImageGenerationConfig): + def get_supported_openai_params( + self, model: str + ) -> list[OpenAIImageGenerationOptionalParams]: # mutable-ok: inherited contract + return list(_SUPPORTED_PARAMS) # mutable-ok: inherited contract + + def map_openai_params( + self, + non_default_params: dict[str, object], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: inherited contract + return {**optional_params, **pick(non_default_params, _SUPPORTED_PARAMS)} # mutable-ok: inherited contract + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + stream: bool | None = None, + ) -> str: + return endpoint_url(api_base, "images/generations") + + def validate_environment( + self, + headers: dict[str, object], # mutable-ok: inherited contract + model: str, + messages: list[AllMessageValues], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, object]: # mutable-ok: inherited contract + return json_headers(headers, api_key, model) + + def transform_image_generation_request( + self, + model: str, + prompt: str, + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + headers: dict[str, object], # mutable-ok: inherited contract + ) -> dict[str, object]: # mutable-ok: inherited contract + return {"model": model, "prompt": prompt, **optional_params} # mutable-ok: inherited contract + + def transform_image_generation_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ImageResponse, + logging_obj: "LiteLLMLoggingObj", + request_data: dict[str, object], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + encoding: "tiktoken.Encoding | None", + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ImageResponse: + body: Final = raw_response.json() + logging_obj.post_call(original_response=body) + response: Final[ImageResponse] = convert_to_model_response_object( + response_object=body, model_response_object=model_response, response_type="image_generation" + ) + set_response_cost_in_hidden_params(response, reported_cost(body)) + return response + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract + ) -> BaseLLMException: + return EdenAIException(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/edenai/messages/transformation.py b/litellm/llms/edenai/messages/transformation.py new file mode 100644 index 00000000000..fbb3ae0c05a --- /dev/null +++ b/litellm/llms/edenai/messages/transformation.py @@ -0,0 +1,79 @@ +""" +Support for Anthropic's `/v1/messages` endpoint on Eden AI. + +Eden AI serves the Anthropic Messages API at `/v3/v1/messages` for every model in its catalog, so +the Anthropic payload is forwarded untranslated and the answer comes back in Anthropic's shape with +Eden's per-request `cost` beside it. Eden does not report a cost inside a Messages stream yet, so +streams fall back to the price map. + +Docs: https://www.edenai.co/docs/api-reference/anthropic-messages/create-anthropic-message +""" + +from typing import TYPE_CHECKING, Final + +import httpx + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.openai_like.json_loader import SimpleProviderConfig +from litellm.llms.openai_like.messages.transformation import JSONProviderAnthropicMessagesConfig +from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +from litellm.types.utils import LlmProviders + +from ..common_utils import EDENAI_API_BASE, EdenAIException, reported_cost, require_api_key + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_EDENAI_PROVIDER_SPEC: Final[dict[str, str]] = { # mutable-ok: SimpleProviderConfig takes a plain dict + "base_url": EDENAI_API_BASE, + "api_key_env": "EDENAI_API_KEY", + "api_base_env": "EDENAI_API_BASE", +} +_EDENAI_PROVIDER: Final = SimpleProviderConfig(LlmProviders.EDENAI.value, _EDENAI_PROVIDER_SPEC) + + +class EdenAIAnthropicMessagesConfig(JSONProviderAnthropicMessagesConfig): + def __init__(self) -> None: + super().__init__(_EDENAI_PROVIDER) + + def validate_anthropic_messages_environment( + self, + headers: dict[str, str], # mutable-ok: inherited contract + model: str, + messages: list[object], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + api_key: str | None = None, + api_base: str | None = None, + ) -> tuple[dict[str, str], str | None]: # mutable-ok: inherited contract + return super().validate_anthropic_messages_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=require_api_key(api_key, model), + api_base=api_base, + ) + + def transform_anthropic_messages_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> AnthropicMessagesResponse: + response: Final = super().transform_anthropic_messages_response( + model=model, raw_response=raw_response, logging_obj=logging_obj + ) + cost: Final = reported_cost(response) + if cost is not None: + logging_obj.model_call_details["response_cost"] = cost # rebind-ok: the per-call record spend logging reads + return response + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract + ) -> BaseLLMException: + return EdenAIException(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/edenai/responses/transformation.py b/litellm/llms/edenai/responses/transformation.py new file mode 100644 index 00000000000..3274e70746c --- /dev/null +++ b/litellm/llms/edenai/responses/transformation.py @@ -0,0 +1,80 @@ +""" +Support for OpenAI's `/v1/responses` endpoint on Eden AI. + +Eden AI serves the Responses API at `/v3/responses` in OpenAI's wire format, so the OpenAI config +does the work; this one points it at Eden and authenticates with the Eden key. Eden reports the +per-request cost on `usage.cost` of every body, the final `response.completed` event included, so +the shared usage-cost lift bills both modes. + +Docs: https://www.edenai.co/docs/v3/llms/responses +""" + +from typing import TYPE_CHECKING, Final + +import httpx + +from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + +from ..common_utils import EdenAIException, authorized_headers, resolve_api_base + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + +class EdenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.EDENAI + + def validate_environment( + self, + headers: dict[str, object], # mutable-ok: inherited contract + model: str, + litellm_params: GenericLiteLLMParams | None, + ) -> dict[str, object]: # mutable-ok: inherited contract + return authorized_headers(headers, litellm_params.api_key if litellm_params else None, model) + + def get_complete_url( + self, + api_base: str | None, + litellm_params: dict[str, object], # mutable-ok: inherited contract + ) -> str: + return super().get_complete_url(api_base=resolve_api_base(api_base), litellm_params=litellm_params) + + def transform_response_api_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> ResponsesAPIResponse: + response: Final = super().transform_response_api_response( + model=model, raw_response=raw_response, logging_obj=logging_obj + ) + set_response_cost_in_hidden_params(response, response.usage.cost if response.usage else None) + return response + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract + ) -> BaseLLMException: + return EdenAIException(message=error_message, status_code=status_code, headers=headers) + + def should_fake_stream( + self, + model: str | None, + stream: bool | None, + custom_llm_provider: str | None = None, + ) -> bool: + """Eden streams every catalog model natively; the base class would fake-stream any model the + price map does not know, which is all of them.""" + return False + + def supports_native_websocket(self) -> bool: + return False diff --git a/litellm/llms/edenai/text_to_speech/transformation.py b/litellm/llms/edenai/text_to_speech/transformation.py new file mode 100644 index 00000000000..50c7ed96725 --- /dev/null +++ b/litellm/llms/edenai/text_to_speech/transformation.py @@ -0,0 +1,85 @@ +""" +Support for OpenAI's `/v1/audio/speech` endpoint on Eden AI, served at `/v3/audio/speech`. The answer +is raw audio, so the real per-request cost travels in the `x-edenai-cost` response header. + +Docs: https://www.edenai.co/docs/api-reference/audio/audio-speech +""" + +from typing import TYPE_CHECKING, Final + +import httpx + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig, TextToSpeechRequestData +from litellm.types.llms.openai import HttpxBinaryResponseContent + +from ..common_utils import EdenAIException, endpoint_url, json_headers, reported_cost + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_SUPPORTED_PARAMS: Final = ("voice", "response_format", "speed", "instructions") + + +class EdenAITextToSpeechConfig(BaseTextToSpeechConfig): + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: inherited contract + return list(_SUPPORTED_PARAMS) # mutable-ok: inherited contract + + def map_openai_params( + self, + model: str, + optional_params: dict[str, object], # mutable-ok: inherited contract + voice: str | dict[str, object] | None = None, # mutable-ok: inherited contract + drop_params: bool = False, + kwargs: dict[str, object] | None = None, # mutable-ok: inherited contract + ) -> tuple[str | None, dict[str, object]]: # mutable-ok: inherited contract + return (voice if isinstance(voice, str) else None), optional_params + + def validate_environment( + self, + headers: dict[str, object], # mutable-ok: inherited contract + model: str, + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, object]: # mutable-ok: inherited contract + return json_headers(headers, api_key, model) + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict[str, object], # mutable-ok: inherited contract + ) -> str: + return endpoint_url(api_base, "audio/speech") + + def transform_text_to_speech_request( + self, + model: str, + input: str, + voice: str | None, + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + headers: dict[str, object], # mutable-ok: inherited contract + ) -> TextToSpeechRequestData: + fields: Final = (("model", model), ("input", input), ("voice", voice), *optional_params.items()) + return TextToSpeechRequestData( + dict_body={key: value for key, value in fields if value is not None} # mutable-ok: TypedDict field + ) + + def transform_text_to_speech_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> HttpxBinaryResponseContent: + response: Final = HttpxBinaryResponseContent(response=raw_response) + response.set_response_cost(reported_cost(raw_response.headers)) + return response + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract + ) -> BaseLLMException: + return EdenAIException(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/edenai/videos/transformation.py b/litellm/llms/edenai/videos/transformation.py new file mode 100644 index 00000000000..25c7bcf24ea --- /dev/null +++ b/litellm/llms/edenai/videos/transformation.py @@ -0,0 +1,146 @@ +""" +Support for OpenAI's `/v1/videos` API on Eden AI, served at `/v3/videos`. A job is created, polled and +downloaded through the OpenAI routes; Eden reports `cost` as 0 on the create response and the settled +amount on the status read once the job completes or fails. + +Docs: https://www.edenai.co/docs/v3/llms/video-generation +""" + +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final + +import httpx +from httpx._types import RequestFiles + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.openai.videos.transformation import OpenAIVideoConfig +from litellm.types.router import GenericLiteLLMParams +from litellm.types.videos.main import VideoObject + +from ..common_utils import EdenAIException, authorized_headers, endpoint_url, reported_cost + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + +def _usage_with_reported_cost( + usage: Mapping[str, object] | None, body: bytes +) -> dict[str, object]: # mutable-ok: VideoObject.usage is a plain dict field + cost: Final = reported_cost(body) + return { # mutable-ok: VideoObject.usage is a plain dict field + key: value + for key, value in (*(usage.items() if usage else ()), ("provider_reported_cost_usd", cost)) + if value is not None + } + + +class EdenAIVideoConfig(OpenAIVideoConfig): + def validate_environment( + self, + headers: dict[str, object], # mutable-ok: inherited contract + model: str, + api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, + ) -> dict[str, object]: # mutable-ok: inherited contract + return authorized_headers(headers, api_key or (litellm_params.api_key if litellm_params else None), model) + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict[str, object], # mutable-ok: inherited contract + ) -> str: + return endpoint_url(api_base, "videos") + + def use_multipart_form_data(self) -> bool: + return False + + def transform_video_create_request( + self, + model: str, + prompt: str, + api_base: str, + video_create_optional_request_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: GenericLiteLLMParams, + headers: dict[str, object], # mutable-ok: inherited contract + ) -> tuple[dict[str, object], RequestFiles, str]: # mutable-ok: inherited contract + """A reference image is a multipart file part, or a JSON `{"file_id"}` / `{"image_url"}` object.""" + reference: Final = video_create_optional_request_params.get("input_reference") + if not isinstance(reference, Mapping): + return super().transform_video_create_request( + model=model, + prompt=prompt, + api_base=api_base, + video_create_optional_request_params=video_create_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + data, files, url = super().transform_video_create_request( + model=model, + prompt=prompt, + api_base=api_base, + video_create_optional_request_params={ # mutable-ok: inherited contract + key: value for key, value in video_create_optional_request_params.items() if key != "input_reference" + }, + litellm_params=litellm_params, + headers=headers, + ) + return {**data, "input_reference": dict(reference)}, files, url # mutable-ok: JSON body + + def transform_video_create_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + custom_llm_provider: str | None = None, + request_data: dict[str, object] | None = None, # mutable-ok: inherited contract + ) -> VideoObject: + video: Final = super().transform_video_create_response( + model=model, + raw_response=raw_response, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + request_data=request_data, + ) + video.usage = _usage_with_reported_cost(video.usage, raw_response.content) + return video + + def transform_video_status_retrieve_response( + self, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + custom_llm_provider: str | None = None, + ) -> VideoObject: + raw_response.raise_for_status() # the shared GET helpers return error bodies instead of raising + video: Final = super().transform_video_status_retrieve_response( + raw_response=raw_response, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider + ) + video.usage = _usage_with_reported_cost(video.usage, raw_response.content) + return video + + def transform_video_content_response( + self, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> bytes: + raw_response.raise_for_status() # the shared GET helpers return error bodies instead of raising + return raw_response.content + + def transform_video_list_response( + self, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + custom_llm_provider: str | None = None, + ) -> dict[str, str]: # mutable-ok: inherited contract + raw_response.raise_for_status() # the shared GET helpers return error bodies instead of raising + return super().transform_video_list_response( + raw_response=raw_response, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider + ) + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract + ) -> BaseLLMException: + return EdenAIException(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index 98528c82f6b..51082a6773b 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -1,6 +1,8 @@ import math +import sys import time -from collections.abc import Mapping +from collections.abc import Callable, Mapping +from dataclasses import dataclass from types import MappingProxyType from typing import Final, TypeAlias @@ -36,11 +38,33 @@ class FalAIVideoError(BaseLLMException): _ALLOWED_ASPECT_RATIOS: Final[frozenset[str]] = frozenset({"auto", "16:9", "9:16", "1:1", "4:3", "3:4", "21:9"}) -_ALLOWED_RESOLUTIONS: Final[frozenset[str]] = frozenset({"480p", "720p", "1080p", "4k"}) -_RESOLUTION_TIERS: Final[tuple[tuple[int, str], ...]] = ( - (480, "480p"), - (720, "720p"), - (1080, "1080p"), + + +@dataclass(frozen=True, slots=True) +class _ModelProfile: + resolutions: frozenset[str] + resolution_tiers: tuple[tuple[int, str], ...] + default_resolution: str + integer_duration: bool + reference_key: str + reference_as_list: bool + + +_SEEDANCE_PROFILE: Final[_ModelProfile] = _ModelProfile( + resolutions=frozenset({"480p", "720p", "1080p", "4k"}), + resolution_tiers=((480, "480p"), (720, "720p"), (1080, "1080p"), (sys.maxsize, "4k")), + default_resolution="720p", + integer_duration=False, + reference_key="image_url", + reference_as_list=False, +) +_H3_PROFILE: Final[_ModelProfile] = _ModelProfile( + resolutions=frozenset({"480P", "768P", "2K", "4K"}), + resolution_tiers=((480, "480P"), (768, "768P"), (1440, "2K"), (sys.maxsize, "4K")), + default_resolution="2K", + integer_duration=True, + reference_key="reference_image_urls", + reference_as_list=True, ) _QUEUE_NAMESPACES: Final[frozenset[str]] = frozenset(("workflows", "comfy")) _STATUS_MAP: Final[Mapping[str, str]] = MappingProxyType( @@ -75,8 +99,12 @@ def _duration_value(value: object) -> str | None: return None -def _resolution_for_short_side(short_side: int) -> str: - return next((resolution for threshold, resolution in _RESOLUTION_TIERS if short_side <= threshold), "4k") +def _profile_for_model(model: str) -> _ModelProfile: + return _H3_PROFILE if model.startswith("minimax/h3/") else _SEEDANCE_PROFILE + + +def _resolution_for_short_side(short_side: int, profile: _ModelProfile) -> str: + return next(resolution for threshold, resolution in profile.resolution_tiers if short_side <= threshold) def _model_path_from_request_url(raw_response: httpx.Response) -> str | None: @@ -97,14 +125,19 @@ def _request_id_from_request_url(raw_response: httpx.Response) -> str | None: return segments[request_id_index] if len(segments) > request_id_index else None -def _size_params(size: object) -> Mapping[str, str]: +def _size_params(size: object, profile: _ModelProfile) -> Mapping[str, str]: if not isinstance(size, str): return MappingProxyType({}) - if size in _ALLOWED_RESOLUTIONS: - return MappingProxyType({"resolution": size}) - if size.count("x") != 1: + normalized_size: Final[str] = size.lower() + canonical_resolution: Final[str | None] = next( + (resolution for resolution in profile.resolutions if resolution.lower() == normalized_size), + None, + ) + if canonical_resolution is not None: + return MappingProxyType({"resolution": canonical_resolution}) + if normalized_size.count("x") != 1: return MappingProxyType({}) - width_text, height_text = size.split("x") + width_text, height_text = normalized_size.split("x") if not (width_text.isdigit() and height_text.isdigit()): return MappingProxyType({}) width: Final[int] = int(width_text) @@ -113,7 +146,7 @@ def _size_params(size: object) -> Mapping[str, str]: return MappingProxyType({}) reduced_gcd: Final[int] = math.gcd(width, height) aspect_ratio: Final[str] = f"{width // reduced_gcd}:{height // reduced_gcd}" - resolution: Final[str] = _resolution_for_short_side(min(width, height)) + resolution: Final[str] = _resolution_for_short_side(min(width, height), profile) if aspect_ratio in _ALLOWED_ASPECT_RATIOS: return MappingProxyType({"resolution": resolution, "aspect_ratio": aspect_ratio}) return MappingProxyType({"resolution": resolution}) @@ -130,12 +163,127 @@ def _response_data(raw_response: httpx.Response) -> Mapping[str, object]: return TypeAdapter(Mapping[str, object]).validate_python(raw_response.json()) +def _response_data_or_none(raw_response: httpx.Response) -> Mapping[str, object] | None: + try: + return _response_data(raw_response) + except ValueError: + return None + + +def _detail_item_text(item: Mapping[str, object]) -> str | None: + message: Final[object] = item.get("msg") + if not isinstance(message, str): + return None + location: Final[object] = item.get("loc") + if isinstance(location, str) and location: + return f"{location}: {message}" + if isinstance(location, (list, tuple)): + location_parts: Final[tuple[str, ...]] = tuple(part for part in location if isinstance(part, str)) + if location_parts: + return f"{'.'.join(location_parts)}: {message}" + return message + + +def _error_text(response_data: Mapping[str, object]) -> str | None: + detail: Final[object] = response_data.get("detail") + if isinstance(detail, str): + return detail + if isinstance(detail, list): + detail_items: Final[tuple[Mapping[str, object], ...]] = tuple( + item for item in detail if isinstance(item, Mapping) + ) + detail_messages: Final[tuple[str, ...]] = tuple( + message for item in detail_items if (message := _detail_item_text(item)) is not None + ) + if detail_messages: + return "; ".join(detail_messages) + error: Final[object] = response_data.get("error") + return error if isinstance(error, str) else None + + +def _result_error(raw_response: httpx.Response) -> str | None: + if raw_response.is_success: + return None + response_data: Final[Mapping[str, object] | None] = _response_data_or_none(raw_response) + error_text: Final[str | None] = _error_text(response_data) if response_data is not None else None + if error_text: + return error_text + response_text: Final[str] = raw_response.text + return response_text or f"fal.ai returned HTTP {raw_response.status_code}" + + +def _terminal_result_error(raw_response: httpx.Response) -> str | None: + if raw_response.status_code == 429 or raw_response.status_code >= 500: + return None + return _result_error(raw_response) + + +def _get_fal_ai_async_httpx_client() -> AsyncHTTPHandler: + return get_async_httpx_client(llm_provider=LlmProviders.FAL_AI) + + def _response_string(response_data: Mapping[str, object], key: str, default: str = "") -> str: value: Final[object] = response_data.get(key) return value if isinstance(value, str) else default +def _result_request( + raw_response: httpx.Response, + response_data: Mapping[str, object], +) -> tuple[str, Mapping[str, str]] | None: + if _response_string(response_data, "status", "IN_QUEUE") != "COMPLETED": + return None + result_url: Final[str] = str(raw_response.request.url).removesuffix("/status") + result_headers: Final[Mapping[str, str]] = MappingProxyType( + { + key: value + for key, value in ( + ("Authorization", raw_response.request.headers.get("Authorization")), + ("Content-Type", raw_response.request.headers.get("Content-Type")), + ) + if value is not None + } + ) + return result_url, result_headers + + +def _status_video_object( + response_data: Mapping[str, object], + raw_response: httpx.Response, + custom_llm_provider: str | None, + result_error: str | None, +) -> VideoObject: + raw_status: Final[str] = _response_string(response_data, "status", "IN_QUEUE") + status: Final[str] = _STATUS_MAP.get(raw_status, "queued") + status_error: Final[str | None] = _error_text(response_data) + error: Final[str | None] = result_error if result_error is not None else status_error + provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER + model_path: Final[str | None] = _model_path_from_request_url(raw_response) + request_id: Final[str] = _response_string(response_data, "request_id") or ( + _request_id_from_request_url(raw_response) or "" + ) + return VideoObject( + id=encode_video_id_with_provider(request_id, provider, model_path), + object="video", + status="failed" if error else status, + created_at=0, + model=model_path, + error=( + {"code": "fal_error", "message": error} if error else None # mutable-ok: VideoObject requires a dict + ), + ) + + class FalAIVideoConfig(BaseVideoConfig): + def __init__( + self, + sync_client_factory: Callable[[], HTTPHandler] = _get_httpx_client, + async_client_factory: Callable[[], AsyncHTTPHandler] = _get_fal_ai_async_httpx_client, + ) -> None: + super().__init__() + self._sync_client_factory: Final = sync_client_factory + self._async_client_factory: Final = async_client_factory + def get_supported_openai_params(self, model: str) -> _SupportedParams: supported_params: Final[_SupportedParams] = [ # mutable-ok: BaseVideoConfig requires a list "model", @@ -158,18 +306,27 @@ class FalAIVideoConfig(BaseVideoConfig): input_reference: Final[object] = video_create_optional_params.get("input_reference") if "input_reference" in video_create_optional_params and not isinstance(input_reference, str): raise ValueError("fal.ai needs a public image URL for input_reference") - input_reference_params: Final[Mapping[str, str]] = ( + profile: Final[_ModelProfile] = _profile_for_model(model) + input_reference_params: Final[Mapping[str, object]] = ( MappingProxyType({}) if not isinstance(input_reference, str) - else MappingProxyType({"image_url": input_reference}) + else MappingProxyType( + { + profile.reference_key: ( + [input_reference] # mutable-ok: fal.ai expects a list for H3 references + if profile.reference_as_list + else input_reference + ), + } + ) ) - duration_params: Final[Mapping[str, str]] = ( + duration_params: Final[Mapping[str, object]] = ( MappingProxyType({}) if "seconds" not in video_create_optional_params - else self._duration_params(video_create_optional_params["seconds"]) + else self._duration_params(video_create_optional_params["seconds"], profile) ) size_params: Final[Mapping[str, str]] = ( - _size_params(video_create_optional_params["size"]) + _size_params(video_create_optional_params["size"], profile) if "size" in video_create_optional_params else MappingProxyType({}) ) @@ -190,11 +347,11 @@ class FalAIVideoConfig(BaseVideoConfig): return mapped_params @staticmethod - def _duration_params(seconds: object) -> Mapping[str, str]: + def _duration_params(seconds: object, profile: _ModelProfile) -> Mapping[str, object]: duration: Final[str | None] = _duration_value(seconds) if duration is None: raise ValueError("fal.ai seconds must be a numeric value") - return MappingProxyType({"duration": duration}) + return MappingProxyType({"duration": int(duration) if profile.integer_duration else duration}) def validate_environment( self, @@ -251,6 +408,7 @@ class FalAIVideoConfig(BaseVideoConfig): request_data: Mapping[str, object] | None = None, ) -> VideoObject: response_data: Final[Mapping[str, object]] = _response_data(raw_response) + profile: Final[_ModelProfile] = _profile_for_model(model) request_params: Final[Mapping[str, object]] = request_data or MappingProxyType({}) request_id: Final[str] = _response_string(response_data, "request_id") provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER @@ -262,7 +420,10 @@ class FalAIVideoConfig(BaseVideoConfig): key: value for key, value in ( ("duration_seconds", duration), - ("video_resolution", resolution if isinstance(resolution, str) else "720p"), + ( + "video_resolution", + resolution if isinstance(resolution, str) else profile.default_resolution, + ), ) if value is not None } @@ -299,25 +460,58 @@ class FalAIVideoConfig(BaseVideoConfig): custom_llm_provider: str | None = None, ) -> VideoObject: response_data: Final[Mapping[str, object]] = _response_data(raw_response) - raw_status: Final[str] = _response_string(response_data, "status", "IN_QUEUE") - status: Final[str] = _STATUS_MAP.get(raw_status, "queued") - error_value: Final[object] = response_data.get("error") - error: Final[str | None] = error_value if isinstance(error_value, str) else None - provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER - model_path: Final[str | None] = _model_path_from_request_url(raw_response) - request_id: Final[str] = _response_string(response_data, "request_id") or ( - _request_id_from_request_url(raw_response) or "" + result_error: Final[str | None] = self._fetch_result_error(raw_response, response_data) + return _status_video_object( + response_data=response_data, + raw_response=raw_response, + custom_llm_provider=custom_llm_provider, + result_error=result_error, ) - return VideoObject( - id=encode_video_id_with_provider(request_id, provider, model_path), - object="video", - status="failed" if error else status, - created_at=0, - model=model_path, - error=( - {"code": "fal_error", "message": error} if error else None # mutable-ok: VideoObject requires a dict - ), + + def _fetch_result_error( + self, + raw_response: httpx.Response, + response_data: Mapping[str, object], + ) -> str | None: + result_request: Final[tuple[str, Mapping[str, str]] | None] = _result_request(raw_response, response_data) + if result_request is None: + return None + result_url, result_headers = result_request + result_response: Final[httpx.Response] = self._sync_client_factory().get( + url=result_url, + headers=result_headers, ) + return _terminal_result_error(result_response) + + async def async_transform_video_status_retrieve_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> VideoObject: + response_data: Final[Mapping[str, object]] = _response_data(raw_response) + result_error: Final[str | None] = await self._fetch_result_error_async(raw_response, response_data) + return _status_video_object( + response_data=response_data, + raw_response=raw_response, + custom_llm_provider=custom_llm_provider, + result_error=result_error, + ) + + async def _fetch_result_error_async( + self, + raw_response: httpx.Response, + response_data: Mapping[str, object], + ) -> str | None: + result_request: Final[tuple[str, Mapping[str, str]] | None] = _result_request(raw_response, response_data) + if result_request is None: + return None + result_url, result_headers = result_request + result_response: Final[httpx.Response] = await self._async_client_factory().get( + url=result_url, + headers=result_headers, + ) + return _terminal_result_error(result_response) @staticmethod def _decode_video_id(video_id: str) -> tuple[str, str]: @@ -355,17 +549,23 @@ class FalAIVideoConfig(BaseVideoConfig): video_url: Final[object] = video_data.get("url") if isinstance(video_url, str) and video_url: return video_url - error_message: Final[str | None] = next( - (value for key in ("error", "detail") if isinstance(value := response_data.get(key), str)), - None, - ) + error_message: Final[str | None] = _error_text(response_data) if error_message: raise ValueError(f"fal.ai video result did not include a video URL: {error_message}") raise ValueError("fal.ai video result did not include a video URL") def transform_video_content_response(self, raw_response: httpx.Response, logging_obj: object) -> bytes: + error: Final[str | None] = _result_error(raw_response) + if error is not None: + raise FalAIVideoError( + status_code=raw_response.status_code, + message=error, + headers=dict(raw_response.headers), # mutable-ok: exception headers require a mutable dictionary + request=raw_response.request, + response=raw_response, + ) video_url: Final[str] = self._extract_video_url(_response_data(raw_response)) - httpx_client: Final[HTTPHandler] = _get_httpx_client() + httpx_client: Final[HTTPHandler] = self._sync_client_factory() video_response: Final[httpx.Response] = httpx_client.get( # pyright: ignore[reportUnknownMemberType] # HTTP handler stubs are untyped video_url ) @@ -373,8 +573,17 @@ class FalAIVideoConfig(BaseVideoConfig): return video_response.content async def async_transform_video_content_response(self, raw_response: httpx.Response, logging_obj: object) -> bytes: + error: Final[str | None] = _result_error(raw_response) + if error is not None: + raise FalAIVideoError( + status_code=raw_response.status_code, + message=error, + headers=dict(raw_response.headers), # mutable-ok: exception headers require a mutable dictionary + request=raw_response.request, + response=raw_response, + ) video_url: Final[str] = self._extract_video_url(_response_data(raw_response)) - async_httpx_client: Final[AsyncHTTPHandler] = get_async_httpx_client(llm_provider=LlmProviders.FAL_AI) + async_httpx_client: Final[AsyncHTTPHandler] = self._async_client_factory() video_response: Final[httpx.Response] = await async_httpx_client.get( # pyright: ignore[reportUnknownMemberType] # HTTP handler stubs are untyped video_url ) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index 978daf119ce..785f4dcefce 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -108,6 +108,9 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert if anthropic_model_info.is_tool_search_used(tools): beta_values.add(get_tool_search_beta_header("vertex_ai")) + if optional_params.get("safeguards") is not None: + beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value) + if beta_values: headers["anthropic-beta"] = ",".join(beta_values) diff --git a/litellm/main.py b/litellm/main.py index 66466f01da4..6704358e3ea 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3568,6 +3568,32 @@ def _complete_vercel_ai_gateway( return response +def _complete_edenai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + api_base: Final = litellm.EdenAIChatConfig.get_api_base(ctx.api_base) + api_key: Final = litellm.EdenAIChatConfig.get_api_key(ctx.api_key or litellm.api_key) + response: Final = base_llm_http_handler.completion( + model=ctx.model, + messages=ctx.messages, + api_base=api_base, + custom_llm_provider="edenai", + model_response=ctx.model_response, + encoding=_get_encoding(), + logging_obj=ctx.logging, + optional_params=ctx.optional_params, + timeout=ctx.timeout, + litellm_params=ctx.litellm_params, + shared_session=ctx.shared_session, + acompletion=ctx.acompletion, + stream=ctx.stream, + api_key=api_key, + headers=ctx.headers or litellm.headers, + client=_dispatch_client_http(ctx), + provider_config=ctx.provider_config, + ) + ctx.logging.post_call(input=ctx.messages, api_key=api_key, original_response=response) + return response + + def _complete_vertex_ai_beta( ctx: _CompletionDispatchContext, ) -> _CompletionDispatchResult: @@ -5771,6 +5797,8 @@ def completion( response = _complete_minimax(_dispatch_ctx) elif custom_llm_provider == "hosted_vllm": response = _complete_hosted_vllm(_dispatch_ctx) + elif custom_llm_provider == "edenai": + response = _complete_edenai(_dispatch_ctx) # rebind-ok: dispatch chain binds response per branch elif ( # A known OpenAI model name only decides the route when nothing else # resolved a provider. get_llm_provider() already maps these names to @@ -6440,6 +6468,22 @@ def embedding( litellm_params=litellm_params_dict, headers=headers or {}, ) + elif custom_llm_provider == "edenai": + response = base_llm_http_handler.embedding( + model=model, + input=input, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, + logging_obj=logging, + timeout=timeout, + model_response=EmbeddingResponse(), + optional_params=optional_params, + client=client, + aembedding=aembedding, + litellm_params=litellm_params_dict, + headers=headers, + ) elif ( custom_llm_provider == "openai_like" or custom_llm_provider == "llamafile" @@ -8142,7 +8186,23 @@ def speech( custom_llm_provider=custom_llm_provider, ) response: HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent] | None = None - if custom_llm_provider == "openai" or ( + if custom_llm_provider == "edenai": + litellm_params_dict["api_base"] = api_base + response = base_llm_http_handler.text_to_speech_handler( + model=model, + input=input, + voice=voice if isinstance(voice, str) else None, + text_to_speech_provider_config=text_to_speech_provider_config or litellm.EdenAITextToSpeechConfig(), + text_to_speech_optional_params=optional_params, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params_dict, + logging_obj=logging_obj, + timeout=timeout, + extra_headers=extra_headers, + client=client, + _is_async=aspeech or False, + ) + elif custom_llm_provider == "openai" or ( custom_llm_provider in litellm.openai_compatible_providers and custom_llm_provider not in AZURE_OPENAI_AUDIO_PROVIDERS ): diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7a2c481397f..4288a87491c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -22889,6 +22889,45 @@ "video" ] }, + "fal_ai/minimax/h3/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.13, + "output_cost_per_second_480p": 0.05, + "output_cost_per_second_768p": 0.06, + "output_cost_per_second_2k": 0.13, + "output_cost_per_second_4k": 0.16, + "source": "https://fal.ai/models/minimax/h3/text-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/minimax/h3/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.13, + "output_cost_per_second_480p": 0.05, + "output_cost_per_second_768p": 0.06, + "output_cost_per_second_2k": 0.13, + "output_cost_per_second_4k": 0.16, + "source": "https://fal.ai/models/minimax/h3/reference-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, "fal_ai/bytedance/seedance-2.0/text-to-video": { "litellm_provider": "fal_ai", "mode": "video_generation", @@ -42972,21 +43011,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.15936e-07, + "input_cost_per_token": 9.00798e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.831872e-06, + "output_cost_per_token": 1.801596e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.6328e-08, + "cache_read_input_token_cost": 7.50665e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -76850,5 +76889,65 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true + }, + "openrouter/xiaomi/mimo-v2.6-flash": { + "cache_read_input_token_cost": 2.8e-09, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": false + }, + "openrouter/xiaomi/mimo-v2.6-pro": { + "cache_read_input_token_cost": 3.6e-09, + "input_cost_per_token": 4.35e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 8.7e-07, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": false + }, + "openrouter/xiaomi/mimo-v2.6-pro-ultraspeed": { + "cache_read_input_token_cost": 3.6e-08, + "input_cost_per_token": 4.35e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 8.7e-06, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": false } } diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 30c1e0b894e..1fcb7600a5e 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -691,7 +691,7 @@ } }, "qwen_ai_platform": { - "display_name": "Qwen AI Platform (`qwen_ai_platform`)", + "display_name": "Qianwen AI Platform (`qwen_ai_platform`)", "url": "https://docs.litellm.ai/docs/providers/qwencloud", "endpoints": { "chat_completions": true, @@ -815,6 +815,24 @@ "interactions": true } }, + "edenai": { + "display_name": "Eden AI (`edenai`)", + "url": "https://docs.litellm.ai/docs/providers/edenai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": true, + "audio_transcriptions": true, + "audio_speech": true, + "moderations": false, + "batches": false, + "rerank": false, + "interactions": false, + "video_generations": true + } + }, "duckduckgo": { "display_name": "DuckDuckGo (`duckduckgo`)", "url": "https://docs.litellm.ai/docs/search/duckduckgo", diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b293ab5a206..e49638e8e9c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2500,9 +2500,8 @@ class MCPServerManager: # Filter blank scopes (e.g. YAML ``scopes: [""]``) the same way the DB-build path does, so # an all-blank list normalizes to None rather than a ``("",)`` tuple that skips the # entra_obo fail-closed scope precondition and POSTs an empty scope to the IdP. - resolved_scopes = self._extract_scopes(server_config.get("scopes")) or ( - gated_oauth_metadata.scopes if gated_oauth_metadata else None - ) + configured_scopes = self._extract_scopes(server_config.get("scopes")) + resolved_scopes = configured_scopes or (gated_oauth_metadata.scopes if gated_oauth_metadata else None) resolved_authorization_url = manual_authorization_url or ( gated_oauth_metadata.authorization_url if gated_oauth_metadata else None ) @@ -2579,6 +2578,7 @@ class MCPServerManager: client_secret=server_config.get("client_secret", None), oauth2_flow=self._explicit_oauth2_flow(config_oauth2_flow), scopes=resolved_scopes, + configured_scopes=tuple(configured_scopes) if configured_scopes else None, issuer=effective_issuer, issuer_is_anchored=use_issuer_anchor, authorization_url=resolved_authorization_url, @@ -3055,6 +3055,18 @@ class MCPServerManager: if scopes_value is not None: scopes = self._extract_scopes(scopes_value) + stored_scopes: Final[object] = credentials_dict.get("scopes") if credentials_dict else None + scopes_as_objects: Final = ( + cast(Sequence[object], stored_scopes) # cast-ok: list shape validated below + if isinstance(stored_scopes, list) + else () + ) + configured_scopes: Final = ( + tuple(scope for scope in scopes_as_objects if isinstance(scope, str)) + if scopes_as_objects and all(isinstance(scope, str) and scope for scope in scopes_as_objects) + else None + ) + name_for_prefix: Final = mcp_server.alias or mcp_server.server_name or mcp_server.server_id mcp_info: Final[MCPInfo] = _mcp_info.copy() @@ -3129,6 +3141,7 @@ class MCPServerManager: client_secret=client_secret_value or getattr(mcp_server, "client_secret", None), oauth2_flow=self._explicit_oauth2_flow(getattr(mcp_server, "oauth2_flow", None)), scopes=resolved_scopes, + configured_scopes=configured_scopes, issuer=effective_issuer, issuer_is_anchored=use_issuer_anchor, authorization_url=manual_authorization_url or getattr(gated_oauth_metadata, "authorization_url", None), @@ -7094,6 +7107,11 @@ class MCPServerManager: spec_path=server.spec_path, transport=server.transport, auth_type=server.auth_type, + credentials=( + {"scopes": list(server.configured_scopes)} # mutable-ok: MCPCredentials requires a JSON-array list + if server.configured_scopes + else None + ), created_at=server.created_at, updated_at=server.updated_at, teams=[], diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3b34440d0bf..3f6ef89fbc6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -641,6 +641,7 @@ class LiteLLMRoutes(enum.Enum): "/v1/models", "/sso/get/ui_settings", "/get/user_banner", + "/get/latest_release_info", ] # NOTE: ROUTES ONLY FOR MASTER KEY - only the Master Key should be able to Reset Spend diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 5326cf3415f..c1388e8bb81 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -22,7 +22,14 @@ import os from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol +from typing import ( + TYPE_CHECKING, + Annotated, + Final, + Literal, + Protocol, + cast, # noqa: TID251 # validated JSON values need explicit narrowing +) from fastapi import ( APIRouter, @@ -628,8 +635,8 @@ if MCP_AVAILABLE: def _preserved_admin_config_credentials( credentials: "MCPCredentials | str | None", - ) -> "dict[str, str] | None": - """Keep only the non-secret admin-config keys, which are stored unencrypted so they lift out + ) -> "dict[str, str | list[str]] | None": # mutable-ok: API response payload + """Keep non-secret admin-config keys and scopes, which are stored unencrypted so they lift out as plaintext; every secret and minted-token key is dropped. Total over every stored shape: a dict is read directly, a JSON-object string is parsed, and @@ -639,15 +646,30 @@ if MCP_AVAILABLE: parsed: object = credentials if isinstance(credentials, str): try: - parsed = json.loads(credentials) + parsed = cast(object, json.loads(credentials)) # cast-ok: JSON parse result is validated below except (ValueError, TypeError): return None if not isinstance(parsed, dict): return None - preserved: Final = { - key: value - for key in MCP_ADMIN_CONFIG_CREDENTIAL_KEYS - if isinstance((value := parsed.get(key)), str) and value + parsed_credentials: Final = cast(Mapping[str, object], parsed) # cast-ok: dict shape validated above + scopes: Final[object] = parsed_credentials.get("scopes") + scopes_as_objects: Final = ( + cast(Sequence[object], scopes) # cast-ok: list shape validated above + if isinstance(scopes, list) + else () + ) + preserved_scopes: Final = ( + {"scopes": cast(list[str], scopes_as_objects)} # cast-ok: every scope is validated below + if scopes_as_objects and all(isinstance(scope, str) and scope for scope in scopes_as_objects) + else {} + ) + preserved: Final = { # mutable-ok: API response payload + **{ + key: value + for key in MCP_ADMIN_CONFIG_CREDENTIAL_KEYS + if isinstance((value := parsed_credentials.get(key)), str) and value + }, + **preserved_scopes, } return preserved or None @@ -827,7 +849,9 @@ if MCP_AVAILABLE: if not credentials: return False as_dict: Final[dict[str, object]] = dict(credentials) - return any(value for key, value in as_dict.items() if key not in MCP_ADMIN_CONFIG_CREDENTIAL_KEYS) + return any( + value for key, value in as_dict.items() if key not in MCP_ADMIN_CONFIG_CREDENTIAL_KEYS and key != "scopes" + ) def _inherit_credentials_from_existing_server( payload: NewMCPServerRequest, diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index ea124776d0b..2ff77b5008b 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -28,6 +28,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.ptu_pricing import ( CUSTOM_PRICING_FIELDS, PTU_EMPTIED_PRICING_FIELDS, @@ -94,6 +95,7 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import ( is_ptu_cost_attribution_enabled, ) from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.credentials_repository import CredentialsRepository from litellm.repositories.model_repository import ModelRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ModelTableRepository @@ -145,7 +147,7 @@ if TYPE_CHECKING: from prisma import types as prisma_types router: Final = APIRouter() -CLEARABLE_LITELLM_PARAMS: Final = frozenset({"cache_control_injection_points"}) +CLEARABLE_LITELLM_PARAMS: Final = frozenset({"cache_control_injection_points", "litellm_credential_name"}) NULL_CLEARABLE_LITELLM_PARAMS: Final = frozenset((*SPECIAL_MODEL_INFO_PARAMS, *CLEARABLE_LITELLM_PARAMS)) @@ -332,6 +334,36 @@ def _raise_on_strategy_router_write_violation( ) +async def _raise_on_invalid_credential_name( + litellm_params: updateLiteLLMParams | None, prisma_client: PrismaClient +) -> None: + if litellm_params is None or "litellm_credential_name" not in litellm_params.model_fields_set: + return + credential_name: Final = litellm_params.litellm_credential_name + if credential_name is None: + return + if credential_name == "": + raise ProxyException( + message="litellm_credential_name cannot be an empty string. Send null to detach the stored credential or omit the field to leave it unchanged.", + type=ProxyErrorTypes.validation_error.value, + code=status.HTTP_400_BAD_REQUEST, + param="litellm_credential_name", + ) + if CredentialAccessor.find_credential(credential_name) is not None: + return + stored_credential: Final = await CredentialsRepository(WriterPinnedClient(prisma_client.db)).find_by_name( + credential_name + ) + if stored_credential is not None: + return + raise ProxyException( + message=f"Credential '{credential_name}' not found. Create it via /credentials before attaching it to a model.", + type=ProxyErrorTypes.validation_error.value, + code=status.HTTP_400_BAD_REQUEST, + param="litellm_credential_name", + ) + + AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY: Final = 5_872_301 _CAPABILITY_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" _STORED_LITELLM_PARAMS_SQL: Final = ( @@ -1110,7 +1142,9 @@ async def patch_model( litellm_params=patch_data.litellm_params, user_api_key_dict=user_api_key_dict, existing_litellm_params=db_model.litellm_params, + null_detaches=True, ) + await _raise_on_invalid_credential_name(patch_data.litellm_params, prisma_client) ModelManagementAuthChecks.can_user_set_aws_session_tags( litellm_params=patch_data.litellm_params, @@ -1920,22 +1954,33 @@ class ModelManagementAuthChecks: litellm_params: GenericLiteLLMParams | None, user_api_key_dict: UserAPIKeyAuth, existing_litellm_params: GenericLiteLLMParams | None = None, + *, + null_detaches: bool = False, ) -> Literal[True]: - if litellm_params is None or litellm_params.litellm_credential_name is None: + if litellm_params is None: return True - if existing_litellm_params is not None and existing_litellm_params.litellm_credential_name is not None: - existing_credential_name: Final = decrypt_value_helper( + if "litellm_credential_name" not in litellm_params.model_fields_set: + return True + if litellm_params.litellm_credential_name is None and not null_detaches: + return True + existing_credential_name: Final = ( + decrypt_value_helper( value=existing_litellm_params.litellm_credential_name, key="litellm_credential_name", exception_type="debug", return_original_value=True, ) - if litellm_params.litellm_credential_name == existing_credential_name: - return True + if existing_litellm_params is not None and existing_litellm_params.litellm_credential_name is not None + else None + ) + requested_credential_name: Final = litellm_params.litellm_credential_name + if requested_credential_name == existing_credential_name: + return True if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return True + action: Final = "detach" if requested_credential_name is None else "attach" raise ProxyException( - message=f"Only a proxy admin can attach a stored credential (litellm_credential_name) to a model. Your role={user_api_key_dict.user_role}.", + message=f"Only a proxy admin can {action} a stored credential (litellm_credential_name) on a model. Your role={user_api_key_dict.user_role}.", type=ProxyErrorTypes.auth_error.value, code=status.HTTP_403_FORBIDDEN, param="litellm_credential_name", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 04e6ee1d23c..1cfce8a917c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -730,6 +730,9 @@ from litellm.proxy.spend_tracking.spend_management_endpoints import ( ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.proxy.types_utils.utils import get_instance_fn +from litellm.proxy.ui_crud_endpoints.latest_release_endpoints import ( + router as latest_release_endpoints_router, +) from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( router as ui_crud_endpoints_router, ) @@ -3546,6 +3549,16 @@ async def increment_spend_counter(counter_key: str, increment: float): return await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) +async def refresh_spend_counter_ttl(counter_key: str) -> bool: + if spend_counter_cache.redis_cache is None: + return False + try: + return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key) + except Exception as e: + verbose_proxy_logger.debug("spend counter TTL refresh skipped for %s: %s", counter_key, e) + return False + + async def _increment_spend_counter_cache(counter_key: str, increment: float): if spend_counter_cache.redis_cache is not None: try: @@ -19267,6 +19280,7 @@ app.include_router(debugging_endpoints_router) app.include_router(rust_control_plane_router) app.include_router(ui_crud_endpoints_router) app.include_router(user_banner_endpoints_router) +app.include_router(latest_release_endpoints_router) app.include_router(team_callback_router) app.include_router(budget_management_router) app.include_router(model_management_router) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index d7e100dd630..0ca08cb7992 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -1129,12 +1129,12 @@ }, { "provider": "Qwen_AI_Platform", - "provider_display_name": "Qwen AI Platform", + "provider_display_name": "Qianwen AI Platform", "litellm_provider": "qwen_ai_platform", "credential_fields": [ { "key": "api_key", - "label": "Qwen AI Platform API Key", + "label": "Qianwen AI Platform API Key", "placeholder": null, "tooltip": null, "required": true, @@ -1146,7 +1146,7 @@ "key": "api_base", "label": "API Base", "placeholder": "https://dashscope.aliyuncs.com/compatible-mode/v1", - "tooltip": "The base URL for Qwen AI Platform. Defaults to https://dashscope.aliyuncs.com/compatible-mode/v1 if not specified.", + "tooltip": "The base URL for Qianwen AI Platform. Defaults to https://dashscope.aliyuncs.com/compatible-mode/v1 if not specified.", "required": true, "field_type": "text", "options": null, @@ -1321,6 +1321,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "EDENAI", + "provider_display_name": "Eden AI", + "litellm_provider": "edenai", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.edenai.run/v3", + "tooltip": "Set to https://api.eu.edenai.run/v3 for the EU endpoint", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "edenai/openai/gpt-mini-latest" + }, { "provider": "ElevenLabs", "provider_display_name": "ElevenLabs", diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 4c4785339c8..f9e5c4ff1e4 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import json import math +import time from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -105,6 +106,48 @@ def get_reserved_counter_keys(budget_reservation: dict | None) -> set: } +_lease_renewals: Final[set[asyncio.Task[None]]] = set() # mutable-ok: asyncio only weak-refs pending tasks + + +def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], counter_keys: frozenset[str]) -> None: + """A reservation lives inside spend counter keys that expire on their Redis TTL. Renew the TTL + while the request is in flight so a request longer than the TTL does not drop its + reservation and admit concurrent requests against the DB floor on any worker.""" + from litellm.proxy.proxy_server import spend_counter_cache + + if spend_counter_cache.redis_cache is None or not counter_keys: + return + task: Final = asyncio.create_task( + _renew_reservation_lease( + budget_reservation=budget_reservation, + counter_keys=counter_keys, + interval=spend_counter_cache.redis_cache.default_ttl / 2, + request_task=asyncio.current_task(), + ) + ) + _lease_renewals.add(task) + task.add_done_callback(_lease_renewals.discard) + + +async def _renew_reservation_lease( + budget_reservation: Mapping[str, object], + counter_keys: frozenset[str], + interval: float, + request_task: asyncio.Task[object] | None, +) -> None: + """Stops on finalization or once the request task that took the reservation is gone, so a + disconnect path that skipped reconciliation falls back to the plain counter TTL.""" + from litellm.proxy.proxy_server import refresh_spend_counter_ttl + + deadline: Final = time.monotonic() + litellm.request_timeout + while time.monotonic() < deadline: + await asyncio.sleep(interval) + if budget_reservation.get("finalized") is True or (request_task is not None and request_task.done()): + return + for counter_key in counter_keys: + await refresh_spend_counter_ttl(counter_key=counter_key) + + def _key_reservation_should_release_for_throttle(counter_key: str, valid_token: UserAPIKeyAuth | None) -> bool: """ Whether an over-budget key's own ``max_budget`` reservation should be @@ -319,13 +362,18 @@ async def reserve_budget_for_request( llm_router=llm_router, input_token_counts=input_token_counts, ) - return { + budget_reservation: Final = { "reserved_cost": reservation_cost, "entries": applied_entries, "finalized": False, "input_cost": min(float(input_cost or 0.0), reservation_cost), "input_tokens": max(input_token_counts.values(), default=None), } + _start_reservation_lease_renewal( + budget_reservation=budget_reservation, + counter_keys=frozenset(get_reserved_counter_keys(budget_reservation=budget_reservation)), + ) + return budget_reservation async def reconcile_budget_reservation( diff --git a/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py new file mode 100644 index 00000000000..ad5cc8efc31 --- /dev/null +++ b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py @@ -0,0 +1,153 @@ +import asyncio +import re +from collections import Counter +from collections.abc import Awaitable, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Annotated, Final, Literal, Protocol, TypeAlias + +import httpx +from fastapi import APIRouter, Depends +from pydantic import BaseModel, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +router: Final = APIRouter() + +LATEST_RELEASE_URL: Final = "https://api.github.com/repos/BerriAI/litellm/releases/latest" +LATEST_RELEASE_FETCH_TIMEOUT_SECONDS: Final = 5 +LATEST_RELEASE_CACHE_TTL_SECONDS: Final = 60 * 60 +LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS: Final = 5 * 60 +LATEST_RELEASE_CACHE_KEY: Final = "latest_release_info" + +_RELEASE_BULLET_PATTERN: Final = re.compile(r"^\*\s+(?:([A-Za-z]+)(?:\([^)]*\))?!?:\s)?\S") +_NEW_CONTRIBUTOR_PATTERN: Final = re.compile(r"^\*\s+@\S+ made their first contribution\b") + +_Bucket: TypeAlias = Literal["new_features", "bug_fixes", "other_updates"] +_PREFIX_BUCKETS: Final[Mapping[str, _Bucket]] = MappingProxyType({"feat": "new_features", "fix": "bug_fixes"}) + + +class LatestReleaseInfo(BaseModel): + version: str + new_features: int + bug_fixes: int + other_updates: int + release_url: str + + +@dataclass(frozen=True, slots=True) +class LatestReleaseUnavailable: + reason: str + + +class _GitHubRelease(BaseModel): + tag_name: str + html_url: str + body: str + + +class _AsyncGetClient(Protocol): + def get(self, url: str, *, timeout: float | None = None) -> Awaitable[httpx.Response]: ... + + +_latest_release_cache: Final = InMemoryCache(max_size_in_memory=1, default_ttl=LATEST_RELEASE_CACHE_TTL_SECONDS) +_latest_release_fetch_lock: Final = asyncio.Lock() + + +def _default_client() -> _AsyncGetClient: + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.types.llms.custom_http import httpxSpecialProvider + + return get_async_httpx_client(llm_provider=httpxSpecialProvider.UI) + + +def _default_cache() -> InMemoryCache: + return _latest_release_cache + + +def _default_fetch_lock() -> asyncio.Lock: + return _latest_release_fetch_lock + + +def _bucket_for(line: str) -> _Bucket | None: + if _NEW_CONTRIBUTOR_PATTERN.match(line) is not None: + return None + match: Final = _RELEASE_BULLET_PATTERN.match(line) + if match is None: + return None + prefix: Final = match.group(1) + return "other_updates" if prefix is None else _PREFIX_BUCKETS.get(prefix.lower(), "other_updates") + + +def count_release_bullets(body: str) -> Mapping[_Bucket, int]: + """Bucket release-note bullets by conventional-commit type or ``other_updates``.""" + return MappingProxyType(Counter(bucket for line in body.splitlines() if (bucket := _bucket_for(line)) is not None)) + + +def parse_latest_release(response: httpx.Response) -> LatestReleaseInfo | LatestReleaseUnavailable: + if response.status_code != 200: + return LatestReleaseUnavailable(reason=f"GitHub responded with status {response.status_code}") + try: + release: Final = _GitHubRelease.model_validate_json(response.content) + except ValidationError as e: + return LatestReleaseUnavailable(reason=f"GitHub release payload was not the expected shape: {e}") + counts: Final = count_release_bullets(release.body) + return LatestReleaseInfo( + version=release.tag_name.removeprefix("v"), + new_features=counts.get("new_features", 0), + bug_fixes=counts.get("bug_fixes", 0), + other_updates=counts.get("other_updates", 0), + release_url=release.html_url, + ) + + +async def fetch_latest_release(client: _AsyncGetClient) -> LatestReleaseInfo | LatestReleaseUnavailable: + try: + response: Final = await client.get(LATEST_RELEASE_URL, timeout=LATEST_RELEASE_FETCH_TIMEOUT_SECONDS) + except httpx.HTTPError as e: + return LatestReleaseUnavailable(reason=f"{type(e).__name__}: {e}") + return parse_latest_release(response) + + +async def get_latest_release_info( + client: _AsyncGetClient, cache: InMemoryCache, fetch_lock: asyncio.Lock +) -> LatestReleaseInfo | LatestReleaseUnavailable: + cached: Final = cache.get_cache(LATEST_RELEASE_CACHE_KEY) + if isinstance(cached, (LatestReleaseInfo, LatestReleaseUnavailable)): + return cached + async with fetch_lock: + cached_after_lock: Final = cache.get_cache(LATEST_RELEASE_CACHE_KEY) + if isinstance(cached_after_lock, (LatestReleaseInfo, LatestReleaseUnavailable)): + return cached_after_lock + result: Final = await fetch_latest_release(client) + ttl: Final = ( + LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS + if isinstance(result, LatestReleaseUnavailable) + else LATEST_RELEASE_CACHE_TTL_SECONDS + ) + cache.set_cache(LATEST_RELEASE_CACHE_KEY, result, ttl=ttl) + return result + + +@router.get( + "/get/latest_release_info", + tags=["UI Settings"], # mutable-ok: FastAPI's route decorator only accepts a list + dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI's route decorator only accepts a list + response_model=LatestReleaseInfo | None, +) +async def latest_release_info( + client: Annotated[_AsyncGetClient, Depends(_default_client)], + cache: Annotated[InMemoryCache, Depends(_default_cache)], + fetch_lock: Annotated[asyncio.Lock, Depends(_default_fetch_lock)], +) -> LatestReleaseInfo | None: + """ + Latest stable LiteLLM GitHub release with its PR count split into new features, bug fixes and other updates. + Returns null when GitHub can't be reached so the dashboard upgrade banner simply doesn't render. + """ + result: Final = await get_latest_release_info(client=client, cache=cache, fetch_lock=fetch_lock) + if isinstance(result, LatestReleaseUnavailable): + verbose_proxy_logger.warning("LiteLLM: latest release info unavailable: %s", result.reason) + return None + return result diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index a22eff79dbb..b38684f1856 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -751,6 +751,7 @@ class ANTHROPIC_BETA_HEADER_VALUES(str, Enum): FAST_MODE_2026_02_01 = "fast-mode-2026-02-01" ADVISOR_TOOL_2026_03_01 = "advisor-tool-2026-03-01" PER_TURN_CONTROL_2026_07_01 = "per-turn-control-2026-07-01" + DANGEROUS_TOOL_USE_2026_09_03 = "dangerous-tool-use-2026-09-03" # Tool search beta header constant (for Anthropic direct API and Microsoft Foundry) diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index f9eaef5e891..3674bb670d5 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1238,6 +1238,7 @@ class BedrockInvokeAnthropicMessagesRequest(TypedDict, total=False): thinking: dict metadata: dict output_config: dict + safeguards: list # `context_management` is allowed for Bedrock InvokeModel only when it # carries `compact_20260112` edits paired with the `compact-2026-01-12` diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 985d31af997..cb32299b143 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -99,6 +99,7 @@ class MCPServer(BaseModel): configured_authorization_url: str | None = None configured_token_url: str | None = None configured_registration_url: str | None = None + configured_scopes: tuple[str, ...] | None = None # How the gateway authenticates to the upstream token endpoint. When # "client_secret_basic" the credentials go in an HTTP Basic Authorization # header (omitted from the body); None defaults to "client_secret_post". diff --git a/litellm/types/router.py b/litellm/types/router.py index 6b43573d0e2..fd426835d65 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -540,6 +540,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): output_cost_per_second: float | None output_cost_per_second_480p: ReadOnly[float | None] output_cost_per_second_720p: ReadOnly[float | None] + output_cost_per_second_768p: ReadOnly[float | None] + output_cost_per_second_2k: ReadOnly[float | None] output_cost_per_second_1080p: float | None output_cost_per_second_4k: ReadOnly[float | None] num_retries: int | None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4075324555b..142304223d0 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -330,6 +330,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): ) # video_generation tier: key output_cost_per_second_ (e.g. 1080p, 720p) output_cost_per_second_480p: ReadOnly[float | None] output_cost_per_second_720p: ReadOnly[float | None] + output_cost_per_second_768p: ReadOnly[float | None] + output_cost_per_second_2k: ReadOnly[float | None] output_cost_per_second_4k: ReadOnly[float | None] ocr_cost_per_page: float | None # for OCR models ocr_cost_per_page_batches: ReadOnly[float | None] @@ -3615,6 +3617,8 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_second_1080p: float | None = None output_cost_per_second_480p: float | None = None output_cost_per_second_720p: float | None = None + output_cost_per_second_768p: float | None = None + output_cost_per_second_2k: float | None = None output_cost_per_second_4k: float | None = None input_cost_per_pixel: float | None = None output_cost_per_pixel: float | None = None @@ -4157,6 +4161,7 @@ class LlmProviders(str, Enum): OCI = "oci" AUTO_ROUTER = "auto_router" VERCEL_AI_GATEWAY = "vercel_ai_gateway" + EDENAI = "edenai" DOTPROMPT = "dotprompt" MANUS = "manus" WANDB = "wandb" diff --git a/litellm/utils.py b/litellm/utils.py index 2bc43f6a756..709f3f6d1dd 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3313,6 +3313,9 @@ def register_model( elif value.get("litellm_provider") == "vercel_ai_gateway": if key not in litellm.vercel_ai_gateway_models: litellm.vercel_ai_gateway_models.add(key) + elif value.get("litellm_provider") == "edenai": + if key not in litellm.edenai_models: + litellm.edenai_models.add(key) elif value.get("litellm_provider") == "vertex_ai-text-models": if key not in litellm.vertex_text_models: litellm.vertex_text_models.add(key) @@ -4895,6 +4898,9 @@ def get_optional_params( return optional_params +EXTRA_BODY_ROUTING_KEYS: Final = frozenset({"model"}) + + def add_provider_specific_params_to_optional_params( optional_params: dict, passed_params: dict, @@ -4920,10 +4926,8 @@ def add_provider_specific_params_to_optional_params( **extra_body, } - if additional_drop_params is not None: - processed_extra_body = {k: v for k, v in initial_extra_body.items() if k not in additional_drop_params} - else: - processed_extra_body = initial_extra_body + dropped_keys: Final = EXTRA_BODY_ROUTING_KEYS | frozenset(additional_drop_params or ()) + processed_extra_body: Final = {k: v for k, v in initial_extra_body.items() if k not in dropped_keys} _ensure_extra_body_is_safe: Final = getattr(sys.modules[__name__], "_ensure_extra_body_is_safe") optional_params["extra_body"] = _ensure_extra_body_is_safe(extra_body=processed_extra_body) @@ -6100,6 +6104,8 @@ def _get_model_info_helper( output_cost_per_second_1080p=_model_info.get("output_cost_per_second_1080p", None), output_cost_per_second_480p=_model_info.get("output_cost_per_second_480p", None), output_cost_per_second_720p=_model_info.get("output_cost_per_second_720p", None), + output_cost_per_second_768p=_model_info.get("output_cost_per_second_768p", None), + output_cost_per_second_2k=_model_info.get("output_cost_per_second_2k", None), output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None), output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None), output_cost_per_image=_model_info.get("output_cost_per_image", None), @@ -6572,6 +6578,11 @@ def validate_environment( keys_in_environment = True else: missing_keys.append("VERCEL_AI_GATEWAY_API_KEY") + elif custom_llm_provider == "edenai": + if "EDENAI_API_KEY" in os.environ: + keys_in_environment = True + else: + missing_keys.append("EDENAI_API_KEY") elif custom_llm_provider == "datarobot": if "DATAROBOT_API_TOKEN" in os.environ: keys_in_environment = True @@ -6822,6 +6833,12 @@ def validate_environment( keys_in_environment = True else: missing_keys.append("VERCEL_AI_GATEWAY_API_KEY") + ## edenai + elif model in litellm.edenai_models: + if "EDENAI_API_KEY" in os.environ: + keys_in_environment = True + else: + missing_keys.append("EDENAI_API_KEY") ## datarobot elif model in litellm.datarobot_models: if "DATAROBOT_API_TOKEN" in os.environ: @@ -8322,6 +8339,7 @@ class ProviderConfigManager: lambda: litellm.VercelAIGatewayConfig(), False, ), + LlmProviders.EDENAI: (litellm.EdenAIChatConfig, False), LlmProviders.COMETAPI: (lambda: litellm.CometAPIConfig(), False), LlmProviders.DATAROBOT: (lambda: litellm.DataRobotConfig(), False), LlmProviders.GEMINI: (lambda: litellm.GoogleAIStudioGeminiConfig(), False), @@ -8624,6 +8642,8 @@ class ProviderConfigManager: return SagemakerEmbeddingConfig.get_model_config(model) elif litellm.LlmProviders.PERPLEXITY == provider: return litellm.PerplexityEmbeddingConfig() + elif litellm.LlmProviders.EDENAI == provider: + return litellm.EdenAIEmbeddingConfig() return None @staticmethod @@ -8744,6 +8764,8 @@ class ProviderConfigManager: ) return GithubCopilotAnthropicMessagesConfig() + elif litellm.LlmProviders.EDENAI == provider: + return litellm.EdenAIAnthropicMessagesConfig() from litellm.llms.openai_like.json_loader import JSONProviderRegistry @@ -8852,6 +8874,8 @@ class ProviderConfigManager: ) return GeminiAudioTranscriptionConfig() + elif litellm.LlmProviders.EDENAI == provider: + return litellm.EdenAIAudioTranscriptionConfig() return None @staticmethod @@ -8954,6 +8978,8 @@ class ProviderConfigManager: return litellm.HostedVLLMResponsesAPIConfig() elif litellm.LlmProviders.FIREWORKS_AI == provider: return litellm.FireworksAIResponsesAPIConfig() + elif litellm.LlmProviders.EDENAI == provider: + return litellm.EdenAIResponsesAPIConfig() elif litellm.LlmProviders.BEDROCK_MANTLE == provider: # Both decisions are data-driven from the model's price-map entry, with # no model-name logic. Capability (can it serve Responses?) comes from @@ -9026,7 +9052,7 @@ class ProviderConfigManager: return litellm.OpenAITextCompletionConfig() @staticmethod - def get_provider_model_info( + def get_provider_model_info( # noqa: C901 # provider dispatch table, one branch per provider model: str | None, provider: LlmProviders, ) -> BaseLLMModelInfo | None: @@ -9063,6 +9089,8 @@ class ProviderConfigManager: return litellm.LemonadeChatConfig() elif LlmProviders.CLARIFAI == provider: return litellm.ClarifaiConfig() + elif LlmProviders.EDENAI == provider: + return litellm.EdenAIChatConfig() elif LlmProviders.BEDROCK == provider: from litellm.llms.bedrock.common_utils import BedrockModelInfo @@ -9411,6 +9439,8 @@ class ProviderConfigManager: ) return get_modelscope_image_generation_config(model) + elif LlmProviders.EDENAI == provider: + return litellm.EdenAIImageGenerationConfig() return None @staticmethod @@ -9446,6 +9476,8 @@ class ProviderConfigManager: from litellm.llms.hosted_vllm.videos import get_hosted_vllm_video_config return get_hosted_vllm_video_config(model) + elif LlmProviders.EDENAI == provider: + return litellm.EdenAIVideoConfig() return None @staticmethod @@ -9766,6 +9798,8 @@ class ProviderConfigManager: ) return AWSPollyTextToSpeechConfig() + elif litellm.LlmProviders.EDENAI == provider: + return litellm.EdenAITextToSpeechConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7a2c481397f..4288a87491c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -22889,6 +22889,45 @@ "video" ] }, + "fal_ai/minimax/h3/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.13, + "output_cost_per_second_480p": 0.05, + "output_cost_per_second_768p": 0.06, + "output_cost_per_second_2k": 0.13, + "output_cost_per_second_4k": 0.16, + "source": "https://fal.ai/models/minimax/h3/text-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/minimax/h3/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.13, + "output_cost_per_second_480p": 0.05, + "output_cost_per_second_768p": 0.06, + "output_cost_per_second_2k": 0.13, + "output_cost_per_second_4k": 0.16, + "source": "https://fal.ai/models/minimax/h3/reference-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, "fal_ai/bytedance/seedance-2.0/text-to-video": { "litellm_provider": "fal_ai", "mode": "video_generation", @@ -42972,21 +43011,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.15936e-07, + "input_cost_per_token": 9.00798e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.831872e-06, + "output_cost_per_token": 1.801596e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.6328e-08, + "cache_read_input_token_cost": 7.50665e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -76850,5 +76889,65 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true + }, + "openrouter/xiaomi/mimo-v2.6-flash": { + "cache_read_input_token_cost": 2.8e-09, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": false + }, + "openrouter/xiaomi/mimo-v2.6-pro": { + "cache_read_input_token_cost": 3.6e-09, + "input_cost_per_token": 4.35e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 8.7e-07, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": false + }, + "openrouter/xiaomi/mimo-v2.6-pro-ultraspeed": { + "cache_read_input_token_cost": 3.6e-08, + "input_cost_per_token": 4.35e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 8.7e-06, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": false } } diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 5b0a23adfea..0509516ac32 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -607,6 +607,10 @@ "type": "number", "minimum": 0 }, + "output_cost_per_second_2k": { + "type": "number", + "minimum": 0 + }, "output_cost_per_second_480p": { "type": "number", "minimum": 0 @@ -619,6 +623,10 @@ "type": "number", "minimum": 0 }, + "output_cost_per_second_768p": { + "type": "number", + "minimum": 0 + }, "output_cost_per_token": { "type": "number", "minimum": 0, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index af9b194bbee..b8d1621cde3 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -744,7 +744,7 @@ } }, "qwen_ai_platform": { - "display_name": "Qwen AI Platform (`qwen_ai_platform`)", + "display_name": "Qianwen AI Platform (`qwen_ai_platform`)", "url": "https://docs.litellm.ai/docs/providers/qwencloud", "endpoints": { "chat_completions": true, @@ -868,6 +868,24 @@ "interactions": true } }, + "edenai": { + "display_name": "Eden AI (`edenai`)", + "url": "https://docs.litellm.ai/docs/providers/edenai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": true, + "audio_transcriptions": true, + "audio_speech": true, + "moderations": false, + "batches": false, + "rerank": false, + "interactions": false, + "video_generations": true + } + }, "duckduckgo": { "display_name": "DuckDuckGo (`duckduckgo`)", "url": "https://docs.litellm.ai/docs/search/duckduckgo", diff --git a/pyproject.toml b/pyproject.toml index 3b17a397f3c..95da93df41e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -236,6 +236,7 @@ e2e-dev = [ "playwright==1.61.0", "websockets>=15.0.1,<16.0", "locust==2.45.0", + "anthropic==0.84.0", "psutil==7.2.2", "mcp>=2.2.0,<3", ] diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index c25e958242f..b00b7dfac95 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -85,6 +85,8 @@ That snippet only conveys intent. What you actually write uses the real harness: Every HTTP call goes through the shared transport, never through `requests.*` in a test. `e2e_http.py` is the only module permitted to call `requests.*`, and that is enforced in CI by `tests/code_coverage_tests/check_e2e_no_raw_requests.py`. A test that imports requests will fail the check +One deliberate exception: LLM-endpoint calls in `llm_translation/` go through the real provider SDKs (OpenAI, Anthropic) via the suite's `sdk` fixture (`llm_translation/sdk_clients.py`), because that is what customers actually run against the proxy (LIT-4577). The SDKs raise their own typed exceptions on failure, which is exactly the customer-observable contract; management routes (model/key CRUD, spend read-back) and endpoints no official SDK covers (e.g. `/v1/rerank`, `/v1/ocr`, custom passthrough paths) stay on the shared transport. Raw HTTP client imports remain banned either way + The shape is layered so tests stay declarative `transport.py` exposes a `Transport` Protocol with `post`, `get`, `delete`, `send`, `stream`, `probe`, plus `bearer(key)` and the `master` header. `HttpTransport` fulfils it, and `SplitTransport` routes each call by path to the data plane or the control plane so a split control-plane/data-plane deployment works without any change in the test diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 6c3dc4d0bd1..2b59e8770b5 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -75,10 +75,10 @@ The suites run against a live proxy, so bring one up first by running the litell Buildkite runs this suite against a Keycloak deployed beside the ephemeral stack by project-releaser. It fetches the realm from the test-runner revision even when it reuses a gateway image from another commit. The GitHub Actions changed-test stack starts the same digest-pinned Keycloak through `.github/e2e-stack/start-idp.sh`, imports the checked-out realm, and exports the IdP URL and credentials in `stack.env`. Both runners configure issuer/audience validation and store the realm, keys and users in a separate schema in the stack's PostgreSQL, so replacing Keycloak preserves token validity. Both wait for realm discovery before running tests. Losing the whole ephemeral database invalidates the stack. Keycloak skips imports into an existing realm, so changes to the realm export require a fresh stack (or deliberately replacing the local data volume). A stack without it fails the JWT tests rather than skipping them -4. Run a suite against it; the harness reads `LITELLM_PROXY_URL` (default `http://localhost:4000`): +4. Run a suite against it; the harness reads `LITELLM_PROXY_URL` (default `http://localhost:4000`). The suites' client dependencies (the provider SDKs, websockets) live in the `e2e-dev` dependency group; `make bootstrap` installs it, and naming the group on the run keeps the command working from any environment state: ```bash - uv run pytest tests/e2e/llm_translation/ -v + uv run --group e2e-dev pytest tests/e2e/llm_translation/ -v ``` The browser tests in the `management/` suite drive the dashboard the proxy serves at `/ui` through playwright, an optional dependency behind `importorskip` (the suite's API tests run without it). It lives in the `e2e-dev` dependency group; install it along with its browser: @@ -206,6 +206,8 @@ That snippet only conveys intent. What you actually write uses the real harness: Every HTTP call goes through the shared transport, never through `requests.*` in a test. `e2e_http.py` is the only module permitted to call `requests.*`, and that is enforced in CI by `tests/code_coverage_tests/check_e2e_no_raw_requests.py`. A test that imports requests will fail the check +One deliberate exception: LLM-endpoint calls in `llm_translation/` go through the real provider SDKs (OpenAI, Anthropic) via the suite's `sdk` fixture (`llm_translation/sdk_clients.py`), because that is what customers actually run against the proxy (LIT-4577). Management routes and endpoints no official SDK covers stay on the shared transport, and raw HTTP client imports remain banned either way + The shape is layered so tests stay declarative `transport.py` exposes a `Transport` Protocol with `post`, `get`, `delete`, `send`, `stream`, `probe`, plus `bearer(key)` and the `master` header. `HttpTransport` fulfils it, and `SplitTransport` routes each call by path to the data plane or the control plane so a split control-plane/data-plane deployment works without any change in the test @@ -230,7 +232,7 @@ Before you push ```bash litellm --config .yml --port 4000 - uv run pytest tests/e2e// -v + uv run --group e2e-dev pytest tests/e2e// -v ``` 4. Capture screenshots of the test run and attach them to the PR as proof diff --git a/tests/e2e/llm_translation/conftest.py b/tests/e2e/llm_translation/conftest.py index f35ecf0760d..9fd45799773 100644 --- a/tests/e2e/llm_translation/conftest.py +++ b/tests/e2e/llm_translation/conftest.py @@ -2,14 +2,16 @@ The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker live in the parent tests/e2e/conftest.py. PassthroughClient holds the shared -ProxyClient, so the `resources` fixture cleans up keys this suite creates. +ProxyClient, so the `resources` fixture cleans up keys this suite creates. The +`sdk` fixture hands tests real provider SDK clients (OpenAI, Anthropic) pointed +at the proxy, the way customers actually call it. """ import pytest -from endpoints_client import EndpointsClient, build_endpoints_client from passthrough_client import PassthroughClient, build_client from proxy_client import ProxyClient +from sdk_clients import SdkClients, build_sdk_clients def pytest_configure(config: pytest.Config) -> None: @@ -25,5 +27,5 @@ def client(proxy: ProxyClient) -> PassthroughClient: @pytest.fixture(scope="session") -def endpoints_client(proxy: ProxyClient) -> EndpointsClient: - return build_endpoints_client(proxy) +def sdk() -> SdkClients: + return build_sdk_clients() diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py deleted file mode 100644 index eb7bae2220c..00000000000 --- a/tests/e2e/llm_translation/endpoints_client.py +++ /dev/null @@ -1,476 +0,0 @@ -"""Client for the non-chat inference endpoints (responses, messages, rerank, -embeddings, audio speech, image generation). - -Each test registers the deployment it needs through /model/new (deleted on -teardown), so nothing is hardcoded into the gateway config, then drives the -endpoint with `send` and parses the provider-native body with a suite-local model -so the assertion is on real content, not just a 200. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Literal - -from e2e_config import SLOW_PROVIDER_TIMEOUT_SECONDS -from e2e_http import BinaryStream, Result, StreamingResponse -from models import CacheControl, ChatMessage, LiteLLMParamsBody, RichMessage, TextBlock -from proxy_client import ProxyClient -from pydantic import BaseModel - -__all__ = [ - "CacheControl", - "ImageEditForm", - "ImagesResult", - "RichMessage", - "TextBlock", - "TranscriptionForm", - "TranscriptionResult", -] - - -class FunctionParameterProperty(BaseModel): - type: str - description: str | None = None - - -class FunctionParameters(BaseModel): - type: Literal["object"] = "object" - properties: dict[str, FunctionParameterProperty] - required: list[str] = [] - - -class ResponsesFunctionTool(BaseModel): - type: Literal["function"] = "function" - name: str - description: str | None = None - parameters: FunctionParameters - - -class ResponsesInputTextPart(BaseModel): - type: Literal["input_text"] = "input_text" - text: str - - -class ResponsesInputImagePart(BaseModel): - type: Literal["input_image"] = "input_image" - image_url: str - - -ResponsesInputContentPart = ResponsesInputTextPart | ResponsesInputImagePart - - -class ResponsesInputMessage(BaseModel): - role: Literal["user", "assistant", "system"] = "user" - content: list[ResponsesInputContentPart] - - -ResponsesInput = str | list[ResponsesInputMessage] - - -class ResponsesRequest(BaseModel): - model: str - input: ResponsesInput - instructions: str | None = None - stream: bool = False - tools: list[ResponsesFunctionTool] | None = None - guardrails: list[str] | None = None - safety_identifier: str | None = None - cache: dict[str, bool] | None = {"no-cache": True} - - -class MessagesRequest(BaseModel): - model: str - max_tokens: int - messages: list[ChatMessage] - cache: dict[str, bool] | None = {"no-cache": True} - - -class RichMessagesRequest(BaseModel): - model: str - max_tokens: int = 64 - system: list[TextBlock] - messages: list[RichMessage] - cache: dict[str, bool] | None = {"no-cache": True} - - -class CompletionsRequest(BaseModel): - model: str - prompt: str - max_tokens: int = 32 - cache: dict[str, bool] | None = {"no-cache": True} - - -class EmbeddingsRequest(BaseModel): - model: str - input: str - cache: dict[str, bool] | None = {"no-cache": True} - - -class RerankRequest(BaseModel): - model: str - query: str - documents: list[str] - top_n: int - cache: dict[str, bool] | None = {"no-cache": True} - - -class SpeechRequest(BaseModel): - model: str - input: str - voice: str - - -class ImageRequest(BaseModel): - model: str - prompt: str - n: int = 1 - size: str = "1024x1024" - - -class ImageEditForm(BaseModel): - model: str - prompt: str - n: int = 1 - - -class TranscriptionForm(BaseModel): - model: str - response_format: str = "json" - - -class ModerationRequest(BaseModel): - model: str - input: str - - -class GenerateContentPart(BaseModel): - text: str - - -class GenerateContentContent(BaseModel): - role: Literal["user"] = "user" - parts: tuple[GenerateContentPart, ...] - - -class GenerateContentBody(BaseModel): - contents: tuple[GenerateContentContent, ...] - - -class ResponsesOutputContent(BaseModel): - type: str | None = None - text: str | None = None - - -class ResponsesOutputItem(BaseModel): - type: str | None = None - content: list[ResponsesOutputContent] = [] - name: str | None = None - arguments: str | None = None - call_id: str | None = None - - -class ResponsesResult(BaseModel): - id: str | None = None - status: str | None = None - model: str | None = None - output: list[ResponsesOutputItem] = [] - - @property - def text(self) -> str: - return "".join( - content.text or "" for item in self.output for content in item.content - ) - - @property - def function_calls(self) -> tuple[ResponsesOutputItem, ...]: - return tuple( - item - for item in self.output - if item.type == "function_call" - and item.name is not None - and item.arguments is not None - ) - - -class ResponsesStreamEvent(BaseModel): - event_id: str | None = None - - -class ResponsesStreamEventType(BaseModel): - type: str - - -class ResponsesOutputTextDeltaEvent(ResponsesStreamEvent): - type: Literal["response.output_text.delta"] - delta: str - - -class AnthropicContentBlock(BaseModel): - type: str | None = None - text: str | None = None - - -class MessagesUsage(BaseModel): - input_tokens: int = 0 - output_tokens: int = 0 - cache_creation_input_tokens: int = 0 - cache_read_input_tokens: int = 0 - - -class MessagesResult(BaseModel): - id: str | None = None - role: str | None = None - model: str | None = None - content: list[AnthropicContentBlock] = [] - usage: MessagesUsage = MessagesUsage() - - @property - def text(self) -> str: - return "".join(block.text or "" for block in self.content) - - -class CompletionChoice(BaseModel): - text: str | None = None - - -class CompletionsResult(BaseModel): - choices: list[CompletionChoice] = [] - - -class EmbeddingItem(BaseModel): - embedding: list[float] = [] - - -class EmbeddingsResult(BaseModel): - data: list[EmbeddingItem] = [] - - @property - def first_vector(self) -> tuple[float, ...]: - return tuple(self.data[0].embedding) if self.data else () - - -class RerankItem(BaseModel): - index: int | None = None - relevance_score: float | None = None - - -class RerankResult(BaseModel): - results: list[RerankItem] = [] - - -class ImageItem(BaseModel): - url: str | None = None - b64_json: str | None = None - - -class ImagesResult(BaseModel): - data: list[ImageItem] = [] - - -class TranscriptionResult(BaseModel): - text: str = "" - - -class ModerationResultItem(BaseModel): - flagged: bool - categories: dict[str, bool] = {} - - @property - def flagged_categories(self) -> tuple[str, ...]: - return tuple(name for name, hit in self.categories.items() if hit) - - -class ModerationResult(BaseModel): - results: list[ModerationResultItem] = [] - - @property - def first(self) -> ModerationResultItem | None: - return self.results[0] if self.results else None - - -@dataclass(frozen=True, slots=True) -class EndpointsClient: - proxy: ProxyClient - - def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str: - return self.proxy.create_model(model_name, litellm_params) - - def delete_model(self, model_id: str) -> None: - self.proxy.delete_model(model_id) - - def _send( - self, path: str, key: str, body: BaseModel, *, stream: bool = False - ) -> StreamingResponse: - return self.proxy.transport.send( - path, - headers=self.proxy.transport.bearer(key), - json=body, - stream=stream, - ) - - def responses( - self, - key: str, - model: str, - text: str, - *, - stream: bool = False, - guardrails: list[str] | None = None, - safety_identifier: str | None = None, - ) -> StreamingResponse: - return self._send( - "/v1/responses", - key, - ResponsesRequest( - model=model, - input=text, - instructions="You are a helpful assistant", - stream=stream, - guardrails=guardrails, - safety_identifier=safety_identifier, - ), - stream=stream, - ) - - def responses_vision( - self, key: str, model: str, text: str, image_url: str - ) -> StreamingResponse: - return self._send( - "/v1/responses", - key, - ResponsesRequest( - model=model, - input=[ - ResponsesInputMessage( - content=[ - ResponsesInputTextPart(text=text), - ResponsesInputImagePart(image_url=image_url), - ] - ) - ], - instructions="You are a helpful assistant", - ), - ) - - def responses_with_tools( - self, key: str, model: str, text: str, tools: list[ResponsesFunctionTool] - ) -> StreamingResponse: - return self._send( - "/v1/responses", - key, - ResponsesRequest( - model=model, - input=text, - instructions="You are a helpful assistant", - tools=tools, - ), - ) - - def messages( - self, key: str, model: str, text: str, *, max_tokens: int = 64 - ) -> StreamingResponse: - return self._send( - "/v1/messages", - key, - MessagesRequest( - model=model, - max_tokens=max_tokens, - messages=[ChatMessage(role="user", content=text)], - ), - ) - - def text_completions( - self, key: str, model: str, prompt: str, *, max_tokens: int = 32 - ) -> StreamingResponse: - return self._send( - "/v1/completions", - key, - CompletionsRequest(model=model, prompt=prompt, max_tokens=max_tokens), - ) - - def embeddings(self, key: str, model: str, text: str) -> StreamingResponse: - return self._send("/embeddings", key, EmbeddingsRequest(model=model, input=text)) - - def rerank( - self, key: str, model: str, query: str, documents: list[str], top_n: int - ) -> StreamingResponse: - return self._send( - "/v1/rerank", - key, - RerankRequest(model=model, query=query, documents=documents, top_n=top_n), - ) - - def audio_speech( - self, key: str, model: str, text: str, *, voice: str = "alloy" - ) -> StreamingResponse: - return self._send( - "/v1/audio/speech", key, SpeechRequest(model=model, input=text, voice=voice) - ) - - def audio_speech_stream( - self, key: str, model: str, text: str, *, voice: str = "alloy" - ) -> BinaryStream: - return self.proxy.transport.stream_binary( - "/v1/audio/speech", - headers=self.proxy.transport.bearer(key), - json=SpeechRequest(model=model, input=text, voice=voice), - ) - - def transcribe( - self, key: str, model: str, *, filename: str, content: bytes - ) -> Result[TranscriptionResult]: - return self.proxy.transport.upload( - "/v1/audio/transcriptions", - headers=self.proxy.transport.bearer(key), - form=TranscriptionForm(model=model), - filename=filename, - content=content, - file_content_type="audio/wav", - response_type=TranscriptionResult, - ) - - def moderations(self, key: str, model: str, text: str) -> Result[ModerationResult]: - return self.proxy.transport.post( - "/v1/moderations", - headers=self.proxy.transport.bearer(key), - json=ModerationRequest(model=model, input=text), - response_type=ModerationResult, - ) - - def images(self, key: str, model: str, prompt: str) -> StreamingResponse: - return self._send( - "/v1/images/generations", key, ImageRequest(model=model, prompt=prompt) - ) - - def image_edit( - self, key: str, model: str, prompt: str, image: bytes, *, filename: str = "image.png" - ) -> Result[ImagesResult]: - return self.proxy.transport.upload( - "/v1/images/edits", - headers=self.proxy.transport.bearer(key), - form=ImageEditForm(model=model, prompt=prompt), - filename=filename, - content=image, - file_content_type="image/png", - file_field="image", - response_type=ImagesResult, - timeout=SLOW_PROVIDER_TIMEOUT_SECONDS, - ) - - def generate_content( - self, key: str, model: str, text: str, *, stream: bool = False - ) -> StreamingResponse: - operation = "streamGenerateContent" if stream else "generateContent" - return self._send( - f"/v1beta/models/{model}:{operation}", - key, - GenerateContentBody( - contents=(GenerateContentContent(parts=(GenerateContentPart(text=text),)),) - ), - stream=stream, - ) - - -def build_endpoints_client(proxy: ProxyClient) -> EndpointsClient: - return EndpointsClient(proxy=proxy) diff --git a/tests/e2e/llm_translation/sdk_clients.py b/tests/e2e/llm_translation/sdk_clients.py new file mode 100644 index 00000000000..145efbdca98 --- /dev/null +++ b/tests/e2e/llm_translation/sdk_clients.py @@ -0,0 +1,62 @@ +"""Real provider SDK clients pointed at the proxy, connected the way customers +connect (LIT-4577). + +The OpenAI SDK drives the OpenAI-compatible surface (/responses, /embeddings, +/images/generations, /moderations, /audio/*) and the Anthropic SDK drives +/v1/messages, each authenticated with a litellm virtual key. Errors surface as +the SDK's own exceptions, exactly what an end user sees. Retries are disabled +so a proxy fault fails the test instead of being papered over, and the timeout +matches the shared transport's request budget. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +from anthropic import Anthropic +from openai import OpenAI + +from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT + +NO_PROXY_CACHE: Final = MappingProxyType({"cache": {"no-cache": True}}) +"""``extra_body`` for every cacheable SDK call (messages, responses, completions, +embeddings): the gateway under test caches those call types, so an identical +re-send would otherwise be served from Redis instead of reaching the provider, +which hides provider-side behavior such as prompt-cache warm-up. The SDKs +themselves cannot bypass it (``Cache-Control`` only sets a TTL on the proxy).""" + + +def response_header(headers: Mapping[str, str], name: str) -> str | None: + """Typed read of an SDK response header: httpx.Headers.get returns Any and + httpx itself is a banned import in suite code, so tests read headers through + the Mapping[str, str] interface Headers fulfils.""" + return headers[name] if name in headers else None + + +@dataclass(frozen=True, slots=True) +class SdkClients: + base_url: str + request_timeout: float + + def openai(self, key: str) -> OpenAI: + return OpenAI( + base_url=self.base_url, + api_key=key, + timeout=self.request_timeout, + max_retries=0, + ) + + def anthropic(self, key: str) -> Anthropic: + return Anthropic( + base_url=self.base_url, + api_key=key, + timeout=self.request_timeout, + max_retries=0, + ) + + +def build_sdk_clients() -> SdkClients: + return SdkClients(base_url=PROXY_BASE_URL, request_timeout=REQUEST_TIMEOUT) diff --git a/tests/e2e/llm_translation/test_audio_speech_e2e.py b/tests/e2e/llm_translation/test_audio_speech_e2e.py index 784007ec789..c3b6fddb632 100644 --- a/tests/e2e/llm_translation/test_audio_speech_e2e.py +++ b/tests/e2e/llm_translation/test_audio_speech_e2e.py @@ -1,20 +1,23 @@ """Live e2e: POST /v1/audio/speech returns audio, non-streamed and streamed. -The non-streamed call asserts an audio (not JSON) body. The streamed call consumes -the response the way a player would and asserts customer-observable streaming: -chunked transfer encoding (a buffered body would carry a content-length) with -non-zero audio bytes. +Both positive calls go through the real OpenAI SDK (LIT-4577). The non-streamed +call asserts an audio (not JSON) body. The streamed call consumes the response +the way a player would and asserts customer-observable streaming: chunked +transfer encoding (a buffered body would carry a content-length) with non-zero +audio bytes. The malformed-body negatives stay on the shared transport because +the SDK refuses to send a request missing its required fields. """ from __future__ import annotations import pytest from e2e_config import unique_marker -from e2e_http import assert_client_error, require_successful_call -from endpoints_client import EndpointsClient +from e2e_http import assert_client_error from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient from pydantic import BaseModel +from sdk_clients import SdkClients, response_header pytestmark = pytest.mark.e2e @@ -25,67 +28,75 @@ class _OptionalSpeechBody(BaseModel): voice: str | None = None -def _register_tts( - endpoints_client: EndpointsClient, resources: ResourceManager -) -> tuple[str, str]: +def _register_tts(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]: model = f"e2e-speech-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody(model="openai/gpt-4o-mini-tts", api_key="os.environ/OPENAI_API_KEY"), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) return model, resources.key() class TestAudioSpeech: @pytest.mark.covers("llm.audio_speech.openai.basic.nonstream.works") def test_audio_speech_returns_audio( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = _register_tts(endpoints_client, resources) - result = endpoints_client.audio_speech(key, model, "Hello!") - require_successful_call(result) - assert "audio" in (result.content_type or ""), ( - f"/audio/speech content-type is not audio: {result.content_type!r}" + model, key = _register_tts(proxy, resources) + client = sdk.openai(key) + + response = client.audio.speech.with_raw_response.create( + model=model, voice="alloy", input="Hello!" ) - assert result.body, "/audio/speech returned an empty body" + content_type = response_header(response.headers, "content-type") + assert "audio" in (content_type or ""), ( + f"/audio/speech content-type is not audio: {content_type!r}" + ) + assert response.content, "/audio/speech returned an empty body" @pytest.mark.covers("llm.audio_speech.openai.basic.stream.works") def test_audio_speech_streams_audio_chunks( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = _register_tts(endpoints_client, resources) - result = endpoints_client.audio_speech_stream( - key, - model, - "Streaming speech should arrive in several audio chunks so a client can " - "begin playback well before the whole clip has finished generating.", + model, key = _register_tts(proxy, resources) + client = sdk.openai(key) + + with client.audio.speech.with_streaming_response.create( + model=model, + voice="alloy", + input=( + "Streaming speech should arrive in several audio chunks so a client can " + "begin playback well before the whole clip has finished generating." + ), + ) as response: + content_type = response_header(response.headers, "content-type") + transfer_encoding = response_header(response.headers, "transfer-encoding") + content_length = response_header(response.headers, "content-length") + total_bytes = sum(len(chunk) for chunk in response.iter_bytes(chunk_size=8192)) + + assert "audio" in (content_type or ""), ( + f"/audio/speech content-type is not audio: {content_type!r}" ) - assert result.ok, ( - f"/audio/speech stream failed (status {result.status_code}); body={result.error_body}" + assert "chunked" in (transfer_encoding or ""), ( + f"/audio/speech did not stream: transfer-encoding={transfer_encoding!r}, " + f"content-length={content_length!r} (a buffered body is not a stream)" ) - assert "audio" in (result.content_type or ""), ( - f"/audio/speech content-type is not audio: {result.content_type!r}" - ) - assert result.chunked, ( - f"/audio/speech did not stream: transfer-encoding={result.transfer_encoding!r}, " - f"content-length={result.content_length!r} (a buffered body is not a stream)" - ) - assert result.content_length is None, ( - f"/audio/speech advertised content-length={result.content_length!r} on a " + assert content_length is None, ( + f"/audio/speech advertised content-length={content_length!r} on a " f"streamed response (a buffered body is not a stream)" ) - assert result.total_bytes > 0, "/audio/speech stream returned no audio bytes" + assert total_bytes > 0, "/audio/speech stream returned no audio bytes" @pytest.mark.skip(reason="stage red: product gap, /v1/audio/speech 500s on missing input instead of 400") @pytest.mark.covers("llm.audio_speech.openai.input_validation.nonstream.works") def test_missing_input_returns_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: - model, key = _register_tts(endpoints_client, resources) - result = endpoints_client.proxy.transport.send( + model, key = _register_tts(proxy, resources) + result = proxy.transport.send( "/v1/audio/speech", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalSpeechBody(model=model, voice="alloy"), ) assert_client_error(result, "speech missing input") @@ -93,12 +104,12 @@ class TestAudioSpeech: @pytest.mark.skip(reason="stage red: product gap, /v1/audio/speech 500s on missing model instead of 400") @pytest.mark.covers("llm.audio_speech.openai.input_validation.nonstream.works") def test_missing_model_returns_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: - _, key = _register_tts(endpoints_client, resources) - result = endpoints_client.proxy.transport.send( + _, key = _register_tts(proxy, resources) + result = proxy.transport.send( "/v1/audio/speech", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalSpeechBody(input="hello", voice="alloy"), ) assert_client_error(result, "speech missing model") @@ -106,12 +117,12 @@ class TestAudioSpeech: @pytest.mark.skip(reason="stage red: product gap, /v1/audio/speech 500s on invalid voice instead of surfacing the provider 4xx") @pytest.mark.covers("llm.audio_speech.openai.input_validation.nonstream.works") def test_invalid_voice_returns_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: - model, key = _register_tts(endpoints_client, resources) - result = endpoints_client.proxy.transport.send( + model, key = _register_tts(proxy, resources) + result = proxy.transport.send( "/v1/audio/speech", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalSpeechBody(model=model, input="hello", voice="invalid_voice_xyz"), ) assert_client_error(result, "speech invalid voice") @@ -119,12 +130,12 @@ class TestAudioSpeech: @pytest.mark.skip(reason="stage red: product gap, /v1/audio/speech 500s on empty input instead of surfacing the provider 4xx") @pytest.mark.covers("llm.audio_speech.openai.input_validation.nonstream.works") def test_empty_input_returns_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: - model, key = _register_tts(endpoints_client, resources) - result = endpoints_client.proxy.transport.send( + model, key = _register_tts(proxy, resources) + result = proxy.transport.send( "/v1/audio/speech", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalSpeechBody(model=model, input="", voice="alloy"), ) assert_client_error(result, "speech empty input") diff --git a/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py b/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py index 735f1a4a703..0ef73653835 100644 --- a/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py +++ b/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py @@ -1,12 +1,13 @@ """Live e2e: POST /v1/audio/transcriptions turns speech into text (vendor §9.7 / LIT-4778). Registers an OpenAI speech-to-text deployment at runtime and uploads a spoken -weather question (the realtime suite's 24kHz WAV fixture) as multipart, asserting -the returned transcript is non-empty and mentions the word it was asked about. -Also pins missing file/model negatives. A model-less request comes back as one of -two 400s depending on whether any wildcard deployment happens to be registered on -the shared proxy, so the assertion accepts either phrasing and holds both to naming -the model as the problem. +weather question (the realtime suite's 24kHz WAV fixture) through the real +OpenAI SDK (LIT-4577), asserting the returned transcript is non-empty and +mentions the word it was asked about. Also pins missing file/model negatives on +the shared multipart transport, since the SDK refuses to send them. A model-less +request comes back as one of two 400s depending on whether any wildcard +deployment happens to be registered on the shared proxy, so the assertion +accepts either phrasing and holds both to naming the model as the problem. """ from __future__ import annotations @@ -16,11 +17,12 @@ from typing import Final import pytest from e2e_config import unique_marker -from e2e_http import UnknownApiError, unwrap -from endpoints_client import EndpointsClient, TranscriptionForm, TranscriptionResult +from e2e_http import UnknownApiError from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient from pydantic import BaseModel +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e @@ -36,32 +38,34 @@ class _OptionalTranscriptionForm(BaseModel): response_format: str = "json" -def _register( - endpoints_client: EndpointsClient, resources: ResourceManager -) -> tuple[str, str]: +class _TranscriptionResult(BaseModel): + text: str = "" + + +def _register(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]: model = f"e2e-transcribe-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="openai/gpt-4o-mini-transcribe", api_key="os.environ/OPENAI_API_KEY" ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) return model, resources.key() class TestAudioTranscriptions: @pytest.mark.covers("llm.audio_transcriptions.openai.basic.nonstream.works") def test_audio_transcriptions_returns_text( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = _register(endpoints_client, resources) - result = unwrap( - endpoints_client.transcribe( - key, model, filename=WEATHER_WAV.name, content=WEATHER_WAV.read_bytes() - ) + model, key = _register(proxy, resources) + client = sdk.openai(key) + + transcription = client.audio.transcriptions.create( + model=model, file=(WEATHER_WAV.name, WEATHER_WAV.read_bytes(), "audio/wav") ) - text = result.text.strip() + text = transcription.text.strip() assert text, "/audio/transcriptions returned an empty transcript" assert "weather" in text.lower(), ( f"transcript of a spoken weather question does not mention weather: {text!r}" @@ -69,17 +73,17 @@ class TestAudioTranscriptions: @pytest.mark.covers("llm.audio_transcriptions.openai.input_validation.nonstream.works") def test_missing_file_returns_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: - model, key = _register(endpoints_client, resources) - result = endpoints_client.proxy.transport.upload( + model, key = _register(proxy, resources) + result = proxy.transport.upload( "/v1/audio/transcriptions", - headers=endpoints_client.proxy.transport.bearer(key), - form=TranscriptionForm(model=model), + headers=proxy.transport.bearer(key), + form=_OptionalTranscriptionForm(model=model), filename="empty.wav", content=b"", file_content_type="audio/wav", - response_type=TranscriptionResult, + response_type=_TranscriptionResult, ) match result: case UnknownApiError(status_code=400, body=body): @@ -95,17 +99,17 @@ class TestAudioTranscriptions: @pytest.mark.covers("llm.audio_transcriptions.openai.input_validation.nonstream.works") def test_missing_model_returns_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: - _, key = _register(endpoints_client, resources) - result = endpoints_client.proxy.transport.upload( + _, key = _register(proxy, resources) + result = proxy.transport.upload( "/v1/audio/transcriptions", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), form=_OptionalTranscriptionForm(), filename=WEATHER_WAV.name, content=WEATHER_WAV.read_bytes(), file_content_type="audio/wav", - response_type=TranscriptionResult, + response_type=_TranscriptionResult, ) match result: case UnknownApiError(status_code=400, body=body): diff --git a/tests/e2e/llm_translation/test_bedrock_web_search_server_tool_e2e.py b/tests/e2e/llm_translation/test_bedrock_web_search_server_tool_e2e.py index 43461239e5f..b4253a82dd8 100644 --- a/tests/e2e/llm_translation/test_bedrock_web_search_server_tool_e2e.py +++ b/tests/e2e/llm_translation/test_bedrock_web_search_server_tool_e2e.py @@ -34,27 +34,22 @@ block alone does not activate it. from __future__ import annotations import pytest - +from anthropic.types import WebSearchTool20250305Param from e2e_config import unique_marker -from e2e_http import unwrap -from endpoints_client import EndpointsClient from lifecycle import ResourceManager -from models import ( - AnthropicMessagesBody, - AnthropicWebSearchTool, - ChatMessage, - LiteLLMParamsBody, -) +from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import NO_PROXY_CACHE, SdkClients pytestmark = pytest.mark.e2e BEDROCK_INVOKE_BACKEND = "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0" -WEB_SEARCH_TOOL = AnthropicWebSearchTool( - type="web_search_20250305", - name="web_search", - max_uses=3, -) +WEB_SEARCH_TOOL: WebSearchTool20250305Param = { + "type": "web_search_20250305", + "name": "web_search", + "max_uses": 3, +} SEARCH_PROMPT = "Use web search to tell me one recent news headline about Anthropic." @@ -68,34 +63,30 @@ class TestBedrockWebSearchServerTool: ) @pytest.mark.covers("llm.messages.bedrock_invoke.web_search_server_tool.nonstream.works") def test_web_search_server_tool_is_served( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: """A bedrock deployment must answer a web_search server-tool request instead of handing the tool to AWS and returning its 400.""" model = f"e2e-bedrock-websearch-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model=BEDROCK_INVOKE_BACKEND, aws_region_name="us-east-1", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + resources.defer(lambda: proxy.delete_model(model_id)) + client = sdk.anthropic(resources.key()) - response = unwrap( - endpoints_client.proxy.messages( - key, - AnthropicMessagesBody( - model=model, - max_tokens=512, - tools=[WEB_SEARCH_TOOL], - messages=[ChatMessage(role="user", content=SEARCH_PROMPT)], - ), - ) + response = client.messages.create( + model=model, + max_tokens=512, + tools=[WEB_SEARCH_TOOL], + messages=[{"role": "user", "content": SEARCH_PROMPT}], + extra_body=NO_PROXY_CACHE, ) - assert response.content, f"no content blocks in response: {response}" + assert response.content, f"no content blocks in response: {response!r}" block_types = [block.type for block in response.content] assert "web_search_tool_result" in block_types, ( "the answer carries no web_search_tool_result block, so the search " diff --git a/tests/e2e/llm_translation/test_cache_control.py b/tests/e2e/llm_translation/test_cache_control.py index 102b3f00698..ceb3620183a 100644 --- a/tests/e2e/llm_translation/test_cache_control.py +++ b/tests/e2e/llm_translation/test_cache_control.py @@ -24,7 +24,7 @@ service_tier lives in test_provider_features_e2e.py. The provider-native cache_control request shape is not expressible with the shared ``ChatBody`` (whose content is a plain string), so the cacheable body is -built from the typed content blocks shared in ``endpoints_client.py``. +built from the typed content blocks shared in ``models.py``. """ from __future__ import annotations @@ -38,9 +38,8 @@ from pydantic import BaseModel from e2e_config import unique_marker from e2e_http import Result, UnknownApiError, unwrap -from endpoints_client import CacheControl, RichMessage, TextBlock from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody, Usage +from models import CacheControl, ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody, RichMessage, TextBlock, Usage from passthrough_client import PassthroughClient import os diff --git a/tests/e2e/llm_translation/test_completions_endpoint_e2e.py b/tests/e2e/llm_translation/test_completions_endpoint_e2e.py index 3fc506e4de9..63fcee3ce36 100644 --- a/tests/e2e/llm_translation/test_completions_endpoint_e2e.py +++ b/tests/e2e/llm_translation/test_completions_endpoint_e2e.py @@ -3,19 +3,19 @@ The legacy text-completion endpoint (prompt-style, non-chat) is the second-busiest route in production yet was previously uncovered; the rest of the "completions" surface is chat only. Registers an OpenAI instruct deployment at runtime (deleted -on teardown), drives /v1/completions through the gateway, and asserts real -generated text came back so a regression that empties the completion fails here. +on teardown), drives /v1/completions through the gateway with the real OpenAI SDK +(LIT-4577), and asserts real generated text came back so a regression that empties +the completion fails here. """ from __future__ import annotations import pytest - from e2e_config import unique_marker -from e2e_http import require_successful_call -from endpoints_client import CompletionsResult, EndpointsClient 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 @@ -23,24 +23,25 @@ pytestmark = pytest.mark.e2e class TestCompletionsEndpoint: @pytest.mark.covers("llm.completions.openai.basic.nonstream.works") def test_text_completion_returns_text( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: model = f"e2e-completions-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="text-completion-openai/gpt-3.5-turbo-instruct", api_key="os.environ/OPENAI_API_KEY", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + resources.defer(lambda: proxy.delete_model(model_id)) + client = sdk.openai(resources.key()) - result = endpoints_client.text_completions( - key, model, "Finish this sentence in a few words: the capital of France is" + completion = client.completions.create( + model=model, + prompt="Finish this sentence in a few words: the capital of France is", + max_tokens=32, + extra_body=NO_PROXY_CACHE, ) - require_successful_call(result) - parsed = CompletionsResult.model_validate_json(result.body) - assert parsed.choices, f"/v1/completions returned no choices: {result.body[:300]}" - completion = (parsed.choices[0].text or "").strip() - assert completion, f"/v1/completions returned an empty completion: {result.body[:300]}" + assert completion.choices, f"/v1/completions returned no choices: {completion!r}" + text = (completion.choices[0].text or "").strip() + assert text, f"/v1/completions returned an empty completion: {completion!r}" diff --git a/tests/e2e/llm_translation/test_credential_messages_e2e.py b/tests/e2e/llm_translation/test_credential_messages_e2e.py index 49ea748430e..52306ce3a7a 100644 --- a/tests/e2e/llm_translation/test_credential_messages_e2e.py +++ b/tests/e2e/llm_translation/test_credential_messages_e2e.py @@ -7,43 +7,47 @@ import os import pytest from e2e_config import unique_marker -from e2e_http import require_successful_call -from endpoints_client import EndpointsClient, MessagesResult from lifecycle import ResourceManager from models import CredentialCreateBody, LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import NO_PROXY_CACHE, SdkClients pytestmark = pytest.mark.e2e class TestCredentialBackedMessages: @pytest.mark.covers("mgmt.credential.new.serves_request") - def test_credential_backed_messages(self, endpoints_client: EndpointsClient, resources: ResourceManager) -> None: + def test_credential_backed_messages(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None: marker = unique_marker() credential_name = f"e2e-cred-{marker}" model = f"e2e-cred-messages-{marker}" anthropic_api_key = os.getenv("ANTHROPIC_API_KEY") assert anthropic_api_key, "ANTHROPIC_API_KEY must be set for this live e2e test" - endpoints_client.proxy.create_credential( + proxy.create_credential( CredentialCreateBody( credential_name=credential_name, credential_values={"api_key": anthropic_api_key}, ) ) - resources.defer(lambda: endpoints_client.proxy.delete_credential(credential_name)) + resources.defer(lambda: proxy.delete_credential(credential_name)) - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="anthropic/claude-haiku-4-5", litellm_credential_name=credential_name, ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) - key = resources.key() - result = endpoints_client.messages(key, model, "reply with one word") - require_successful_call(result) - parsed = MessagesResult.model_validate_json(result.body) - assert parsed.role == "assistant", f"unexpected role: {result.body[:300]}" - assert parsed.text.strip(), f"/v1/messages returned no text: {result.body[:300]}" + client = sdk.anthropic(resources.key()) + message = client.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": "reply with one word"}], + extra_body=NO_PROXY_CACHE, + ) + assert message.role == "assistant", f"unexpected role: {message.role!r}" + text = "".join(block.text for block in message.content if block.type == "text") + assert text.strip(), f"/v1/messages returned no text: {message.content!r}" diff --git a/tests/e2e/llm_translation/test_custom_pricing_e2e.py b/tests/e2e/llm_translation/test_custom_pricing_e2e.py index b4ff631a56b..1cebf90fa21 100644 --- a/tests/e2e/llm_translation/test_custom_pricing_e2e.py +++ b/tests/e2e/llm_translation/test_custom_pricing_e2e.py @@ -25,7 +25,6 @@ from pydantic import BaseModel, RootModel from e2e_config import unique_marker from proxy_client import ProxyClient from e2e_http import Success, unwrap -from endpoints_client import EndpointsClient from lifecycle import ResourceManager from models import ( ChatBody, @@ -71,7 +70,7 @@ def _approx_equal(actual: float, expected: float) -> bool: def _provision( - endpoints_client: EndpointsClient, + proxy: ProxyClient, resources: ResourceManager, prefix: str, *, @@ -84,7 +83,7 @@ def _provision( marker keeps the name unique so concurrent runs on the shared proxy never collide.""" model_name = f"{prefix}-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model_name, LiteLLMParamsBody( model=BACKEND_MODEL, @@ -93,15 +92,15 @@ def _provision( output_cost_per_token=output_cost_per_token, ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) return model_name def _provision_custom_priced( - endpoints_client: EndpointsClient, resources: ResourceManager + proxy: ProxyClient, resources: ResourceManager ) -> str: return _provision( - endpoints_client, + proxy, resources, "custom-priced-flash", input_cost_per_token=CUSTOM_INPUT_RATE, @@ -151,14 +150,14 @@ def _poll_breakdown_row(proxy: ProxyClient, key: str, response_id: str | None) - class TestCustomPricing: def test_custom_pricing_is_billed_at_configured_rate( self, - endpoints_client: EndpointsClient, + proxy: ProxyClient, resources: ResourceManager, scoped_key: str, ) -> None: - model = _provision_custom_priced(endpoints_client, resources) + model = _provision_custom_priced(proxy, resources) chat = unwrap( - endpoints_client.proxy.chat( + proxy.chat( scoped_key, ChatBody( model=model, @@ -172,7 +171,7 @@ class TestCustomPricing: ) ) - row = _poll_breakdown_row(endpoints_client.proxy, scoped_key, chat.id) + row = _poll_breakdown_row(proxy, scoped_key, chat.id) assert row.metadata and row.metadata.cost_breakdown # guaranteed by the poll breakdown = row.metadata.cost_breakdown @@ -195,10 +194,10 @@ class TestCustomPricing: ) def test_model_info_reports_custom_pricing( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: - model = _provision_custom_priced(endpoints_client, resources) - entry = _model_info_entry(endpoints_client.proxy.model_info(), model) + model = _provision_custom_priced(proxy, resources) + entry = _model_info_entry(proxy.model_info(), model) assert entry.litellm_params.input_cost_per_token == CUSTOM_INPUT_RATE, ( f"/model/info litellm_params input rate " @@ -210,20 +209,20 @@ class TestCustomPricing: ) def test_custom_pricing_is_isolated_from_sibling_deployment( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: # Register the override first so its rate is in the backend cost map before # the sibling resolves; a leak (LIT-3897) would then poison the sibling. - custom = _provision_custom_priced(endpoints_client, resources) + custom = _provision_custom_priced(proxy, resources) sibling = _provision( - endpoints_client, + proxy, resources, "base-flash", input_cost_per_token=None, output_cost_per_token=None, ) - entries = {entry.model_name: entry for entry in endpoints_client.proxy.model_info()} + entries = {entry.model_name: entry for entry in proxy.model_info()} custom_entry = entries.get(custom) sibling_entry = entries.get(sibling) assert custom_entry is not None, f"{custom} absent from /model/info" diff --git a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py index 5520ca0cee5..41282260b7e 100644 --- a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py +++ b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py @@ -1,23 +1,23 @@ """Live e2e: POST /embeddings returns a real vector across OpenAI, Bedrock, Vertex, Cohere. -Each test registers the deployment it needs at runtime (deleted on teardown) and -asserts a non-empty, non-zero vector came back. The LIT-3167 guard in -tests/e2e/embeddings/ covers the Gemini embedding path; embeddings cost tracking is -covered by tests/e2e/quota_management/spend_tracking/. +Each test registers the deployment it needs at runtime (deleted on teardown), +drives the endpoint with the real OpenAI SDK (LIT-4577), and asserts a +non-empty, non-zero vector came back. The LIT-3167 guard in +tests/e2e/embeddings/ covers the Gemini embedding path; embeddings cost tracking +is covered by tests/e2e/quota_management/spend_tracking/. Malformed bodies the +SDK refuses to build stay on the shared transport. """ from __future__ import annotations import pytest from e2e_config import provider_edge_base, unique_marker -from e2e_http import ( - assert_client_error, - require_successful_call, -) -from endpoints_client import EmbeddingsResult, EndpointsClient +from e2e_http import assert_client_error from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient from pydantic import BaseModel +from sdk_clients import NO_PROXY_CACHE, SdkClients pytestmark = pytest.mark.e2e @@ -39,35 +39,47 @@ def _openai_embeddings_params() -> LiteLLMParamsBody: ) +def _register( + proxy: ProxyClient, resources: ResourceManager, prefix: str, params: LiteLLMParamsBody +) -> tuple[str, str]: + model = f"{prefix}-{unique_marker()}" + model_id = proxy.create_model(model, params) + resources.defer(lambda: proxy.delete_model(model_id)) + return model, resources.key() + + +def _assert_embedding_vector( + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + prefix: str, + params: LiteLLMParamsBody, +) -> None: + model, key = _register(proxy, resources, prefix, params) + client = sdk.openai(key) + + embeddings = client.embeddings.create(model=model, input="Say this is a test!", extra_body=NO_PROXY_CACHE) + assert embeddings.data, f"/embeddings returned no data: {embeddings!r}" + vector = embeddings.data[0].embedding + assert vector, f"/embeddings returned no vector: {embeddings!r}" + assert any(component != 0.0 for component in vector), "embedding vector is all zeros" + + class TestEmbeddingsEndpoint: @pytest.mark.replayable @pytest.mark.covers("llm.embeddings.openai.basic.nonstream.works") - def test_embeddings_returns_vector( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model = f"e2e-embeddings-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - _openai_embeddings_params(), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.embeddings(key, model, "Say this is a test!") - require_successful_call(result) - parsed = EmbeddingsResult.model_validate_json(result.body) - assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" - assert any(component != 0.0 for component in parsed.first_vector), ( - f"embedding vector is all zeros: {result.body[:300]}" - ) + def test_embeddings_returns_vector(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None: + _assert_embedding_vector(proxy, resources, sdk, "e2e-embeddings", _openai_embeddings_params()) @pytest.mark.covers("llm.embeddings.bedrock.basic.nonstream.works") def test_bedrock_embeddings_returns_vector( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-embeddings-bedrock-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + _assert_embedding_vector( + proxy, + resources, + sdk, + "e2e-embeddings-bedrock", LiteLLMParamsBody( model="bedrock/amazon.titan-embed-text-v2:0", aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", @@ -75,110 +87,62 @@ class TestEmbeddingsEndpoint: aws_region_name="os.environ/AWS_REGION", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.embeddings(key, model, "Say this is a test!") - require_successful_call(result) - parsed = EmbeddingsResult.model_validate_json(result.body) - assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" - assert any(component != 0.0 for component in parsed.first_vector), ( - f"embedding vector is all zeros: {result.body[:300]}" - ) @pytest.mark.covers("llm.embeddings.cohere.basic.nonstream.works") def test_cohere_embeddings_returns_vector( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-embeddings-cohere-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + _assert_embedding_vector( + proxy, + resources, + sdk, + "e2e-embeddings-cohere", LiteLLMParamsBody(model="cohere/embed-v4.0", api_key="os.environ/COHERE_API_KEY"), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.embeddings(key, model, "Say this is a test!") - require_successful_call(result) - parsed = EmbeddingsResult.model_validate_json(result.body) - assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" - assert any(component != 0.0 for component in parsed.first_vector), ( - f"embedding vector is all zeros: {result.body[:300]}" - ) @pytest.mark.covers("llm.embeddings.vertex.basic.nonstream.works") def test_vertex_embeddings_returns_vector( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-embeddings-vertex-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + _assert_embedding_vector( + proxy, + resources, + sdk, + "e2e-embeddings-vertex", LiteLLMParamsBody( model="vertex_ai/text-embedding-005", vertex_project="os.environ/VERTEXAI_PROJECT", vertex_location="us-central1", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.embeddings(key, model, "Say this is a test!") - require_successful_call(result) - parsed = EmbeddingsResult.model_validate_json(result.body) - assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" - assert any(component != 0.0 for component in parsed.first_vector), ( - f"embedding vector is all zeros: {result.body[:300]}" - ) @pytest.mark.replayable @pytest.mark.covers("llm.embeddings.openai.basic.nonstream.works") - def test_array_input_returns_vectors( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model = f"e2e-embeddings-array-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - _openai_embeddings_params(), + def test_array_input_returns_vectors(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None: + model, key = _register(proxy, resources, "e2e-embeddings-array", _openai_embeddings_params()) + embeddings = sdk.openai(key).embeddings.create( + model=model, input=["Hello", "World", "Test"], extra_body=NO_PROXY_CACHE ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - result = endpoints_client.proxy.transport.send( - "/embeddings", - headers=endpoints_client.proxy.transport.bearer(key), - json=_OptionalEmbeddingsBody(model=model, input=["Hello", "World", "Test"]), - ) - require_successful_call(result) - parsed = EmbeddingsResult.model_validate_json(result.body) - assert len(parsed.data) == 3, f"expected 3 vectors: {result.body[:300]}" + assert len(embeddings.data) == 3, f"expected 3 vectors: {embeddings!r}" @pytest.mark.replayable @pytest.mark.covers("llm.embeddings.openai.input_validation.nonstream.works") - def test_missing_model_returns_client_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: + def test_missing_model_returns_client_error(self, proxy: ProxyClient, resources: ResourceManager) -> None: key = resources.key() - result = endpoints_client.proxy.transport.send( + result = proxy.transport.send( "/embeddings", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalEmbeddingsBody(input="hello"), ) assert_client_error(result, "embeddings missing model") @pytest.mark.replayable @pytest.mark.covers("llm.embeddings.openai.input_validation.nonstream.works") - def test_missing_input_returns_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model = f"e2e-embeddings-missin-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - _openai_embeddings_params(), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - result = endpoints_client.proxy.transport.send( + def test_missing_input_returns_error(self, proxy: ProxyClient, resources: ResourceManager) -> None: + model, key = _register(proxy, resources, "e2e-embeddings-missin", _openai_embeddings_params()) + result = proxy.transport.send( "/embeddings", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalEmbeddingsBody(model=model), ) assert_client_error(result, "embeddings missing input") diff --git a/tests/e2e/llm_translation/test_google_native_e2e.py b/tests/e2e/llm_translation/test_google_native_e2e.py index 40fd6eca765..6910519c6df 100644 --- a/tests/e2e/llm_translation/test_google_native_e2e.py +++ b/tests/e2e/llm_translation/test_google_native_e2e.py @@ -1,19 +1,41 @@ +"""Live e2e: the Gemini-native generateContent routes through the gateway. + +Google's own SDKs read these routes, and the streaming test asserts the exact SSE +framing they expect (no doubled ``data:`` prefix, no bytes literal, no OpenAI +``[DONE]`` sentinel), which an SDK would hide, so this passthrough surface stays on +the shared transport. +""" + from __future__ import annotations -import pytest -from pydantic import BaseModel +from typing import Literal +import pytest from e2e_config import unique_marker from e2e_http import StreamingResponse, require_successful_call -from endpoints_client import EndpointsClient from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from pydantic import BaseModel pytestmark = pytest.mark.e2e UPSTREAM_MODEL = "gemini/gemini-2.5-flash" +class _GenerateContentPart(BaseModel): + text: str + + +class _GenerateContentContent(BaseModel): + role: Literal["user"] = "user" + parts: tuple[_GenerateContentPart, ...] + + +class _GenerateContentBody(BaseModel): + contents: tuple[_GenerateContentContent, ...] + + class _StreamPart(BaseModel): text: str | None = None @@ -30,16 +52,27 @@ class _StreamEvent(BaseModel): candidates: tuple[_StreamCandidate, ...] = () -def _managed_deployment(client: EndpointsClient, resources: ResourceManager) -> str: +def _managed_deployment(proxy: ProxyClient, resources: ResourceManager) -> str: model = f"e2e-google-native-{unique_marker()}" - model_id = client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody(model=UPSTREAM_MODEL, api_key="os.environ/GEMINI_API_KEY"), ) - resources.defer(lambda: client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) return model +def _generate_content(proxy: ProxyClient, key: str, model: str, text: str, *, stream: bool = False) -> StreamingResponse: + operation = "streamGenerateContent" if stream else "generateContent" + body = _GenerateContentBody(contents=(_GenerateContentContent(parts=(_GenerateContentPart(text=text),)),)) + return proxy.transport.send( + f"/v1beta/models/{model}:{operation}", + headers=proxy.transport.bearer(key), + json=body, + stream=stream, + ) + + def _streamed_text(result: StreamingResponse) -> str: return "".join( part.text @@ -54,15 +87,13 @@ class TestGoogleNativeGenerateContent: @pytest.mark.covers("llm.google_native.gemini.basic.nonstream.cost_logged") def test_generate_content_returns_response_cost_header( self, - endpoints_client: EndpointsClient, + proxy: ProxyClient, resources: ResourceManager, scoped_key: str, ) -> None: - model = _managed_deployment(endpoints_client, resources) + model = _managed_deployment(proxy, resources) - result = endpoints_client.generate_content( - scoped_key, model, f"Reply with the single word ok. {unique_marker()}" - ) + result = _generate_content(proxy, scoped_key, model, f"Reply with the single word ok. {unique_marker()}") require_successful_call(result) assert result.call_id, "generateContent must stamp x-litellm-call-id" @@ -75,13 +106,14 @@ class TestGoogleNativeGenerateContent: @pytest.mark.covers("llm.google_native.gemini.basic.stream.works") def test_stream_generate_content_frames_sse_the_way_google_sdks_expect( self, - endpoints_client: EndpointsClient, + proxy: ProxyClient, resources: ResourceManager, scoped_key: str, ) -> None: - model = _managed_deployment(endpoints_client, resources) + model = _managed_deployment(proxy, resources) - result = endpoints_client.generate_content( + result = _generate_content( + proxy, scoped_key, model, f"Count from one to five, one number per line. {unique_marker()}", diff --git a/tests/e2e/llm_translation/test_image_edits_e2e.py b/tests/e2e/llm_translation/test_image_edits_e2e.py index 0197c8739fd..e95b054862e 100644 --- a/tests/e2e/llm_translation/test_image_edits_e2e.py +++ b/tests/e2e/llm_translation/test_image_edits_e2e.py @@ -1,23 +1,24 @@ """Live e2e: POST /v1/images/edits returns an edited image. -Registers an OpenAI image model, then sends a small PNG plus an edit prompt as a -multipart request to /v1/images/edits and asserts the response carries an image -(url or base64). /images/edits is a distinct native route from -/images/generations: it is multipart file upload with the image sent as the -`image` part, not a JSON body. The fixture image is a small generated 64x64 PNG, -so no external asset is needed. +Registers an OpenAI image model, then sends a small PNG plus an edit prompt +through the real OpenAI SDK (LIT-4577) to /v1/images/edits and asserts the +response carries an image (url or base64). /images/edits is a distinct native +route from /images/generations: it is multipart file upload with the image sent +as the `image` part, not a JSON body. The fixture image is a small generated +64x64 PNG, so no external asset is needed. """ from __future__ import annotations import base64 +import openai import pytest -from e2e_config import unique_marker -from e2e_http import Result, UnknownApiError, unwrap -from endpoints_client import EndpointsClient, ImageEditForm, ImagesResult +from e2e_config import SLOW_PROVIDER_TIMEOUT_SECONDS, unique_marker from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e @@ -28,51 +29,54 @@ _TEST_PNG = base64.b64decode( ) -def _register_image_model(endpoints_client: EndpointsClient, resources: ResourceManager) -> tuple[str, str]: +def _register_image_model(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]: model = f"e2e-image-edit-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody(model="openai/gpt-image-1", api_key="os.environ/OPENAI_API_KEY"), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) return model, resources.key() -def _assert_client_error(result: Result[ImagesResult], context: str) -> None: - match result: - case UnknownApiError(status_code=status) if 400 <= status < 500: - return - case other: - pytest.fail(f"{context}: expected 4xx, got {other!r}") +def _image_part(content: bytes) -> tuple[str, bytes, str]: + return ("image.png", content, "image/png") + + +def _assert_client_error(error: openai.APIStatusError, context: str) -> None: + assert 400 <= error.status_code < 500, f"{context}: expected 4xx, got {error.status_code}: {error.message}" class TestImageEdit: @pytest.mark.covers("llm.images_edits.openai.basic.nonstream.works") - def test_image_edit_returns_image(self, endpoints_client: EndpointsClient, resources: ResourceManager) -> None: - model, key = _register_image_model(endpoints_client, resources) + def test_image_edit_returns_image(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None: + model, key = _register_image_model(proxy, resources) + client = sdk.openai(key) - edited = unwrap(endpoints_client.image_edit(key, model, "Add a small red circle in the center", _TEST_PNG)) - assert edited.data, f"/images/edits returned no data: {edited}" - first = edited.data[0] - assert first.b64_json or first.url, f"edited image has neither b64_json nor url: {first}" - - @pytest.mark.covers("llm.images_edits.openai.input_validation.nonstream.works") - def test_empty_prompt_returns_error(self, endpoints_client: EndpointsClient, resources: ResourceManager) -> None: - model, key = _register_image_model(endpoints_client, resources) - result = endpoints_client.image_edit(key, model, "", _TEST_PNG) - _assert_client_error(result, "empty image-edit prompt") - - @pytest.mark.covers("llm.images_edits.openai.input_validation.nonstream.works") - def test_empty_image_returns_error(self, endpoints_client: EndpointsClient, resources: ResourceManager) -> None: - model, key = _register_image_model(endpoints_client, resources) - result = endpoints_client.proxy.transport.upload( - "/v1/images/edits", - headers=endpoints_client.proxy.transport.bearer(key), - form=ImageEditForm(model=model, prompt="add a red circle"), - filename="image.png", - content=b"", - file_content_type="image/png", - file_field="image", - response_type=ImagesResult, + edited = client.images.edit( + model=model, + image=_image_part(_TEST_PNG), + prompt="Add a small red circle in the center", + timeout=SLOW_PROVIDER_TIMEOUT_SECONDS, ) - _assert_client_error(result, "empty image-edit file") + assert edited.data, f"/images/edits returned no data: {edited!r}" + first = edited.data[0] + assert first.b64_json or first.url, f"edited image has neither b64_json nor url: {first!r}" + + @pytest.mark.covers("llm.images_edits.openai.input_validation.nonstream.works") + def test_empty_prompt_returns_error(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None: + model, key = _register_image_model(proxy, resources) + client = sdk.openai(key) + + with pytest.raises(openai.APIStatusError) as raised: + client.images.edit(model=model, image=_image_part(_TEST_PNG), prompt="") + _assert_client_error(raised.value, "empty image-edit prompt") + + @pytest.mark.covers("llm.images_edits.openai.input_validation.nonstream.works") + def test_empty_image_returns_error(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None: + model, key = _register_image_model(proxy, resources) + client = sdk.openai(key) + + with pytest.raises(openai.APIStatusError) as raised: + client.images.edit(model=model, image=_image_part(b""), prompt="add a red circle") + _assert_client_error(raised.value, "empty image-edit file") diff --git a/tests/e2e/llm_translation/test_image_generation_e2e.py b/tests/e2e/llm_translation/test_image_generation_e2e.py index 3b0d7da635f..1db40e7e15a 100644 --- a/tests/e2e/llm_translation/test_image_generation_e2e.py +++ b/tests/e2e/llm_translation/test_image_generation_e2e.py @@ -1,22 +1,22 @@ """Live e2e: POST /v1/images/generations returns an image. -Registers an OpenAI image deployment at runtime and asserts the response carries a -generated image (url or base64). Migrated from -litellm-regression-tests/tests/test_inference_endpoints.py. +Registers an image deployment at runtime, drives it through the real OpenAI SDK +(LIT-4577), and asserts the response carries a generated image (url or base64). +Malformed bodies the SDK refuses to build stay on the shared transport. Migrated +from litellm-regression-tests/tests/test_inference_endpoints.py. """ from __future__ import annotations import pytest from e2e_config import unique_marker -from e2e_http import ( - assert_client_error, - require_successful_call, -) -from endpoints_client import EndpointsClient, ImagesResult +from e2e_http import assert_client_error from lifecycle import ResourceManager from models import LiteLLMParamsBody +from openai.types import ImagesResponse +from proxy_client import ProxyClient from pydantic import BaseModel +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e @@ -28,44 +28,46 @@ class _OptionalImageBody(BaseModel): size: str | None = None -def _assert_image_returned(body: str) -> None: - parsed = ImagesResult.model_validate_json(body) - assert parsed.data, f"/images/generations returned no data: {body[:300]}" - first = parsed.data[0] - assert first.b64_json or first.url, ( - f"generated image has neither b64_json nor url: {body[:300]}" - ) +def _assert_image_returned(images: ImagesResponse) -> None: + data = images.data or [] + assert data, f"/images/generations returned no data: {images!r}" + first = data[0] + assert first.b64_json or first.url, f"generated image has neither b64_json nor url: {first!r}" -def _register_openai_image( - endpoints_client: EndpointsClient, resources: ResourceManager -) -> tuple[str, str]: - model = f"e2e-image-{unique_marker()}" - model_id = endpoints_client.create_model( - model, +def _register(proxy: ProxyClient, resources: ResourceManager, prefix: str, params: LiteLLMParamsBody) -> tuple[str, str]: + model = f"{prefix}-{unique_marker()}" + model_id = proxy.create_model(model, params) + resources.defer(lambda: proxy.delete_model(model_id)) + return model, resources.key() + + +def _register_openai_image(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]: + return _register( + proxy, + resources, + "e2e-image", LiteLLMParamsBody(model="openai/gpt-image-1-mini", api_key="os.environ/OPENAI_API_KEY"), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - return model, resources.key() class TestImageGeneration: @pytest.mark.covers("llm.images_generations.openai.basic.nonstream.works") def test_image_generation_returns_image( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = _register_openai_image(endpoints_client, resources) - result = endpoints_client.images(key, model, "Draw a cute cat") - require_successful_call(result) - _assert_image_returned(result.body) + model, key = _register_openai_image(proxy, resources) + images = sdk.openai(key).images.generate(model=model, prompt="Draw a cute cat", n=1, size="1024x1024") + _assert_image_returned(images) @pytest.mark.covers("llm.images_generations.bedrock.basic.nonstream.works", exercised_on=["images_generations"]) def test_bedrock_image_generation_returns_image( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-bedrock-image-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + model, key = _register( + proxy, + resources, + "e2e-bedrock-image", LiteLLMParamsBody( model="bedrock/amazon.nova-canvas-v1:0", aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", @@ -73,58 +75,46 @@ class TestImageGeneration: aws_region_name="os.environ/AWS_REGION", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.images(key, model, "Draw a cute cat") - require_successful_call(result) - _assert_image_returned(result.body) + images = sdk.openai(key).images.generate(model=model, prompt="Draw a cute cat", n=1, size="1024x1024") + _assert_image_returned(images) @pytest.mark.skip(reason="stage red: product gap, /v1/images/generations 500s (aimage_generation TypeError) on missing prompt instead of 400") @pytest.mark.covers("llm.images_generations.openai.input_validation.nonstream.works") - def test_missing_prompt_returns_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model, key = _register_openai_image(endpoints_client, resources) - result = endpoints_client.proxy.transport.send( + def test_missing_prompt_returns_error(self, proxy: ProxyClient, resources: ResourceManager) -> None: + model, key = _register_openai_image(proxy, resources) + result = proxy.transport.send( "/v1/images/generations", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalImageBody(model=model), ) assert_client_error(result, "images missing prompt") @pytest.mark.covers("llm.images_generations.openai.input_validation.nonstream.works") - def test_empty_prompt_returns_client_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model, key = _register_openai_image(endpoints_client, resources) - result = endpoints_client.proxy.transport.send( + def test_empty_prompt_returns_client_error(self, proxy: ProxyClient, resources: ResourceManager) -> None: + model, key = _register_openai_image(proxy, resources) + result = proxy.transport.send( "/v1/images/generations", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalImageBody(model=model, prompt=""), ) assert_client_error(result, "images empty prompt") @pytest.mark.covers("llm.images_generations.openai.input_validation.nonstream.works") - def test_invalid_size_returns_client_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model, key = _register_openai_image(endpoints_client, resources) - result = endpoints_client.proxy.transport.send( + def test_invalid_size_returns_client_error(self, proxy: ProxyClient, resources: ResourceManager) -> None: + model, key = _register_openai_image(proxy, resources) + result = proxy.transport.send( "/v1/images/generations", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalImageBody(model=model, prompt="a blue square", size="999x999"), ) assert_client_error(result, "images invalid size") @pytest.mark.covers("llm.images_generations.openai.input_validation.nonstream.works") - def test_invalid_n_returns_client_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model, key = _register_openai_image(endpoints_client, resources) - result = endpoints_client.proxy.transport.send( + def test_invalid_n_returns_client_error(self, proxy: ProxyClient, resources: ResourceManager) -> None: + model, key = _register_openai_image(proxy, resources) + result = proxy.transport.send( "/v1/images/generations", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalImageBody(model=model, prompt="a blue square", n=0), ) assert_client_error(result, "images invalid n") diff --git a/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py b/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py index 07be68a964b..8629cf12013 100644 --- a/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py +++ b/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py @@ -1,9 +1,9 @@ """Live e2e: POST /v1/messages routed to Azure AI Foundry Anthropic deployments. Registers `azure_ai/` deployments at runtime and drives the Messages -endpoint through the gateway across the behaviors an Anthropic client relies on: -a basic completion, a streamed completion, and tool use (non-streaming and -streaming). Auth is the Azure API key (`x-api-key`); the deployment reads +endpoint through the gateway with the real Anthropic SDK (LIT-4577) across the +behaviors an Anthropic client relies on: a basic completion, a streamed +completion, and tool use (non-streaming and streaming). The deployment reads `AZURE_AI_API_BASE` / `AZURE_AI_API_KEY` from the proxy env, so no secret is sent in the request. """ @@ -11,52 +11,39 @@ sent in the request. from __future__ import annotations import pytest +from anthropic.types import RawMessageStreamEvent, ToolParam + from e2e_config import unique_marker -from e2e_http import StreamingResponse, require_successful_call, unwrap -from endpoints_client import EndpointsClient from lifecycle import ResourceManager -from models import ( - AnthropicCustomTool, - AnthropicMessagesBody, - ChatMessage, - JsonSchemaProperty, - LiteLLMParamsBody, - ToolInputSchema, -) +from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import NO_PROXY_CACHE, SdkClients pytestmark = pytest.mark.e2e AZURE_FOUNDRY_MODEL = "azure_ai/claude-haiku-4-5" -WEATHER_TOOL = AnthropicCustomTool( - name="get_weather", - description="Get the current weather for a city.", - input_schema=ToolInputSchema( - properties={"city": JsonSchemaProperty(type="string")}, - required=["city"], - ), -) +WEATHER_TOOL: ToolParam = { + "name": "get_weather", + "description": "Get the current weather for a city.", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, +} -def _assert_streamed_ok(result: StreamingResponse) -> None: - require_successful_call(result) - assert result.is_streaming, f"response was not streamed: {result.headers}" - assert not result.stream_error, f"stream errored: {result.stream_error}" - assert result.stream_events, "stream produced no SSE events" - assert any("content_block_delta" in event for event in result.stream_events), ( - "stream carried no content deltas" - ) - assert any("message_stop" in event for event in result.stream_events), ( - "stream never reached message_stop" - ) +def _assert_streamed_ok(event_types: list[str]) -> None: + assert event_types, "stream produced no SSE events" + assert "content_block_delta" in event_types, "stream carried no content deltas" + assert "message_stop" in event_types, "stream never reached message_stop" class TestAzureFoundryMessages: - def _register( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> tuple[str, str]: + def _register(self, proxy: ProxyClient, resources: ResourceManager) -> str: model = f"e2e-azure-foundry-messages-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model=AZURE_FOUNDRY_MODEL, @@ -64,91 +51,72 @@ class TestAzureFoundryMessages: api_key="os.environ/AZURE_AI_API_KEY", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - return model, resources.key(models=[model]) + resources.defer(lambda: proxy.delete_model(model_id)) + return model @pytest.mark.covers("llm.messages.azure_foundry.basic.nonstream.works") - def test_basic_nonstream( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model, key = self._register(endpoints_client, resources) - response = unwrap( - endpoints_client.proxy.messages( - key, - AnthropicMessagesBody( - model=model, - max_tokens=64, - messages=[ChatMessage(role="user", content="Reply with one word.")], - ), - ) + def test_basic_nonstream(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None: + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key(models=[model])) + + message = client.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": "Reply with one word."}], + extra_body=NO_PROXY_CACHE, ) - assert response.content, f"no content blocks in response: {response}" - text = "".join(block.text or "" for block in response.content if block.type == "text") - assert text.strip(), f"/v1/messages returned no text: {response}" + assert message.content, f"no content blocks in response: {message!r}" + text = "".join(block.text for block in message.content if block.type == "text") + assert text.strip(), f"/v1/messages returned no text: {message.content!r}" @pytest.mark.covers("llm.messages.azure_foundry.basic.stream.works") - def test_basic_stream( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model, key = self._register(endpoints_client, resources) - result = endpoints_client.proxy.messages_stream( - key, - AnthropicMessagesBody( - model=model, - max_tokens=64, - stream=True, - messages=[ChatMessage(role="user", content="Count from one to three.")], - ), + def test_basic_stream(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None: + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key(models=[model])) + + stream = client.messages.create( + model=model, + max_tokens=64, + stream=True, + messages=[{"role": "user", "content": "Count from one to three."}], + extra_body=NO_PROXY_CACHE, ) - _assert_streamed_ok(result) + _assert_streamed_ok([event.type for event in stream]) @pytest.mark.covers("llm.messages.azure_foundry.tool_use.nonstream.works") - def test_tool_use_nonstream( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model, key = self._register(endpoints_client, resources) - response = unwrap( - endpoints_client.proxy.messages( - key, - AnthropicMessagesBody( - model=model, - max_tokens=256, - tools=[WEATHER_TOOL], - messages=[ - ChatMessage(role="user", content="What is the weather in Paris? Use the tool.") - ], - ), - ) + def test_tool_use_nonstream(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None: + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key(models=[model])) + + message = client.messages.create( + model=model, + max_tokens=256, + tools=[WEATHER_TOOL], + messages=[{"role": "user", "content": "What is the weather in Paris? Use the tool."}], + extra_body=NO_PROXY_CACHE, ) - assert response.content, f"no content blocks in response: {response}" - assert any(block.type == "tool_use" for block in response.content), ( - f"model did not call the tool: {response}" + assert message.content, f"no content blocks in response: {message!r}" + assert any(block.type == "tool_use" for block in message.content), ( + f"model did not call the tool: {message.content!r}" ) @pytest.mark.covers("llm.messages.azure_foundry.tool_use.stream.works") - def test_tool_use_stream( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model, key = self._register(endpoints_client, resources) - result = endpoints_client.proxy.messages_stream( - key, - AnthropicMessagesBody( - model=model, - max_tokens=256, - stream=True, - tools=[WEATHER_TOOL], - messages=[ - ChatMessage(role="user", content="What is the weather in Paris? Use the tool.") - ], - ), - ) - require_successful_call(result) - assert result.is_streaming, f"response was not streamed: {result.headers}" - assert not result.stream_error, f"stream errored: {result.stream_error}" - assert result.stream_events, "stream produced no SSE events" - assert any("tool_use" in event for event in result.stream_events), ( - "stream carried no tool_use block" - ) - assert any("message_stop" in event for event in result.stream_events), ( - "stream never reached message_stop" + def test_tool_use_stream(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None: + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key(models=[model])) + + stream = client.messages.create( + model=model, + max_tokens=256, + stream=True, + tools=[WEATHER_TOOL], + messages=[{"role": "user", "content": "What is the weather in Paris? Use the tool."}], + extra_body=NO_PROXY_CACHE, ) + events: list[RawMessageStreamEvent] = list(stream) + event_types = [event.type for event in events] + assert event_types, "stream produced no SSE events" + assert any( + event.type == "content_block_start" and event.content_block.type == "tool_use" for event in events + ), "stream carried no tool_use block" + assert "message_stop" in event_types, "stream never reached message_stop" diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index 09ec48daa2f..d048d1343eb 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -1,40 +1,42 @@ """Live e2e: POST /v1/messages (Anthropic Messages API) returns a real completion. Registers an Anthropic deployment at runtime, drives the Messages endpoint through -the gateway, and asserts an assistant message with text came back, both -non-streaming and streamed. Migrated from +the gateway with the real Anthropic SDK, the client customers actually use +(LIT-4577), and asserts an assistant message with text came back, both +non-streaming and streamed. Malformed bodies the SDK refuses to build stay on the +shared transport. Migrated from litellm-regression-tests/tests/test_inference_endpoints.py. """ from __future__ import annotations +import time from typing import Final import pytest -from e2e_config import ( - STREAM_MIN_LEAD_SECONDS, - provider_edge_base, - provider_paces_stream, - unique_marker, +from anthropic import Anthropic +from anthropic.types import ( + InputJSONDelta, + Message, + MessageParam, + RawContentBlockDeltaEvent, + RawContentBlockStartEvent, + RawContentBlockStopEvent, + RawMessageDeltaEvent, + RawMessageStreamEvent, + TextBlock, + TextDelta, + ToolChoiceParam, + ToolParam, + ToolUseBlock, ) -from e2e_http import assert_client_error, require_successful_call, unwrap -from endpoints_client import EndpointsClient, MessagesResult +from e2e_config import STREAM_MIN_LEAD_SECONDS, provider_edge_base, provider_paces_stream, unique_marker +from e2e_http import assert_client_error from lifecycle import ResourceManager -from models import ( - AnthropicAssistantTurn, - AnthropicContentBlock, - AnthropicCustomTool, - AnthropicMessagesBody, - AnthropicToolChoice, - AnthropicToolResultBlock, - AnthropicToolResultTurn, - ChatMessage, - JsonSchemaProperty, - LiteLLMParamsBody, - SpendLogRow, - ToolInputSchema, -) +from models import ChatMessage, LiteLLMParamsBody, SpendLogRow +from proxy_client import ProxyClient from pydantic import BaseModel, ConfigDict +from sdk_clients import NO_PROXY_CACHE, SdkClients, response_header pytestmark = [pytest.mark.e2e, pytest.mark.replayable] @@ -45,35 +47,17 @@ class _OptionalMessagesBody(BaseModel): max_tokens: int | None = None -class _MessagesEventDelta(BaseModel): - text: str = "" - - -class _MessagesEventUsage(BaseModel): - output_tokens: int | None = None - - -class _MessagesStreamEvent(BaseModel): - """One Anthropic SSE event, keeping only what the stream's shape is asserted on. - - ``delta.text`` is populated on ``content_block_delta`` and absent on the - ``message_delta`` that closes the turn, which is the event carrying ``usage``.""" - - type: str - delta: _MessagesEventDelta | None = None - usage: _MessagesEventUsage | None = None - - ANTHROPIC_BACKEND = "anthropic/claude-haiku-4-5" -WEATHER_TOOL = AnthropicCustomTool( - name="get_weather", - description="Get the current weather for a city.", - input_schema=ToolInputSchema( - properties={"city": JsonSchemaProperty(type="string")}, - required=["city"], - ), -) +WEATHER_TOOL: ToolParam = { + "name": "get_weather", + "description": "Get the current weather for a city.", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, +} def _approx_equal(actual: float, expected: float) -> bool: @@ -87,60 +71,67 @@ def _anthropic_params() -> LiteLLMParamsBody: handler appends ``/v1/messages`` to ``api_base`` itself, where the OpenAI handler appends only ``/chat/completions``.""" base = provider_edge_base("anthropic") - return LiteLLMParamsBody( - model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY", api_base=base - ) + return LiteLLMParamsBody(model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY", api_base=base) + + +def _register( + proxy: ProxyClient, + resources: ResourceManager, + params: LiteLLMParamsBody | None = None, + prefix: str = "e2e-messages", +) -> tuple[str, str]: + model = f"{prefix}-{unique_marker()}" + model_id = proxy.create_model(model, _anthropic_params() if params is None else params) + resources.defer(lambda: proxy.delete_model(model_id)) + return model, resources.key() + + +def _text(message: Message) -> str: + return "".join(block.text for block in message.content if isinstance(block, TextBlock)) + + +def _user_turn(text: str) -> MessageParam: + return {"role": "user", "content": text} class TestAnthropicMessages: - def _register( - self, - endpoints_client: EndpointsClient, - resources: ResourceManager, - params: LiteLLMParamsBody | None = None, - ) -> tuple[str, str]: - model = f"e2e-messages-{unique_marker()}" - model_id = endpoints_client.create_model( - model, _anthropic_params() if params is None else params - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - return model, resources.key() - @pytest.mark.covers("llm.messages.anthropic.basic.nonstream.works") - def test_messages_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model, key = self._register(endpoints_client, resources) + def test_messages_returns_completion(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None: + model, key = _register(proxy, resources) + client = sdk.anthropic(key) - result = endpoints_client.messages(key, model, "reply with one word") - require_successful_call(result) - parsed = MessagesResult.model_validate_json(result.body) - assert parsed.role == "assistant", f"unexpected role: {result.body[:300]}" - assert parsed.text.strip(), f"/v1/messages returned no text: {result.body[:300]}" + message = client.messages.create( + model=model, max_tokens=64, messages=[_user_turn("reply with one word")], extra_body=NO_PROXY_CACHE + ) + assert message.role == "assistant", f"unexpected role: {message.role!r}" + assert _text(message).strip(), f"/v1/messages returned no text: {message.content!r}" @pytest.mark.covers("llm.messages.anthropic.basic.nonstream.cost_logged") def test_messages_logs_cost_matching_the_response_header( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-messages-cost-{unique_marker()}" - model_id = endpoints_client.create_model(model, _anthropic_params()) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model, key = _register(proxy, resources, prefix="e2e-messages-cost") + client = sdk.anthropic(key) - result = endpoints_client.messages(key, model, f"reply with one word {unique_marker()}") - require_successful_call(result) - parsed = MessagesResult.model_validate_json(result.body) - assert parsed.role == "assistant" and parsed.text.strip(), ( - f"/v1/messages returned no assistant text: {result.body[:300]}" + raw = client.messages.with_raw_response.create( + model=model, + max_tokens=64, + messages=[_user_turn(f"reply with one word {unique_marker()}")], + extra_body=NO_PROXY_CACHE, + ) + message = raw.parse() + assert message.role == "assistant" and _text(message).strip(), ( + f"/v1/messages returned no assistant text: {message.content!r}" ) # The customer reads per-request cost off the response header (LIT-4076), so # it must be present and positive on /v1/messages, not only /chat/completions. - header_cost = result.response_cost - assert header_cost is not None and header_cost > 0, ( - "x-litellm-response-cost header missing or non-positive on /v1/messages; " - f"headers={result.headers}" + raw_header_cost = response_header(raw.headers, "x-litellm-response-cost") + assert raw_header_cost is not None, ( + f"x-litellm-response-cost header missing on /v1/messages; headers={dict(raw.headers)}" ) + header_cost = float(raw_header_cost) + assert header_cost > 0, f"x-litellm-response-cost header non-positive on /v1/messages: {header_cost}" # Correlate the spend row by the unique scoped key, not the Anthropic response # id: on /v1/messages the spend-log request_id is the proxy's own call id, which @@ -150,11 +141,9 @@ class TestAnthropicMessages: def _priced(rows: list[SpendLogRow]) -> bool: return any(r.spend is not None and r.spend > 0 for r in rows) - rows = endpoints_client.proxy.poll_logs_for_key(key, predicate=_priced) + rows = proxy.poll_logs_for_key(key, predicate=_priced) priced = [r for r in rows if r.spend is not None and r.spend > 0] - assert priced, ( - f"no priced /spend/logs row landed for key {key} within the poll window; got {rows}" - ) + assert priced, f"no priced /spend/logs row landed for key {key} within the poll window; got {rows}" row = priced[0] assert (row.prompt_tokens or 0) > 0 and (row.completion_tokens or 0) > 0, ( f"messages spend row missing token counts, so the cost is not real usage: {row}" @@ -166,9 +155,7 @@ class TestAnthropicMessages: @pytest.mark.covers("llm.messages.anthropic.basic.stream.works") @pytest.mark.provider_live - def test_messages_streams_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: + def test_messages_streams_completion(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None: """Edge-wired like its non-streaming siblings, so record and replay both carry the streamed response. @@ -178,51 +165,45 @@ class TestAnthropicMessages: the first content delta must instead reach the client well before ``message_stop``, which a buffered response cannot do. Replay serves chunks back to back, so only live and record runs judge the timing.""" - model, key = self._register(endpoints_client, resources) + model, key = _register(proxy, resources) + client = sdk.anthropic(key) - result = endpoints_client.proxy.messages_stream( - key, - AnthropicMessagesBody( - model=model, - max_tokens=800, - stream=True, - messages=[ChatMessage(role="user", content="Count from 1 to 200, one number per line.")], - ), + started: Final = time.monotonic() + stream = client.messages.create( + model=model, + max_tokens=800, + stream=True, + messages=[_user_turn("Count from 1 to 200, one number per line.")], + extra_body=NO_PROXY_CACHE, ) - require_successful_call(result) - assert result.is_streaming, f"response was not streamed: {result.headers}" - assert not result.stream_error, f"stream errored: {result.stream_error}" - assert result.stream_events, "stream produced no SSE events" + arrivals: Final = tuple((event, time.monotonic() - started) for event in stream) + assert arrivals, "stream produced no SSE events" - events = [ - _MessagesStreamEvent.model_validate_json(event) for event in result.stream_events - ] - types = [event.type for event in events] - delta_positions = [ + events: Final = tuple(event for event, _ in arrivals) + types: Final = tuple(event.type for event in events) + delta_positions: Final = tuple( index for index, event in enumerate(events) if event.type == "content_block_delta" - ] + ) assert delta_positions, f"stream carried no content deltas: {types}" - text = "".join( + text: Final = "".join( event.delta.text for event in events - if event.type == "content_block_delta" and event.delta is not None + if isinstance(event, RawContentBlockDeltaEvent) and isinstance(event.delta, TextDelta) ) - assert text.strip(), f"content deltas assembled to no text: {result.stream_events[:5]}" + assert text.strip(), f"content deltas assembled to no text: {events[:5]}" - usage_positions = [ - index - for index, event in enumerate(events) - if event.type == "message_delta" and event.usage is not None - ] + usage_positions: Final = tuple( + index for index, event in enumerate(events) if isinstance(event, RawMessageDeltaEvent) + ) assert usage_positions, f"stream never reported usage: {types}" assert "message_stop" in types, f"stream never reached message_stop: {types}" - stop_position = types.index("message_stop") + stop_position: Final = types.index("message_stop") assert delta_positions[-1] < usage_positions[0] < stop_position, ( f"usage did not land between the last content delta and message_stop: {types}" ) - first_delta_at: Final = result.stream_event_arrivals[delta_positions[0]] - stop_at: Final = result.stream_event_arrivals[stop_position] + first_delta_at: Final = arrivals[delta_positions[0]][1] + stop_at: Final = arrivals[stop_position][1] if provider_paces_stream(): assert stop_at - first_delta_at >= STREAM_MIN_LEAD_SECONDS, ( f"first content delta reached the client {first_delta_at:.2f}s after the request " @@ -231,142 +212,125 @@ class TestAnthropicMessages: ) @pytest.mark.covers("llm.messages.anthropic.tool_use.nonstream.works") - def test_messages_tool_use( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model, key = self._register(endpoints_client, resources) + def test_messages_tool_use(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None: + model, key = _register(proxy, resources) + client = sdk.anthropic(key) - response = unwrap( - endpoints_client.proxy.messages( - key, - AnthropicMessagesBody( - model=model, - max_tokens=256, - tools=[WEATHER_TOOL], - messages=[ - ChatMessage(role="user", content="What is the weather in Paris? Use the tool.") - ], - ), - ) + message = client.messages.create( + model=model, + max_tokens=256, + tools=[WEATHER_TOOL], + messages=[_user_turn("What is the weather in Paris? Use the tool.")], + extra_body=NO_PROXY_CACHE, ) - assert response.content, f"no content blocks in response: {response}" - assert any(block.type == "tool_use" for block in response.content), ( - f"model did not call the tool: {response}" + assert message.content, f"no content blocks in response: {message!r}" + assert any(isinstance(block, ToolUseBlock) for block in message.content), ( + f"model did not call the tool: {message.content!r}" ) - @pytest.mark.skip(reason="stage red: product gap, /v1/messages 500s (anthropic_messages TypeError) on missing messages instead of 400") + @pytest.mark.skip( + reason="stage red: product gap, /v1/messages 500s (anthropic_messages TypeError) on missing messages instead of 400" + ) @pytest.mark.covers("llm.messages.anthropic.input_validation.nonstream.works") - def test_missing_messages_returns_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model, key = self._register(endpoints_client, resources) - result = endpoints_client.proxy.transport.send( + def test_missing_messages_returns_error(self, proxy: ProxyClient, resources: ResourceManager) -> None: + model, key = _register(proxy, resources) + result = proxy.transport.send( "/v1/messages", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalMessagesBody(model=model, max_tokens=50), ) assert_client_error(result, "messages missing messages") - @pytest.mark.skip(reason="stage red: product gap, /v1/messages 500s (anthropic_messages TypeError) on missing max_tokens instead of 400") + @pytest.mark.skip( + reason="stage red: product gap, /v1/messages 500s (anthropic_messages TypeError) on missing max_tokens instead of 400" + ) @pytest.mark.covers("llm.messages.anthropic.input_validation.nonstream.works") - def test_missing_max_tokens_returns_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model, key = self._register(endpoints_client, resources) - result = endpoints_client.proxy.transport.send( + def test_missing_max_tokens_returns_error(self, proxy: ProxyClient, resources: ResourceManager) -> None: + model, key = _register(proxy, resources) + result = proxy.transport.send( "/v1/messages", - headers=endpoints_client.proxy.transport.bearer(key), - json=_OptionalMessagesBody( - model=model, messages=[ChatMessage(role="user", content="hi")] - ), + headers=proxy.transport.bearer(key), + json=_OptionalMessagesBody(model=model, messages=[ChatMessage(role="user", content="hi")]), ) assert_client_error(result, "messages missing max_tokens") @pytest.mark.covers("llm.messages.anthropic.input_validation.nonstream.works") - def test_missing_model_returns_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - _, key = self._register(endpoints_client, resources) - result = endpoints_client.proxy.transport.send( + def test_missing_model_returns_error(self, proxy: ProxyClient, resources: ResourceManager) -> None: + _, key = _register(proxy, resources) + result = proxy.transport.send( "/v1/messages", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalMessagesBody(messages=[ChatMessage(role="user", content="hi")], max_tokens=50), ) assert_client_error(result, "messages missing model") -class _BridgeDelta(BaseModel): - type: str | None = None - partial_json: str | None = None - stop_reason: str | None = None - - -class _BridgeEvent(BaseModel): - type: str - index: int | None = None - content_block: AnthropicContentBlock | None = None - delta: _BridgeDelta | None = None - - class _ParcelInput(BaseModel): model_config = ConfigDict(extra="forbid", strict=True) parcel: str shelf: int -def _tool_from_stream(events: tuple[_BridgeEvent, ...]) -> AnthropicContentBlock: +def _tool_from_stream(events: tuple[RawMessageStreamEvent, ...]) -> ToolUseBlock: starts: Final = tuple( - event - for event in events - if event.type == "content_block_start" - and event.content_block is not None - and event.content_block.type == "tool_use" + (index, event.index, event.content_block) + for index, event in enumerate(events) + if isinstance(event, RawContentBlockStartEvent) and isinstance(event.content_block, ToolUseBlock) ) assert len(starts) == 1, "expected exactly one tool call" - start: Final = starts[0] - block: Final = start.content_block - assert block is not None and block.id and start.index is not None + start_position, block_index, block = starts[0] + assert block.id fragments: Final = tuple( - event - for event in events - if event.type == "content_block_delta" and event.delta is not None and event.delta.type == "input_json_delta" + (index, event.index, event.delta.partial_json) + for index, event in enumerate(events) + if isinstance(event, RawContentBlockDeltaEvent) and isinstance(event.delta, InputJSONDelta) ) assert fragments, "tool stream contained no argument fragments" - assert all(event.index == start.index for event in fragments), "tool fragments changed index" - positions: Final = tuple(i for i, event in enumerate(events) if event in fragments) + assert all(fragment_block == block_index for _, fragment_block, _ in fragments), "tool fragments changed index" + positions: Final = tuple(index for index, _, _ in fragments) stops: Final = tuple( - i for i, event in enumerate(events) if event.type == "content_block_stop" and event.index == start.index + index + for index, event in enumerate(events) + if isinstance(event, RawContentBlockStopEvent) and event.index == block_index ) - assert len(stops) == 1 and events.index(start) < positions[0] <= positions[-1] < stops[0] - assert tuple( - event.delta.stop_reason for event in events if event.type == "message_delta" and event.delta is not None - ) == ("tool_use",) - terminal_positions: Final = tuple(i for i, event in enumerate(events) if event.type == "message_delta") + assert len(stops) == 1 and start_position < positions[0] <= positions[-1] < stops[0] + terminal_positions: Final = tuple( + index for index, event in enumerate(events) if isinstance(event, RawMessageDeltaEvent) + ) + stop_reasons: Final = tuple(event.delta.stop_reason for event in events if isinstance(event, RawMessageDeltaEvent)) + assert stop_reasons == ("tool_use",) assert len(terminal_positions) == 1 and stops[0] < terminal_positions[0] < len(events) - 1 - assert tuple(i for i, event in enumerate(events) if event.type == "message_stop") == (len(events) - 1,), ( + assert tuple(index for index, event in enumerate(events) if event.type == "message_stop") == (len(events) - 1,), ( "tool stream did not terminate exactly once" ) - arguments: Final = _ParcelInput.model_validate_json( - "".join(event.delta.partial_json or "" for event in fragments if event.delta is not None) - ) - return AnthropicContentBlock(type="tool_use", id=block.id, name=block.name, input=arguments.model_dump()) + arguments: Final = _ParcelInput.model_validate_json("".join(partial for _, _, partial in fragments)) + return ToolUseBlock(type="tool_use", id=block.id, name=block.name, input=arguments.model_dump()) -def _parcel_result(tool: AnthropicContentBlock, result: AnthropicToolResultBlock) -> AnthropicToolResultTurn: - assert tool.id and result.tool_use_id == tool.id, "tool result ID does not match the emitted call" - return AnthropicToolResultTurn(content=[result]) - - -def _request_tool( - client: EndpointsClient, key: str, request: AnthropicMessagesBody, stream: bool -) -> AnthropicContentBlock: +def _request_tool(client: Anthropic, model: str, question: MessageParam, tool: ToolParam, stream: bool) -> ToolUseBlock: + tool_choice: Final[ToolChoiceParam] = {"type": "tool", "name": tool["name"]} if stream: - response: Final = client.proxy.messages_stream(key, request) - require_successful_call(response) - assert response.is_streaming and not response.stream_error - return _tool_from_stream(tuple(_BridgeEvent.model_validate_json(event) for event in response.stream_events)) - response_body: Final = unwrap(client.proxy.messages(key, request)) - blocks: Final = tuple(block for block in response_body.content or () if block.type == "tool_use") + events: Final = tuple( + client.messages.create( + model=model, + max_tokens=2048, + messages=[question], + tools=[tool], + tool_choice=tool_choice, + stream=True, + extra_body=NO_PROXY_CACHE, + ) + ) + return _tool_from_stream(events) + message: Final = client.messages.create( + model=model, + max_tokens=2048, + messages=[question], + tools=[tool], + tool_choice=tool_choice, + extra_body=NO_PROXY_CACHE, + ) + blocks: Final = tuple(block for block in message.content if isinstance(block, ToolUseBlock)) assert len(blocks) == 1 return blocks[0] @@ -375,55 +339,49 @@ class TestOpenAIMessagesToolContinuation: @pytest.mark.provider_live @pytest.mark.parametrize("stream", [True, False], ids=["stream", "nonstream"]) def test_required_tool_arguments_and_correlated_result( - self, endpoints_client: EndpointsClient, resources: ResourceManager, stream: bool + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, stream: bool ) -> None: model: Final = f"e2e-bridge-tool-{unique_marker()}" base: Final = provider_edge_base("openai") - model_id: Final = endpoints_client.create_model( + model_id: Final = proxy.create_model( model, LiteLLMParamsBody( model="openai/gpt-5.6", api_key="os.environ/OPENAI_API_KEY", api_base=f"{base}/v1" if base else None ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key: Final = resources.key(models=[model]) - tool: Final = AnthropicCustomTool( - name="locate_parcel", - description="Look up the receipt for a parcel on a shelf. Return the receipt verbatim.", - input_schema=ToolInputSchema( - properties={"parcel": JsonSchemaProperty(type="string"), "shelf": JsonSchemaProperty(type="integer")}, - required=["parcel", "shelf"], - ), + resources.defer(lambda: proxy.delete_model(model_id)) + client: Final = sdk.anthropic(resources.key(models=[model])) + tool: Final[ToolParam] = { + "name": "locate_parcel", + "description": "Look up the receipt for a parcel on a shelf. Return the receipt verbatim.", + "input_schema": { + "type": "object", + "properties": {"parcel": {"type": "string"}, "shelf": {"type": "integer"}}, + "required": ["parcel", "shelf"], + }, + } + question: Final = _user_turn( + "Call locate_parcel with parcel exactly amber-kite and shelf exactly 7. " + "After the tool result, reply with only the receipt returned by the tool." ) - question: Final = ChatMessage( - role="user", - content="Call locate_parcel with parcel exactly amber-kite and shelf exactly 7. After the tool result, reply with only the receipt returned by the tool.", - ) - request: Final = AnthropicMessagesBody( - model=model, - max_tokens=2048, - messages=[question], - tools=[tool], - tool_choice=AnthropicToolChoice(type="tool", name=tool.name), - stream=stream, - ) - emitted: Final = _request_tool(endpoints_client, key, request, stream) + emitted: Final = _request_tool(client, model, question, tool, stream) assert emitted.id and emitted.name == "locate_parcel" assert emitted.input == {"parcel": "amber-kite", "shelf": 7}, "required tool arguments were lost or changed" receipt: Final = f"receipt-{unique_marker()}" - result_turn: Final = _parcel_result(emitted, AnthropicToolResultBlock(tool_use_id=emitted.id, content=receipt)) - continuation: Final = unwrap( - endpoints_client.proxy.messages( - key, - AnthropicMessagesBody( - model=model, - max_tokens=2048, - tools=[tool], - tool_choice=AnthropicToolChoice(type="none"), - messages=[question, AnthropicAssistantTurn(content=[emitted]), result_turn], - ), - ) + continuation: Final = client.messages.create( + model=model, + max_tokens=2048, + tools=[tool], + tool_choice={"type": "none"}, + messages=[ + question, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": emitted.id, "name": emitted.name, "input": emitted.input}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": emitted.id, "content": receipt}]}, + ], + extra_body=NO_PROXY_CACHE, ) - answer: Final = "".join(block.text or "" for block in continuation.content or ()) - assert answer.strip() == receipt, "continuation did not consume the correlated tool result" - assert all(block.type != "tool_use" for block in continuation.content or ()) + assert _text(continuation).strip() == receipt, "continuation did not consume the correlated tool result" + assert all(not isinstance(block, ToolUseBlock) for block in continuation.content) diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py index 557a2cb64e9..e9b4b394996 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py @@ -17,27 +17,28 @@ entry whose prefix spans ``system`` plus message turns is invalidated when the reminder is hoisted (the ``system`` field mutates and a turn disappears from ``messages``), while an entry ending at the system block itself would survive the hoist and mask the regression. + +Calls go through the real Anthropic SDK (LIT-4577). The SDK's ``MessageParam`` +type only admits user/assistant roles, so the system reminder turn is cast to +it; the SDK serializes the dict verbatim, which is exactly the wire shape under +test. """ from __future__ import annotations import time +from collections.abc import Sequence +from typing import cast import pytest -from pydantic import BaseModel - +from anthropic import Anthropic +from anthropic.types import Message, MessageParam, TextBlockParam from e2e_config import unique_marker -from e2e_http import Result, unwrap -from endpoints_client import ( - CacheControl, - EndpointsClient, - MessagesResult, - RichMessage, - RichMessagesRequest, - TextBlock, -) from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from pydantic import BaseModel +from sdk_clients import NO_PROXY_CACHE, SdkClients pytestmark = pytest.mark.e2e @@ -49,54 +50,54 @@ CACHE_PRIMING_INTERVAL_SECONDS = 3.0 CACHE_WARM_CONSECUTIVE_READS = 3 -def _cacheable_system_block(marker: str) -> TextBlock: +def _cacheable_system_block(marker: str) -> TextBlockParam: """A system prompt at roughly twice the 4096-token minimum cacheable size of Haiku 4.5 (the smallest model here), unique per run so no other run's cache entry can satisfy the read. The marker appears once instead of in every paragraph: repeating it swung the block's size by ~1800 tokens with the marker's own tokenization and left it under the minimum on ~15% of runs, so the system breakpoint went uncached and the priming loop never saw a read.""" - text = f"Run {marker}.\n" + " ".join( - f"Reference paragraph {index}." for index in range(1500) + text = f"Run {marker}.\n" + " ".join(f"Reference paragraph {index}." for index in range(1500)) + return {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}} + + +def _user_turn(text: str, *, cached: bool = False) -> MessageParam: + block: TextBlockParam = ( + {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}} + if cached + else {"type": "text", "text": text} ) - return TextBlock(text=text, cache_control=CacheControl()) + return {"role": "user", "content": [block]} -def _user_turn(text: str, *, cached: bool = False) -> RichMessage: - block = TextBlock(text=text, cache_control=CacheControl() if cached else None) - return RichMessage(role="user", content=[block]) - - -def _system_reminder_turn() -> RichMessage: - return RichMessage( - role="system", - content=[ - TextBlock( - text="Answer with exactly one word." - ) - ], +def _system_reminder_turn() -> MessageParam: + return cast( + "MessageParam", + { + "role": "system", + "content": [{"type": "text", "text": "Answer with exactly one word."}], + }, ) -def _post_messages( - client: EndpointsClient, key: str, body: RichMessagesRequest -) -> Result[MessagesResult]: - return client.proxy.transport.post( - "/v1/messages", - headers=client.proxy.transport.bearer(key), - json=body, - response_type=MessagesResult, +def _assistant_turn(text: str) -> MessageParam: + return {"role": "assistant", "content": [{"type": "text", "text": text}]} + + +def _text(message: Message) -> str: + return "".join(block.text for block in message.content if block.type == "text") + + +def _send(client: Anthropic, model: str, system_block: TextBlockParam, messages: Sequence[MessageParam]) -> Message: + return client.messages.create( + model=model, max_tokens=64, system=[system_block], messages=messages, extra_body=NO_PROXY_CACHE ) -def _register_invoke_deployment( - client: EndpointsClient, resources: ResourceManager, bedrock_model: str -) -> str: +def _register_invoke_deployment(proxy: ProxyClient, resources: ResourceManager, bedrock_model: str) -> str: model = f"e2e-midsys-{unique_marker()}" - model_id = client.create_model( - model, LiteLLMParamsBody(model=bedrock_model, aws_region_name=AWS_REGION) - ) - resources.defer(lambda: client.delete_model(model_id)) + model_id = proxy.create_model(model, LiteLLMParamsBody(model=bedrock_model, aws_region_name=AWS_REGION)) + resources.defer(lambda: proxy.delete_model(model_id)) return model @@ -118,9 +119,7 @@ class PrimedCache(BaseModel): return self.prefix_read_tokens + self.first_turn_creation_tokens -def _prime_prompt_cache( - client: EndpointsClient, key: str, model: str, system_block: TextBlock -) -> PrimedCache: +def _prime_prompt_cache(client: Anthropic, model: str, system_block: TextBlockParam) -> PrimedCache: """Send first-turn calls (fresh cache-marked user turn each attempt, identical system prefix) until one both reads the system prefix back from cache and writes its own user-turn chunk, then re-send that exact turn until @@ -132,19 +131,17 @@ def _prime_prompt_cache( deadline = time.monotonic() + CACHE_PRIMING_DEADLINE_SECONDS while True: user_text = _first_turn_user_text(unique_marker()) - body = RichMessagesRequest( - model=model, - system=[system_block], - messages=[_user_turn(user_text, cached=True)], - ) - usage = unwrap(_post_messages(client, key, body)).usage - if usage.cache_read_input_tokens > 0 and usage.cache_creation_input_tokens > 0: + first_turn = (_user_turn(user_text, cached=True),) + usage = _send(client, model, system_block, first_turn).usage + read_tokens = usage.cache_read_input_tokens or 0 + creation_tokens = usage.cache_creation_input_tokens or 0 + if read_tokens > 0 and creation_tokens > 0: primed = PrimedCache( first_user_text=user_text, - prefix_read_tokens=usage.cache_read_input_tokens, - first_turn_creation_tokens=usage.cache_creation_input_tokens, + prefix_read_tokens=read_tokens, + first_turn_creation_tokens=creation_tokens, ) - if _first_turn_reads_back(client, key, body, primed.full_prefix_tokens, deadline): + if _first_turn_reads_back(client, model, system_block, first_turn, primed.full_prefix_tokens, deadline): return primed if time.monotonic() >= deadline: pytest.fail( @@ -155,15 +152,20 @@ def _prime_prompt_cache( def _reads_full_prefix( - client: EndpointsClient, key: str, body: RichMessagesRequest, full_prefix_tokens: int + client: Anthropic, + model: str, + system_block: TextBlockParam, + messages: Sequence[MessageParam], + full_prefix_tokens: int, ) -> bool: - return unwrap(_post_messages(client, key, body)).usage.cache_read_input_tokens >= full_prefix_tokens + return (_send(client, model, system_block, messages).usage.cache_read_input_tokens or 0) >= full_prefix_tokens def _first_turn_reads_back( - client: EndpointsClient, - key: str, - body: RichMessagesRequest, + client: Anthropic, + model: str, + system_block: TextBlockParam, + messages: Sequence[MessageParam], full_prefix_tokens: int, deadline: float, ) -> bool: @@ -172,12 +174,24 @@ def _first_turn_reads_back( fresh entry can be missing from the region the next request lands on; each miss re-creates the entry there, so the streak converges as the regions warm up.""" while time.monotonic() < deadline: - if all(_reads_full_prefix(client, key, body, full_prefix_tokens) for _ in range(CACHE_WARM_CONSECUTIVE_READS)): + if all( + _reads_full_prefix(client, model, system_block, messages, full_prefix_tokens) + for _ in range(CACHE_WARM_CONSECUTIVE_READS) + ): return True time.sleep(CACHE_PRIMING_INTERVAL_SECONDS) return False +def _reminder_turn_messages(primed: PrimedCache) -> tuple[MessageParam, ...]: + return ( + _user_turn(primed.first_user_text, cached=True), + _system_reminder_turn(), + _assistant_turn("OK."), + _user_turn("Reply with one word again.", cached=True), + ) + + #: Kept in sync with the copy in test_messages_mid_conversation_system_native_providers_e2e.py; #: the e2e suites stay self-contained rather than importing across test modules. MID_CONVERSATION_CACHE_SKIP_REASON = ( @@ -195,32 +209,18 @@ class TestBedrockInvokeMidConversationSystem: exercised_on=[], ) def test_flagged_model_keeps_prompt_cache_across_system_reminder( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = _register_invoke_deployment( - endpoints_client, resources, FLAGGED_INVOKE_MODEL - ) - key = resources.key(models=[model]) + model = _register_invoke_deployment(proxy, resources, FLAGGED_INVOKE_MODEL) + client = sdk.anthropic(resources.key(models=[model])) system_block = _cacheable_system_block(unique_marker()) - primed = _prime_prompt_cache(endpoints_client, key, model, system_block) + primed = _prime_prompt_cache(client, model, system_block) - reminder_turn_body = RichMessagesRequest( - model=model, - system=[system_block], - messages=[ - _user_turn(primed.first_user_text, cached=True), - _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="OK.")]), - _user_turn("Reply with one word again.", cached=True), - ], - ) - second = unwrap(_post_messages(endpoints_client, key, reminder_turn_body)) + second = _send(client, model, system_block, _reminder_turn_messages(primed)) - assert second.text.strip(), ( - f"{model}: reminder turn returned no completion text" - ) - assert second.usage.cache_read_input_tokens >= primed.full_prefix_tokens, ( + assert _text(second).strip(), f"{model}: reminder turn returned no completion text" + assert (second.usage.cache_read_input_tokens or 0) >= primed.full_prefix_tokens, ( f"{model}: turn with a mid-conversation system reminder read " f"{second.usage.cache_read_input_tokens} cached tokens, expected at " f"least the {primed.full_prefix_tokens} cached on turn one " @@ -235,37 +235,23 @@ class TestBedrockInvokeMidConversationSystem: exercised_on=[], ) def test_unflagged_model_converts_system_reminder_and_succeeds( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = _register_invoke_deployment( - endpoints_client, resources, UNFLAGGED_INVOKE_MODEL - ) - key = resources.key(models=[model]) + model = _register_invoke_deployment(proxy, resources, UNFLAGGED_INVOKE_MODEL) + client = sdk.anthropic(resources.key(models=[model])) system_block = _cacheable_system_block(unique_marker()) - primed = _prime_prompt_cache(endpoints_client, key, model, system_block) + primed = _prime_prompt_cache(client, model, system_block) - reminder_turn_body = RichMessagesRequest( - model=model, - system=[system_block], - messages=[ - _user_turn(primed.first_user_text, cached=True), - _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="OK.")]), - _user_turn("Reply with one word again.", cached=True), - ], - ) - second = unwrap(_post_messages(endpoints_client, key, reminder_turn_body)) + second = _send(client, model, system_block, _reminder_turn_messages(primed)) - assert second.role == "assistant", ( - f"{model}: unexpected role {second.role!r}" - ) - assert second.text.strip(), ( + assert second.role == "assistant", f"{model}: unexpected role {second.role!r}" + assert _text(second).strip(), ( f"{model}: conversation with a mid-conversation system reminder " f"returned no text; the reminder was forwarded in place to a model " f"that rejects role 'system' inside messages instead of being converted to a user turn" ) - assert second.usage.cache_read_input_tokens >= primed.full_prefix_tokens, ( + assert (second.usage.cache_read_input_tokens or 0) >= primed.full_prefix_tokens, ( f"{model}: reminder turn read {second.usage.cache_read_input_tokens} " f"cached tokens, expected at least the {primed.full_prefix_tokens} " f"cached on turn one ({primed.prefix_read_tokens} system prefix + " diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py index 8c448399be1..9f5ed8b05da 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py @@ -24,27 +24,28 @@ entry whose prefix spans ``system`` plus message turns is invalidated when the reminder is hoisted (the ``system`` field mutates and a turn disappears from ``messages``), while an entry ending at the system block itself would survive the hoist and mask the regression. + +Calls go through the real Anthropic SDK (LIT-4577). The SDK's ``MessageParam`` +type only admits user/assistant roles, so the system reminder turn is cast to +it; the SDK serializes the dict verbatim, which is exactly the wire shape under +test. """ from __future__ import annotations import time +from collections.abc import Sequence +from typing import cast import pytest -from pydantic import BaseModel - +from anthropic import Anthropic +from anthropic.types import Message, MessageParam, TextBlockParam from e2e_config import unique_marker -from e2e_http import Result, unwrap -from endpoints_client import ( - CacheControl, - EndpointsClient, - MessagesResult, - RichMessage, - RichMessagesRequest, - TextBlock, -) from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from pydantic import BaseModel +from sdk_clients import NO_PROXY_CACHE, SdkClients pytestmark = pytest.mark.e2e @@ -69,46 +70,54 @@ def _vertex_params(model: str, location: str) -> LiteLLMParamsBody: ) -def _cacheable_system_block(marker: str) -> TextBlock: +def _cacheable_system_block(marker: str) -> TextBlockParam: """A system prompt at roughly twice the 4096-token minimum cacheable size of Haiku 4.5 (the smallest model here), unique per run so no other run's cache entry can satisfy the read. The marker appears once instead of in every paragraph: repeating it swung the block's size by ~1800 tokens with the marker's own tokenization and left it under the minimum on ~15% of runs, so the system breakpoint went uncached and the priming loop never saw a read.""" - text = f"Run {marker}.\n" + " ".join( - f"Reference paragraph {index}." for index in range(1500) + text = f"Run {marker}.\n" + " ".join(f"Reference paragraph {index}." for index in range(1500)) + return {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}} + + +def _user_turn(text: str, *, cached: bool = False) -> MessageParam: + block: TextBlockParam = ( + {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}} + if cached + else {"type": "text", "text": text} ) - return TextBlock(text=text, cache_control=CacheControl()) + return {"role": "user", "content": [block]} -def _user_turn(text: str, *, cached: bool = False) -> RichMessage: - block = TextBlock(text=text, cache_control=CacheControl() if cached else None) - return RichMessage(role="user", content=[block]) - - -def _system_reminder_turn() -> RichMessage: - return RichMessage( - role="system", - content=[TextBlock(text="Answer with exactly one word.")], +def _system_reminder_turn() -> MessageParam: + return cast( + "MessageParam", + { + "role": "system", + "content": [{"type": "text", "text": "Answer with exactly one word."}], + }, ) -def _post_messages(client: EndpointsClient, key: str, body: RichMessagesRequest) -> Result[MessagesResult]: - return client.proxy.transport.post( - "/v1/messages", - headers=client.proxy.transport.bearer(key), - json=body, - response_type=MessagesResult, +def _assistant_turn(text: str) -> MessageParam: + return {"role": "assistant", "content": [{"type": "text", "text": text}]} + + +def _text(message: Message) -> str: + return "".join(block.text for block in message.content if block.type == "text") + + +def _send(client: Anthropic, model: str, system_block: TextBlockParam, messages: Sequence[MessageParam]) -> Message: + return client.messages.create( + model=model, max_tokens=64, system=[system_block], messages=messages, extra_body=NO_PROXY_CACHE ) -def _register_deployment( - client: EndpointsClient, resources: ResourceManager, params: LiteLLMParamsBody -) -> str: +def _register_deployment(proxy: ProxyClient, resources: ResourceManager, params: LiteLLMParamsBody) -> str: model = f"e2e-midsys-{unique_marker()}" - model_id = client.create_model(model, params) - resources.defer(lambda: client.delete_model(model_id)) + model_id = proxy.create_model(model, params) + resources.defer(lambda: proxy.delete_model(model_id)) return model @@ -130,9 +139,7 @@ class PrimedCache(BaseModel): return self.prefix_read_tokens + self.first_turn_creation_tokens -def _prime_prompt_cache( - client: EndpointsClient, key: str, model: str, system_block: TextBlock -) -> PrimedCache: +def _prime_prompt_cache(client: Anthropic, model: str, system_block: TextBlockParam) -> PrimedCache: """Send first-turn calls (fresh cache-marked user turn each attempt, identical system prefix) until one both reads the system prefix back from cache and writes its own user-turn chunk, then re-send that exact turn until @@ -144,19 +151,17 @@ def _prime_prompt_cache( deadline = time.monotonic() + CACHE_PRIMING_DEADLINE_SECONDS while True: user_text = _first_turn_user_text(unique_marker()) - body = RichMessagesRequest( - model=model, - system=[system_block], - messages=[_user_turn(user_text, cached=True)], - ) - usage = unwrap(_post_messages(client, key, body)).usage - if usage.cache_read_input_tokens > 0 and usage.cache_creation_input_tokens > 0: + first_turn = (_user_turn(user_text, cached=True),) + usage = _send(client, model, system_block, first_turn).usage + read_tokens = usage.cache_read_input_tokens or 0 + creation_tokens = usage.cache_creation_input_tokens or 0 + if read_tokens > 0 and creation_tokens > 0: primed = PrimedCache( first_user_text=user_text, - prefix_read_tokens=usage.cache_read_input_tokens, - first_turn_creation_tokens=usage.cache_creation_input_tokens, + prefix_read_tokens=read_tokens, + first_turn_creation_tokens=creation_tokens, ) - if _first_turn_reads_back(client, key, body, primed.full_prefix_tokens, deadline): + if _first_turn_reads_back(client, model, system_block, first_turn, primed.full_prefix_tokens, deadline): return primed if time.monotonic() >= deadline: pytest.fail( @@ -167,15 +172,20 @@ def _prime_prompt_cache( def _reads_full_prefix( - client: EndpointsClient, key: str, body: RichMessagesRequest, full_prefix_tokens: int + client: Anthropic, + model: str, + system_block: TextBlockParam, + messages: Sequence[MessageParam], + full_prefix_tokens: int, ) -> bool: - return unwrap(_post_messages(client, key, body)).usage.cache_read_input_tokens >= full_prefix_tokens + return (_send(client, model, system_block, messages).usage.cache_read_input_tokens or 0) >= full_prefix_tokens def _first_turn_reads_back( - client: EndpointsClient, - key: str, - body: RichMessagesRequest, + client: Anthropic, + model: str, + system_block: TextBlockParam, + messages: Sequence[MessageParam], full_prefix_tokens: int, deadline: float, ) -> bool: @@ -184,12 +194,24 @@ def _first_turn_reads_back( fresh entry can be missing from the region the next request lands on; each miss re-creates the entry there, so the streak converges as the regions warm up.""" while time.monotonic() < deadline: - if all(_reads_full_prefix(client, key, body, full_prefix_tokens) for _ in range(CACHE_WARM_CONSECUTIVE_READS)): + if all( + _reads_full_prefix(client, model, system_block, messages, full_prefix_tokens) + for _ in range(CACHE_WARM_CONSECUTIVE_READS) + ): return True time.sleep(CACHE_PRIMING_INTERVAL_SECONDS) return False +def _reminder_turn_messages(primed: PrimedCache) -> tuple[MessageParam, ...]: + return ( + _user_turn(primed.first_user_text, cached=True), + _system_reminder_turn(), + _assistant_turn("OK."), + _user_turn("Reply with one word again.", cached=True), + ) + + #: Why the flagged-model cache checks are skipped rather than failing. The #: assertions below are correct and must be restored unchanged when the bug is #: fixed; they are the regression guard for a real billing cost. @@ -209,28 +231,18 @@ MID_CONVERSATION_CACHE_SKIP_REASON = ( def _assert_flagged_model_keeps_cache( - client: EndpointsClient, resources: ResourceManager, params: LiteLLMParamsBody + proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, params: LiteLLMParamsBody ) -> None: - model = _register_deployment(client, resources, params) - key = resources.key(models=[model]) + model = _register_deployment(proxy, resources, params) + client = sdk.anthropic(resources.key(models=[model])) system_block = _cacheable_system_block(unique_marker()) - primed = _prime_prompt_cache(client, key, model, system_block) + primed = _prime_prompt_cache(client, model, system_block) - reminder_turn_body = RichMessagesRequest( - model=model, - system=[system_block], - messages=[ - _user_turn(primed.first_user_text, cached=True), - _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="OK.")]), - _user_turn("Reply with one word again.", cached=True), - ], - ) - second = unwrap(_post_messages(client, key, reminder_turn_body)) + second = _send(client, model, system_block, _reminder_turn_messages(primed)) - assert second.text.strip(), f"{model}: reminder turn returned no completion text" - assert second.usage.cache_read_input_tokens >= primed.full_prefix_tokens, ( + assert _text(second).strip(), f"{model}: reminder turn returned no completion text" + assert (second.usage.cache_read_input_tokens or 0) >= primed.full_prefix_tokens, ( f"{model}: turn with a mid-conversation system reminder read " f"{second.usage.cache_read_input_tokens} cached tokens, expected at " f"least the {primed.full_prefix_tokens} cached on turn one " @@ -242,33 +254,23 @@ def _assert_flagged_model_keeps_cache( def _assert_unflagged_model_converts_and_succeeds( - client: EndpointsClient, resources: ResourceManager, params: LiteLLMParamsBody + proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, params: LiteLLMParamsBody ) -> None: - model = _register_deployment(client, resources, params) - key = resources.key(models=[model]) + model = _register_deployment(proxy, resources, params) + client = sdk.anthropic(resources.key(models=[model])) system_block = _cacheable_system_block(unique_marker()) - primed = _prime_prompt_cache(client, key, model, system_block) + primed = _prime_prompt_cache(client, model, system_block) - reminder_turn_body = RichMessagesRequest( - model=model, - system=[system_block], - messages=[ - _user_turn(primed.first_user_text, cached=True), - _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="OK.")]), - _user_turn("Reply with one word again.", cached=True), - ], - ) - second = unwrap(_post_messages(client, key, reminder_turn_body)) + second = _send(client, model, system_block, _reminder_turn_messages(primed)) assert second.role == "assistant", f"{model}: unexpected role {second.role!r}" - assert second.text.strip(), ( + assert _text(second).strip(), ( f"{model}: conversation with a mid-conversation system reminder returned " f"no text; the reminder was forwarded in place to a model that rejects " f"role 'system' inside messages instead of being converted to a user turn" ) - assert second.usage.cache_read_input_tokens >= primed.full_prefix_tokens, ( + assert (second.usage.cache_read_input_tokens or 0) >= primed.full_prefix_tokens, ( f"{model}: reminder turn read {second.usage.cache_read_input_tokens} cached " f"tokens, expected at least the {primed.full_prefix_tokens} cached on turn " f"one ({primed.prefix_read_tokens} system prefix + " @@ -289,20 +291,18 @@ class TestAzureFoundryMidConversationSystem: exercised_on=[], ) def test_flagged_model_keeps_prompt_cache_across_system_reminder( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - _assert_flagged_model_keeps_cache(endpoints_client, resources, _azure_params(self.FLAGGED_MODEL)) + _assert_flagged_model_keeps_cache(proxy, resources, sdk, _azure_params(self.FLAGGED_MODEL)) @pytest.mark.covers( "llm.messages.azure_foundry.mid_conversation_system.nonstream.works", exercised_on=[], ) def test_unflagged_model_converts_system_reminder_and_succeeds( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - _assert_unflagged_model_converts_and_succeeds( - endpoints_client, resources, _azure_params(self.UNFLAGGED_MODEL) - ) + _assert_unflagged_model_converts_and_succeeds(proxy, resources, sdk, _azure_params(self.UNFLAGGED_MODEL)) class TestVertexMidConversationSystem: @@ -323,10 +323,10 @@ class TestVertexMidConversationSystem: exercised_on=[], ) def test_flagged_model_keeps_prompt_cache_across_system_reminder( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: _assert_flagged_model_keeps_cache( - endpoints_client, resources, _vertex_params(self.FLAGGED_MODEL, self.FLAGGED_LOCATION) + proxy, resources, sdk, _vertex_params(self.FLAGGED_MODEL, self.FLAGGED_LOCATION) ) @pytest.mark.covers( @@ -334,8 +334,8 @@ class TestVertexMidConversationSystem: exercised_on=[], ) def test_unflagged_model_converts_system_reminder_and_succeeds( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: _assert_unflagged_model_converts_and_succeeds( - endpoints_client, resources, _vertex_params(self.UNFLAGGED_MODEL, self.UNFLAGGED_LOCATION) + proxy, resources, sdk, _vertex_params(self.UNFLAGGED_MODEL, self.UNFLAGGED_LOCATION) ) diff --git a/tests/e2e/llm_translation/test_moderations_e2e.py b/tests/e2e/llm_translation/test_moderations_e2e.py index 0395a4b2848..e936f7b335a 100644 --- a/tests/e2e/llm_translation/test_moderations_e2e.py +++ b/tests/e2e/llm_translation/test_moderations_e2e.py @@ -1,19 +1,23 @@ """Live e2e: POST /v1/moderations classifies content against the provider policy. -Registers OpenAI's omni moderation model at runtime and asserts the product -promise on both sides of the decision: clearly violent text comes back flagged -with at least one policy category tripped, and benign text comes back not flagged. +Registers OpenAI's omni moderation model at runtime, drives it through the real +OpenAI SDK (LIT-4577), and asserts the product promise on both sides of the +decision: clearly violent text comes back flagged with at least one policy +category tripped, and benign text comes back not flagged. The malformed-body +negative stays on the shared transport, since the SDK refuses to send it. """ from __future__ import annotations import pytest from e2e_config import unique_marker -from e2e_http import assert_client_error, unwrap -from endpoints_client import EndpointsClient +from e2e_http import assert_client_error from lifecycle import ResourceManager from models import LiteLLMParamsBody -from pydantic import BaseModel +from openai.types import Moderation +from proxy_client import ProxyClient +from pydantic import BaseModel, TypeAdapter +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e @@ -26,59 +30,63 @@ class _OptionalModerationBody(BaseModel): input: str | None = None -def _register_moderation_model( - endpoints_client: EndpointsClient, resources: ResourceManager -) -> str: +def _register_moderation_model(proxy: ProxyClient, resources: ResourceManager) -> str: model = f"e2e-moderation-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="openai/omni-moderation-latest", api_key="os.environ/OPENAI_API_KEY" ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) return model +_CATEGORY_FLAGS = TypeAdapter(dict[str, bool | None]) + + +def _flagged_categories(item: Moderation) -> tuple[str, ...]: + flags = _CATEGORY_FLAGS.validate_python(item.categories.model_dump()) + return tuple(name for name, hit in flags.items() if hit) + + class TestModerations: @pytest.mark.covers("llm.moderations.openai.basic.nonstream.works") def test_moderations_flags_violent_content( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = _register_moderation_model(endpoints_client, resources) - key = resources.key() + model = _register_moderation_model(proxy, resources) + client = sdk.openai(resources.key()) - result = unwrap(endpoints_client.moderations(key, model, VIOLENT_TEXT)) - item = result.first - assert item is not None, f"/moderations returned no results: {result}" - assert item.flagged, f"violent text was not flagged: {item}" - assert item.flagged_categories, ( - f"flagged result reported no true category: {item}" - ) + moderation = client.moderations.create(model=model, input=VIOLENT_TEXT) + assert moderation.results, f"/moderations returned no results: {moderation!r}" + item = moderation.results[0] + assert item.flagged, f"violent text was not flagged: {item!r}" + assert _flagged_categories(item), f"flagged result reported no true category: {item!r}" def test_moderations_passes_benign_content( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = _register_moderation_model(endpoints_client, resources) - key = resources.key() + model = _register_moderation_model(proxy, resources) + client = sdk.openai(resources.key()) - result = unwrap(endpoints_client.moderations(key, model, BENIGN_TEXT)) - item = result.first - assert item is not None, f"/moderations returned no results: {result}" + moderation = client.moderations.create(model=model, input=BENIGN_TEXT) + assert moderation.results, f"/moderations returned no results: {moderation!r}" + item = moderation.results[0] assert not item.flagged, ( - f"benign text was flagged as {item.flagged_categories}: {item}" + f"benign text was flagged as {_flagged_categories(item)}: {item!r}" ) @pytest.mark.skip(reason="stage red: product gap, /v1/moderations 500s (KeyError 'input') on missing input instead of 400") @pytest.mark.covers("llm.moderations.openai.input_validation.nonstream.works") def test_missing_input_returns_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: - model = _register_moderation_model(endpoints_client, resources) + model = _register_moderation_model(proxy, resources) key = resources.key() - result = endpoints_client.proxy.transport.send( + result = proxy.transport.send( "/v1/moderations", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalModerationBody(model=model), ) assert_client_error(result, "moderations missing input") diff --git a/tests/e2e/llm_translation/test_ocr_rust_e2e.py b/tests/e2e/llm_translation/test_ocr_rust_e2e.py index e83920111c7..c2560b199af 100644 --- a/tests/e2e/llm_translation/test_ocr_rust_e2e.py +++ b/tests/e2e/llm_translation/test_ocr_rust_e2e.py @@ -21,9 +21,9 @@ from typing import Protocol import pytest from e2e_config import unique_marker from e2e_http import assert_client_error, unwrap -from endpoints_client import EndpointsClient from lifecycle import ResourceManager from models import LiteLLMParamsBody, OcrBody, OcrDocument, OcrResponse +from proxy_client import ProxyClient from pydantic import BaseModel pytestmark = pytest.mark.e2e @@ -149,28 +149,28 @@ def _assert_ocr_document(response: OcrResponse) -> None: class TestRustOcrGateway: @pytest.mark.parametrize("case", RUST_OCR_CASES, ids=_CASE_IDS) def test_rust_ocr_response( - self, endpoints_client: EndpointsClient, resources: ResourceManager, case: _OcrCase + self, proxy: ProxyClient, resources: ResourceManager, case: _OcrCase ) -> None: model = f"rust-ocr-{case.suffix}-{unique_marker()}" - model_id = endpoints_client.create_model(model, case.provider.litellm_params()) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + model_id = proxy.create_model(model, case.provider.litellm_params()) + resources.defer(lambda: proxy.delete_model(model_id)) key = resources.key() - response = unwrap(endpoints_client.proxy.ocr(key, OcrBody(model=model, document=case.document))) + response = unwrap(proxy.ocr(key, OcrBody(model=model, document=case.document))) _assert_ocr_document(response) @pytest.mark.skip(reason="stage red: product gap, /v1/ocr 500s (aocr TypeError) on missing document instead of 400") @pytest.mark.covers("llm.ocr.openai.input_validation.nonstream.works") def test_missing_document_returns_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: model = f"rust-ocr-val-{unique_marker()}" - model_id = endpoints_client.create_model(model, MistralOcr().litellm_params()) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + model_id = proxy.create_model(model, MistralOcr().litellm_params()) + resources.defer(lambda: proxy.delete_model(model_id)) key = resources.key() - result = endpoints_client.proxy.transport.send( + result = proxy.transport.send( "/v1/ocr", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalOcrBody(model=model), ) assert_client_error(result, "ocr missing document") diff --git a/tests/e2e/llm_translation/test_passthrough_headers_e2e.py b/tests/e2e/llm_translation/test_passthrough_headers_e2e.py index d6f832afb53..26f98774c63 100644 --- a/tests/e2e/llm_translation/test_passthrough_headers_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_headers_e2e.py @@ -20,9 +20,8 @@ from pydantic import BaseModel, Field from e2e_config import unique_marker from e2e_http import AuthHeaders, NoBody, require_successful_call, unwrap -from endpoints_client import MessagesResult from lifecycle import ResourceManager -from models import ChatMessage, KeyGenerateBody +from models import AnthropicMessagesResponse, ChatMessage, KeyGenerateBody from passthrough_client import PassthroughClient pytestmark = pytest.mark.e2e @@ -165,8 +164,9 @@ class TestPassthroughHeaders: json=_messages_body(), ) require_successful_call(result) - completion = MessagesResult.model_validate_json(result.body) - assert completion.text.strip(), ( + completion = AnthropicMessagesResponse.model_validate_json(result.body) + text = "".join(block.text or "" for block in (completion.content or [])) + assert text.strip(), ( f"static x-api-key must reach Anthropic for the call to succeed at all; got {result.body[:300]}" ) diff --git a/tests/e2e/llm_translation/test_rerank_e2e.py b/tests/e2e/llm_translation/test_rerank_e2e.py index c9f58b2c03c..87b8618e6fb 100644 --- a/tests/e2e/llm_translation/test_rerank_e2e.py +++ b/tests/e2e/llm_translation/test_rerank_e2e.py @@ -1,19 +1,20 @@ """Live e2e: POST /v1/rerank ranks documents by relevance. -Registers a Cohere rerank deployment at runtime and asserts the endpoint returns -scored results within the requested top_n. Migrated from +Registers Cohere and Bedrock rerank deployments at runtime and asserts the +endpoint returns scored results within the requested top_n. No official +OpenAI/Anthropic SDK covers /v1/rerank, so the call rides the shared typed +transport via ProxyClient.rerank. Migrated from litellm-regression-tests/tests/test_inference_endpoints.py. """ from __future__ import annotations import pytest - from e2e_config import unique_marker -from e2e_http import require_successful_call -from endpoints_client import EndpointsClient, RerankResult +from e2e_http import unwrap from lifecycle import ResourceManager -from models import LiteLLMParamsBody +from models import LiteLLMParamsBody, RerankBody, RerankResponse +from proxy_client import ProxyClient pytestmark = pytest.mark.e2e @@ -26,38 +27,39 @@ DOCUMENTS = [ QUERY = "What is the capital of the United States?" -def _assert_top_n_scored(body: str) -> None: - parsed = RerankResult.model_validate_json(body) - assert parsed.results, f"/rerank returned no results: {body[:300]}" - assert len(parsed.results) <= 3, f"top_n=3 not honored: {body[:300]}" - assert parsed.results[0].relevance_score is not None, ( - f"top rerank result has no relevance_score: {body[:300]}" +def _assert_top_n_scored(response: RerankResponse) -> None: + assert response.results, f"/rerank returned no results: {response!r}" + assert len(response.results) <= 3, f"top_n=3 not honored: {response!r}" + assert response.results[0].relevance_score is not None, ( + f"top rerank result has no relevance_score: {response!r}" + ) + + +def _rerank_top_3(proxy: ProxyClient, key: str, model: str) -> RerankResponse: + return unwrap( + proxy.rerank(key, RerankBody(model=model, query=QUERY, documents=DOCUMENTS, top_n=3)) ) class TestRerank: @pytest.mark.covers("llm.rerank.cohere.basic.nonstream.works") - def test_rerank_scores_top_n( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: + def test_rerank_scores_top_n(self, proxy: ProxyClient, resources: ResourceManager) -> None: model = f"e2e-rerank-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody(model="cohere/rerank-v3.5", api_key="os.environ/COHERE_API_KEY"), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) key = resources.key() - result = endpoints_client.rerank(key, model, QUERY, DOCUMENTS, top_n=3) - require_successful_call(result) - _assert_top_n_scored(result.body) + _assert_top_n_scored(_rerank_top_3(proxy, key, model)) @pytest.mark.covers("llm.rerank.bedrock.basic.nonstream.works", exercised_on=["rerank"]) def test_bedrock_rerank_scores_top_n( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: model = f"e2e-bedrock-rerank-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0", @@ -66,9 +68,7 @@ class TestRerank: aws_region_name="os.environ/AWS_REGION", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) key = resources.key() - result = endpoints_client.rerank(key, model, QUERY, DOCUMENTS, top_n=3) - require_successful_call(result) - _assert_top_n_scored(result.body) + _assert_top_n_scored(_rerank_top_3(proxy, key, model)) diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index 525231de917..9cc70da63b0 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -1,12 +1,15 @@ """Live e2e: POST /v1/responses returns a real completion. -Registers an OpenAI deployment at runtime, drives the Responses API through the -gateway, and asserts output text came back. Migrated from +Registers an OpenAI deployment at runtime and drives the Responses API through +the gateway with the real OpenAI SDK, the client customers actually use +(LIT-4577), asserting output text came back. Malformed bodies the SDK refuses +to build stay on the shared transport. Migrated from litellm-regression-tests/tests/test_inference_endpoints.py. """ from __future__ import annotations +import contextlib import json import threading from collections.abc import Mapping @@ -14,26 +17,23 @@ from dataclasses import dataclass, field from types import MappingProxyType from typing import Final, cast +import openai import pytest from e2e_config import PROVIDER_EDGE_ADVERTISE_HOST, PROVIDER_EDGE_BIND_HOST, unique_marker -from e2e_http import ( - assert_client_error, - require_successful_call, -) -from endpoints_client import ( - EndpointsClient, - FunctionParameterProperty, - FunctionParameters, - ResponsesFunctionTool, - ResponsesOutputTextDeltaEvent, - ResponsesResult, - ResponsesStreamEventType, -) +from e2e_http import assert_client_error from lifecycle import ResourceManager from models import ChatBody, ChatMessage, LiteLLMParamsBody +from openai.types.responses import ( + FunctionToolParam, + Response, + ResponseFunctionToolCall, + ResponseInputParam, +) from provider_edge import LiveEdge, start_provider_edge from provider_edge_bedrock import bedrock_signer -from pydantic import BaseModel, ValidationError +from proxy_client import ProxyClient +from pydantic import BaseModel +from sdk_clients import NO_PROXY_CACHE, SdkClients pytestmark = pytest.mark.e2e @@ -45,6 +45,8 @@ class _OptionalResponsesBody(BaseModel): BEDROCK_CONVERSE_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" +INSTRUCTIONS = "You are a helpful assistant" +CAT_IMAGE_URL = "https://upload.wikimedia.org/wikipedia/commons/3/3a/Cat03.jpg" BEDROCK_EDGE_REGION: Final = "us-east-1" BEDROCK_EDGE_MOUNT: Final = f"bedrock/{BEDROCK_EDGE_REGION}" @@ -73,14 +75,25 @@ class ConverseRequestCapture: return tuple(self._bodies) -WEATHER_TOOL = ResponsesFunctionTool( - name="get_weather", - description="Get the weather for a location", - parameters=FunctionParameters( - properties={"location": FunctionParameterProperty(type="string")}, - required=["location"], - ), -) +WEATHER_TOOL: FunctionToolParam = { + "type": "function", + "name": "get_weather", + "description": "Get the weather for a location", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + "strict": False, +} + + +def _openai_params() -> LiteLLMParamsBody: + return LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY") + + +def _anthropic_params() -> LiteLLMParamsBody: + return LiteLLMParamsBody(model="anthropic/claude-haiku-4-5", api_key="os.environ/ANTHROPIC_API_KEY") def _bedrock_params() -> LiteLLMParamsBody: @@ -92,6 +105,27 @@ def _bedrock_params() -> LiteLLMParamsBody: ) +def _register( + proxy: ProxyClient, resources: ResourceManager, params: LiteLLMParamsBody, prefix: str = "e2e-responses" +) -> str: + model = f"{prefix}-{unique_marker()}" + model_id = proxy.create_model(model, params) + resources.defer(lambda: proxy.delete_model(model_id)) + return model + + +def _function_calls(response: Response) -> tuple[ResponseFunctionToolCall, ...]: + return tuple(item for item in response.output if isinstance(item, ResponseFunctionToolCall)) + + +def _assert_weather_call(response: Response) -> None: + function_call = next((call for call in _function_calls(response) if call.name == "get_weather"), None) + assert function_call is not None, f"no get_weather function call: {response.output!r}" + raw_arguments = cast(object, json.loads(function_call.arguments)) + arguments = WeatherArguments.model_validate(raw_arguments) + assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + + class WeatherArguments(BaseModel): location: str @@ -99,250 +133,184 @@ class WeatherArguments(BaseModel): class TestResponses: @pytest.mark.covers("llm.responses.openai.basic.nonstream.works") def test_responses_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _openai_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses(key, model, "reply with one word") - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - assert parsed.text.strip(), f"/responses returned no output text: {result.body[:300]}" + response = client.responses.create( + model=model, input="reply with one word", instructions=INSTRUCTIONS, extra_body=NO_PROXY_CACHE + ) + assert response.output_text.strip(), f"/responses returned no output text: {response.output!r}" @pytest.mark.covers("llm.responses.openai.basic.stream.works") def test_responses_streaming_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _openai_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses(key, model, "reply with one word", stream=True) - require_successful_call(result) - delta_events = tuple( - parsed - for event in result.stream_events - if (parsed := _parse_stream_event(event)) is not None + stream = client.responses.create( + model=model, + input="reply with one word", + instructions=INSTRUCTIONS, + stream=True, + extra_body=NO_PROXY_CACHE, + ) + events = tuple(stream) + assert events, "responses stream returned no events" + deltas = tuple(event.delta for event in events if event.type == "response.output_text.delta") + assert any(delta for delta in deltas), "responses stream returned no text deltas" + assert events[-1].type == "response.completed", ( + f"responses stream did not terminate with response.completed: {events[-1].type}" ) - - assert any(event.delta for event in delta_events), "responses stream returned no text deltas" - assert result.stream_events, "responses stream returned no events" - assert ( - ResponsesStreamEventType.model_validate_json(result.stream_events[-1]).type - == "response.completed" - ), "responses stream did not terminate with response.completed" @pytest.mark.covers("llm.responses.openai.basic.nonstream.cost_logged") - def test_responses_logs_cost( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), + def test_responses_logs_cost(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None: + model = _register(proxy, resources, _openai_params()) + client = sdk.openai(resources.key()) + + raw = client.responses.with_raw_response.create( + model=model, + input=f"reply with one word {unique_marker()}", + instructions=INSTRUCTIONS, + extra_body=NO_PROXY_CACHE, + ) + response = raw.parse() + assert response.output_text.strip(), f"/responses returned no output text: {response.output!r}" + assert raw.headers.get("x-litellm-call-id") and response.id, ( + f"missing response identifiers: id={response.id!r}, headers={dict(raw.headers)}" ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - result = endpoints_client.responses(key, model, f"reply with one word {unique_marker()}") - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - assert parsed.text.strip(), f"/responses returned no output text: {result.body[:300]}" - assert result.call_id and parsed.id, f"missing response identifiers: {result.body[:300]}" - - rows = endpoints_client.proxy.poll_logs_for_request_id( - parsed.id, + rows = proxy.poll_logs_for_request_id( + response.id, predicate=lambda logged_rows: any((row.spend or 0) > 0 for row in logged_rows), ) row = next((logged_row for logged_row in rows if (logged_row.spend or 0) > 0), None) - assert row is not None, f"no costed spend row for response id {parsed.id}" + assert row is not None, f"no costed spend row for response id {response.id}" assert "gpt-4o-mini" in (row.model or ""), f"unexpected spend row model: {row.model}" @pytest.mark.covers("llm.responses.openai.tool_use.nonstream.works") def test_responses_returns_function_call( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _openai_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses_with_tools( - key, - model, - "What is the weather in San Francisco? Use the get_weather tool.", - [ - ResponsesFunctionTool( - name="get_weather", - description="Get the weather for a location", - parameters=FunctionParameters( - properties={"location": FunctionParameterProperty(type="string")}, - required=["location"], - ), - ) - ], + response = client.responses.create( + model=model, + input="What is the weather in San Francisco? Use the get_weather tool.", + instructions=INSTRUCTIONS, + tools=[WEATHER_TOOL], + extra_body=NO_PROXY_CACHE, ) - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - function_call = next( - (call for call in parsed.function_calls if call.name == "get_weather"), - None, - ) - assert function_call is not None, f"no get_weather function call: {result.body[:500]}" - assert function_call.arguments is not None - raw_arguments = cast(object, json.loads(function_call.arguments)) - arguments = WeatherArguments.model_validate(raw_arguments) - assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + _assert_weather_call(response) @pytest.mark.covers("llm.responses.openai.vision.nonstream.works") def test_responses_vision_describes_image( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + model = _register( + proxy, + resources, LiteLLMParamsBody(model="openai/gpt-4o", api_key="os.environ/OPENAI_API_KEY"), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + client = sdk.openai(resources.key()) - result = endpoints_client.responses_vision( - key, - model, - "What animal is shown in this image? Answer in one word", - "https://upload.wikimedia.org/wikipedia/commons/3/3a/Cat03.jpg", + vision_input: ResponseInputParam = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "What animal is shown in this image? Answer in one word"}, + {"type": "input_image", "image_url": CAT_IMAGE_URL, "detail": "auto"}, + ], + } + ] + response = client.responses.create( + model=model, input=vision_input, instructions=INSTRUCTIONS, extra_body=NO_PROXY_CACHE + ) + text = response.output_text.strip().lower() + assert text, f"/responses vision returned no output text: {response.output!r}" + assert any(keyword in text for keyword in ("cat", "feline")), ( + f"vision response did not describe the image: {text[:300]}" ) - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - text = parsed.text.strip().lower() - assert text, f"/responses vision returned no output text: {result.body[:300]}" - assert any( - keyword in text - for keyword in ("cat", "feline") - ), f"vision response did not describe the image: {parsed.text[:300]}" @pytest.mark.covers("llm.responses.anthropic.basic.nonstream.works") def test_responses_anthropic_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="anthropic/claude-haiku-4-5", api_key="os.environ/ANTHROPIC_API_KEY" - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _anthropic_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses(key, model, "reply with one word") - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - assert parsed.text.strip(), f"/responses returned no output text: {result.body[:300]}" + response = client.responses.create( + model=model, input="reply with one word", instructions=INSTRUCTIONS, extra_body=NO_PROXY_CACHE + ) + assert response.output_text.strip(), f"/responses returned no output text: {response.output!r}" @pytest.mark.covers("llm.responses.anthropic.tool_use.nonstream.works") def test_responses_anthropic_returns_function_call( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="anthropic/claude-haiku-4-5", api_key="os.environ/ANTHROPIC_API_KEY" - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _anthropic_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses_with_tools( - key, - model, - "What is the weather in San Francisco? Use the get_weather tool.", - [ - ResponsesFunctionTool( - name="get_weather", - description="Get the weather for a location", - parameters=FunctionParameters( - properties={"location": FunctionParameterProperty(type="string")}, - required=["location"], - ), - ) - ], + response = client.responses.create( + model=model, + input="What is the weather in San Francisco? Use the get_weather tool.", + instructions=INSTRUCTIONS, + tools=[WEATHER_TOOL], + extra_body=NO_PROXY_CACHE, ) - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - function_call = next( - (call for call in parsed.function_calls if call.name == "get_weather"), - None, - ) - assert function_call is not None, f"no get_weather function call: {result.body[:500]}" - assert function_call.arguments is not None - raw_arguments = cast(object, json.loads(function_call.arguments)) - arguments = WeatherArguments.model_validate(raw_arguments) - assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + _assert_weather_call(response) @pytest.mark.covers("llm.responses.bedrock_converse.basic.nonstream.works") def test_responses_bedrock_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model(model, _bedrock_params()) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _bedrock_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses(key, model, "reply with one word") - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - assert parsed.text.strip(), f"/responses over bedrock returned no output text: {result.body[:300]}" + response = client.responses.create( + model=model, input="reply with one word", instructions=INSTRUCTIONS, extra_body=NO_PROXY_CACHE + ) + assert response.output_text.strip(), f"/responses over bedrock returned no output text: {response.output!r}" @pytest.mark.covers("llm.responses.bedrock_converse.tool_use.nonstream.works") def test_responses_bedrock_returns_function_call( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model(model, _bedrock_params()) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _bedrock_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses_with_tools( - key, model, "What is the weather in San Francisco? Use the get_weather tool.", [WEATHER_TOOL] + response = client.responses.create( + model=model, + input="What is the weather in San Francisco? Use the get_weather tool.", + instructions=INSTRUCTIONS, + tools=[WEATHER_TOOL], + extra_body=NO_PROXY_CACHE, ) - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - function_call = next((call for call in parsed.function_calls if call.name == "get_weather"), None) - assert function_call is not None, f"no get_weather function call over bedrock: {result.body[:500]}" - assert function_call.arguments is not None - raw_arguments = cast(object, json.loads(function_call.arguments)) - arguments = WeatherArguments.model_validate(raw_arguments) - assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + _assert_weather_call(response) @pytest.mark.provider_edge_host @pytest.mark.parametrize("endpoint", ["/v1/responses", "/v1/chat/completions"]) def test_bedrock_forwards_allowed_safety_identifier_as_additional_model_request_field( - self, endpoints_client: EndpointsClient, resources: ResourceManager, endpoint: str + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, endpoint: str ) -> None: + """Judges the Converse bodies the edge captured, not the reply: Claude on + Bedrock rejects the forwarded field with a 400, which the chat leg's + ``Result`` carries as a value and the OpenAI SDK raises.""" capture: Final = ConverseRequestCapture() edge: Final = start_provider_edge( LiveEdge(observe_request=capture.observe, sign=bedrock_signer(BEDROCK_EDGE_REGION)), - mounts=MappingProxyType({BEDROCK_EDGE_MOUNT: f"https://bedrock-runtime.{BEDROCK_EDGE_REGION}.amazonaws.com"}), + mounts=MappingProxyType( + {BEDROCK_EDGE_MOUNT: f"https://bedrock-runtime.{BEDROCK_EDGE_REGION}.amazonaws.com"} + ), bind_host=PROVIDER_EDGE_BIND_HOST, advertise_host=PROVIDER_EDGE_ADVERTISE_HOST, ) resources.defer(edge.shutdown) model: Final = f"e2e-responses-{unique_marker()}" - model_id: Final = endpoints_client.create_model( + model_id: Final = proxy.create_model( model, LiteLLMParamsBody( model=BEDROCK_CONVERSE_BACKEND, @@ -353,14 +321,21 @@ class TestResponses: allowed_openai_params=["safety_identifier"], ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) key: Final = resources.key() safety_identifier: Final = f"end-user-{unique_marker()}" if endpoint == "/v1/responses": - endpoints_client.responses(key, model, "reply with one word", safety_identifier=safety_identifier) + with contextlib.suppress(openai.BadRequestError): + sdk.openai(key).responses.create( + model=model, + input="reply with one word", + instructions=INSTRUCTIONS, + safety_identifier=safety_identifier, + extra_body=NO_PROXY_CACHE, + ) else: - endpoints_client.proxy.chat( + proxy.chat( key, ChatBody( model=model, @@ -375,59 +350,37 @@ class TestResponses: f"{endpoint} did not forward safety_identifier to Bedrock Converse on every attempt: {capture.bodies}" ) - @pytest.mark.skip(reason="stage red: product gap, /v1/responses 500s (aresponses TypeError) on missing input instead of 400") + @pytest.mark.skip( + reason="stage red: product gap, /v1/responses 500s (aresponses TypeError) on missing input instead of 400" + ) @pytest.mark.covers("llm.responses.openai.input_validation.nonstream.works") - def test_missing_input_returns_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model = f"e2e-responses-val-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + def test_missing_input_returns_error(self, proxy: ProxyClient, resources: ResourceManager) -> None: + model = _register(proxy, resources, _openai_params(), prefix="e2e-responses-val") key = resources.key() - result = endpoints_client.proxy.transport.send( + result = proxy.transport.send( "/v1/responses", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalResponsesBody(model=model), ) assert_client_error(result, "responses missing input") @pytest.mark.covers("llm.responses.openai.input_validation.nonstream.works") - def test_missing_model_returns_client_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: + def test_missing_model_returns_client_error(self, proxy: ProxyClient, resources: ResourceManager) -> None: key = resources.key() - result = endpoints_client.proxy.transport.send( + result = proxy.transport.send( "/v1/responses", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalResponsesBody(input="ping"), ) assert_client_error(result, "responses missing model") @pytest.mark.covers("llm.responses.openai.input_validation.nonstream.works") - def test_empty_input_returns_client_error( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model = f"e2e-responses-val-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + def test_empty_input_returns_client_error(self, proxy: ProxyClient, resources: ResourceManager) -> None: + model = _register(proxy, resources, _openai_params(), prefix="e2e-responses-val") key = resources.key() - result = endpoints_client.proxy.transport.send( + result = proxy.transport.send( "/v1/responses", - headers=endpoints_client.proxy.transport.bearer(key), + headers=proxy.transport.bearer(key), json=_OptionalResponsesBody(model=model, input=""), ) assert_client_error(result, "responses empty input") - -def _parse_stream_event( - event: str, -) -> ResponsesOutputTextDeltaEvent | None: - try: - return ResponsesOutputTextDeltaEvent.model_validate_json(event) - except ValidationError: - return None diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 47ef672ebec..9440505e58c 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -697,6 +697,26 @@ class EmbedResponse(BaseModel): model: str | None = None +# ---------- rerank ---------- + + +class RerankBody(BaseModel): + model: str + query: str + documents: list[str] + top_n: int + cache: dict[str, bool] | None = {"no-cache": True} + + +class RerankItem(BaseModel): + index: int | None = None + relevance_score: float | None = None + + +class RerankResponse(BaseModel): + results: list[RerankItem] = [] + + # ---------- ocr ---------- diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index c6ede240c3b..2f32361e083 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -79,6 +79,8 @@ from models import ( ModelUpdateBody, OcrBody, OcrResponse, + RerankBody, + RerankResponse, RouterCurrentValues, RouterSettingsResponse, SpendLogRow, @@ -940,6 +942,16 @@ class ProxyClient: timeout=SLOW_PROVIDER_TIMEOUT_SECONDS, ) + def rerank(self, key: str, body: RerankBody) -> Result[RerankResponse]: + """POST /v1/rerank (Cohere-format). No official OpenAI/Anthropic SDK + covers this route, so it stays on the shared typed transport.""" + return self.transport.post( + "/v1/rerank", + headers=self.transport.bearer(key), + json=body, + response_type=RerankResponse, + ) + def count_tokens(self, key: str, body: CountTokensBody) -> Result[CountTokensResponse]: """POST /v1/messages/count_tokens (Anthropic-native). Sends the anthropic-version header so the native path accepts it; harmless on the diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index 9f1118ab1e3..e07cbe6b2a3 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -34,12 +34,19 @@ def delete_key_if_present(candidate: Gateway, key: str) -> None: assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == [] -def eventually(read: Callable[[], T], satisfied: Callable[[T], bool], seconds: float = 10) -> T: +def eventually( + read: Callable[[], T], + satisfied: Callable[[T], bool], + seconds: float = 10, + return_last_on_timeout: bool = False, +) -> T: deadline: Final = time.monotonic() + seconds while True: observed: Final = read() if satisfied(observed): return observed + if return_last_on_timeout and time.monotonic() >= deadline: + return observed assert time.monotonic() < deadline, f"State did not converge: {observed!r}" time.sleep(0.1) @@ -58,12 +65,32 @@ class Gateway: *, key: str | None = None, params: Mapping[str, str] | None = None, + headers: Mapping[str, str] | None = None, ) -> httpx.Response: + request_headers: Final = { + "Authorization": f"Bearer {self.key if key is None else key}", + **(headers or {}), + } return self.client.request( method, path, json=body, params=params, + headers=request_headers, + ) + + def request_multipart( + self, + path: str, + fields: Mapping[str, str], + files: Mapping[str, tuple[str, bytes, str]], + *, + key: str | None = None, + ) -> httpx.Response: + return self.client.post( + path, + data=fields, + files=files, headers={"Authorization": f"Bearer {self.key if key is None else key}"}, ) diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index 1ad02b6a3f2..759df7003e4 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -1,8 +1,10 @@ from __future__ import annotations import argparse +import asyncio +import base64 from collections import deque -from collections.abc import Mapping +from collections.abc import AsyncIterator, Mapping import json from dataclasses import dataclass, field import os @@ -10,6 +12,7 @@ from pathlib import Path from queue import SimpleQueue import struct from typing import Final, cast +import uuid import zlib import httpx @@ -17,11 +20,13 @@ import uvicorn from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from starlette.applications import Starlette from starlette.requests import Request -from starlette.responses import JSONResponse, Response +from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.routing import Route from _fake_openai_endpoint_server import chat_completions, completions, embeddings, health, moderations from integration.cost_calculation.cost_tracking_case import ( + BinaryResponse, + EventStreamEvent, EventStreamResponse, JsonResponse, SseResponse, @@ -75,10 +80,15 @@ def _aws_str_header(name: str, value: str) -> bytes: ) -def _aws_event_frame(event_type: str, payload: Mapping[str, JsonValue], scenario_id: str) -> bytes: +def _aws_event_frame( + event_type: str, + payload: Mapping[str, JsonValue], + scenario_id: str, + unique_id: str, +) -> bytes: payload_bytes: Final = json.dumps(payload, separators=(",", ":")).replace( "$REQUEST_ID", scenario_id - ).encode() + ).replace("$UNIQUE_ID", unique_id).encode() headers_bytes: Final = ( _aws_str_header(":event-type", event_type) + _aws_str_header(":content-type", "application/json") @@ -193,9 +203,11 @@ class Provider: async def scripted(self, request: Request) -> Response: segments: Final = tuple(segment for segment in cast(str, request.path_params["path"]).split("/") if segment) - if not segments: - return JSONResponse({"error": "Unknown scenario"}, status_code=404) - scenario_id: Final = segments[0].split(":", 1)[0] + scenario_id: Final = ( + segments[0].split(":", 1)[0] + if segments and self.scenario_store.get(segments[0].split(":", 1)[0]) is not None + else request.headers.get("x-scripted-scenario", "") + ) response: Final = self.scenario_store.get(scenario_id) if response is None: return JSONResponse({"error": "Unknown scenario"}, status_code=404) @@ -203,22 +215,58 @@ class Provider: @staticmethod def _response(response: StoredResponse, scenario_id: str) -> Response: + unique_id: Final = f"{scenario_id}-{uuid.uuid4().hex[:8]}" match response: case JsonResponse(): return Response( content=json.dumps(response.body, separators=(",", ":")).replace( "$REQUEST_ID", scenario_id + ).replace( + "$UNIQUE_ID", unique_id ).encode(), media_type=response.content_type, + status_code=response.status, + ) + case BinaryResponse(): + return Response( + content=b"\x00" * response.length, + media_type=response.content_type, ) case SseResponse(): + if response.frame_delay_ms > 0: + async def stream() -> AsyncIterator[bytes]: + for frame in response.frames: + yield ( + f"{frame.replace('$REQUEST_ID', scenario_id).replace('$UNIQUE_ID', unique_id)}\n\n" + ).encode() + await asyncio.sleep(response.frame_delay_ms / 1000) + + return StreamingResponse(stream(), media_type=response.content_type) stream_body: Final = ("\n\n".join(response.frames) + "\n\n").replace( "$REQUEST_ID", scenario_id - ) + ).replace("$UNIQUE_ID", unique_id) return Response(content=stream_body.encode(), media_type=response.content_type) case EventStreamResponse(): + events: Final = ( + tuple( + EventStreamEvent( + event_type="chunk", + payload={ + "bytes": base64.b64encode( + json.dumps(event.payload, separators=(",", ":")) + .replace("$REQUEST_ID", scenario_id) + .replace("$UNIQUE_ID", unique_id) + .encode() + ).decode(), + }, + ) + for event in response.events + ) + if response.framing == "invoke" + else response.events + ) event_body: Final = b"".join( - _aws_event_frame(event.event_type, event.payload, scenario_id) for event in response.events + _aws_event_frame(event.event_type, event.payload, scenario_id, unique_id) for event in events ) return Response(content=event_body, media_type=response.content_type) diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index cd7e84f81b6..331e89c8918 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -166,6 +166,9 @@ "tests/integration/providers/test_fal_ai_video_wire.py::test_fal_video_create_status_and_content_follow_queue_wire_contract": [ "other.provider_wire.fal_ai.video_queue_create_status_and_content_download" ], + "tests/integration/providers/test_fal_ai_video_wire.py::test_fal_video_failed_result_reports_failed_status_and_fal_error": [ + "other.provider_wire.fal_ai.video_failed_result_surfaces_fal_error" + ], "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_gpt_image_25_generation_sends_quality_and_size_and_charges_keyed_row": [ "other.provider_wire.fal_ai.gpt_image_generation_quality_size_wire_and_keyed_pricing" ], @@ -175,6 +178,9 @@ "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_gpt_image_25_edit_inlines_upload_as_data_url_and_charges_keyed_row": [ "other.provider_wire.fal_ai.image_edit_json_data_urls_and_keyed_pricing" ], + "tests/integration/providers/test_fal_ai_video_wire.py::test_fal_h3_video_create_uses_canonical_body_and_status_path": [ + "other.provider_wire.fal_ai.video_queue_create_status_and_content_download" + ], "tests/integration/mcp/test_mcp_lifecycle.py::test_saved_headers_reach_real_mcp_tool_and_survive_unrelated_edit": [ "mcp.call_tool.saved_headers.reach_actual_transport" ], @@ -394,6 +400,15 @@ "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-gpt-5.6-stream_full_usage]": [ "quota_management.spend_tracking.scripted_wire.logs_cost" ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_native_json]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-upstream_500_zero_spend]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-upstream_429_zero_spend]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-haiku-4-5-input_text]": [ "quota_management.spend_tracking.cost_matrix.logs_cost" ], @@ -1324,6 +1339,333 @@ "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[us.anthropic.claude-opus-5-v1:0-stream_full_usage]": [ "quota_management.spend_tracking.scripted_wire.logs_cost" ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[whisper-next-transcriptions-per-second]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[whisper-verbose-next-transcriptions-duration]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-4o-transcribe-next-transcriptions-tokens]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[nova-next-transcriptions-per-second]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-whisper-next-transcriptions-deployment]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[tts-next-speech-per-character]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[tts-next-hd-speech-per-character]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-tts-next-speech-deployment]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dall-e-3-next-images-standard]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dall-e-3-next-images-hd]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dall-e-3-next-images-wide]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dall-e-3-next-images-two]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-image-next-images-low]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[imagen-next-images-one]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[amazon-nova-canvas-next-images-one]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-image-next-images-edit]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-text-embeddings-4-large-deployment]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-cohere-embeddings-v4]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-cohere-rerank-v4]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-embeddings-titan-v2]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[cohere-embeddings-v5]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[cohere-rerank-v4-one]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[cohere-rerank-v4-three]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[cohere-rerank-v4-total-tokens-fallback]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[fireworks-embeddings-v1]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gemini-embeddings-002]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[omni-moderations-next-list]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[omni-moderations-next-single]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[text-completions-openai-basic]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[text-completions-openai-n-best]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[text-completions-openai-stream-usage]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[text-embeddings-3-large-dimensions]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[text-embeddings-4-small-batch]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[text-embeddings-4-small-single]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[text-embeddings-4-small-token-array]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[together-completions-v1]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[together-embeddings-v1]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[vertex-embeddings-text-006]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_cache_read]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_reasoning]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_stream]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_stream_cache_read]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_incomplete]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_previous_response_id]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_web_search_medium]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.3-codex-responses_file_search]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_service_tier_flex]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_service_tier_priority]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_input_text]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_cache_read]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_cache_write_5m]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_cache_write_1h]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_web_search]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_stream]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_stream_cache_read]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_tiered_input_above_200k]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-haiku-4-5-messages_input_text]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[us.anthropic.claude-opus-5-v1:0-messages_input_text]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[anthropic.claude-sonnet-5-v1:0-messages_cache_read]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gemini-3.1-pro-passthrough-generate_content_priced_via_gemini_key]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gemini-3.1-pro-passthrough-stream_generate_content_priced_via_vertex_key]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-passthrough-messages]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-passthrough-messages_cache_read]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[anthropic.claude-sonnet-5-v1:0-passthrough-converse]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[anthropic.claude-sonnet-5-v1:0-passthrough-converse_stream]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dashscope-qwen4-max-tiered_input]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dashscope-qwen4-max-tiered_boundary_stays_lower_tier]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dashscope-qwen4-max-tiered_second_tier]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dashscope-qwen4-max-tiered_above_top_range]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gemini-gemini-3.8-flash-lite-input_below_128k]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gemini-gemini-3.8-flash-lite-input_above_128k]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-cache_creation_1h_above_200k]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[openrouter-anthropic-claude-sonnet-5-provider_reported_cost]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[openrouter-anthropic-claude-sonnet-5-token_priced]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[perplexity-sonar-next-no_search]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[deepseek-deepseek-v4-chat-prompt_cache_hit]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[deepseek-deepseek-v4-chat-no_cache_fields_bills_zero_cache]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[xai-grok-5-reasoning_folded_into_completion]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[xai-grok-5-live_search]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[xai-grok-5-provider_reported_cost]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-invoke-haiku-json]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-invoke-haiku-stream]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-converse-profile-base-model]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-converse-eu-regional-key]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-converse-apac-bare-fallback]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-converse-nova-2-pro]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-converse-mistral-large-3-stream]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-ai-gpt-5.4-mini-latest]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-ai-gpt-5.4-mini-latest-stream]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-pinned-gpt-5.4-mini-stream]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[groq-qwen-3.8-json]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[groq-qwen-3.8-stream_x_groq_recount]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[cohere-command-a-v2-tokens]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[mistral-medium-2604-json]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[openai-deployment-pricing-override]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-upstream_400_zero_spend]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-upstream_401_zero_spend]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-upstream_500_stream_request_zero_spend]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_upstream_500_zero_spend]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_upstream_500_zero_spend]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-fallback_billed_to_answering_deployment]": [ + "quota_management.spend_tracking.routing.fallback_billing" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-n_2_choices]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-finish_reason_length]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-stream_usage_in_empty_choices_chunk]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-stream_usage_in_last_delta_chunk]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-unknown_model_response_model_unknown]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-unknown_model_response_model_known]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-chat_request_to_embedding_entry]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-client_disconnect_mid_stream]": [ + "quota_management.spend_tracking.scripted_wire.client_disconnect" + ], "tests/integration/mcp/test_mcp_lifecycle.py::test_health_intersects_route_restricted_key_grants_in_both_management_modes": [ "other.mcp.health.restricted_keys_intersect_grants_in_both_modes" ], diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index f1b8901d626..bc68419f01d 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -4,6 +4,7 @@ import functools import json import os from collections.abc import Mapping +from dataclasses import dataclass from hashlib import sha256 from typing import Final @@ -13,8 +14,8 @@ from pydantic import BaseModel, ConfigDict from integration._support.client import JSON_OBJECT, Scenario, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.upstream import delete_scenario, register_scenario -from integration.cost_calculation.cost_tracking_case import CostTrackingTestCase +from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import CostTrackingTestCase, StoredResponse class CostBreakdown(BaseModel): @@ -40,22 +41,51 @@ class CostRow(BaseModel): model_config = ConfigDict(extra="ignore") spend: float | None = None + status: str | None = None prompt_tokens: int | None = None completion_tokens: int | None = None + model_id: str | None = None metadata: CostMetadata | None = None @property - def breakdown(self) -> CostBreakdown: - assert self.metadata is not None and self.metadata.cost_breakdown is not None - return self.metadata.cost_breakdown + def breakdown(self) -> CostBreakdown | None: + return self.metadata.cost_breakdown if self.metadata is not None else None + + +class FailureRow(BaseModel): + model_config = ConfigDict(extra="ignore") + + spend: float + status: str + prompt_tokens: int | None = None + completion_tokens: int | None = None + + +class DailySpend(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + spend: float + prompt_tokens: int + completion_tokens: int + api_requests: int + + +class Rollups(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + key_spend: float + team_spend: float + user_spend: float + end_user_spend: float + daily_user: DailySpend + daily_team: DailySpend def approx_equal(actual: float, expected: float) -> bool: return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) -def assert_total_is_sum_of_components(row: CostRow, context: str) -> None: - breakdown: Final = row.breakdown +def assert_total_is_sum_of_components(row: CostRow, breakdown: CostBreakdown, context: str) -> None: total: Final = sum( cost or 0.0 for cost in (breakdown.input_cost, breakdown.output_cost, breakdown.tool_usage_cost) @@ -74,7 +104,7 @@ def _row(value: Mapping[str, object]) -> CostRow | None: metadata_value: Final = value.get("metadata") metadata: Final = json.loads(metadata_value) if isinstance(metadata_value, str) else metadata_value parsed: Final = CostRow.model_validate({**value, "metadata": metadata}) - return parsed if parsed.metadata and parsed.metadata.cost_breakdown else None + return parsed if parsed.metadata is not None or (parsed.spend is not None and parsed.status is not None) else None def poll_cost_row(key: str) -> CostRow: @@ -82,7 +112,8 @@ def poll_cost_row(key: str) -> CostRow: def read() -> CostRow | None: rows: Final = read_rows( - 'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + 'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id ' + 'FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,), ) return next((parsed for row in rows if (parsed := _row(row)) is not None), None) @@ -92,6 +123,120 @@ def poll_cost_row(key: str) -> CostRow: return result +def read_rows_now(key: str) -> tuple[CostRow, ...]: + digest: Final = sha256(key.encode()).hexdigest() + rows: Final = read_rows( + 'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id ' + 'FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"', + (digest,), + ) + return tuple(parsed for row in rows if (parsed := _row(row)) is not None) + + +def poll_rows(key: str, count: int) -> tuple[CostRow, ...]: + result: Final = eventually( + lambda: read_rows_now(key), + lambda rows: len(rows) >= count, + seconds=60, + ) + return result + + +def poll_rollups( + key: str, + team_id: str, + user_id: str, + end_user_id: str, + target_spend: float, + target_requests: int, +) -> Rollups: + digest: Final = sha256(key.encode()).hexdigest() + + def read() -> Rollups | None: + key_rows: Final = read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', + (digest,), + ) + team_rows: Final = read_rows( + 'SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id=%s', + (team_id,), + ) + user_rows: Final = read_rows( + 'SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=%s', + (user_id,), + ) + end_user_rows: Final = read_rows( + 'SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', + (end_user_id,), + ) + daily_user_rows: Final = read_rows( + 'SELECT spend, prompt_tokens, completion_tokens, api_requests ' + 'FROM "LiteLLM_DailyUserSpend" WHERE user_id=%s AND api_key=%s AND date=CURRENT_DATE::text', + (user_id, digest), + ) + daily_team_rows: Final = read_rows( + 'SELECT spend, prompt_tokens, completion_tokens, api_requests ' + 'FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s AND api_key=%s AND date=CURRENT_DATE::text', + (team_id, digest), + ) + if not all((key_rows, team_rows, user_rows, end_user_rows, daily_user_rows, daily_team_rows)): + return None + rollups: Final = Rollups( + key_spend=float(key_rows[0]["spend"]), + team_spend=float(team_rows[0]["spend"]), + user_spend=float(user_rows[0]["spend"]), + end_user_spend=float(end_user_rows[0]["spend"]), + daily_user=DailySpend.model_validate(daily_user_rows[0]), + daily_team=DailySpend.model_validate(daily_team_rows[0]), + ) + return rollups + + def settled(value: Rollups | None) -> bool: + return value is not None and all( + ( + approx_equal(value.key_spend, target_spend), + approx_equal(value.team_spend, target_spend), + approx_equal(value.user_spend, target_spend), + approx_equal(value.end_user_spend, target_spend), + approx_equal(value.daily_user.spend, target_spend), + approx_equal(value.daily_team.spend, target_spend), + value.daily_user.api_requests == target_requests, + value.daily_team.api_requests == target_requests, + ) + ) + + result: Final = eventually( + read, + settled, + seconds=20, + return_last_on_timeout=True, + ) + assert result is not None + return result + + +def poll_failure_row(key: str) -> FailureRow: + digest: Final = sha256(key.encode()).hexdigest() + + def read() -> FailureRow | None: + rows: Final = read_rows( + 'SELECT spend, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (digest,), + ) + return next( + ( + parsed + for row in rows + if (parsed := FailureRow.model_validate(row)).status == "failure" + ), + None, + ) + + result: Final = eventually(read, lambda row: row is not None, seconds=60) + assert result is not None + return result + + @functools.cache def _vertex_private_key_pem() -> str: return rsa.generate_private_key(public_exponent=65537, key_size=2048).private_bytes( @@ -116,32 +261,57 @@ def _vertex_service_account_json(url: str) -> str: ) +@dataclass(frozen=True, slots=True) +class RegisteredDeployment: + model_name: str + identity: str + handle: ScenarioHandle + + def register_scenario_deployment( scenario: Scenario, case: CostTrackingTestCase, marker: str, key: str, -) -> str: + *, + response: StoredResponse | None = None, + marker_suffix: str = "", +) -> RegisteredDeployment: control_url: Final = os.environ["INTEGRATION_UPSTREAM_URL"].rstrip("/") run_marker: Final = sha256(key.encode()).hexdigest()[:12] - handle: Final = register_scenario(f"sc-{marker}-{run_marker}", case.response) + handle: Final = register_scenario( + f"sc-{marker}{marker_suffix}-{run_marker}", + case.response if response is None else response, + ) scenario.cleanups.callback(delete_scenario, handle) - model_name: Final = f"cost-{marker}-{run_marker}" + registered_model_name: Final = f"cost-{marker}{marker_suffix}-{run_marker}" parameters: Final = { "model": case.litellm_model, "api_key": case.api_key, "api_base": handle.api_base(), **case.litellm_params, + **( + { + key: value + for key, value in ( + ("input_cost_per_token", case.deployment.input_cost_per_token), + ("output_cost_per_token", case.deployment.output_cost_per_token), + ) + if value is not None + } + if case.deployment is not None + else {} + ), **( {"vertex_credentials": _vertex_service_account_json(control_url)} - if case.rates.litellm_provider == "vertex_ai-language-models" + if case.rates.litellm_provider.startswith("vertex_ai") else {} ), } created: Final = scenario.gateway.post( "/model/new", JSON_OBJECT.validate_python({ - "model_name": model_name, + "model_name": registered_model_name, "litellm_params": parameters, "model_info": ( {"base_model": case.base_model} @@ -152,4 +322,4 @@ def register_scenario_deployment( ) identity: Final = string_value(object_value(created["model_info"])["id"]) scenario.cleanups.callback(scenario.delete_model, identity) - return model_name + return RegisteredDeployment(model_name=registered_model_name, identity=identity, handle=handle) diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index 6af95f995ff..7c156937518 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -25,6 +25,14 @@ class ProviderSpecificEntry(BaseModel): us: float | None = None +class TieredPrice(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + range: tuple[float, float] + input_cost_per_token: float + output_cost_per_token: float + + class CostMapEntry(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") @@ -35,19 +43,33 @@ class CostMapEntry(BaseModel): max_output_tokens: int | None = None supports_function_calling: bool | None = None input_cost_per_token: float | None = None + input_cost_per_query: float | None = None output_cost_per_token: float | None = None + input_cost_per_token_above_128k_tokens: float | None = None + output_cost_per_token_above_128k_tokens: float | None = None + output_vector_size: int | None = None + input_cost_per_token_batches: float | None = None cache_read_input_token_cost: float | None = None cache_creation_input_token_cost: float | None = None cache_creation_input_token_cost_above_1hr: float | None = None + cache_creation_input_token_cost_above_1hr_above_200k_tokens: float | None = None cache_read_input_token_cost_above_200k_tokens: float | None = None cache_creation_input_token_cost_above_200k_tokens: float | None = None - output_cost_per_reasoning_token: float | None = None - input_cost_per_audio_token: float | None = None - output_cost_per_audio_token: float | None = None - input_cost_per_image_token: float | None = None - input_cost_per_video_token: float | None = None input_cost_per_token_above_200k_tokens: float | None = None output_cost_per_token_above_200k_tokens: float | None = None + tiered_pricing: tuple[TieredPrice, ...] | None = None + output_cost_per_reasoning_token: float | None = None + input_cost_per_audio_token: float | None = None + input_cost_per_second: float | None = None + output_cost_per_second: float | None = None + input_cost_per_character: float | None = None + output_cost_per_character: float | None = None + input_cost_per_image: float | None = None + output_cost_per_image: float | None = None + output_cost_per_audio_token: float | None = None + input_cost_per_image_token: float | None = None + output_cost_per_image_token: float | None = None + input_cost_per_video_token: float | None = None input_cost_per_token_flex: float | None = None output_cost_per_token_flex: float | None = None input_cost_per_token_priority: float | None = None @@ -64,6 +86,24 @@ class Deployment(BaseModel): model: str | None = None base_model: str | None = None + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + + +class WavUpload(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + kind: Literal["wav"] + seconds: float + + +class PngUpload(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + kind: Literal["png"] + + +Upload: TypeAlias = Annotated[WavUpload | PngUpload, Field(discriminator="kind")] class JsonResponse(BaseModel): @@ -71,6 +111,7 @@ class JsonResponse(BaseModel): content_type: Literal["application/json"] body: dict[str, JsonValue] + status: int = 200 class SseResponse(BaseModel): @@ -78,6 +119,7 @@ class SseResponse(BaseModel): content_type: Literal["text/event-stream"] frames: tuple[str, ...] + frame_delay_ms: int = Field(default=0, ge=0) class EventStreamEvent(BaseModel): @@ -92,10 +134,18 @@ class EventStreamResponse(BaseModel): content_type: Literal["application/vnd.amazon.eventstream"] events: tuple[EventStreamEvent, ...] + framing: Literal["converse", "invoke"] = "converse" + + +class BinaryResponse(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + content_type: Literal["audio/mpeg"] + length: int StoredResponse: TypeAlias = Annotated[ - JsonResponse | SseResponse | EventStreamResponse, + JsonResponse | SseResponse | EventStreamResponse | BinaryResponse, Field(discriminator="content_type"), ] @@ -108,6 +158,13 @@ class ExactExpected(BaseModel): output_cost: float prompt_tokens: int completion_tokens: int + cache_read_cost: float | None = None + cache_creation_cost: float | None = None + reasoning_cost: float | None = None + tool_usage_cost: float | None = None + breakdown_persisted: bool = True + cost_header: bool = True + rollups: bool = False class RecountRates(BaseModel): @@ -121,9 +178,25 @@ class RecountExpected(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") recount: RecountRates + prompt_tokens: int | None = None + completion_tokens: int | None = None + min_completion_tokens: int | None = None + max_completion_tokens: int | None = None -Expected: TypeAlias = ExactExpected | RecountExpected +class FailureDetails(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + status: int + + +class FailureExpected(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + failure: FailureDetails + + +Expected: TypeAlias = ExactExpected | RecountExpected | FailureExpected class CostTrackingTestCase(BaseModel): @@ -132,10 +205,29 @@ class CostTrackingTestCase(BaseModel): name: str covers: str model: str + endpoint: ( + Literal[ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages", + "/v1/embeddings", + "/v1/rerank", + "/v1/completions", + "/v1/moderations", + "/v1/audio/transcriptions", + "/v1/audio/speech", + "/v1/images/generations", + "/v1/images/edits", + ] + | Annotated[str, Field(pattern=r"^/(gemini|anthropic|bedrock)/")] + ) = "/v1/chat/completions" deployment: Deployment | None = None + upload: Upload | None = None request: dict[str, JsonValue] response: StoredResponse expected: Expected + fallback_from: StoredResponse | None = None + disconnect_after_frames: int | None = Field(default=None, ge=1) @property def rates(self) -> CostMapEntry: @@ -146,16 +238,23 @@ class CostTrackingTestCase(BaseModel): provider: Final = self.rates.litellm_provider prefix: Final = ( "openai" - if provider == "openai" and self.rates.mode == "chat" + if provider == "openai" + and ( + self.endpoint == "/v1/responses" + or self.rates.mode + in {"chat", "embedding", "moderation", "audio_transcription", "audio_speech", "image_generation"} + ) else "openai/responses" if provider == "openai" else _PROVIDER_PREFIXES.get(provider) ) if prefix is None: raise ValueError(f"unsupported cost-map provider {provider} for {self.model}") - return self.deployment.model if self.deployment and self.deployment.model is not None else ( - self.model if prefix == "" else f"{prefix}/{self.model}" - ) + if self.deployment and self.deployment.model is not None: + return self.deployment.model + if prefix == "" or self.model.startswith(f"{prefix}/"): + return self.model + return f"{prefix}/{self.model}" @property def litellm_params(self) -> Mapping[str, str]: @@ -169,6 +268,24 @@ class CostTrackingTestCase(BaseModel): def base_model(self) -> str | None: return self.deployment.base_model if self.deployment else None + @property + def passthrough_provider(self) -> Literal["gemini", "anthropic", "bedrock"] | None: + provider: Final = self.endpoint.removeprefix("/").split("/", 1)[0] + if provider == "gemini": + return "gemini" + if provider == "anthropic": + return "anthropic" + if provider == "bedrock": + return "bedrock" + return None + + @property + def reports_provider_cost(self) -> bool: + if not isinstance(self.response, JsonResponse): + return False + usage: Final = self.response.body.get("usage") + return isinstance(usage, dict) and isinstance(usage.get("cost"), (int, float)) + class _CasesFile(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") @@ -180,17 +297,39 @@ class _CasesFile(BaseModel): _PROVIDER_PREFIXES: Final[Mapping[str, str]] = MappingProxyType( { "anthropic": "anthropic", + "bedrock": "bedrock", "bedrock_converse": "bedrock/converse", + "deepgram": "deepgram", + "text-completion-openai": "text-completion-openai", + "cohere": "cohere", "vertex_ai-language-models": "vertex_ai", + "vertex_ai-image-models": "vertex_ai", + "vertex_ai-embedding-models": "vertex_ai", "gemini": "", "together_ai": "", "fireworks_ai": "", "azure": "", + "dashscope": "", + "openrouter": "", + "perplexity": "", + "deepseek": "", + "xai": "", + "azure_ai": "azure_ai", + "groq": "groq", + "mistral": "mistral", + "cohere_chat": "cohere_chat", } ) _LITELLM_PARAMS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType( { "anthropic": MappingProxyType({}), + "bedrock": MappingProxyType( + { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + } + ), "bedrock_converse": MappingProxyType( { "aws_access_key_id": "AKIASCRIPTEDPROVIDER", @@ -198,14 +337,32 @@ _LITELLM_PARAMS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType( "aws_region_name": "us-east-1", } ), + "deepgram": MappingProxyType({}), + "text-completion-openai": MappingProxyType({}), + "cohere": MappingProxyType({}), "vertex_ai-language-models": MappingProxyType( {"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"} ), + "vertex_ai-image-models": MappingProxyType( + {"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"} + ), + "vertex_ai-embedding-models": MappingProxyType( + {"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"} + ), "gemini": MappingProxyType({}), "together_ai": MappingProxyType({}), "fireworks_ai": MappingProxyType({}), "azure": MappingProxyType({"api_version": "2025-04-01-preview"}), "openai": MappingProxyType({}), + "dashscope": MappingProxyType({}), + "openrouter": MappingProxyType({}), + "perplexity": MappingProxyType({}), + "deepseek": MappingProxyType({}), + "xai": MappingProxyType({}), + "azure_ai": MappingProxyType({}), + "groq": MappingProxyType({}), + "mistral": MappingProxyType({}), + "cohere_chat": MappingProxyType({}), } ) @@ -240,6 +397,103 @@ def data_errors() -> tuple[str, ...]: or case.expected.recount.output_cost_per_token != (COST_MAP[case.model].output_cost_per_token or 0.0) ) ) + component_mismatches: Final = sorted( + case.name + for case in CASES + if isinstance(case.expected, ExactExpected) + and any( + component is not None + for component in ( + case.expected.cache_read_cost, + case.expected.cache_creation_cost, + case.expected.reasoning_cost, + case.expected.tool_usage_cost, + ) + ) + and ( + (case.expected.cache_read_cost or 0.0) + (case.expected.cache_creation_cost or 0.0) + > case.expected.input_cost + or (case.expected.reasoning_cost or 0.0) > case.expected.output_cost + or not _approx_equal( + case.expected.input_cost + + case.expected.output_cost + + (case.expected.tool_usage_cost or 0.0), + case.expected.spend, + ) + ) + ) + failure_response_mismatches: Final = sorted( + case.name + for case in CASES + if ( + isinstance(case.expected, FailureExpected) + and ( + not isinstance(case.response, JsonResponse) + or not 400 <= case.response.status <= 599 + or not 400 <= case.expected.failure.status <= 599 + ) + ) + or ( + not isinstance(case.expected, FailureExpected) + and isinstance(case.response, JsonResponse) + and case.response.status != 200 + ) + ) + invalid_opt_outs: Final = sorted( + case.name + for case in CASES + if isinstance(case.expected, ExactExpected) + and ( + ( + not case.expected.breakdown_persisted + and case.passthrough_provider is None + and case.rates.mode != "image_generation" + and not case.reports_provider_cost + ) + or ( + not case.expected.cost_header + and case.passthrough_provider is None + and not isinstance(case.response, SseResponse) + and case.expected.spend != 0.0 + ) + ) + ) + invalid_fallbacks: Final = sorted( + case.name + for case in CASES + if case.fallback_from is not None + and ( + not isinstance(case.fallback_from, JsonResponse) + or not 400 <= case.fallback_from.status <= 599 + ) + ) + invalid_disconnects: Final = sorted( + case.name + for case in CASES + if case.disconnect_after_frames is not None + and ( + not isinstance(case.response, SseResponse) + or case.response.frame_delay_ms <= 0 + or not isinstance(case.expected, RecountExpected) + ) + ) + invalid_rollup_ids: Final = sorted( + case.name + for case in CASES + if isinstance(case.expected, ExactExpected) + and case.expected.rollups + and "$UNIQUE_ID" not in case.response.model_dump_json() + ) + invalid_pinned_tool_ids: Final = sorted( + case.name + for case in CASES + if isinstance(case.expected, RecountExpected) + and (case.expected.prompt_tokens is not None or case.expected.completion_tokens is not None) + and any( + marker in case.response.model_dump_json() + for marker in ('"id": "call_$REQUEST_ID"', '"id": "toolu_$REQUEST_ID"') + ) + ) return tuple( message for message in ( @@ -248,6 +502,21 @@ def data_errors() -> tuple[str, ...]: f"duplicate case names: {duplicate_names}" if duplicate_names else None, f"cost-map entries share input_cost_per_token: {shared_input_rates}" if shared_input_rates else None, f"recount rates differ from cost-map rates: {recount_mismatches}" if recount_mismatches else None, + f"breakdown components are inconsistent: {component_mismatches}" if component_mismatches else None, + f"failure response statuses are inconsistent: {failure_response_mismatches}" + if failure_response_mismatches + else None, + f"invalid passthrough opt-outs: {invalid_opt_outs}" if invalid_opt_outs else None, + f"invalid fallback responses: {invalid_fallbacks}" if invalid_fallbacks else None, + f"invalid disconnect cases: {invalid_disconnects}" if invalid_disconnects else None, + f"rollup responses lack $UNIQUE_ID: {invalid_rollup_ids}" if invalid_rollup_ids else None, + f"pinned tool IDs contain $REQUEST_ID: {invalid_pinned_tool_ids}" + if invalid_pinned_tool_ids + else None, ) if message is not None ) + + +def _approx_equal(actual: float, expected: float) -> bool: + return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index d8b9be3a558..63b47dcba14 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -1,5 +1,82 @@ { "cost_map": { + "dashscope/qwen4-max": { + "litellm_provider": "dashscope", + "mode": "chat", + "max_input_tokens": 252000, + "max_output_tokens": 65536, + "tiered_pricing": [ + { + "range": [ + 0, + 32000 + ], + "input_cost_per_token": 1.3e-06, + "output_cost_per_token": 6.5e-06 + }, + { + "range": [ + 32000, + 128000 + ], + "input_cost_per_token": 2.6e-06, + "output_cost_per_token": 1.3e-05 + }, + { + "range": [ + 128000, + 252000 + ], + "input_cost_per_token": 3.1e-06, + "output_cost_per_token": 1.55e-05 + } + ] + }, + "gemini/gemini-3.8-flash-lite": { + "litellm_provider": "gemini", + "mode": "chat", + "input_cost_per_token": 1.1e-07, + "output_cost_per_token": 4.4e-07, + "input_cost_per_token_above_128k_tokens": 2.2e-07, + "output_cost_per_token_above_128k_tokens": 8.8e-07 + }, + "openrouter/anthropic/claude-sonnet-5": { + "litellm_provider": "openrouter", + "mode": "chat", + "input_cost_per_token": 3.2e-06, + "output_cost_per_token": 1.6e-05 + }, + "perplexity/sonar-next": { + "litellm_provider": "perplexity", + "mode": "chat", + "input_cost_per_token": 1.13e-06, + "output_cost_per_token": 1.05e-06, + "search_context_cost_per_query": { + "search_context_size_low": 0.005, + "search_context_size_medium": 0.008, + "search_context_size_high": 0.012 + } + }, + "deepseek/deepseek-v4-chat": { + "litellm_provider": "deepseek", + "mode": "chat", + "input_cost_per_token": 2.9e-07, + "output_cost_per_token": 4.3e-07, + "cache_read_input_token_cost": 2.9e-08, + "cache_creation_input_token_cost": 0.0 + }, + "xai/grok-5": { + "litellm_provider": "xai", + "mode": "chat", + "input_cost_per_token": 1.35e-06, + "output_cost_per_token": 2.7e-06, + "cache_read_input_token_cost": 2.1e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.005, + "search_context_size_medium": 0.005, + "search_context_size_high": 0.005 + } + }, "gpt-5.6": { "cache_read_input_token_cost": 1.75e-07, "input_cost_per_audio_token": 4e-05, @@ -165,6 +242,7 @@ "claude-sonnet-5": { "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, @@ -408,6 +486,234 @@ "mode": "chat", "output_cost_per_token": 3.6e-06, "supports_function_calling": true + }, + "whisper-next": { + "litellm_provider": "openai", + "mode": "audio_transcription", + "input_cost_per_second": 0.0001 + }, + "whisper-verbose-next": { + "litellm_provider": "openai", + "mode": "audio_transcription", + "input_cost_per_second": 0.0002 + }, + "gpt-4o-transcribe-next": { + "litellm_provider": "openai", + "mode": "audio_transcription", + "input_cost_per_token": 2.11e-06, + "output_cost_per_token": 3.11e-06, + "input_cost_per_audio_token": 1e-05 + }, + "nova-next": { + "litellm_provider": "deepgram", + "mode": "audio_transcription", + "input_cost_per_second": 0.0003 + }, + "azure/whisper-next": { + "litellm_provider": "azure", + "mode": "audio_transcription", + "input_cost_per_second": 0.00011 + }, + "tts-next": { + "litellm_provider": "openai", + "mode": "audio_speech", + "input_cost_per_character": 1e-05 + }, + "tts-next-hd": { + "litellm_provider": "openai", + "mode": "audio_speech", + "input_cost_per_character": 2e-05 + }, + "azure/tts-next": { + "litellm_provider": "azure", + "mode": "audio_speech", + "input_cost_per_character": 1.1e-05 + }, + "gpt-image-next": { + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_token": 1.71e-06, + "output_cost_per_token": 4.3e-06, + "input_cost_per_image_token": 2.2e-06, + "output_cost_per_image_token": 5.1e-06 + }, + "1024-x-1024/dall-e-3-next": { + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image": 0.04 + }, + "hd/1024-x-1024/dall-e-3-next": { + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image": 0.08 + }, + "1792-x-1024/dall-e-3-next": { + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image": 0.06 + }, + "low/1024-x-1024/gpt-image-next": { + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_token": 1.7e-06, + "output_cost_per_token": 4.3e-06, + "input_cost_per_image_token": 2.2e-06, + "output_cost_per_image_token": 5.1e-06 + }, + "1024-x-1024/imagen-next": { + "litellm_provider": "vertex_ai-image-models", + "mode": "image_generation", + "output_cost_per_image": 0.05 + }, + "amazon.nova-canvas-next": { + "litellm_provider": "bedrock", + "mode": "image_generation", + "output_cost_per_image": 0.045 + }, + "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0": { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 1.19e-06, + "output_cost_per_token": 5.01e-06 + }, + "eu.anthropic.claude-sonnet-5-v1:0": { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3.4e-06, + "output_cost_per_token": 1.7e-05 + }, + "amazon.nova-2-pro-preview-20251202-v1:0": { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 2.1875e-06, + "output_cost_per_token": 1.75e-05 + }, + "mistral.mistral-large-3-675b-instruct": { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 5.1e-07, + "output_cost_per_token": 1.51e-06 + }, + "azure_ai/gpt-5.4-mini-2026-03-17": { + "litellm_provider": "azure_ai", + "mode": "chat", + "input_cost_per_token": 7.5e-07, + "output_cost_per_token": 4.5e-06 + }, + "groq/qwen/qwen3.8-27b": { + "litellm_provider": "groq", + "mode": "chat", + "input_cost_per_token": 8e-07, + "output_cost_per_token": 4e-06 + }, + "cohere_chat/v2/command-a-03-2025": { + "litellm_provider": "cohere_chat", + "mode": "chat", + "input_cost_per_token": 2.51e-06, + "output_cost_per_token": 1.001e-05 + }, + "mistral/mistral-medium-2604": { + "litellm_provider": "mistral", + "mode": "chat", + "input_cost_per_token": 1.51e-06, + "output_cost_per_token": 7.51e-06 + }, + "text-embedding-3-large": { + "litellm_provider": "openai", + "mode": "embedding", + "input_cost_per_token": 1.3e-07 + }, + "text-embedding-4-small": { + "input_cost_per_token": 1.01e-06, + "output_cost_per_token": 0, + "litellm_provider": "openai", + "mode": "embedding" + }, + "text-embedding-3-large-next": { + "input_cost_per_token": 1.02e-06, + "output_cost_per_token": 0, + "litellm_provider": "openai", + "mode": "embedding" + }, + "azure/text-embedding-4-large": { + "input_cost_per_token": 1.03e-06, + "output_cost_per_token": 0, + "litellm_provider": "azure", + "mode": "embedding" + }, + "embed-v5": { + "input_cost_per_token": 1.04e-06, + "output_cost_per_token": 0, + "litellm_provider": "cohere", + "mode": "embedding" + }, + "amazon.titan-embed-text-v2:0": { + "input_cost_per_token": 1.05e-06, + "output_cost_per_token": 0, + "litellm_provider": "bedrock", + "mode": "embedding" + }, + "cohere.embed-english-v4": { + "input_cost_per_token": 1.06e-06, + "output_cost_per_token": 0, + "litellm_provider": "bedrock", + "mode": "embedding" + }, + "text-embedding-006": { + "input_cost_per_token": 1.07e-06, + "output_cost_per_token": 0, + "litellm_provider": "vertex_ai-embedding-models", + "mode": "embedding" + }, + "gemini/gemini-embedding-002": { + "input_cost_per_token": 1.08e-06, + "output_cost_per_token": 0, + "litellm_provider": "gemini", + "mode": "embedding" + }, + "together_ai/together-embed-v1": { + "input_cost_per_token": 1.09e-06, + "output_cost_per_token": 0, + "litellm_provider": "together_ai", + "mode": "embedding" + }, + "fireworks_ai/fireworks-embed-v1": { + "input_cost_per_token": 1.1e-06, + "output_cost_per_token": 0, + "litellm_provider": "fireworks_ai", + "mode": "embedding" + }, + "rerank-v4": { + "input_cost_per_token": 1.11e-06, + "output_cost_per_token": 0, + "input_cost_per_query": 0.0021, + "litellm_provider": "cohere", + "mode": "rerank" + }, + "cohere.rerank-v4:0": { + "input_cost_per_token": 1.12e-06, + "output_cost_per_token": 0, + "input_cost_per_query": 0.0022, + "litellm_provider": "bedrock", + "mode": "rerank" + }, + "gpt-3.5-turbo-instruct-next": { + "input_cost_per_token": 1.14e-06, + "output_cost_per_token": 2.14e-06, + "litellm_provider": "text-completion-openai", + "mode": "completion" + }, + "omni-moderation-next": { + "input_cost_per_token": null, + "output_cost_per_token": 0, + "litellm_provider": "openai", + "mode": "moderation" + }, + "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo": { + "input_cost_per_token": 1.16e-06, + "output_cost_per_token": 2.16e-06, + "litellm_provider": "together_ai", + "mode": "completion" } }, "cases": [ @@ -534,7 +840,8 @@ "input_cost": 0.00616704, "output_cost": 0.00627, "prompt_tokens": 12928, - "completion_tokens": 380 + "completion_tokens": 380, + "cache_read_cost": 0.00405504 } }, { @@ -605,7 +912,8 @@ "input_cost": 0.0397056, "output_cost": 0.005775, "prompt_tokens": 9728, - "completion_tokens": 350 + "completion_tokens": 350, + "cache_creation_cost": 0.038016 } }, { @@ -681,7 +989,8 @@ "input_cost": 0.0574464, "output_cost": 0.005775, "prompt_tokens": 9728, - "completion_tokens": 350 + "completion_tokens": 350, + "cache_creation_cost": 0.0557568 } }, { @@ -3417,7 +3726,8 @@ "input_cost": 0.002232, "output_cost": 0.065484, "prompt_tokens": 1240, - "completion_tokens": 4040 + "completion_tokens": 4040, + "reasoning_cost": 0.05742 } }, { @@ -3638,7 +3948,8 @@ "input_cost": 0.003312, "output_cost": 0.0059328, "prompt_tokens": 1840, - "completion_tokens": 412 + "completion_tokens": 412, + "tool_usage_cost": 0.0125 } }, { @@ -4494,7 +4805,8 @@ "input_cost": 0.0018688, "output_cost": 0.0019, "prompt_tokens": 12928, - "completion_tokens": 380 + "completion_tokens": 380, + "cache_read_cost": 0.0012288 } }, { @@ -6710,7 +7022,7 @@ "response": { "content_type": "application/json", "body": { - "id": "msg_$REQUEST_ID", + "id": "msg_$UNIQUE_ID", "type": "message", "role": "assistant", "model": "claude-sonnet-5", @@ -6732,7 +7044,8 @@ "input_cost": 0.00552, "output_cost": 0.00618, "prompt_tokens": 1840, - "completion_tokens": 412 + "completion_tokens": 412, + "rollups": true } }, { @@ -7393,7 +7706,9 @@ "recount": { "input_cost_per_token": 3e-06, "output_cost_per_token": 1.5e-05 - } + }, + "prompt_tokens": 47, + "completion_tokens": 10 } }, { @@ -7467,7 +7782,7 @@ "content_type": "text/event-stream", "frames": [ "event: message_start\ndata: {\"type\": \"message_start\", \"message\": {\"id\": \"msg_$REQUEST_ID\", \"type\": \"message\", \"role\": \"assistant\", \"model\": \"claude-sonnet-5\", \"content\": [], \"stop_reason\": null}}", - "event: content_block_start\ndata: {\"type\": \"content_block_start\", \"index\": 0, \"content_block\": {\"type\": \"tool_use\", \"id\": \"toolu_$REQUEST_ID\", \"name\": \"get_weather\", \"input\": {}}}", + "event: content_block_start\ndata: {\"type\": \"content_block_start\", \"index\": 0, \"content_block\": {\"type\": \"tool_use\", \"id\": \"call_fixture_0001\", \"name\": \"get_weather\", \"input\": {}}}", "event: content_block_delta\ndata: {\"type\": \"content_block_delta\", \"index\": 0, \"delta\": {\"type\": \"input_json_delta\", \"partial_json\": \"{\\\"city\\\": \\\"Berlin\\\", \\\"days\\\": 7, \\\"units\\\": \\\"metric\\\", \\\"notes\\\": \\\"filler filler filler filler fil\"}}", "event: content_block_delta\ndata: {\"type\": \"content_block_delta\", \"index\": 0, \"delta\": {\"type\": \"input_json_delta\", \"partial_json\": \"ler filler filler filler filler filler filler filler filler filler filler filler filler fi\"}}", "event: content_block_delta\ndata: {\"type\": \"content_block_delta\", \"index\": 0, \"delta\": {\"type\": \"input_json_delta\", \"partial_json\": \"ller filler filler filler filler filler filler filler filler filler filler filler filler \\\"}\"}}", @@ -7480,7 +7795,8 @@ "recount": { "input_cost_per_token": 3e-06, "output_cost_per_token": 1.5e-05 - } + }, + "min_completion_tokens": 60 } }, { @@ -7537,7 +7853,9 @@ "recount": { "input_cost_per_token": 3e-06, "output_cost_per_token": 1.5e-05 - } + }, + "prompt_tokens": 301, + "completion_tokens": 9 } }, { @@ -8011,6 +8329,7 @@ "spend": 0.0012456, "input_cost": 0.0010176, "output_cost": 0.000228, + "cache_read_cost": 0.0009216, "prompt_tokens": 12928, "completion_tokens": 380 } @@ -12566,7 +12885,9 @@ "recount": { "input_cost_per_token": 5.2e-07, "output_cost_per_token": 3.12e-06 - } + }, + "prompt_tokens": 48, + "completion_tokens": 12 } }, { @@ -12646,7 +12967,8 @@ "recount": { "input_cost_per_token": 5.2e-07, "output_cost_per_token": 3.12e-06 - } + }, + "min_completion_tokens": 60 } }, { @@ -12698,7 +13020,9 @@ "recount": { "input_cost_per_token": 5.2e-07, "output_cost_per_token": 3.12e-06 - } + }, + "prompt_tokens": 302, + "completion_tokens": 10 } }, { @@ -16955,7 +17279,8 @@ "input_cost": 0.00276, "output_cost": 0.004944, "prompt_tokens": 1840, - "completion_tokens": 412 + "completion_tokens": 412, + "tool_usage_cost": 0.0025 } }, { @@ -20517,7 +20842,7 @@ "response": { "content_type": "application/json", "body": { - "id": "chatcmpl-$REQUEST_ID", + "id": "chatcmpl-$UNIQUE_ID", "object": "chat.completion", "created": 1789788262, "model": "gpt-5.6", @@ -20543,7 +20868,8 @@ "input_cost": 0.00322, "output_cost": 0.005768, "prompt_tokens": 1840, - "completion_tokens": 412 + "completion_tokens": 412, + "rollups": true } }, { @@ -21298,7 +21624,9 @@ "recount": { "input_cost_per_token": 1.75e-06, "output_cost_per_token": 1.4e-05 - } + }, + "prompt_tokens": 49, + "completion_tokens": 12 } }, { @@ -21372,7 +21700,7 @@ "content_type": "text/event-stream", "frames": [ "data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"role\": \"assistant\"}, \"finish_reason\": null}], \"usage\": null}", - "data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"role\": \"assistant\", \"tool_calls\": [{\"index\": 0, \"id\": \"call_$REQUEST_ID\", \"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"arguments\": \"\"}}]}, \"finish_reason\": null}], \"usage\": null}", + "data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"role\": \"assistant\", \"tool_calls\": [{\"index\": 0, \"id\": \"call_fixture_0001\", \"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"arguments\": \"\"}}]}, \"finish_reason\": null}], \"usage\": null}", "data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"tool_calls\": [{\"index\": 0, \"function\": {\"arguments\": \"{\\\"city\\\": \\\"Berlin\\\", \\\"days\\\": 7, \\\"units\\\": \\\"metric\\\", \\\"notes\\\": \\\"filler filler filler filler fil\"}}]}, \"finish_reason\": null}], \"usage\": null}", "data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"tool_calls\": [{\"index\": 0, \"function\": {\"arguments\": \"ler filler filler filler filler filler filler filler filler filler filler filler filler fi\"}}]}, \"finish_reason\": null}], \"usage\": null}", "data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"tool_calls\": [{\"index\": 0, \"function\": {\"arguments\": \"ller filler filler filler filler filler filler filler filler filler filler filler filler \\\"}\"}}]}, \"finish_reason\": null}], \"usage\": null}", @@ -21384,7 +21712,8 @@ "recount": { "input_cost_per_token": 1.75e-06, "output_cost_per_token": 1.4e-05 - } + }, + "min_completion_tokens": 60 } }, { @@ -21439,7 +21768,9 @@ "recount": { "input_cost_per_token": 1.75e-06, "output_cost_per_token": 1.4e-05 - } + }, + "prompt_tokens": 302, + "completion_tokens": 11 } }, { @@ -21797,6 +22128,122 @@ "completion_tokens": 1592 } }, + { + "name": "gpt-5.6-responses_native_json", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/responses", + "request": { + "model": "$MODEL", + "input": "responses native fixture", + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "resp_$REQUEST_ID", + "object": "response", + "status": "completed", + "created_at": 1700000000, + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "scripted response", + "annotations": [] + } + ] + } + ], + "usage": { + "input_tokens": 11, + "output_tokens": 7, + "total_tokens": 18, + "input_tokens_details": { + "cached_tokens": 0 + }, + "output_tokens_details": { + "reasoning_tokens": 0 + } + } + } + }, + "expected": { + "spend": 0.00011725, + "input_cost": 1.925e-05, + "output_cost": 9.8e-05, + "prompt_tokens": 11, + "completion_tokens": 7 + } + }, + { + "name": "gpt-5.6-upstream_500_zero_spend", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "scripted upstream failure 500" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "status": 500, + "body": { + "error": { + "message": "scripted upstream failure", + "type": "server_error", + "code": "500" + } + } + }, + "expected": { + "failure": { + "status": 500 + } + } + }, + { + "name": "gpt-5.6-upstream_429_zero_spend", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "scripted upstream failure 429" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "status": 429, + "body": { + "error": { + "message": "scripted upstream failure", + "type": "rate_limit_error", + "code": "429" + } + } + }, + "expected": { + "failure": { + "status": 429 + } + } + }, { "name": "meta.llama4-maverick-17b-instruct-v1:0-input_text", "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", @@ -25653,6 +26100,4856 @@ "prompt_tokens": 11056, "completion_tokens": 412 } + }, + { + "name": "whisper-next-transcriptions-per-second", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "whisper-next", + "endpoint": "/v1/audio/transcriptions", + "upload": { + "kind": "wav", + "seconds": 3.5 + }, + "request": { + "language": "en", + "response_format": "json" + }, + "response": { + "content_type": "application/json", + "body": { + "text": "hello" + } + }, + "expected": { + "spend": 0.00035, + "input_cost": 0.00035, + "output_cost": 0, + "prompt_tokens": 0, + "completion_tokens": 0 + } + }, + { + "name": "whisper-verbose-next-transcriptions-duration", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "whisper-verbose-next", + "endpoint": "/v1/audio/transcriptions", + "upload": { + "kind": "wav", + "seconds": 3.5 + }, + "request": { + "response_format": "verbose_json" + }, + "response": { + "content_type": "application/json", + "body": { + "text": "hello", + "duration": 12.25 + } + }, + "expected": { + "spend": 0.00245, + "input_cost": 0.00245, + "output_cost": 0, + "prompt_tokens": 0, + "completion_tokens": 0 + } + }, + { + "name": "gpt-4o-transcribe-next-transcriptions-tokens", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-4o-transcribe-next", + "endpoint": "/v1/audio/transcriptions", + "upload": { + "kind": "wav", + "seconds": 1.0 + }, + "request": { + "response_format": "json" + }, + "response": { + "content_type": "application/json", + "body": { + "text": "hello", + "usage": { + "type": "tokens", + "input_tokens": 10, + "output_tokens": 2, + "total_tokens": 12, + "input_token_details": { + "text_tokens": 2, + "audio_tokens": 8 + } + } + } + }, + "expected": { + "spend": 9.044e-05, + "input_cost": 8.422e-05, + "output_cost": 6.22e-06, + "prompt_tokens": 10, + "completion_tokens": 2 + } + }, + { + "name": "nova-next-transcriptions-per-second", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "nova-next", + "endpoint": "/v1/audio/transcriptions", + "upload": { + "kind": "wav", + "seconds": 4.0 + }, + "request": {}, + "response": { + "content_type": "application/json", + "body": { + "results": { + "channels": [ + { + "alternatives": [ + { + "transcript": "hello", + "confidence": 0.9 + } + ] + } + ] + }, + "metadata": { + "duration": 4.0, + "channels": 1 + } + } + }, + "expected": { + "spend": 0.0012, + "input_cost": 0.0012, + "output_cost": 0, + "prompt_tokens": 0, + "completion_tokens": 0 + } + }, + { + "name": "azure-whisper-next-transcriptions-deployment", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "azure/whisper-next", + "endpoint": "/v1/audio/transcriptions", + "deployment": { + "model": "azure/cc-whisper-deployment", + "base_model": "azure/whisper-next" + }, + "upload": { + "kind": "wav", + "seconds": 3.5 + }, + "request": { + "response_format": "json" + }, + "response": { + "content_type": "application/json", + "body": { + "text": "hello" + } + }, + "expected": { + "spend": 0.000385, + "input_cost": 0.000385, + "output_cost": 0, + "prompt_tokens": 0, + "completion_tokens": 0 + } + }, + { + "name": "tts-next-speech-per-character", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "tts-next", + "endpoint": "/v1/audio/speech", + "request": { + "input": "hello world", + "voice": "alloy", + "response_format": "mp3" + }, + "response": { + "content_type": "audio/mpeg", + "length": 2048 + }, + "expected": { + "spend": 0.0001, + "input_cost": 0.0001, + "output_cost": 0, + "prompt_tokens": 0, + "completion_tokens": 0 + } + }, + { + "name": "tts-next-hd-speech-per-character", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "tts-next-hd", + "endpoint": "/v1/audio/speech", + "request": { + "input": "hello world", + "voice": "alloy", + "response_format": "mp3" + }, + "response": { + "content_type": "audio/mpeg", + "length": 2048 + }, + "expected": { + "spend": 0.0002, + "input_cost": 0.0002, + "output_cost": 0, + "prompt_tokens": 0, + "completion_tokens": 0 + } + }, + { + "name": "azure-tts-next-speech-deployment", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "azure/tts-next", + "endpoint": "/v1/audio/speech", + "deployment": { + "model": "azure/cc-tts-deployment", + "base_model": "azure/tts-next" + }, + "request": { + "input": "hello world", + "voice": "alloy", + "response_format": "mp3" + }, + "response": { + "content_type": "audio/mpeg", + "length": 2048 + }, + "expected": { + "spend": 0.00011, + "input_cost": 0.00011, + "output_cost": 0, + "prompt_tokens": 0, + "completion_tokens": 0 + } + }, + { + "name": "dall-e-3-next-images-standard", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "1024-x-1024/dall-e-3-next", + "endpoint": "/v1/images/generations", + "deployment": { + "model": "openai/dall-e-3-next" + }, + "request": { + "prompt": "a deterministic square", + "size": "1024x1024", + "quality": "standard", + "n": 1 + }, + "response": { + "content_type": "application/json", + "body": { + "created": 1700000000, + "data": [ + { + "url": "https://x/1.png" + } + ] + } + }, + "expected": { + "spend": 0.04, + "input_cost": 0.04, + "output_cost": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "breakdown_persisted": false + } + }, + { + "name": "dall-e-3-next-images-hd", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "hd/1024-x-1024/dall-e-3-next", + "endpoint": "/v1/images/generations", + "deployment": { + "model": "openai/dall-e-3-next" + }, + "request": { + "prompt": "a deterministic square", + "size": "1024x1024", + "quality": "hd", + "n": 1 + }, + "response": { + "content_type": "application/json", + "body": { + "created": 1700000001, + "data": [ + { + "url": "https://x/1.png" + } + ] + } + }, + "expected": { + "spend": 0.08, + "input_cost": 0.08, + "output_cost": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "breakdown_persisted": false + } + }, + { + "name": "dall-e-3-next-images-wide", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "1792-x-1024/dall-e-3-next", + "endpoint": "/v1/images/generations", + "deployment": { + "model": "openai/dall-e-3-next" + }, + "request": { + "prompt": "a deterministic wide image", + "size": "1792x1024", + "quality": "standard", + "n": 1 + }, + "response": { + "content_type": "application/json", + "body": { + "created": 1700000002, + "data": [ + { + "url": "https://x/1.png" + } + ] + } + }, + "expected": { + "spend": 0.06, + "input_cost": 0.06, + "output_cost": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "breakdown_persisted": false + } + }, + { + "name": "dall-e-3-next-images-two", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "1024-x-1024/dall-e-3-next", + "endpoint": "/v1/images/generations", + "deployment": { + "model": "openai/dall-e-3-next" + }, + "request": { + "prompt": "two deterministic squares", + "size": "1024x1024", + "quality": "standard", + "n": 2 + }, + "response": { + "content_type": "application/json", + "body": { + "created": 1700000003, + "data": [ + { + "url": "https://x/1.png" + }, + { + "url": "https://x/2.png" + } + ] + } + }, + "expected": { + "spend": 0.08, + "input_cost": 0.08, + "output_cost": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "breakdown_persisted": false + } + }, + { + "name": "gpt-image-next-images-low", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-image-next", + "endpoint": "/v1/images/generations", + "deployment": { + "model": "openai/gpt-image-next" + }, + "request": { + "prompt": "a deterministic generated image", + "size": "1024x1024", + "quality": "low", + "n": 1 + }, + "response": { + "content_type": "application/json", + "body": { + "created": 1700000004, + "data": [ + { + "b64_json": "AA==" + } + ], + "usage": { + "total_tokens": 30, + "input_tokens": 10, + "output_tokens": 20, + "input_tokens_details": { + "text_tokens": 10, + "image_tokens": 0 + } + } + } + }, + "expected": { + "spend": 0.0001191, + "input_cost": 1.71e-05, + "output_cost": 0.000102, + "prompt_tokens": 10, + "completion_tokens": 20, + "breakdown_persisted": false + } + }, + { + "name": "imagen-next-images-one", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "1024-x-1024/imagen-next", + "endpoint": "/v1/images/generations", + "request": { + "prompt": "a deterministic vertex image", + "sampleCount": 1 + }, + "response": { + "content_type": "application/json", + "body": { + "predictions": [ + { + "bytesBase64Encoded": "AA==", + "mimeType": "image/png" + } + ] + } + }, + "expected": { + "spend": 0.05, + "input_cost": 0.05, + "output_cost": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "breakdown_persisted": false + } + }, + { + "name": "amazon-nova-canvas-next-images-one", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "amazon.nova-canvas-next", + "endpoint": "/v1/images/generations", + "deployment": { + "model": "amazon.nova-canvas-next" + }, + "request": { + "prompt": "a deterministic bedrock image" + }, + "response": { + "content_type": "application/json", + "body": { + "images": [ + "AA==" + ] + } + }, + "expected": { + "spend": 0.045, + "input_cost": 0.045, + "output_cost": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "breakdown_persisted": false + } + }, + { + "name": "gpt-image-next-images-edit", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "low/1024-x-1024/gpt-image-next", + "endpoint": "/v1/images/edits", + "deployment": { + "model": "openai/gpt-image-next" + }, + "upload": { + "kind": "png" + }, + "request": { + "prompt": "edit this deterministic image", + "size": "1024x1024", + "quality": "low", + "n": 1 + }, + "response": { + "content_type": "application/json", + "body": { + "created": 1700000005, + "data": [ + { + "b64_json": "AA==" + } + ], + "usage": { + "total_tokens": 30, + "input_tokens": 10, + "output_tokens": 20, + "input_tokens_details": { + "text_tokens": 10, + "image_tokens": 0 + } + } + } + }, + "expected": { + "spend": 0.000119, + "input_cost": 1.7e-05, + "output_cost": 0.000102, + "prompt_tokens": 10, + "completion_tokens": 20, + "breakdown_persisted": false + } + }, + { + "name": "text-embeddings-4-small-single", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "text-embedding-4-small", + "endpoint": "/v1/embeddings", + "request": { + "model": "$MODEL", + "input": "one embedding" + }, + "response": { + "content_type": "application/json", + "body": { + "object": "list", + "data": [ + { + "object": "embedding", + "embedding": [ + 0.1, + 0.2, + 0.3 + ], + "index": 0 + } + ], + "model": "text-embedding-4-small", + "usage": { + "prompt_tokens": 7, + "total_tokens": 7 + } + } + }, + "expected": { + "spend": 7.07e-06, + "input_cost": 7.07e-06, + "output_cost": 0.0, + "prompt_tokens": 7, + "completion_tokens": 0 + } + }, + { + "name": "text-embeddings-4-small-batch", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "text-embedding-4-small", + "endpoint": "/v1/embeddings", + "request": { + "model": "$MODEL", + "input": [ + "one", + "two", + "three" + ] + }, + "response": { + "content_type": "application/json", + "body": { + "object": "list", + "data": [ + { + "object": "embedding", + "embedding": [ + 0.1, + 0.2, + 0.3 + ], + "index": 0 + }, + { + "object": "embedding", + "embedding": [ + 0.1, + 0.2, + 0.3 + ], + "index": 1 + }, + { + "object": "embedding", + "embedding": [ + 0.1, + 0.2, + 0.3 + ], + "index": 2 + } + ], + "model": "text-embedding-4-small", + "usage": { + "prompt_tokens": 21, + "total_tokens": 21 + } + } + }, + "expected": { + "spend": 2.1210000000000002e-05, + "input_cost": 2.1210000000000002e-05, + "output_cost": 0.0, + "prompt_tokens": 21, + "completion_tokens": 0 + } + }, + { + "name": "text-embeddings-4-small-token-array", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "text-embedding-4-small", + "endpoint": "/v1/embeddings", + "request": { + "model": "$MODEL", + "input": [ + 1, + 2, + 3, + 4 + ] + }, + "response": { + "content_type": "application/json", + "body": { + "object": "list", + "data": [ + { + "object": "embedding", + "embedding": [ + 0.1, + 0.2, + 0.3 + ], + "index": 0 + } + ], + "model": "text-embedding-4-small", + "usage": { + "prompt_tokens": 9, + "total_tokens": 9 + } + } + }, + "expected": { + "spend": 9.090000000000001e-06, + "input_cost": 9.090000000000001e-06, + "output_cost": 0.0, + "prompt_tokens": 9, + "completion_tokens": 0 + } + }, + { + "name": "text-embeddings-3-large-dimensions", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "text-embedding-3-large-next", + "endpoint": "/v1/embeddings", + "request": { + "model": "$MODEL", + "input": "large embedding", + "dimensions": 3 + }, + "response": { + "content_type": "application/json", + "body": { + "object": "list", + "data": [ + { + "object": "embedding", + "embedding": [ + 0.1, + 0.2, + 0.3 + ], + "index": 0 + } + ], + "model": "text-embedding-3-large-next", + "usage": { + "prompt_tokens": 8, + "total_tokens": 8 + } + } + }, + "expected": { + "spend": 8.16e-06, + "input_cost": 8.16e-06, + "output_cost": 0.0, + "prompt_tokens": 8, + "completion_tokens": 0 + } + }, + { + "name": "azure-text-embeddings-4-large-deployment", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "azure/text-embedding-4-large", + "endpoint": "/v1/embeddings", + "request": { + "model": "$MODEL", + "input": "azure embedding" + }, + "response": { + "content_type": "application/json", + "body": { + "object": "list", + "data": [ + { + "object": "embedding", + "embedding": [ + 0.1, + 0.2, + 0.3 + ], + "index": 0 + } + ], + "model": "azure/text-embedding-4-large", + "usage": { + "prompt_tokens": 8, + "total_tokens": 8 + } + } + }, + "expected": { + "spend": 8.24e-06, + "input_cost": 8.24e-06, + "output_cost": 0.0, + "prompt_tokens": 8, + "completion_tokens": 0 + }, + "deployment": { + "model": "azure/cc-pinned-embedding-deployment", + "base_model": "azure/text-embedding-4-large" + } + }, + { + "name": "cohere-embeddings-v5", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "embed-v5", + "endpoint": "/v1/embeddings", + "request": { + "model": "$MODEL", + "input": "cohere embedding", + "input_type": "search_query" + }, + "response": { + "content_type": "application/json", + "body": { + "id": "emb-1", + "embeddings": { + "float": [ + [ + 0.1, + 0.2, + 0.3 + ] + ] + }, + "meta": { + "billed_units": { + "input_tokens": 11 + } + } + } + }, + "expected": { + "spend": 1.144e-05, + "input_cost": 1.144e-05, + "output_cost": 0.0, + "prompt_tokens": 11, + "completion_tokens": 0 + } + }, + { + "name": "bedrock-embeddings-titan-v2", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "amazon.titan-embed-text-v2:0", + "endpoint": "/v1/embeddings", + "request": { + "model": "$MODEL", + "input": "titan embedding" + }, + "response": { + "content_type": "application/json", + "body": { + "embedding": [ + 0.1, + 0.2, + 0.3 + ], + "inputTextTokenCount": 10 + } + }, + "expected": { + "spend": 1.05e-05, + "input_cost": 1.05e-05, + "output_cost": 0.0, + "prompt_tokens": 10, + "completion_tokens": 0 + } + }, + { + "name": "bedrock-cohere-embeddings-v4", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "cohere.embed-english-v4", + "endpoint": "/v1/embeddings", + "request": { + "model": "$MODEL", + "input": "bedrock cohere embedding" + }, + "response": { + "content_type": "application/json", + "body": { + "embeddings": [ + [ + 0.1, + 0.2, + 0.3 + ] + ], + "id": "emb-bedrock-cohere-1", + "response_type": "embeddings_floats", + "texts": [ + "bedrock cohere embedding" + ] + } + }, + "expected": { + "spend": 5.3e-06, + "input_cost": 5.3e-06, + "output_cost": 0.0, + "prompt_tokens": 5, + "completion_tokens": 0 + } + }, + { + "name": "vertex-embeddings-text-006", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "text-embedding-006", + "endpoint": "/v1/embeddings", + "request": { + "model": "$MODEL", + "input": "vertex embedding" + }, + "response": { + "content_type": "application/json", + "body": { + "predictions": [ + { + "embeddings": { + "values": [ + 0.1, + 0.2, + 0.3 + ], + "statistics": { + "token_count": 7, + "truncated": false + } + } + } + ] + } + }, + "expected": { + "spend": 7.4899999999999994e-06, + "input_cost": 7.4899999999999994e-06, + "output_cost": 0.0, + "prompt_tokens": 7, + "completion_tokens": 0 + } + }, + { + "name": "gemini-embeddings-002", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gemini/gemini-embedding-002", + "endpoint": "/v1/embeddings", + "request": { + "model": "$MODEL", + "input": "gemini embedding" + }, + "response": { + "content_type": "application/json", + "body": { + "embeddings": [ + { + "values": [ + 0.1, + 0.2, + 0.3 + ] + } + ], + "usageMetadata": { + "promptTokenCount": 7, + "totalTokenCount": 7 + } + } + }, + "expected": { + "spend": 3.24e-06, + "input_cost": 3.24e-06, + "output_cost": 0.0, + "prompt_tokens": 3, + "completion_tokens": 0 + } + }, + { + "name": "together-embeddings-v1", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "together_ai/together-embed-v1", + "endpoint": "/v1/embeddings", + "request": { + "model": "$MODEL", + "input": "together embedding" + }, + "response": { + "content_type": "application/json", + "body": { + "object": "list", + "data": [ + { + "object": "embedding", + "embedding": [ + 0.1, + 0.2, + 0.3 + ], + "index": 0 + } + ], + "model": "together-embed-v1", + "usage": { + "prompt_tokens": 7, + "total_tokens": 7 + } + } + }, + "expected": { + "spend": 7.63e-06, + "input_cost": 7.63e-06, + "output_cost": 0.0, + "prompt_tokens": 7, + "completion_tokens": 0 + } + }, + { + "name": "fireworks-embeddings-v1", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "fireworks_ai/fireworks-embed-v1", + "endpoint": "/v1/embeddings", + "request": { + "model": "$MODEL", + "input": "fireworks embedding" + }, + "response": { + "content_type": "application/json", + "body": { + "object": "list", + "data": [ + { + "object": "embedding", + "embedding": [ + 0.1, + 0.2, + 0.3 + ], + "index": 0 + } + ], + "model": "fireworks-embed-v1", + "usage": { + "prompt_tokens": 7, + "total_tokens": 7 + } + } + }, + "expected": { + "spend": 7.7e-06, + "input_cost": 7.7e-06, + "output_cost": 0.0, + "prompt_tokens": 7, + "completion_tokens": 0 + } + }, + { + "name": "cohere-rerank-v4-one", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "rerank-v4", + "endpoint": "/v1/rerank", + "request": { + "model": "$MODEL", + "query": "rank this", + "documents": [ + "a", + "b" + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "rr-$REQUEST_ID", + "results": [ + { + "index": 0, + "relevance_score": 0.9 + } + ], + "meta": { + "api_version": { + "version": "2" + }, + "billed_units": { + "search_units": 1 + } + } + } + }, + "expected": { + "spend": 0.0021, + "input_cost": 0.0021, + "output_cost": 0.0, + "prompt_tokens": 0, + "completion_tokens": 0 + } + }, + { + "name": "cohere-rerank-v4-three", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "rerank-v4", + "endpoint": "/v1/rerank", + "request": { + "model": "$MODEL", + "query": "rank this", + "documents": [ + "a long document", + "another long document", + "third long document" + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "rr-three-$REQUEST_ID", + "results": [ + { + "index": 0, + "relevance_score": 0.9 + } + ], + "meta": { + "api_version": { + "version": "2" + }, + "billed_units": { + "search_units": 3 + } + } + } + }, + "expected": { + "spend": 0.0063, + "input_cost": 0.0063, + "output_cost": 0.0, + "prompt_tokens": 0, + "completion_tokens": 0 + } + }, + { + "name": "cohere-rerank-v4-total-tokens-fallback", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "rerank-v4", + "endpoint": "/v1/rerank", + "request": { + "model": "$MODEL", + "query": "rank this", + "documents": [ + "fallback a", + "fallback b" + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "rr-fallback-$REQUEST_ID", + "results": [ + { + "index": 0, + "relevance_score": 0.8 + } + ], + "meta": { + "billed_units": { + "total_tokens": 99 + } + } + } + }, + "expected": { + "spend": 0.0, + "input_cost": 0.0, + "output_cost": 0.0, + "prompt_tokens": 0, + "completion_tokens": 0 + } + }, + { + "name": "bedrock-cohere-rerank-v4", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "cohere.rerank-v4:0", + "endpoint": "/v1/rerank", + "request": { + "model": "$MODEL", + "query": "rank this", + "documents": [ + "a", + "b" + ] + }, + "response": { + "content_type": "application/json", + "body": { + "results": [ + { + "index": 0, + "relevanceScore": 0.9 + } + ], + "response_id": "rr-3", + "token_count": 1 + } + }, + "expected": { + "spend": 0.0022, + "input_cost": 0.0022, + "output_cost": 0.0, + "prompt_tokens": 0, + "completion_tokens": 0 + } + }, + { + "name": "text-completions-openai-basic", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-3.5-turbo-instruct-next", + "endpoint": "/v1/completions", + "request": { + "model": "$MODEL", + "prompt": "complete this" + }, + "response": { + "content_type": "application/json", + "body": { + "id": "cmpl-basic-$REQUEST_ID", + "object": "text_completion", + "choices": [ + { + "text": "done", + "index": 0, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 4, + "total_tokens": 13 + } + } + }, + "expected": { + "spend": 1.882e-05, + "input_cost": 1.026e-05, + "output_cost": 8.56e-06, + "prompt_tokens": 9, + "completion_tokens": 4 + } + }, + { + "name": "text-completions-openai-stream-usage", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-3.5-turbo-instruct-next", + "endpoint": "/v1/completions", + "request": { + "model": "$MODEL", + "prompt": "complete this", + "stream": true, + "stream_options": { + "include_usage": true + } + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "data: {\"id\": \"cmpl-$REQUEST_ID\", \"object\": \"text_completion\", \"created\": 1789789000, \"model\": \"gpt-3.5-turbo-instruct-next\", \"choices\": [{\"text\": \"done\", \"index\": 0, \"finish_reason\": null}], \"usage\": null}", + "data: {\"id\": \"cmpl-$REQUEST_ID\", \"object\": \"text_completion\", \"created\": 1789789000, \"model\": \"gpt-3.5-turbo-instruct-next\", \"choices\": [], \"usage\": {\"prompt_tokens\": 9, \"completion_tokens\": 4, \"total_tokens\": 13}}", + "data: [DONE]" + ] + }, + "expected": { + "spend": 1.882e-05, + "input_cost": 1.026e-05, + "output_cost": 8.56e-06, + "prompt_tokens": 9, + "completion_tokens": 4 + } + }, + { + "name": "text-completions-openai-n-best", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-3.5-turbo-instruct-next", + "endpoint": "/v1/completions", + "request": { + "model": "$MODEL", + "prompt": "complete this twice", + "n": 2 + }, + "response": { + "content_type": "application/json", + "body": { + "id": "cmpl-n-best-$REQUEST_ID", + "object": "text_completion", + "choices": [ + { + "text": "done", + "index": 0, + "finish_reason": "stop" + }, + { + "text": "also done", + "index": 1, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 8, + "total_tokens": 17 + } + } + }, + "expected": { + "spend": 2.738e-05, + "input_cost": 1.026e-05, + "output_cost": 1.712e-05, + "prompt_tokens": 9, + "completion_tokens": 8 + } + }, + { + "name": "together-completions-v1", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", + "endpoint": "/v1/completions", + "request": { + "model": "$MODEL", + "prompt": "together complete" + }, + "response": { + "content_type": "application/json", + "body": { + "id": "cmpl-together-$REQUEST_ID", + "object": "text_completion", + "choices": [ + { + "text": "done", + "index": 0, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 4, + "total_tokens": 13 + } + } + }, + "expected": { + "spend": 1.908e-05, + "input_cost": 1.0439999999999998e-05, + "output_cost": 8.64e-06, + "prompt_tokens": 9, + "completion_tokens": 4 + } + }, + { + "name": "omni-moderations-next-single", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "omni-moderation-next", + "endpoint": "/v1/moderations", + "request": { + "model": "$MODEL", + "input": "safe text" + }, + "response": { + "content_type": "application/json", + "body": { + "id": "modr-single-$REQUEST_ID", + "model": "omni-moderation-next", + "results": [ + { + "flagged": false, + "categories": {}, + "category_scores": {} + } + ] + } + }, + "expected": { + "spend": 0.0, + "input_cost": 0.0, + "output_cost": 0.0, + "prompt_tokens": 0, + "completion_tokens": 0 + } + }, + { + "name": "omni-moderations-next-list", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "omni-moderation-next", + "endpoint": "/v1/moderations", + "request": { + "model": "$MODEL", + "input": [ + "safe text", + "more safe text" + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "modr-list-$REQUEST_ID", + "model": "omni-moderation-next", + "results": [ + { + "flagged": false, + "categories": {}, + "category_scores": {} + }, + { + "flagged": false, + "categories": {}, + "category_scores": {} + } + ] + } + }, + "expected": { + "spend": 0.0, + "input_cost": 0.0, + "output_cost": 0.0, + "prompt_tokens": 0, + "completion_tokens": 0 + } + }, + { + "name": "gpt-5.6-responses_cache_read", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/responses", + "request": { + "model": "$MODEL", + "input": "summarize this text" + }, + "response": { + "content_type": "application/json", + "body": { + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "scripted response", + "annotations": [] + } + ] + } + ], + "usage": { + "input_tokens": 12928, + "output_tokens": 380, + "total_tokens": 13308, + "input_tokens_details": { + "cached_tokens": 12288 + }, + "output_tokens_details": { + "reasoning_tokens": 0 + } + } + } + }, + "expected": { + "spend": 0.0085904, + "input_cost": 0.0032704, + "output_cost": 0.00532, + "prompt_tokens": 12928, + "completion_tokens": 380, + "cache_read_cost": 0.0021504 + } + }, + { + "name": "gpt-5.6-responses_reasoning", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/responses", + "request": { + "model": "$MODEL", + "input": "reason about this text" + }, + "response": { + "content_type": "application/json", + "body": { + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "reasoning", + "id": "rs_$REQUEST_ID", + "status": "completed", + "summary": [] + }, + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "scripted response", + "annotations": [] + } + ] + } + ], + "usage": { + "input_tokens": 1240, + "output_tokens": 4040, + "total_tokens": 5280, + "input_tokens_details": { + "cached_tokens": 0 + }, + "output_tokens_details": { + "reasoning_tokens": 3480 + } + } + } + }, + "expected": { + "spend": 0.06569, + "input_cost": 0.00217, + "output_cost": 0.06352, + "prompt_tokens": 1240, + "completion_tokens": 4040, + "reasoning_cost": 0.05568 + } + }, + { + "name": "gpt-5.6-responses_stream", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/responses", + "request": { + "model": "$MODEL", + "input": "stream this text", + "stream": true + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_$REQUEST_ID\",\"object\":\"response\",\"created_at\":1700000000,\"status\":\"in_progress\",\"model\":\"gpt-5.6\",\"output\":[{\"type\":\"message\",\"id\":\"msg_$REQUEST_ID\",\"status\":\"completed\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"scripted response\",\"annotations\":[]}]}],\"usage\":null}}", + "event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"message\",\"id\":\"msg_$REQUEST_ID\",\"status\":\"in_progress\",\"role\":\"assistant\",\"content\":[]}}", + "event: response.content_part.added\ndata: {\"type\":\"response.content_part.added\",\"item_id\":\"msg_$REQUEST_ID\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"content_part\",\"text\":\"\"}}", + "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"item_id\":\"msg_$REQUEST_ID\",\"output_index\":0,\"content_index\":0,\"delta\":\"scripted \"}", + "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"item_id\":\"msg_$REQUEST_ID\",\"output_index\":0,\"content_index\":0,\"delta\":\"response\"}", + "event: response.output_text.done\ndata: {\"type\":\"response.output_text.done\",\"item_id\":\"msg_$REQUEST_ID\",\"output_index\":0,\"content_index\":0,\"text\":\"scripted response\"}", + "event: response.content_part.done\ndata: {\"type\":\"response.content_part.done\",\"item_id\":\"msg_$REQUEST_ID\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"output_text\",\"text\":\"scripted response\",\"annotations\":[]}}", + "event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"message\",\"id\":\"msg_$REQUEST_ID\",\"status\":\"completed\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"scripted response\",\"annotations\":[]}]}}", + "event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_$REQUEST_ID\",\"object\":\"response\",\"created_at\":1700000000,\"status\":\"completed\",\"model\":\"gpt-5.6\",\"output\":[{\"type\":\"message\",\"id\":\"msg_$REQUEST_ID\",\"status\":\"completed\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"scripted response\",\"annotations\":[]}]}],\"usage\":{\"input_tokens\":1840,\"output_tokens\":412,\"total_tokens\":2252,\"input_tokens_details\":{\"cached_tokens\":0},\"output_tokens_details\":{\"reasoning_tokens\":0}}}}" + ] + }, + "expected": { + "spend": 0.008988, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gpt-5.6-responses_stream_cache_read", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/responses", + "request": { + "model": "$MODEL", + "input": "stream cached text", + "stream": true + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_$REQUEST_ID\",\"object\":\"response\",\"created_at\":1700000000,\"status\":\"in_progress\",\"model\":\"gpt-5.6\",\"output\":[{\"type\":\"message\",\"id\":\"msg_$REQUEST_ID\",\"status\":\"completed\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"scripted response\",\"annotations\":[]}]}],\"usage\":null}}", + "event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"message\",\"id\":\"msg_$REQUEST_ID\",\"status\":\"in_progress\",\"role\":\"assistant\",\"content\":[]}}", + "event: response.content_part.added\ndata: {\"type\":\"response.content_part.added\",\"item_id\":\"msg_$REQUEST_ID\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"content_part\",\"text\":\"\"}}", + "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"item_id\":\"msg_$REQUEST_ID\",\"output_index\":0,\"content_index\":0,\"delta\":\"scripted \"}", + "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"item_id\":\"msg_$REQUEST_ID\",\"output_index\":0,\"content_index\":0,\"delta\":\"response\"}", + "event: response.output_text.done\ndata: {\"type\":\"response.output_text.done\",\"item_id\":\"msg_$REQUEST_ID\",\"output_index\":0,\"content_index\":0,\"text\":\"scripted response\"}", + "event: response.content_part.done\ndata: {\"type\":\"response.content_part.done\",\"item_id\":\"msg_$REQUEST_ID\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"output_text\",\"text\":\"scripted response\",\"annotations\":[]}}", + "event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"message\",\"id\":\"msg_$REQUEST_ID\",\"status\":\"completed\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"scripted response\",\"annotations\":[]}]}}", + "event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_$REQUEST_ID\",\"object\":\"response\",\"created_at\":1700000000,\"status\":\"completed\",\"model\":\"gpt-5.6\",\"output\":[{\"type\":\"message\",\"id\":\"msg_$REQUEST_ID\",\"status\":\"completed\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"scripted response\",\"annotations\":[]}]}],\"usage\":{\"input_tokens\":12928,\"output_tokens\":380,\"total_tokens\":13308,\"input_tokens_details\":{\"cached_tokens\":12288},\"output_tokens_details\":{\"reasoning_tokens\":0}}}}" + ] + }, + "expected": { + "spend": 0.0085904, + "input_cost": 0.0032704, + "output_cost": 0.00532, + "prompt_tokens": 12928, + "completion_tokens": 380, + "cache_read_cost": 0.0021504 + } + }, + { + "name": "gpt-5.6-responses_incomplete", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/responses", + "request": { + "model": "$MODEL", + "input": "truncate this text", + "max_output_tokens": 100 + }, + "response": { + "content_type": "application/json", + "body": { + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "incomplete", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "scripted response", + "annotations": [] + } + ] + } + ], + "usage": { + "input_tokens": 1840, + "output_tokens": 100, + "total_tokens": 1940, + "input_tokens_details": { + "cached_tokens": 0 + }, + "output_tokens_details": { + "reasoning_tokens": 0 + } + }, + "incomplete_details": { + "reason": "max_output_tokens" + } + } + }, + "expected": { + "spend": 0.00462, + "input_cost": 0.00322, + "output_cost": 0.0014, + "prompt_tokens": 1840, + "completion_tokens": 100 + } + }, + { + "name": "gpt-5.6-responses_previous_response_id", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/responses", + "request": { + "model": "$MODEL", + "input": "continue this text", + "previous_response_id": "resp_scripted_prior" + }, + "response": { + "content_type": "application/json", + "body": { + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "scripted response", + "annotations": [] + } + ] + } + ], + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "total_tokens": 2252, + "input_tokens_details": { + "cached_tokens": 0 + }, + "output_tokens_details": { + "reasoning_tokens": 0 + } + } + } + }, + "expected": { + "spend": 0.008988, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gpt-5.6-responses_web_search_medium", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/responses", + "request": { + "model": "$MODEL", + "input": "search this text", + "tools": [ + { + "type": "web_search_preview", + "search_context_size": "medium" + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "web_search_call", + "id": "ws_$REQUEST_ID", + "status": "completed", + "action": { + "type": "search", + "query": "scripted query" + } + }, + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "scripted response", + "annotations": [] + } + ] + } + ], + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "total_tokens": 2252, + "input_tokens_details": { + "cached_tokens": 0 + }, + "output_tokens_details": { + "reasoning_tokens": 0 + } + } + } + }, + "expected": { + "spend": 0.021488, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412, + "tool_usage_cost": 0.0125 + } + }, + { + "name": "gpt-5.3-codex-responses_file_search", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.3-codex", + "endpoint": "/v1/responses", + "request": { + "model": "$MODEL", + "input": "search files", + "tools": [ + { + "type": "file_search", + "vector_store_ids": [ + "vs_scripted" + ] + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.3-codex", + "output": [ + { + "type": "file_search_call", + "id": "fs_$REQUEST_ID", + "status": "completed", + "queries": [ + "scripted query" + ], + "results": [] + }, + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "scripted response", + "annotations": [] + } + ] + } + ], + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "total_tokens": 2252, + "input_tokens_details": { + "cached_tokens": 0 + }, + "output_tokens_details": { + "reasoning_tokens": 0 + } + } + } + }, + "expected": { + "spend": 0.010204, + "input_cost": 0.00276, + "output_cost": 0.004944, + "prompt_tokens": 1840, + "completion_tokens": 412, + "tool_usage_cost": 0.0025 + } + }, + { + "name": "gpt-5.6-responses_service_tier_flex", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/responses", + "request": { + "model": "$MODEL", + "input": "flex text", + "service_tier": "flex" + }, + "response": { + "content_type": "application/json", + "body": { + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "scripted response", + "annotations": [] + } + ] + } + ], + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "total_tokens": 2252, + "input_tokens_details": { + "cached_tokens": 0 + }, + "output_tokens_details": { + "reasoning_tokens": 0 + }, + "service_tier": "flex" + }, + "service_tier": "flex" + } + }, + "expected": { + "spend": 0.004494, + "input_cost": 0.00161, + "output_cost": 0.002884, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gpt-5.6-responses_service_tier_priority", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/responses", + "request": { + "model": "$MODEL", + "input": "priority text", + "service_tier": "priority" + }, + "response": { + "content_type": "application/json", + "body": { + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "scripted response", + "annotations": [] + } + ] + } + ], + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "total_tokens": 2252, + "input_tokens_details": { + "cached_tokens": 0 + }, + "output_tokens_details": { + "reasoning_tokens": 0 + }, + "service_tier": "priority" + }, + "service_tier": "priority" + } + }, + "expected": { + "spend": 0.017976, + "input_cost": 0.00644, + "output_cost": 0.011536, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "claude-sonnet-5-messages_input_text", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "claude-sonnet-5", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [ + { + "role": "user", + "content": "summarize the attached material in one line" + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [ + { + "type": "text", + "text": "scripted response" + } + ], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 1840, + "output_tokens": 412 + } + } + }, + "expected": { + "spend": 0.0117, + "input_cost": 0.00552, + "output_cost": 0.00618, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "claude-sonnet-5-messages_cache_read", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "claude-sonnet-5", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "cached text", + "cache_control": { + "type": "ephemeral" + } + } + ] + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [ + { + "type": "text", + "text": "scripted response" + } + ], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 640, + "output_tokens": 380, + "cache_read_input_tokens": 12288 + } + } + }, + "expected": { + "spend": 0.0113064, + "input_cost": 0.0056064, + "output_cost": 0.0057, + "prompt_tokens": 12928, + "completion_tokens": 380, + "cache_read_cost": 0.0036864 + } + }, + { + "name": "claude-sonnet-5-messages_cache_write_5m", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "claude-sonnet-5", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 350, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "cache this text", + "cache_control": { + "type": "ephemeral" + } + } + ] + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [ + { + "type": "text", + "text": "scripted response" + } + ], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 512, + "output_tokens": 350, + "cache_creation_input_tokens": 9216, + "cache_creation": { + "ephemeral_5m_input_tokens": 9216, + "ephemeral_1h_input_tokens": 0 + } + } + } + }, + "expected": { + "spend": 0.041346, + "input_cost": 0.036096, + "output_cost": 0.00525, + "prompt_tokens": 9728, + "completion_tokens": 350, + "cache_creation_cost": 0.03456 + } + }, + { + "name": "claude-sonnet-5-messages_cache_write_1h", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "claude-sonnet-5", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 350, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "cache this text for an hour", + "cache_control": { + "type": "ephemeral", + "ttl": "1h" + } + } + ] + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [ + { + "type": "text", + "text": "scripted response" + } + ], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 512, + "output_tokens": 350, + "cache_creation_input_tokens": 9216, + "cache_creation": { + "ephemeral_5m_input_tokens": 2048, + "ephemeral_1h_input_tokens": 7168 + } + } + } + }, + "expected": { + "spend": 0.057474, + "input_cost": 0.052224, + "output_cost": 0.00525, + "prompt_tokens": 9728, + "completion_tokens": 350, + "cache_creation_cost": 0.050688 + } + }, + { + "name": "claude-sonnet-5-messages_web_search", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "claude-sonnet-5", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [ + { + "role": "user", + "content": "summarize the attached material in one line" + } + ], + "tools": [ + { + "type": "web_search_20250305", + "name": "web_search" + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [ + { + "type": "server_tool_use", + "id": "srv_$REQUEST_ID", + "name": "web_search", + "input": { + "query": "scripted query" + } + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srv_$REQUEST_ID", + "content": [ + { + "type": "web_search_result", + "title": "scripted result", + "url": "https://scripted.example" + } + ] + }, + { + "type": "text", + "text": "scripted response" + } + ], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "server_tool_use": { + "web_search_requests": 2 + } + } + } + }, + "expected": { + "spend": 0.0317, + "input_cost": 0.00552, + "output_cost": 0.00618, + "prompt_tokens": 1840, + "completion_tokens": 412, + "tool_usage_cost": 0.02 + } + }, + { + "name": "claude-sonnet-5-messages_stream", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "claude-sonnet-5", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [ + { + "role": "user", + "content": "summarize the attached material in one line" + } + ], + "stream": true + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_$REQUEST_ID\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-sonnet-5\",\"content\":[],\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":1840}}}", + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"scripted \"}}", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"response\"}}", + "event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}", + "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":412}}", + "event: message_stop\ndata: {\"type\":\"message_stop\"}" + ] + }, + "expected": { + "spend": 0.0117, + "input_cost": 0.00552, + "output_cost": 0.00618, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "claude-sonnet-5-messages_stream_cache_read", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "claude-sonnet-5", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 380, + "messages": [ + { + "role": "user", + "content": "summarize the attached material in one line" + } + ], + "stream": true + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_$REQUEST_ID\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-sonnet-5\",\"content\":[],\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":640,\"cache_read_input_tokens\":12288}}}", + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"scripted \"}}", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"response\"}}", + "event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}", + "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":380}}", + "event: message_stop\ndata: {\"type\":\"message_stop\"}" + ] + }, + "expected": { + "spend": 0.0113064, + "input_cost": 0.0056064, + "output_cost": 0.0057, + "prompt_tokens": 12928, + "completion_tokens": 380, + "cache_read_cost": 0.0036864 + } + }, + { + "name": "claude-sonnet-5-messages_tiered_input_above_200k", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "claude-sonnet-5", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 620, + "messages": [ + { + "role": "user", + "content": "summarize the attached material in one line" + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [ + { + "type": "text", + "text": "scripted response" + } + ], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 210000, + "output_tokens": 620 + } + } + }, + "expected": { + "spend": 1.27395, + "input_cost": 1.26, + "output_cost": 0.01395, + "prompt_tokens": 210000, + "completion_tokens": 620 + } + }, + { + "name": "claude-haiku-4-5-messages_input_text", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "claude-haiku-4-5", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [ + { + "role": "user", + "content": "summarize the attached material in one line" + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [ + { + "type": "text", + "text": "scripted response" + } + ], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 1840, + "output_tokens": 412 + } + } + }, + "expected": { + "spend": 0.0039, + "input_cost": 0.00184, + "output_cost": 0.00206, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "us.anthropic.claude-opus-5-v1:0-messages_input_text", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "us.anthropic.claude-opus-5-v1:0", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [ + { + "role": "user", + "content": "summarize the attached material in one line" + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "text": "scripted response" + } + ] + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1840, + "outputTokens": 412 + }, + "metrics": { + "latencyMs": 42 + } + } + }, + "expected": { + "spend": 0.02145, + "input_cost": 0.01012, + "output_cost": 0.01133, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "anthropic.claude-sonnet-5-v1:0-messages_cache_read", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "anthropic.claude-sonnet-5-v1:0", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [ + { + "role": "user", + "content": "summarize the attached material in one line" + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "text": "scripted response" + } + ] + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 640, + "outputTokens": 380, + "totalTokens": 13308, + "cacheReadInputTokens": 12288 + }, + "metrics": { + "latencyMs": 42 + } + } + }, + "expected": { + "spend": 0.01243704, + "input_cost": 0.00616704, + "output_cost": 0.00627, + "prompt_tokens": 12928, + "completion_tokens": 380, + "cache_read_cost": 0.00405504 + } + }, + { + "name": "gemini-3.1-pro-passthrough-generate_content_priced_via_gemini_key", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gemini/gemini-3.1-pro", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line." + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "5fdf6b7dd9b9 summarize the attached material in one line and name the city weather" + } + ] + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "scripted answer 5fdf6b7dd9b9" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0 + } + ], + "usageMetadata": { + "promptTokenCount": 1840, + "candidatesTokenCount": 412, + "totalTokenCount": 2252, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 1840 + } + ] + }, + "modelVersion": "gemini-3.1-pro" + } + }, + "expected": { + "spend": 0.008624, + "input_cost": 0.00368, + "output_cost": 0.004944, + "prompt_tokens": 1840, + "completion_tokens": 412, + "breakdown_persisted": false, + "cost_header": false + }, + "endpoint": "/gemini/v1beta/models/$MODEL:generateContent" + }, + { + "name": "gemini-3.1-pro-passthrough-stream_generate_content_priced_via_vertex_key", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "gemini-3.1-pro", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line." + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "52b6a80ff038 summarize the attached material in one line and name the city weather" + } + ] + } + ], + "stream": true, + "stream_options": { + "include_usage": true + }, + "allowed_openai_params": [] + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "data: {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \"scripted answer 52b6a80ff038\"}], \"role\": \"model\"}, \"finishReason\": \"STOP\", \"index\": 0}], \"modelVersion\": \"gemini-3.1-pro\"}", + "data: {\"candidates\": [], \"usageMetadata\": {\"promptTokenCount\": 1840, \"candidatesTokenCount\": 412, \"totalTokenCount\": 2252, \"promptTokensDetails\": [{\"modality\": \"TEXT\", \"tokenCount\": 1840}]}, \"modelVersion\": \"gemini-3.1-pro\"}" + ] + }, + "expected": { + "spend": 0.0090552, + "input_cost": 0.003864, + "output_cost": 0.0051912, + "prompt_tokens": 1840, + "completion_tokens": 412, + "breakdown_persisted": false, + "cost_header": false + }, + "endpoint": "/gemini/v1beta/models/$MODEL:streamGenerateContent?alt=sse" + }, + { + "name": "claude-sonnet-5-passthrough-messages", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "claude-sonnet-5", + "endpoint": "/anthropic/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [ + { + "role": "user", + "content": "summarize the attached material in one line" + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [ + { + "type": "text", + "text": "scripted response" + } + ], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 1840, + "output_tokens": 412 + } + } + }, + "expected": { + "spend": 0.0117, + "input_cost": 0.00552, + "output_cost": 0.00618, + "prompt_tokens": 1840, + "completion_tokens": 412, + "breakdown_persisted": false, + "cost_header": false + } + }, + { + "name": "claude-sonnet-5-passthrough-messages_cache_read", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "claude-sonnet-5", + "endpoint": "/anthropic/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "cached text", + "cache_control": { + "type": "ephemeral" + } + } + ] + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [ + { + "type": "text", + "text": "scripted response" + } + ], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 640, + "output_tokens": 380, + "cache_read_input_tokens": 12288 + } + } + }, + "expected": { + "spend": 0.0113064, + "input_cost": 0.0056064, + "output_cost": 0.0057, + "prompt_tokens": 12928, + "completion_tokens": 380, + "cache_read_cost": 0.0036864, + "breakdown_persisted": false, + "cost_header": false + } + }, + { + "name": "anthropic.claude-sonnet-5-v1:0-passthrough-converse", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "anthropic.claude-sonnet-5-v1:0", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line." + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "9aad4de0556c summarize the attached material in one line and name the city weather" + } + ] + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "text": "scripted answer 9aad4de0556c" + } + ] + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252 + }, + "metrics": { + "latencyMs": 42 + } + } + }, + "expected": { + "spend": 0.01287, + "input_cost": 0.006072, + "output_cost": 0.006798, + "prompt_tokens": 1840, + "completion_tokens": 412, + "cost_header": false + }, + "endpoint": "/bedrock/model/$MODEL/converse" + }, + { + "name": "anthropic.claude-sonnet-5-v1:0-passthrough-converse_stream", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "anthropic.claude-sonnet-5-v1:0", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line." + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "a9257967d38a summarize the attached material in one line and name the city weather" + } + ] + } + ], + "stream": true, + "stream_options": { + "include_usage": true + }, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/vnd.amazon.eventstream", + "events": [ + { + "event_type": "messageStart", + "payload": { + "role": "assistant" + } + }, + { + "event_type": "contentBlockDelta", + "payload": { + "delta": { + "text": "scripted answer a9257967d38a" + }, + "contentBlockIndex": 0 + } + }, + { + "event_type": "contentBlockStop", + "payload": { + "contentBlockIndex": 0 + } + }, + { + "event_type": "messageStop", + "payload": { + "stopReason": "end_turn" + } + }, + { + "event_type": "metadata", + "payload": { + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252 + }, + "metrics": { + "latencyMs": 42 + } + } + } + ] + }, + "expected": { + "spend": 0.01287, + "input_cost": 0.006072, + "output_cost": 0.006798, + "prompt_tokens": 1840, + "completion_tokens": 412, + "cost_header": false + }, + "endpoint": "/bedrock/model/$MODEL/converse-stream" + }, + { + "name": "dashscope-qwen4-max-tiered_input", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "dashscope/qwen4-max", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "tiered input" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "$REQUEST_ID", + "object": "chat.completion", + "model": "qwen4-max", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "ok" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "expected": { + "spend": 0.00507, + "input_cost": 0.002392, + "output_cost": 0.002678, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "dashscope-qwen4-max-tiered_boundary_stays_lower_tier", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "dashscope/qwen4-max", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "tier boundary" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "$REQUEST_ID", + "object": "chat.completion", + "model": "qwen4-max", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "ok" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 32000, + "completion_tokens": 412, + "total_tokens": 32412 + } + } + }, + "expected": { + "spend": 0.044278, + "input_cost": 0.0416, + "output_cost": 0.002678, + "prompt_tokens": 32000, + "completion_tokens": 412 + } + }, + { + "name": "dashscope-qwen4-max-tiered_second_tier", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "dashscope/qwen4-max", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "tier two" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "$REQUEST_ID", + "object": "chat.completion", + "model": "qwen4-max", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "ok" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 40000, + "completion_tokens": 412, + "total_tokens": 40412 + } + } + }, + "expected": { + "spend": 0.109356, + "input_cost": 0.104, + "output_cost": 0.005356, + "prompt_tokens": 40000, + "completion_tokens": 412 + } + }, + { + "name": "dashscope-qwen4-max-tiered_above_top_range", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "dashscope/qwen4-max", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "top tier" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "$REQUEST_ID", + "object": "chat.completion", + "model": "qwen4-max", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "ok" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 300000, + "completion_tokens": 412, + "total_tokens": 300412 + } + } + }, + "expected": { + "spend": 0.936386, + "input_cost": 0.93, + "output_cost": 0.006386, + "prompt_tokens": 300000, + "completion_tokens": 412 + } + }, + { + "name": "gemini-gemini-3.8-flash-lite-input_below_128k", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gemini/gemini-3.8-flash-lite", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "base pricing" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "ok" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0 + } + ], + "usageMetadata": { + "promptTokenCount": 1840, + "candidatesTokenCount": 412, + "totalTokenCount": 2252 + }, + "modelVersion": "gemini-3.8-flash-lite" + } + }, + "expected": { + "spend": 0.00038368, + "input_cost": 0.0002024, + "output_cost": 0.00018128, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gemini-gemini-3.8-flash-lite-input_above_128k", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gemini/gemini-3.8-flash-lite", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "above threshold" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "ok" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0 + } + ], + "usageMetadata": { + "promptTokenCount": 130000, + "candidatesTokenCount": 412, + "totalTokenCount": 130412 + }, + "modelVersion": "gemini-3.8-flash-lite" + } + }, + "expected": { + "spend": 0.02896256, + "input_cost": 0.0286, + "output_cost": 0.00036256, + "prompt_tokens": 130000, + "completion_tokens": 412 + } + }, + { + "name": "claude-sonnet-5-cache_creation_1h_above_200k", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "claude-sonnet-5", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "one hour cache" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [ + { + "type": "text", + "text": "ok" + } + ], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 150000, + "cache_creation_input_tokens": 60000, + "cache_creation": { + "ephemeral_5m_input_tokens": 0, + "ephemeral_1h_input_tokens": 60000 + }, + "cache_read_input_tokens": 0, + "output_tokens": 412 + } + } + }, + "expected": { + "spend": 1.62927, + "input_cost": 1.62, + "output_cost": 0.00927, + "cache_creation_cost": 0.72, + "prompt_tokens": 210000, + "completion_tokens": 412 + } + }, + { + "name": "openrouter-anthropic-claude-sonnet-5-provider_reported_cost", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "openrouter/anthropic/claude-sonnet-5", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "reported cost" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "anthropic/claude-sonnet-5", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "ok" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "cost": 0.0421 + } + } + }, + "expected": { + "spend": 0.0421, + "input_cost": 0.0, + "output_cost": 0.0421, + "prompt_tokens": 1840, + "completion_tokens": 412, + "breakdown_persisted": false + } + }, + { + "name": "openrouter-anthropic-claude-sonnet-5-token_priced", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "openrouter/anthropic/claude-sonnet-5", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "token pricing" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "anthropic/claude-sonnet-5", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "ok" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "expected": { + "spend": 0.01248, + "input_cost": 0.005888, + "output_cost": 0.006592, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "perplexity-sonar-next-no_search", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "perplexity/sonar-next", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "no search" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "sonar-next", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "ok" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "expected": {"spend": 0.0025118, "input_cost": 0.0020792, "output_cost": 0.0004326, "prompt_tokens": 1840, "completion_tokens": 412} + }, + { + "name": "deepseek-deepseek-v4-chat-prompt_cache_hit", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "deepseek/deepseek-v4-chat", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "cache hit" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "deepseek-v4-chat", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "ok" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_cache_hit_tokens": 1200, + "prompt_cache_miss_tokens": 640, + "prompt_tokens_details": { + "cached_tokens": 1200 + } + } + } + }, + "expected": { + "spend": 0.00039756, + "input_cost": 0.0002204, + "output_cost": 0.00017716, + "cache_read_cost": 3.48e-05, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "deepseek-deepseek-v4-chat-no_cache_fields_bills_zero_cache", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "deepseek/deepseek-v4-chat", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "no cache" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "deepseek-v4-chat", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "ok" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "expected": { + "spend": 0.00071076, + "input_cost": 0.0005336, + "output_cost": 0.00017716, + "cache_read_cost": 0.0, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "xai-grok-5-reasoning_folded_into_completion", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "xai/grok-5", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "reasoning" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "grok-5", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "ok" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2552, + "completion_tokens_details": { + "reasoning_tokens": 300 + } + } + } + }, + "expected": { + "spend": 0.0044064, + "input_cost": 0.002484, + "output_cost": 0.0019224, + "prompt_tokens": 1840, + "completion_tokens": 712 + } + }, + { + "name": "xai-grok-5-live_search", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "xai/grok-5", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "live search" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "grok-5", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "ok" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "server_side_tool_usage_details": { + "web_search_calls": 2 + } + } + } + }, + "expected": { + "spend": 0.0135964, + "input_cost": 0.002484, + "output_cost": 0.0011124, + "tool_usage_cost": 0.01, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "xai-grok-5-provider_reported_cost", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "xai/grok-5", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "reported xai cost" + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "grok-5", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "ok" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "cost": 0.0421 + } + } + }, + "expected": { + "spend": 0.0421, + "input_cost": 0.0, + "output_cost": 0.0421, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "bedrock-invoke-haiku-json", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "bedrock-invoke-haiku-json" + } + ], + "stream": false, + "max_tokens": 412 + }, + "response": { + "content_type": "application/json", + "body": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "anthropic.claude-haiku-4-5-20251001-v1:0", + "content": [ + { + "type": "text", + "text": "scripted response" + } + ], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 1840, + "output_tokens": 412 + } + } + }, + "expected": { + "spend": 0.00425372, + "input_cost": 0.0021896, + "output_cost": 0.00206412, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "bedrock-invoke-haiku-stream", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "bedrock-invoke-haiku-stream" + } + ], + "stream": true, + "max_tokens": 412 + }, + "response": { + "content_type": "application/vnd.amazon.eventstream", + "framing": "invoke", + "events": [ + { + "event_type": "message_start", + "payload": { + "type": "message_start", + "message": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "anthropic.claude-haiku-4-5-20251001-v1:0", + "content": [], + "stop_reason": null, + "stop_sequence": null, + "usage": { + "input_tokens": 1840, + "output_tokens": 0 + } + } + } + }, + { + "event_type": "content_block_delta", + "payload": { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "text_delta", + "text": "scripted response" + } + } + }, + { + "event_type": "message_delta", + "payload": { + "type": "message_delta", + "delta": { + "stop_reason": "end_turn" + }, + "usage": { + "output_tokens": 412 + } + } + }, + { + "event_type": "message_stop", + "payload": { + "type": "message_stop" + } + } + ] + }, + "expected": { + "spend": 0.00425372, + "input_cost": 0.0021896, + "output_cost": 0.00206412, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "bedrock-converse-profile-base-model", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "anthropic.claude-sonnet-5-v1:0", + "deployment": { + "model": "bedrock/converse/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + "base_model": "anthropic.claude-sonnet-5-v1:0" + }, + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "bedrock-converse-profile-base-model" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "text": "scripted response" + } + ] + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252 + }, + "metrics": { + "latencyMs": 1 + } + } + }, + "expected": { + "spend": 0.01287, + "input_cost": 0.006072, + "output_cost": 0.006798, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "bedrock-converse-eu-regional-key", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "eu.anthropic.claude-sonnet-5-v1:0", + "deployment": { + "model": "bedrock/converse/eu.anthropic.claude-sonnet-5-v1:0" + }, + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "bedrock-converse-eu-regional-key" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "text": "scripted response" + } + ] + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252 + }, + "metrics": { + "latencyMs": 1 + } + } + }, + "expected": { + "spend": 0.01326, + "input_cost": 0.006256, + "output_cost": 0.007004, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "bedrock-converse-apac-bare-fallback", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "anthropic.claude-sonnet-5-v1:0", + "deployment": { + "model": "bedrock/converse/apac.anthropic.claude-sonnet-5-v1:0" + }, + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "bedrock-converse-apac-bare-fallback" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "text": "scripted response" + } + ] + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252 + }, + "metrics": { + "latencyMs": 1 + } + } + }, + "expected": { + "spend": 0.01287, + "input_cost": 0.006072, + "output_cost": 0.006798, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "bedrock-converse-nova-2-pro", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "amazon.nova-2-pro-preview-20251202-v1:0", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "bedrock-converse-nova-2-pro" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "text": "scripted response" + } + ] + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252 + }, + "metrics": { + "latencyMs": 1 + } + } + }, + "expected": { + "spend": 0.011235, + "input_cost": 0.004025, + "output_cost": 0.00721, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "bedrock-converse-mistral-large-3-stream", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "mistral.mistral-large-3-675b-instruct", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line." + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "bedrock-converse-mistral-large-3-stream" + } + ] + } + ], + "stream": true, + "stream_options": { + "include_usage": true + }, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/vnd.amazon.eventstream", + "events": [ + { + "event_type": "messageStart", + "payload": { + "role": "assistant" + } + }, + { + "event_type": "contentBlockDelta", + "payload": { + "delta": { + "text": "scripted response" + }, + "contentBlockIndex": 0 + } + }, + { + "event_type": "contentBlockStop", + "payload": { + "contentBlockIndex": 0 + } + }, + { + "event_type": "messageStop", + "payload": { + "stopReason": "end_turn" + } + }, + { + "event_type": "metadata", + "payload": { + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252 + }, + "metrics": { + "latencyMs": 1 + } + } + } + ] + }, + "expected": { + "spend": 0.00156052, + "input_cost": 0.0009384, + "output_cost": 0.00062212, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "azure-ai-gpt-5.4-mini-latest", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "azure_ai/gpt-5.4-mini-2026-03-17", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "azure-ai-gpt-5.4-mini-latest" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "$MODEL", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted response" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + }, + "service_tier": "default" + } + }, + "expected": { + "spend": 0.003234, + "input_cost": 0.00138, + "output_cost": 0.001854, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "azure-ai-gpt-5.4-mini-latest-stream", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "azure_ai/gpt-5.4-mini-2026-03-17", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "azure-ai-gpt-5.4-mini-latest-stream" + } + ], + "stream": true + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"$MODEL\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"ok\"},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"$MODEL\",\"choices\":[],\"usage\":{\"prompt_tokens\":1840,\"completion_tokens\":412,\"total_tokens\":2252}}\n\n", + "data: [DONE]\n\n" + ] + }, + "expected": { + "spend": 0.003234, + "input_cost": 0.00138, + "output_cost": 0.001854, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "azure-pinned-gpt-5.4-mini-stream", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "azure/gpt-5.4-mini", + "deployment": { + "model": "azure/cc-pinned-deployment", + "base_model": "azure/gpt-5.4-mini" + }, + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "azure-pinned-gpt-5.4-mini-stream" + } + ], + "stream": true + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"$MODEL\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"ok\"},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"$MODEL\",\"choices\":[],\"usage\":{\"prompt_tokens\":1840,\"completion_tokens\":412,\"total_tokens\":2252}}\n\n", + "data: [DONE]\n\n" + ] + }, + "expected": { + "spend": 0.00184896, + "input_cost": 0.0006624, + "output_cost": 0.00118656, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "groq-qwen-3.8-json", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "groq/qwen/qwen3.8-27b", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "groq-qwen-3.8-json" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "$MODEL", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted response" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + }, + "service_tier": "default", + "x_groq": { + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + } + }, + "expected": { + "spend": 0.00312, + "input_cost": 0.001472, + "output_cost": 0.001648, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "groq-qwen-3.8-stream_x_groq_recount", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "groq/qwen/qwen3.8-27b", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "groq-qwen-3.8-stream_x_groq_recount" + } + ], + "stream": true + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"$MODEL\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"ok\"},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"$MODEL\",\"choices\":[],\"x_groq\":{\"usage\":{\"prompt_tokens\":1840,\"completion_tokens\":412,\"total_tokens\":2252}}}\n\n", + "data: [DONE]\n\n" + ] + }, + "expected": { + "recount": { + "input_cost_per_token": 8e-07, + "output_cost_per_token": 4e-06 + } + } + }, + { + "name": "cohere-command-a-v2-tokens", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "cohere_chat/v2/command-a-03-2025", + "deployment": { + "model": "cohere_chat/v2/command-a-03-2025" + }, + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "cohere-command-a-v2-tokens" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "$REQUEST_ID", + "message": { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "scripted response" + } + ] + }, + "finish_reason": "COMPLETE", + "usage": { + "tokens": { + "input_tokens": 1840, + "output_tokens": 412, + "total_tokens": 2252 + }, + "billed_units": { + "input_tokens": 1800, + "output_tokens": 400 + } + } + } + }, + "expected": { + "spend": 0.00874252, + "input_cost": 0.0046184, + "output_cost": 0.00412412, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "mistral-medium-2604-json", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "mistral/mistral-medium-2604", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "mistral-medium-2604-json" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "$MODEL", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted response" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + }, + "service_tier": "default" + } + }, + "expected": { + "spend": 0.00587252, + "input_cost": 0.0027784, + "output_cost": 0.00309412, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "openai-deployment-pricing-override", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "deployment": { + "model": "openai/cc-custom-model", + "input_cost_per_token": 7e-06, + "output_cost_per_token": 2.1e-05 + }, + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "openai-deployment-pricing-override" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "$MODEL", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted response" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + }, + "service_tier": "default" + } + }, + "expected": { + "spend": 0.021532, + "input_cost": 0.01288, + "output_cost": 0.008652, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gpt-5.6-upstream_400_zero_spend", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "status": 400, + "body": { + "error": { + "message": "scripted upstream failure 400", + "type": "server_error", + "code": "400" + } + } + }, + "expected": { + "failure": { + "status": 400 + } + } + }, + { + "name": "gpt-5.6-upstream_401_zero_spend", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "status": 401, + "body": { + "error": { + "message": "scripted upstream failure 401", + "type": "server_error", + "code": "401" + } + } + }, + "expected": { + "failure": { + "status": 401 + } + } + }, + { + "name": "gpt-5.6-upstream_500_stream_request_zero_spend", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": true + }, + "response": { + "content_type": "application/json", + "status": 500, + "body": { + "error": { + "message": "scripted upstream failure 500", + "type": "server_error", + "code": "500" + } + } + }, + "expected": { + "failure": { + "status": 500 + } + } + }, + { + "name": "gpt-5.6-responses_upstream_500_zero_spend", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/responses", + "request": { + "model": "$MODEL", + "input": "proxy behaviour probe", + "stream": false + }, + "response": { + "content_type": "application/json", + "status": 500, + "body": { + "error": { + "message": "scripted upstream failure 500", + "type": "server_error", + "code": "500" + } + } + }, + "expected": { + "failure": { + "status": 500 + } + } + }, + { + "name": "claude-sonnet-5-messages_upstream_500_zero_spend", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ] + }, + "response": { + "content_type": "application/json", + "status": 500, + "body": { + "error": { + "message": "scripted upstream failure 500", + "type": "server_error", + "code": "500" + } + } + }, + "expected": { + "failure": { + "status": 500 + } + } + }, + { + "name": "gpt-5.6-fallback_billed_to_answering_deployment", + "covers": "quota_management.spend_tracking.routing.fallback_billing", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1789788262, + "model": "gpt-5.6", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "fallback_from": { + "content_type": "application/json", + "status": 500, + "body": { + "error": { + "message": "scripted upstream failure 500", + "type": "server_error", + "code": "500" + } + } + }, + "expected": { + "spend": 0.008988, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gpt-5.6-n_2_choices", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1789788262, + "model": "gpt-5.6", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer" + }, + "finish_reason": "stop" + }, + { + "index": 1, + "message": { + "role": "assistant", + "content": "second choice" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "expected": { + "spend": 0.008988, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gpt-5.6-finish_reason_length", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1789788262, + "model": "gpt-5.6", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "truncated" + }, + "finish_reason": "length" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "expected": { + "spend": 0.008988, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gpt-5.6-stream_usage_in_empty_choices_chunk", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": true, + "stream_options": { + "include_usage": true + } + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"scripted answer\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[],\"usage\":{\"prompt_tokens\":1840,\"completion_tokens\":412,\"total_tokens\":2252}}", + "data: [DONE]" + ] + }, + "expected": { + "spend": 0.008988, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412, + "cost_header": false + } + }, + { + "name": "gpt-5.6-stream_usage_in_last_delta_chunk", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": true, + "stream_options": { + "include_usage": true + } + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"scripted answer\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":1840,\"completion_tokens\":412,\"total_tokens\":2252}}", + "data: [DONE]" + ] + }, + "expected": { + "spend": 0.008988, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412, + "cost_header": false + } + }, + { + "name": "gpt-5.6-unknown_model_response_model_unknown", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "deployment": { + "model": "openai/not-in-any-map-xyz" + }, + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1789788262, + "model": "not-in-any-map-xyz", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "expected": { + "spend": 0.0, + "input_cost": 0.0, + "output_cost": 0.0, + "prompt_tokens": 1840, + "completion_tokens": 412, + "cost_header": false + } + }, + { + "name": "gpt-5.6-unknown_model_response_model_known", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "deployment": { + "model": "openai/not-in-any-map-xyz" + }, + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1789788262, + "model": "gpt-5.6", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "expected": { + "spend": 0.008988, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gpt-5.6-chat_request_to_embedding_entry", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "text-embedding-3-large", + "deployment": { + "model": "openai/text-embedding-3-large" + }, + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1789788262, + "model": "text-embedding-3-large", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "expected": { + "spend": 0.0002392, + "input_cost": 0.0002392, + "output_cost": 0.0, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gpt-5.6-client_disconnect_mid_stream", + "covers": "quota_management.spend_tracking.scripted_wire.client_disconnect", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": true, + "stream_options": { + "include_usage": true + } + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-0\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-1\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-2\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-3\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-4\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-5\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-6\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-7\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-8\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-9\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-10\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-11\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-12\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-13\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-14\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-15\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-16\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-17\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-18\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-19\"},\"finish_reason\":null}],\"usage\":null}", + "data: [DONE]" + ], + "frame_delay_ms": 200 + }, + "disconnect_after_frames": 3, + "expected": { + "recount": { + "input_cost_per_token": 1.75e-06, + "output_cost_per_token": 1.4e-05 + }, + "prompt_tokens": 10, + "min_completion_tokens": 9, + "max_completion_tokens": 30 + } } ] } diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index a8a56fbfbbd..d646278efff 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -2,25 +2,43 @@ from __future__ import annotations +import io +import json +import struct +import time +import uuid +import wave +import zlib from hashlib import sha256 +from itertools import islice from typing import Final, cast +import httpx import pytest - from integration._support.client import JSON_OBJECT, Gateway +from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.conftest import ( + CostBreakdown, + CostRow, approx_equal, assert_total_is_sum_of_components, poll_cost_row, + poll_failure_row, + poll_rollups, + poll_rows, + read_rows_now, register_scenario_deployment, ) from integration.cost_calculation.cost_tracking_case import ( CASES, + BinaryResponse, CostTrackingTestCase, ExactExpected, + FailureExpected, RecountExpected, data_errors, ) +from pydantic import JsonValue if _data_errors := data_errors(): raise ValueError("\n".join(_data_errors)) @@ -32,6 +50,47 @@ _CASES: Final = tuple( ) +def _wav_bytes(seconds: float) -> bytes: + frame_count: Final = round(16000 * seconds) + output: Final = io.BytesIO() + with wave.open(output, "wb") as wav: + wav.setnchannels(1) + wav.setsampwidth(2) + wav.setframerate(16000) + wav.writeframes(b"\x00\x00" * frame_count) + return output.getvalue() + + +def _png_bytes() -> bytes: + def chunk(kind: bytes, payload: bytes) -> bytes: + return ( + struct.pack(">I", len(payload)) + + kind + + payload + + struct.pack(">I", zlib.crc32(kind + payload) & 0xFFFFFFFF) + ) + + return ( + b"\x89PNG\r\n\x1a\n" + + chunk(b"IHDR", struct.pack(">IIBBBBB", 1, 1, 8, 6, 0, 0, 0)) + + chunk(b"IDAT", zlib.compress(b"\x00\x00\x00\x00\x00")) + + chunk(b"IEND", b"") + ) + + +def _multipart_request(gateway: Gateway, case: CostTrackingTestCase, model_name: str, key: str) -> httpx.Response: + assert case.upload is not None + fields: Final = { + field: value if isinstance(value, str) else json.dumps(value, separators=(",", ":")) + for field, value in {**case.request, "model": model_name}.items() + } + if case.upload.kind == "wav": + files: Final = {"file": ("audio.wav", _wav_bytes(case.upload.seconds), "audio/wav")} + else: + files = {"image": ("image.png", _png_bytes(), "image/png")} + return gateway.request_multipart(case.endpoint, fields, files, key=key) + + def _assert_stream_has_no_error(response_text: str) -> None: for line in response_text.splitlines(): if not line.startswith("data:"): @@ -40,62 +99,341 @@ def _assert_stream_has_no_error(response_text: str) -> None: if payload == "[DONE]": continue parsed = JSON_OBJECT.validate_json(payload) - assert "error" not in parsed, f"stream carried an error event: {parsed}" + assert ( + "error" not in parsed and parsed.get("type") not in {"error", "response.failed"} + ), f"stream carried an error event: {parsed}" + + +def _replace_model(value: JsonValue, model_name: str) -> JsonValue: + if isinstance(value, str): + return value.replace("$MODEL", model_name) + if isinstance(value, list): + return [_replace_model(item, model_name) for item in value] + if isinstance(value, dict): + return {key: _replace_model(item, model_name) for key, item in value.items()} + return value + + +def _assert_breakdown( + case: CostTrackingTestCase, + expected: ExactExpected, + breakdown: CostBreakdown, + response: httpx.Response, +) -> None: + assert breakdown.input_cost is not None and approx_equal(breakdown.input_cost, expected.input_cost), ( + f"{case.name}: input_cost {breakdown.input_cost} != expected {expected.input_cost}" + ) + assert breakdown.output_cost is not None and approx_equal(breakdown.output_cost, expected.output_cost), ( + f"{case.name}: output_cost {breakdown.output_cost} != expected {expected.output_cost}" + ) + for field, header_name, actual_component, expected_component in ( + ( + "cache_read_cost", + "x-litellm-response-cost-cache-read", + breakdown.cache_read_cost, + expected.cache_read_cost, + ), + ( + "cache_creation_cost", + "x-litellm-response-cost-cache-creation", + breakdown.cache_creation_cost, + expected.cache_creation_cost, + ), + ( + "reasoning_cost", + "x-litellm-response-cost-reasoning", + breakdown.reasoning_cost, + expected.reasoning_cost, + ), + ( + "tool_usage_cost", + "x-litellm-response-cost-tool-usage", + breakdown.tool_usage_cost, + expected.tool_usage_cost, + ), + ): + if expected_component is None: + continue + omitted_component_allowed: Final = expected_component == 0.0 + assert (actual_component is None and omitted_component_allowed) or ( + actual_component is not None and approx_equal(actual_component, expected_component) + ), f"{case.name}: {field} {actual_component} != expected {expected_component}" + if expected.cost_header and case.response.content_type == "application/json": + header: Final = response.headers.get(header_name) + assert (header is None and omitted_component_allowed) or ( + header is not None and approx_equal(float(header), expected_component) + ), f"{case.name}: {header_name} {header} != expected {expected_component}" + if expected.cost_header and case.response.content_type == "application/json" and any( + component is not None + for component in ( + expected.cache_read_cost, + expected.cache_creation_cost, + expected.reasoning_cost, + expected.tool_usage_cost, + ) + ): + input_header: Final = response.headers.get("x-litellm-response-cost-input") + output_header: Final = response.headers.get("x-litellm-response-cost-output") + expected_input_header: Final = expected.input_cost - ( + expected.cache_read_cost or 0.0 + ) - (expected.cache_creation_cost or 0.0) + assert input_header is not None and approx_equal(float(input_header), expected_input_header), ( + f"{case.name}: x-litellm-response-cost-input {input_header} != expected {expected_input_header}" + ) + assert output_header is not None and approx_equal(float(output_header), expected.output_cost), ( + f"{case.name}: x-litellm-response-cost-output {output_header} != expected {expected.output_cost}" + ) + + +def _assert_exact( + case: CostTrackingTestCase, + expected: ExactExpected, + row: CostRow, + response: httpx.Response, +) -> None: + assert row.spend is not None and approx_equal(row.spend, expected.spend), ( + f"{case.name}: spend {row.spend} != expected {expected.spend} " + f"(breakdown {row.breakdown.model_dump() if row.breakdown is not None else None})" + ) + breakdown: Final = row.breakdown + if expected.breakdown_persisted: + assert breakdown is not None, f"{case.name}: no cost_breakdown persisted" + if breakdown is not None: + _assert_breakdown(case, expected, breakdown, response) + assert row.prompt_tokens == expected.prompt_tokens, ( + f"{case.name}: prompt_tokens {row.prompt_tokens} != expected {expected.prompt_tokens}" + ) + assert row.completion_tokens == expected.completion_tokens, ( + f"{case.name}: completion_tokens {row.completion_tokens} != expected {expected.completion_tokens}" + ) + if breakdown is not None: + assert_total_is_sum_of_components(row, breakdown, case.name) + + +def _assert_recount(case: CostTrackingTestCase, expected: RecountExpected, row: CostRow) -> None: + assert row.prompt_tokens is not None and row.prompt_tokens > 0, ( + f"{case.name}: recount case counted no input tokens: prompt_tokens={row.prompt_tokens}" + ) + assert row.completion_tokens is not None and row.completion_tokens > 0, ( + f"{case.name}: recount case counted no output tokens: completion_tokens={row.completion_tokens}" + ) + if expected.prompt_tokens is not None: + assert row.prompt_tokens == expected.prompt_tokens, ( + f"{case.name}: prompt_tokens {row.prompt_tokens} != pinned {expected.prompt_tokens}" + ) + if expected.completion_tokens is not None: + assert row.completion_tokens == expected.completion_tokens, ( + f"{case.name}: completion_tokens {row.completion_tokens} != pinned {expected.completion_tokens}" + ) + if expected.min_completion_tokens is not None: + assert row.completion_tokens >= expected.min_completion_tokens, ( + f"{case.name}: completion_tokens {row.completion_tokens} < minimum {expected.min_completion_tokens}" + ) + if expected.max_completion_tokens is not None: + assert row.completion_tokens <= expected.max_completion_tokens, ( + f"{case.name}: completion_tokens {row.completion_tokens} > maximum {expected.max_completion_tokens}" + ) + recount: Final = row.prompt_tokens * expected.recount.input_cost_per_token + ( + row.completion_tokens * expected.recount.output_cost_per_token + ) + assert row.spend is not None and approx_equal(row.spend, recount), ( + f"{case.name}: spend {row.spend} != recount {recount} at map rates" + ) + assert row.breakdown is not None, f"{case.name}: no cost_breakdown persisted" + assert_total_is_sum_of_components(row, row.breakdown, case.name) @pytest.mark.parametrize("case", _CASES) def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) -> None: marker: Final = sha256(case.name.encode()).hexdigest()[:12] with gateway.scenario() as scenario: - key: Final = scenario.key() - model_name: Final = register_scenario_deployment(scenario, case, marker, key) - response: Final = gateway.request( - "POST", - "/v1/chat/completions", - {**case.request, "model": model_name}, - key=key, + expected: Final = case.expected + team_id: Final = scenario.team() if isinstance(expected, ExactExpected) and expected.rollups else None + user_id: Final = ( + scenario.user(team_id=team_id) + if team_id is not None + else None ) + key: Final = ( + scenario.key(team_id=team_id, user_id=user_id) + if team_id is not None and user_id is not None + else scenario.key() + ) + passthrough_provider: Final = case.passthrough_provider + scenario_id: Final = f"sc-{marker}-{sha256(key.encode()).hexdigest()[:12]}" + scenario_handle: Final = ( + register_scenario(scenario_id, case.response) + if passthrough_provider in {"gemini", "anthropic"} + else None + ) + if scenario_handle is not None: + scenario.cleanups.callback(delete_scenario, scenario_handle) + deployment: Final = ( + register_scenario_deployment(scenario, case, marker, key) + if passthrough_provider not in {"gemini", "anthropic"} + else None + ) + fallback_deployment: Final = ( + register_scenario_deployment( + scenario, + case, + marker, + key, + response=case.fallback_from, + marker_suffix="-fb", + ) + if case.fallback_from is not None + else None + ) + model_name: Final = ( + case.model + if passthrough_provider in {"gemini", "anthropic"} + else deployment.model_name if deployment is not None else None + ) + assert model_name is not None + request_model: Final = ( + case.model.rsplit("/", 1)[-1] + if passthrough_provider in {"gemini", "anthropic"} + else fallback_deployment.model_name if fallback_deployment is not None else model_name + ) + base_request_values: Final = ( + _replace_model(case.request, request_model) + if passthrough_provider is not None + else {**case.request, "model": model_name} + ) + end_user_id: Final = ( + f"end-user-{uuid.uuid4()}" + if isinstance(expected, ExactExpected) and expected.rollups + else None + ) + request_body: Final = JSON_OBJECT.validate_python( + { + **base_request_values, + **( + {"model": fallback_deployment.model_name, "fallbacks": [model_name]} + if fallback_deployment is not None + else {} + ), + **( + {"user": end_user_id, "cache": {"no-cache": True}} + if end_user_id is not None + else {} + ), + } + ) + request_headers: Final = ( + { + "x-pass-x-scripted-scenario": scenario_id, + **( + {"x-goog-api-key": key} + if passthrough_provider == "gemini" + else {} + ), + } + if passthrough_provider is not None + else {} + ) + request_path: Final = ( + case.endpoint.replace("$MODEL", request_model) + if passthrough_provider is not None + else case.endpoint + ) + if case.disconnect_after_frames is not None: + with gateway.client.stream( + "POST", + request_path, + json=request_body, + headers={"Authorization": f"Bearer {key}", **request_headers}, + ) as stream_response: + frames: Final = tuple( + islice( + (line for line in stream_response.iter_lines() if line.startswith("data:")), + case.disconnect_after_frames, + ) + ) + assert len(frames) == case.disconnect_after_frames + row: Final = poll_cost_row(key) + assert isinstance(expected, RecountExpected) + assert row.status == "success", f"{case.name}: disconnect row status was {row.status}" + _assert_recount(case, expected, row) + return + responses: Final = tuple( + ( + _multipart_request(gateway, case, model_name, key) + if case.upload is not None + else gateway.request("POST", request_path, request_body, key=key, headers=request_headers) + ) + for _ in range(3 if isinstance(expected, ExactExpected) and expected.rollups else 1) + ) + response: Final = responses[0] + if isinstance(expected, FailureExpected): + assert response.status_code == case.expected.failure.status, ( + f"{case.name}: proxy returned {response.status_code}, expected {case.expected.failure.status}: " + f"{response.text[:400]}" + ) + response_cost: Final = response.headers.get("x-litellm-response-cost") + assert response_cost is None or approx_equal(float(response_cost), 0.0), ( + f"{case.name}: failure response cost was {response_cost}" + ) + row: Final = poll_failure_row(key) + assert row.spend == 0, f"{case.name}: failure spend was {row.spend}" + return assert response.is_success, f"{case.name}: proxy returned {response.status_code}: {response.text[:400]}" if case.response.content_type == "text/event-stream": _assert_stream_has_no_error(response.text) - row: Final = poll_cost_row(key) - if isinstance(case.expected, RecountExpected): - assert row.prompt_tokens is not None and row.prompt_tokens > 0, ( - f"{case.name}: recount case counted no input tokens: prompt_tokens={row.prompt_tokens}" - ) - assert row.completion_tokens is not None and row.completion_tokens > 0, ( - f"{case.name}: recount case counted no output tokens: completion_tokens={row.completion_tokens}" - ) - recount: Final = row.prompt_tokens * case.expected.recount.input_cost_per_token + ( - row.completion_tokens * case.expected.recount.output_cost_per_token - ) - assert row.spend is not None and approx_equal(row.spend, recount), ( - f"{case.name}: spend {row.spend} != recount {recount} at map rates" - ) - assert_total_is_sum_of_components(row, case.name) + rows: Final = poll_rows(key, len(responses)) + if isinstance(expected, RecountExpected): + row: Final = rows[0] + _assert_recount(case, expected, row) return - expected: Final = case.expected assert isinstance(expected, ExactExpected) - if case.response.content_type == "application/json": + if fallback_deployment is not None: + assert deployment is not None + time.sleep(3) + settled_rows: Final = read_rows_now(key) + assert len(settled_rows) == 1 + assert settled_rows[0].status == "success" + assert settled_rows[0].model_id == deployment.identity + if isinstance(case.response, BinaryResponse): + header: Final = response.headers.get("x-litellm-response-cost") + if header is not None: + assert approx_equal(float(header), expected.spend), ( + f"{case.name}: x-litellm-response-cost {header} != expected {expected.spend}" + ) + elif case.response.content_type == "application/json": header: Final = cast(str | None, response.headers.get("x-litellm-response-cost")) - assert header is not None and approx_equal(float(header), expected.spend), ( - f"{case.name}: x-litellm-response-cost {header} != expected {expected.spend}" + if expected.cost_header and expected.spend != 0: + assert header is not None and approx_equal(float(header), expected.spend), ( + f"{case.name}: x-litellm-response-cost {header} != expected {expected.spend}" + ) + elif header is not None: + assert approx_equal(float(header), expected.spend), ( + f"{case.name}: x-litellm-response-cost {header} != expected {expected.spend}" + ) + for row in rows: + _assert_exact(case, expected, row, response) + if expected.rollups: + assert deployment is not None and team_id is not None and user_id is not None + assert end_user_id is not None + target_spend: Final = expected.spend * 3 + target_requests: Final = 3 + rollups: Final = poll_rollups( + key, + team_id, + user_id, + end_user_id, + target_spend, + target_requests, ) - assert row.spend is not None and approx_equal(row.spend, expected.spend), ( - f"{case.name}: spend {row.spend} != expected {expected.spend} " - f"(breakdown {row.breakdown.model_dump()})" - ) - breakdown: Final = row.breakdown - assert breakdown.input_cost is not None and approx_equal(breakdown.input_cost, expected.input_cost), ( - f"{case.name}: input_cost {breakdown.input_cost} != expected {expected.input_cost}" - ) - assert breakdown.output_cost is not None and approx_equal(breakdown.output_cost, expected.output_cost), ( - f"{case.name}: output_cost {breakdown.output_cost} != expected {expected.output_cost}" - ) - assert row.prompt_tokens == expected.prompt_tokens, ( - f"{case.name}: prompt_tokens {row.prompt_tokens} != expected {expected.prompt_tokens}" - ) - assert row.completion_tokens == expected.completion_tokens, ( - f"{case.name}: completion_tokens {row.completion_tokens} != expected {expected.completion_tokens}" - ) - assert_total_is_sum_of_components(row, case.name) + assert approx_equal(rollups.key_spend, target_spend) + assert approx_equal(rollups.team_spend, target_spend) + assert approx_equal(rollups.user_spend, target_spend) + assert approx_equal(rollups.end_user_spend, target_spend) + assert approx_equal(rollups.daily_user.spend, target_spend) + assert approx_equal(rollups.daily_team.spend, target_spend) + assert rollups.daily_user.prompt_tokens == expected.prompt_tokens * 3 + assert rollups.daily_user.completion_tokens == expected.completion_tokens * 3 + assert rollups.daily_user.api_requests == 3 + assert rollups.daily_team.prompt_tokens == expected.prompt_tokens * 3 + assert rollups.daily_team.completion_tokens == expected.completion_tokens * 3 + assert rollups.daily_team.api_requests == 3 diff --git a/tests/integration/providers/test_fal_ai_video_wire.py b/tests/integration/providers/test_fal_ai_video_wire.py index 8c72810ffb6..827818c6780 100644 --- a/tests/integration/providers/test_fal_ai_video_wire.py +++ b/tests/integration/providers/test_fal_ai_video_wire.py @@ -7,6 +7,7 @@ from integration._support.client import Gateway from integration._support.wire import Reply, Request, wire_server _MODEL: Final = "bytedance/seedance-2.5/text-to-video" +_H3_MODEL: Final = "minimax/h3/text-to-video" _MP4: Final = b"\x00\x00\x00\x18ftypmp42" + uuid.uuid4().bytes * 4 @@ -65,5 +66,111 @@ def test_fal_video_create_status_and_content_follow_queue_wire_contract(gateway: ("POST", f"/{_MODEL}"), ("GET", f"/bytedance/seedance-2.5/requests/{request_id}/status"), ("GET", f"/bytedance/seedance-2.5/requests/{request_id}"), + ("GET", f"/bytedance/seedance-2.5/requests/{request_id}"), ("GET", f"/files/{request_id}.mp4"), ] + + +@pytest.mark.covers("other.provider_wire.fal_ai.video_queue_create_status_and_content_download") +def test_fal_h3_video_create_uses_canonical_body_and_status_path(gateway: Gateway) -> None: + request_id: Final = "fal-h3-req-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + assert request.headers["authorization"] == "Key synthetic-fal-key" + if request.method == "POST": + assert request.target == f"/{_H3_MODEL}" + assert json.loads(request.body) == { + "prompt": "a cat playing volleyball on a beach", + "duration": 6, + "resolution": "2K", + } + return Reply( + body=json.dumps({"status": "IN_QUEUE", "request_id": request_id, "queue_position": 0}).encode() + ) + assert request.method == "GET" + if request.target == f"/minimax/h3/requests/{request_id}/status": + return Reply(body=json.dumps({"status": "COMPLETED", "request_id": request_id}).encode()) + assert request.target == f"/minimax/h3/requests/{request_id}" + return Reply(body=json.dumps({"video": {"url": f"{wire_url}/files/{request_id}.mp4"}}).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + wire_url: Final = wire.url + model: Final = scenario.model( + model=f"fal_ai/{_H3_MODEL}", + api_base=wire.url, + api_key="synthetic-fal-key", + ) + created: Final = gateway.post( + "/v1/videos", + { + "model": model, + "prompt": "a cat playing volleyball on a beach", + "seconds": 6, + "size": "2k", + }, + ) + assert created["status"] == "queued" + video_id: Final = created["id"] + status: Final = gateway.get(f"/v1/videos/{video_id}") + assert status["status"] == "completed" + content: Final = gateway.request("GET", f"/v1/videos/{video_id}/content") + assert content.status_code == 200, content.text + assert content.content == _MP4 + assert [(request.method, request.target) for request in wire.drain()] == [ + ("POST", f"/{_H3_MODEL}"), + ("GET", f"/minimax/h3/requests/{request_id}/status"), + ("GET", f"/minimax/h3/requests/{request_id}"), + ("GET", f"/minimax/h3/requests/{request_id}"), + ("GET", f"/files/{request_id}.mp4"), + ] + + +@pytest.mark.covers("other.provider_wire.fal_ai.video_failed_result_surfaces_fal_error") +def test_fal_video_failed_result_reports_failed_status_and_fal_error(gateway: Gateway) -> None: + request_id: Final = "fal-failed-req-" + uuid.uuid4().hex + error_body: Final = { + "detail": [ + { + "loc": ["body", "input.reference_image_urls"], + "msg": "Failed to download the file. Please check if the URL is accessible and try again.", + "type": "file_download_error", + } + ] + } + + def respond(request: Request) -> Reply: + assert request.headers["authorization"] == "Key synthetic-fal-key" + if request.method == "POST": + assert request.target == f"/{_MODEL}" + return Reply( + body=json.dumps({"status": "IN_QUEUE", "request_id": request_id, "queue_position": 0}).encode() + ) + assert request.method == "GET" + if request.target == f"/bytedance/seedance-2.5/requests/{request_id}/status": + return Reply(body=json.dumps({"status": "COMPLETED", "request_id": request_id}).encode()) + assert request.target == f"/bytedance/seedance-2.5/requests/{request_id}" + return Reply(status=422, body=json.dumps(error_body).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"fal_ai/{_MODEL}", + api_base=wire.url, + api_key="synthetic-fal-key", + ) + created: Final = gateway.post( + "/v1/videos", + { + "model": model, + "prompt": "a cat playing volleyball on a beach", + "seconds": "4", + "size": "1280x720", + }, + ) + assert created["status"] == "queued" + video_id: Final = created["id"] + status: Final = gateway.get(f"/v1/videos/{video_id}") + assert status["status"] == "failed" + assert "input.reference_image_urls: Failed to download the file" in status["error"]["message"] + content: Final = gateway.request("GET", f"/v1/videos/{video_id}/content") + assert content.status_code == 422, content.text + assert "Failed to download the file" in content.text diff --git a/tests/local_testing/test_basic_python_version.py b/tests/local_testing/test_basic_python_version.py index ef500fdff42..f032486debd 100644 --- a/tests/local_testing/test_basic_python_version.py +++ b/tests/local_testing/test_basic_python_version.py @@ -3,6 +3,7 @@ import os import subprocess import time import traceback +from typing import Final import pytest @@ -253,6 +254,7 @@ def _run_proxy_server_smoke_test(extra_proxy_args=None): raise filepath = os.path.dirname(os.path.abspath(__file__)) config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" + proxy_env: Final = {**os.environ, "LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY": "true"} server_process = subprocess.Popen( [ "uv", @@ -266,6 +268,7 @@ def _run_proxy_server_smoke_test(extra_proxy_args=None): *extra_proxy_args, ], cwd=PROJECT_ROOT, + env=proxy_env, ) # Allow some time for the server to start (increased for CI environments) diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index 17de3cf1e8a..01f2a13d252 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -872,6 +872,45 @@ def test_chat_choices_win_over_a_responses_output_list(): assert data.finish_reasons == ("stop",) +def _ocr_payload(pages: list[object]): + return _sample_payload( + call_type="aocr", + custom_llm_provider="mistral", + model="mistral-ocr-latest", + messages=None, + response={"object": "ocr", "model": "mistral-ocr-latest", "pages": pages, "usage_info": {"pages_processed": 2}}, + ) + + +def test_ocr_pages_become_one_assistant_choice_joined_in_page_order(): + data = LLMCallSpanData.from_standard_logging_payload( + _ocr_payload([{"index": 0, "markdown": "# Invoice"}, {"index": 1, "markdown": "Total: 42"}]), + capture_content=True, + ) + + assert data.choices_out == ( + { + "message": {"role": "assistant", "content": "# Invoice\n\nTotal: 42", "refusal": None, "tool_calls": None}, + "finish_reason": None, + }, + ) + assert data.finish_reasons == () + + +def test_ocr_output_follows_the_content_capture_gate(): + data = LLMCallSpanData.from_standard_logging_payload(_ocr_payload([{"index": 0, "markdown": "# Invoice"}])) + + assert data.choices_out == () + + +def test_ocr_pages_without_markdown_stay_empty(): + data = LLMCallSpanData.from_standard_logging_payload( + _ocr_payload([{"index": 0, "images": []}, "not-a-page"]), capture_content=True + ) + + assert data.choices_out == () + + def test_request_identity_prefers_canonical_team_keys(): from litellm.integrations.otel.model.payloads import RequestIdentity diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py index 4e375de0494..9fa198c4ec5 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py @@ -227,6 +227,28 @@ def test_langfuse_mapper_renders_a_responses_api_call_from_the_standard_logging_ assert attrs["langfuse.observation.type"] == "generation" +def test_langfuse_mapper_renders_an_ocr_call_with_the_page_markdown_as_output(): + payload = { + "call_type": "aocr", + "custom_llm_provider": "mistral", + "model": "mistral-ocr-latest", + "messages": None, + "response": { + "object": "ocr", + "model": "mistral-ocr-latest", + "pages": [{"index": 0, "markdown": "# Invoice"}, {"index": 1, "markdown": "Total: 42"}], + "usage_info": {"pages_processed": 2}, + }, + } + data = LLMCallSpanData.from_standard_logging_payload(payload, capture_content=True) + attrs = LangfuseMapper().map(data) + + assert json.loads(attrs["langfuse.observation.output"]) == [ + {"role": "assistant", "content": "# Invoice\n\nTotal: 42", "refusal": None, "tool_calls": None} + ] + assert attrs["langfuse.observation.type"] == "generation" + + # --------------------------------------------------------------------------- # # Weave # --------------------------------------------------------------------------- # diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 325052ebda9..277ae33a076 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -25,7 +25,7 @@ from litellm.litellm_core_utils.litellm_logging import ( set_callbacks, ) from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo -from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.utils import ( CallTypes, LiteLLMRealtimeStreamLoggingObject, @@ -7415,3 +7415,55 @@ class TestAzurePTUSpilloverCost: finally: litellm.model_cost.pop(custom_model_id, None) self._unregister_models() + + +def _completed_responses_event(usage: ResponseAPIUsage) -> ResponseCompletedEvent: + return ResponseCompletedEvent( + type="response.completed", + response=ResponsesAPIResponse( + id="resp-1", created_at=1, object="response", status="completed", model="codex-mini-latest", output=[], usage=usage + ), + ) + + +def _responses_stream_logging_obj() -> LitellmLogging: + logging_obj = _make_logging_obj(stream=True) + logging_obj.update_environment_variables( + model="openai/codex-mini-latest", user="", optional_params={}, litellm_params={"api_base": ""} + ) + return logging_obj + + +def test_get_assembled_streaming_response_bills_a_provider_reported_usage_cost(): + """A Responses stream whose completed event carries ``usage.cost`` is billed that number, + the way an assembled chat stream already is, instead of a price-map estimate.""" + logging_obj = _responses_stream_logging_obj() + now = datetime.datetime.now() + + assembled = logging_obj._get_assembled_streaming_response( + result=_completed_responses_event(ResponseAPIUsage(input_tokens=12, output_tokens=2, total_tokens=14, cost=0.0042)), + start_time=now, + end_time=now, + is_async=True, + streaming_chunks=[], + ) + + assert assembled._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == 0.0042 + assert logging_obj._response_cost_calculator(result=assembled) == 0.0042 + + +def test_get_assembled_streaming_response_without_usage_cost_leaves_pricing_to_the_price_map(): + logging_obj = _responses_stream_logging_obj() + now = datetime.datetime.now() + + assembled = logging_obj._get_assembled_streaming_response( + result=_completed_responses_event(ResponseAPIUsage(input_tokens=12, output_tokens=2, total_tokens=14)), + start_time=now, + end_time=now, + is_async=True, + streaming_chunks=[], + ) + + assert "additional_headers" not in assembled._hidden_params + price_map_cost = logging_obj._response_cost_calculator(result=assembled) + assert price_map_cost is not None and 0 < price_map_cost != 0.0042 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index cc4eb1d4136..507467b721f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -2,7 +2,7 @@ import asyncio import json import os import uuid -from typing import Any, Dict, List +from typing import Any, Dict, Final, List import httpx import pytest @@ -1584,3 +1584,104 @@ async def test_anthropic_messages_streaming_forwards_safeguards_and_keeps_safegu assert captured["body"]["safeguards"] == safeguards assert events[0]["message"]["safeguard_results"] == safeguard_results assert [e for e in events if e["type"] == "message_delta"][0]["delta"]["safeguard_results"] == safeguard_results + + +def _claude_code_auto_mode_request() -> tuple[list[dict[str, object]], list[dict[str, object]]]: + """Shapes are what Claude Code 2.1.278 sends and Bedrock Invoke / Vertex rawPredict return, captured 2026-09-21.""" + safeguards = [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}] + tool_verdicts = {"toolu_01": {"type": "evaluated", "outcome": "not_flagged"}} + safeguard_results = [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": tool_verdicts}}] + return safeguards, safeguard_results + + +def _upstream_answering_with(safeguard_results: list[dict[str, object]], captured: dict[str, object]) -> AsyncHTTPHandler: + def upstream_records_the_request(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + captured["anthropic-beta"] = request.headers.get("anthropic-beta") + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + "safeguard_results": safeguard_results, + }, + request=request, + ) + + upstream = AsyncHTTPHandler() + upstream.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream_records_the_request)) + return upstream + + +_CLIENT_BETA_HEADERS: Final = ( + pytest.param({"anthropic-beta": "dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14"}, id="client_sends_beta"), + pytest.param({"anthropic-beta": "interleaved-thinking-2025-05-14"}, id="client_omits_beta"), + pytest.param({}, id="client_sends_no_beta_header"), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("client_headers", _CLIENT_BETA_HEADERS) +async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_beta_to_bedrock_invoke( + local_beta_headers_config, client_headers +): + """Bedrock Invoke takes betas in the body's `anthropic_beta` and 400s on `safeguards` without the beta, so the beta rides along with the field.""" + from litellm.llms.anthropic.experimental_pass_through.messages import handler + + safeguards, safeguard_results = _claude_code_auto_mode_request() + captured: dict[str, object] = {} + + response = await handler.anthropic_messages( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="bedrock/us.anthropic.claude-sonnet-5", + custom_llm_provider="bedrock", + aws_access_key_id="test-access-key", + aws_secret_access_key="test-secret-key", + aws_region_name="us-east-1", + client=_upstream_answering_with(safeguard_results, captured), + safeguards=safeguards, + extra_headers=client_headers, + ) + + assert captured["body"]["safeguards"] == safeguards + assert captured["body"]["anthropic_beta"] == ["dangerous-tool-use-2026-09-03"] + assert response["safeguard_results"] == safeguard_results + + +@pytest.mark.asyncio +@pytest.mark.parametrize("client_headers", _CLIENT_BETA_HEADERS) +async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_beta_to_vertex( + local_beta_headers_config, client_headers +): + """Vertex rawPredict takes the beta as the `anthropic-beta` header and 400s on `safeguards` without it, so the beta rides along with the field.""" + from litellm.llms.anthropic.experimental_pass_through.messages import handler + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + safeguards, safeguard_results = _claude_code_auto_mode_request() + captured: dict[str, object] = {} + + with patch.object(VertexBase, "_ensure_access_token", return_value=("test-token", "test-project")): + response = await handler.anthropic_messages( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="vertex_ai/claude-sonnet-5", + custom_llm_provider="vertex_ai", + vertex_project="test-project", + vertex_location="global", + vertex_credentials="{}", + client=_upstream_answering_with(safeguard_results, captured), + safeguards=safeguards, + extra_headers=client_headers, + ) + + assert captured["body"]["safeguards"] == safeguards + assert "anthropic_beta" not in captured["body"] + assert captured["anthropic-beta"].split(",").count("dangerous-tool-use-2026-09-03") == 1 + assert response["safeguard_results"] == safeguard_results diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 7be005c0efe..ea8b722b849 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -1651,6 +1651,92 @@ def test_bedrock_messages_allowlist_filters_anthropic_only_fields(): assert set(result).issubset(cfg.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS) +@pytest.mark.parametrize( + "client_beta_header", + ["dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14", "interleaved-thinking-2025-05-14"], + ids=["client_sends_beta", "client_omits_beta"], +) +def test_bedrock_messages_forwards_safeguards_with_dangerous_tool_use_beta(local_beta_headers_config, client_beta_header): + """ + Claude Code's server-side auto-mode classifier sends `safeguards` alongside the + dangerous-tool-use-2026-09-03 beta. Bedrock Invoke accepts the pair, answers + "safeguards: Extra inputs are not permitted" for the field alone, and returns + `safeguard_results: []` for the beta alone, so the field reaches it unchanged + and the beta rides along whether or not the client sent it, as every other + body-driven beta does here. + """ + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + safeguards = [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}] + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": [{"type": "text", "text": "Hello"}]}], + anthropic_messages_optional_request_params={"max_tokens": 64, "safeguards": safeguards}, + litellm_params=GenericLiteLLMParams(), + headers={"anthropic-beta": client_beta_header}, + ) + + assert result["safeguards"] == safeguards + assert result["anthropic_beta"].count("dangerous-tool-use-2026-09-03") == 1 + + +def test_bedrock_messages_does_not_add_dangerous_tool_use_beta_without_safeguards(local_beta_headers_config): + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": [{"type": "text", "text": "Hello"}]}], + anthropic_messages_optional_request_params={"max_tokens": 64}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert "safeguards" not in result + assert "dangerous-tool-use-2026-09-03" not in result.get("anthropic_beta", []) + + +def test_bedrock_messages_stream_decoder_keeps_safeguard_results(): + """Bedrock streams the classifier verdicts on message_start and on the final message_delta, exactly as api.anthropic.com does.""" + decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="us.anthropic.claude-sonnet-5") + tool_verdicts = {"toolu_01": {"type": "evaluated", "outcome": "not_flagged"}} + safeguard_results = [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": tool_verdicts}}] + + message_start = decoder._chunk_parser( + { + "type": "message_start", + "message": { + "id": "msg_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 3, "output_tokens": 0}, + "safeguard_results": safeguard_results, + }, + } + ) + + assert isinstance(message_start, dict) + assert message_start["message"]["safeguard_results"] == safeguard_results + + message_delta = decoder._chunk_parser( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, "safeguard_results": safeguard_results}, + "usage": {"output_tokens": 1}, + "amazon-bedrock-invocationMetrics": {"inputTokenCount": 3, "outputTokenCount": 1}, + } + ) + + assert isinstance(message_delta, dict) + assert message_delta["delta"]["safeguard_results"] == safeguard_results + + def test_bedrock_messages_filters_user_provided_unsupported_beta_header(): """ In proxy deployments the client (e.g. Claude Code) doesn't know the backend diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py index 6bacf8f3d94..5f69b36c87a 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py @@ -439,6 +439,19 @@ class TestBetaHeadersOnTheWire: assert _sent_betas(route) == ["context-1m-2025-08-07", "context-management-2025-06-27"] assert _sent_body(route)["context_management"] == {"edits": [{"type": "clear_tool_uses_20250919"}]} + @pytest.mark.asyncio + @respx.mock + async def test_safeguards_reach_mantle_with_the_dangerous_tool_use_beta(self): + """Mantle answers 400 "safeguards: Extra inputs are not permitted" when the field + arrives without dangerous-tool-use-2026-09-03 (probed 2026-09-21), so the beta + has to ride along even when the client never sent the header.""" + safeguards = [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}] + + route = await self._send(safeguards=safeguards) + + assert _sent_betas(route) == ["dangerous-tool-use-2026-09-03"] + assert _sent_body(route)["safeguards"] == safeguards + @pytest.mark.asyncio @respx.mock async def test_betas_and_version_never_travel_in_the_body(self): diff --git a/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py b/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py index 7862297bcd5..69ae66d3818 100644 --- a/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py +++ b/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py @@ -1,10 +1,14 @@ +import json import math +import re +from pathlib import Path import pytest import litellm from litellm import completion, get_llm_provider from litellm.llms.dashscope.chat.transformation import DashScopeChatConfig +from litellm.llms.dashscope.common_utils import missing_dashscope_family_key_message from litellm.llms.dashscope.cost_calculator import ( cost_per_token as dashscope_cost_per_token, ) @@ -53,6 +57,7 @@ BRAND_CASES = [ pytest.param( { "provider": "qwencloud", + "display_name": "QwenCloud", "enum": LlmProviders.QWENCLOUD, "key_env": "QWENCLOUD_API_KEY", "base_env": "QWENCLOUD_API_BASE", @@ -69,6 +74,7 @@ BRAND_CASES = [ pytest.param( { "provider": "qwen_ai_platform", + "display_name": "Qianwen AI Platform", "enum": LlmProviders.QWEN_AI_PLATFORM, "key_env": "QWEN_AI_PLATFORM_API_KEY", "base_env": "QWEN_AI_PLATFORM_API_BASE", @@ -89,6 +95,13 @@ BRAND_CASES = [ def clear_dashscope_family_env(monkeypatch): for env_var in DASHSCOPE_FAMILY_ENV_VARS: monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + + +@pytest.fixture +def no_provider_traffic(respx_mock, monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + return respx_mock class TestQwenBrandProviderResolution: @@ -250,6 +263,51 @@ class TestQwenBrandDefaultUrls: ) +class TestQwenBrandUserFacingNames: + RETIRED_MAINLAND_NAME = "Qwen AI Platform" + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_missing_key_message_names_brand(self, brand): + message = missing_dashscope_family_key_message(brand["provider"]) + assert brand["display_name"] in message + assert brand["key_env"] in message + assert self.RETIRED_MAINLAND_NAME not in message + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_embedding_without_key_names_brand(self, brand, no_provider_traffic): + with pytest.raises(litellm.APIConnectionError, match=re.escape(brand["display_name"])) as exc_info: + litellm.embedding(model=f"{brand['provider']}/text-embedding-v4", input=["hello"]) + assert self.RETIRED_MAINLAND_NAME not in str(exc_info.value) + assert no_provider_traffic.calls.call_count == 0 + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_rerank_without_key_names_brand(self, brand, no_provider_traffic): + with pytest.raises(litellm.APIConnectionError, match=re.escape(brand["display_name"])) as exc_info: + litellm.rerank(model=f"{brand['provider']}/gte-rerank-v2", query="q", documents=["a", "b"]) + assert self.RETIRED_MAINLAND_NAME not in str(exc_info.value) + assert no_provider_traffic.calls.call_count == 0 + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_image_generation_without_key_names_brand(self, brand, no_provider_traffic): + with pytest.raises(litellm.APIConnectionError, match=re.escape(brand["display_name"])) as exc_info: + litellm.image_generation(model=f"{brand['provider']}/qwen-image", prompt="a cup of coffee") + assert self.RETIRED_MAINLAND_NAME not in str(exc_info.value) + assert no_provider_traffic.calls.call_count == 0 + + @pytest.mark.parametrize("brand", BRAND_CASES) + @pytest.mark.parametrize( + "matrix_path", + [ + Path(litellm.__file__).parent / "provider_endpoints_support_backup.json", + Path(litellm.__file__).parent.parent / "provider_endpoints_support.json", + ], + ids=["backup", "root"], + ) + def test_supported_endpoints_matrix_display_name(self, brand, matrix_path): + matrix = json.loads(matrix_path.read_text()) + assert matrix["providers"][brand["provider"]]["display_name"] == f"{brand['display_name']} (`{brand['provider']}`)" + + class TestQwenBrandCostParity: @pytest.fixture(autouse=True) def setup_model_cost_map(self, monkeypatch): diff --git a/tests/test_litellm/llms/edenai/audio_transcription/test_edenai_audio_transcription_transformation.py b/tests/test_litellm/llms/edenai/audio_transcription/test_edenai_audio_transcription_transformation.py new file mode 100644 index 00000000000..3f6cba8a91b --- /dev/null +++ b/tests/test_litellm/llms/edenai/audio_transcription/test_edenai_audio_transcription_transformation.py @@ -0,0 +1,155 @@ +"""Eden AI `/v3/audio/transcriptions`: OpenAI's speech-to-text API served by Eden's gateway, which +reports the real per-request cost at the top level of the JSON body.""" + +import httpx +import pytest + +import litellm +from litellm.cost_calculator import get_response_cost_from_hidden_params +from litellm.llms.edenai.audio_transcription.transformation import EdenAIAudioTranscriptionConfig +from litellm.llms.edenai.common_utils import EdenAIException +from litellm.types.utils import LlmProviders, TranscriptionResponse +from litellm.utils import ProviderConfigManager + +EDEN_BASE = "https://api.edenai.run/v3" +EDEN_TRANSCRIPTIONS_URL = f"{EDEN_BASE}/audio/transcriptions" +EDEN_REPORTED_COST = 0.0042 +MODEL = "edenai/openai/whisper-1" +SELLER_MODEL = "openai/whisper-1" +AUDIO_FILE = ("hello.mp3", b"ID3\x04\x00fake-mp3-bytes", "audio/mpeg") + + +def _eden_transcription(cost: float | None = EDEN_REPORTED_COST) -> dict: + """Live `/v3/audio/transcriptions` body: Whisper's verbose shape plus Eden's top-level `cost` + and `provider`, with `duration` present whatever `response_format` was asked for.""" + body = { + "text": "Hello there.", + "usage": {"type": "duration", "seconds": 1.0}, + "language": "english", + "task": "transcribe", + "duration": 0.62, + "words": None, + "segments": [{"id": 0, "start": 0.0, "end": 0.8, "text": " Hello there."}], + "provider": "openai", + } + return body if cost is None else {**body, "cost": cost} + + +def _multipart_body(respx_mock) -> str: + return respx_mock.calls.last.request.content.decode(errors="replace") + + +class TestRegistration: + def test_eden_is_a_native_transcription_provider(self): + config = ProviderConfigManager.get_provider_audio_transcription_config( + model=SELLER_MODEL, provider=LlmProviders.EDENAI + ) + + assert isinstance(config, EdenAIAudioTranscriptionConfig) + + +class TestRequestTransformation: + def test_sends_the_file_as_multipart_without_forcing_verbose_json(self): + request = EdenAIAudioTranscriptionConfig().transform_audio_transcription_request( + model=SELLER_MODEL, audio_file=AUDIO_FILE, optional_params={"language": "en"}, litellm_params={} + ) + + assert request.data == {"model": SELLER_MODEL, "language": "en"} + assert request.files == {"file": AUDIO_FILE} + + def test_sdk_style_extra_body_is_flattened_into_form_fields(self): + """LiteLLM parks `model` and any non-OpenAI kwarg under `extra_body` for the OpenAI SDK, and a + nested dict cannot ride in a multipart form.""" + request = EdenAIAudioTranscriptionConfig().transform_audio_transcription_request( + model=SELLER_MODEL, + audio_file=AUDIO_FILE, + optional_params={"language": "en", "extra_body": {"model": SELLER_MODEL, "user": "u-1"}}, + litellm_params={}, + ) + + assert request.data == {"model": SELLER_MODEL, "language": "en", "user": "u-1"} + + def test_missing_key_is_an_authentication_error_before_any_request(self, no_eden_key, respx_mock): + with pytest.raises(litellm.AuthenticationError, match="EDENAI_API_KEY"): + litellm.transcription(model=MODEL, file=AUDIO_FILE) + assert not respx_mock.calls + + +class TestTranscription: + def test_posts_multipart_to_eden_with_the_bearer_key_and_the_seller_model_id(self, eden_key, respx_mock): + respx_mock.post(EDEN_TRANSCRIPTIONS_URL).mock(return_value=httpx.Response(200, json=_eden_transcription())) + + response = litellm.transcription(model=MODEL, file=AUDIO_FILE, language="en", temperature=0) + + assert isinstance(response, TranscriptionResponse) + assert response.text == "Hello there." + request = respx_mock.calls.last.request + assert request.headers["Authorization"] == f"Bearer {eden_key}" + assert request.headers["Content-Type"].startswith("multipart/form-data") + body = _multipart_body(respx_mock) + assert f'name="model"\r\n\r\n{SELLER_MODEL}' in body + assert 'name="language"\r\n\r\nen' in body + assert 'name="temperature"\r\n\r\n0' in body + assert 'name="file"; filename="hello.mp3"' in body + assert "verbose_json" not in body + + def test_eden_reported_cost_beats_the_price_map(self, eden_key, respx_mock): + respx_mock.post(EDEN_TRANSCRIPTIONS_URL).mock(return_value=httpx.Response(200, json=_eden_transcription())) + + response = litellm.transcription(model=MODEL, file=AUDIO_FILE) + + assert get_response_cost_from_hidden_params(response._hidden_params) == EDEN_REPORTED_COST + assert response._hidden_params["response_cost"] == EDEN_REPORTED_COST + + def test_a_body_without_cost_leaves_pricing_to_the_price_map(self, eden_key, respx_mock): + respx_mock.post(EDEN_TRANSCRIPTIONS_URL).mock( + return_value=httpx.Response(200, json=_eden_transcription(cost=None)) + ) + + response = litellm.transcription(model=MODEL, file=AUDIO_FILE) + + assert get_response_cost_from_hidden_params(response._hidden_params) is None + assert response.duration == 0.62 + assert response.usage is not None + assert response.usage.seconds == 1.0 + + def test_a_plain_text_answer_is_the_transcript(self, eden_key, respx_mock): + respx_mock.post(EDEN_TRANSCRIPTIONS_URL).mock( + return_value=httpx.Response(200, text="Hello there.", headers={"content-type": "text/plain"}) + ) + + response = litellm.transcription(model=MODEL, file=AUDIO_FILE, response_format="text") + + assert response.text == "Hello there." + assert 'name="response_format"\r\n\r\ntext' in _multipart_body(respx_mock) + + @pytest.mark.asyncio + async def test_async_call_tracks_the_same_cost(self, eden_key, httpx_transport, respx_mock): + respx_mock.post(EDEN_TRANSCRIPTIONS_URL).mock(return_value=httpx.Response(200, json=_eden_transcription())) + + response = await litellm.atranscription(model=MODEL, file=AUDIO_FILE) + + assert response.text == "Hello there." + assert response._hidden_params["response_cost"] == EDEN_REPORTED_COST + + +class TestErrors: + def test_sync_401_surfaces_as_an_eden_error_with_the_status_code(self, eden_key, respx_mock): + """`litellm.transcription` does not map provider errors onto the OpenAI exception classes the + way its async twin does, so the proxy relies on the status code the provider exception carries.""" + respx_mock.post(EDEN_TRANSCRIPTIONS_URL).mock( + return_value=httpx.Response(401, json={"detail": "Invalid token."}) + ) + + with pytest.raises(EdenAIException, match="Invalid token") as excinfo: + litellm.transcription(model=MODEL, file=AUDIO_FILE) + assert excinfo.value.status_code == 401 + + @pytest.mark.asyncio + async def test_async_401_maps_to_authentication_error(self, eden_key, httpx_transport, respx_mock): + respx_mock.post(EDEN_TRANSCRIPTIONS_URL).mock( + return_value=httpx.Response(401, json={"detail": "Invalid token."}) + ) + + with pytest.raises(litellm.AuthenticationError, match="Invalid token"): + await litellm.atranscription(model=MODEL, file=AUDIO_FILE) diff --git a/tests/test_litellm/llms/edenai/chat/test_edenai_chat_transformation.py b/tests/test_litellm/llms/edenai/chat/test_edenai_chat_transformation.py new file mode 100644 index 00000000000..4f1e2c11a51 --- /dev/null +++ b/tests/test_litellm/llms/edenai/chat/test_edenai_chat_transformation.py @@ -0,0 +1,453 @@ +"""Eden AI (`edenai/...`) chat provider: an OpenAI-compatible gateway that reports the real +per-request cost at the top level of every response instead of leaving it to the price map.""" + +import json +from pathlib import Path + +import httpx +import pytest + +import litellm +from litellm.cost_calculator import get_response_cost_from_hidden_params, response_cost_calculator +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.edenai.chat.transformation import EdenAIChatCompletionStreamingHandler, EdenAIChatConfig +from litellm.llms.edenai.common_utils import EdenAIException +from litellm.proxy.auth.model_checks import get_provider_models +from litellm.types.router import LiteLLM_Params +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +REPO_ROOT = Path(__file__).resolve().parents[5] +EDEN_BASE = "https://api.edenai.run/v3" +EDEN_EU_BASE = "https://api.eu.edenai.run/v3" +EDEN_CHAT_URL = f"{EDEN_BASE}/chat/completions" +EDEN_REPORTED_COST = 0.0042 +EDEN_USAGE = {"completion_tokens": 1, "prompt_tokens": 9, "total_tokens": 10} +MESSAGES = [{"role": "user", "content": "Say OK"}] + + +def _eden_chat_completion(cost: float | None = EDEN_REPORTED_COST) -> dict: + """Live `/v3/chat/completions` body: OpenAI shape plus Eden's top-level `cost`, `provider` + and `status`, with `model` echoing the seller's bare model name.""" + body = { + "status": "success", + "id": "chatcmpl-eden-1", + "created": 1788347376, + "model": "gpt-4.1-nano", + "object": "chat.completion", + "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "OK", "role": "assistant"}}], + "usage": EDEN_USAGE, + "provider": "openai", + } + return body if cost is None else {**body, "cost": cost} + + +def _eden_stream_chunk( + delta: dict, finish_reason: str | None = None, usage: dict | None = None, cost: float | None = None +) -> dict: + chunk = { + "id": "chatcmpl-eden-stream", + "created": 1788347377, + "model": "openai/gpt-4.1-nano", + "object": "chat.completion.chunk", + "choices": [{"finish_reason": finish_reason, "index": 0, "delta": delta, "logprobs": None}], + } + if usage is not None: + chunk["usage"] = usage + if cost is not None: + chunk["cost"] = cost + return chunk + + +def _eden_stream_frames(cost: float | None = EDEN_REPORTED_COST) -> tuple[dict, ...]: + """Live stream with `stream_options.include_usage`: the usage frame comes after the + finish_reason frame, keeps one empty choice, and carries Eden's `cost` at the top level.""" + return ( + _eden_stream_chunk({"role": "assistant", "content": ""}), + _eden_stream_chunk({"content": "OK"}), + _eden_stream_chunk({"content": None}, finish_reason="stop"), + _eden_stream_chunk({"content": None, "role": None}, usage=EDEN_USAGE, cost=cost), + ) + + +def _sse(frames: tuple[dict, ...]) -> httpx.Response: + body = "".join(f"data: {json.dumps(frame)}\n\n" for frame in frames) + "data: [DONE]\n\n" + return httpx.Response(200, content=body.encode(), headers={"content-type": "text/event-stream"}) + + +def _request_body(respx_mock) -> dict: + return json.loads(respx_mock.calls.last.request.content) + + +class TestProviderResolution: + @pytest.mark.parametrize( + "requested, sent_to_eden", + [ + ("edenai/openai/gpt-4.1-nano", "openai/gpt-4.1-nano"), + ("edenai/gpt-4o", "gpt-4o"), + ("edenai/vertex/gemini-3.7-flash@eu", "vertex/gemini-3.7-flash@eu"), + ("edenai/fireworks_ai/accounts/fireworks/models/glm-5p3", "fireworks_ai/accounts/fireworks/models/glm-5p3"), + ("edenai/cloudflare/@cf/qwen/qwen3.8-27b", "cloudflare/@cf/qwen/qwen3.8-27b"), + ], + ) + def test_strips_only_the_edenai_prefix(self, eden_key, requested, sent_to_eden): + model, provider, api_key, api_base = get_llm_provider(requested) + + assert (model, provider, api_key, api_base) == (sent_to_eden, "edenai", eden_key, EDEN_BASE) + + def test_env_api_base_moves_the_key_to_the_eu_endpoint(self, eden_key, monkeypatch): + monkeypatch.setenv("EDENAI_API_BASE", EDEN_EU_BASE) + + _, provider, api_key, api_base = get_llm_provider("edenai/openai/gpt-4.1-nano") + + assert (provider, api_key, api_base) == ("edenai", eden_key, EDEN_EU_BASE) + + def test_explicit_credentials_win_over_env(self, eden_key): + _, _, api_key, api_base = get_llm_provider( + "edenai/openai/gpt-4.1-nano", api_key="explicit-key", api_base="https://eden.internal/v3" + ) + + assert (api_key, api_base) == ("explicit-key", "https://eden.internal/v3") + + def test_eden_api_base_is_recognised_without_the_prefix(self, eden_key): + model, provider, api_key, api_base = get_llm_provider("gpt-4.1-nano", api_base=EDEN_BASE) + + assert (model, provider, api_key, api_base) == ("gpt-4.1-nano", "edenai", eden_key, EDEN_BASE) + + +class TestRegistration: + def test_provider_is_registered_everywhere_routing_looks(self): + assert LlmProviders.EDENAI.value == "edenai" + assert "edenai" in litellm.provider_list + assert "edenai" in litellm.openai_compatible_providers + assert EDEN_BASE in litellm.openai_compatible_endpoints + assert isinstance( + ProviderConfigManager.get_provider_chat_config(model="openai/gpt-4.1-nano", provider=LlmProviders.EDENAI), + EdenAIChatConfig, + ) + + def test_supported_params_are_the_openai_chat_params(self): + supported = litellm.get_supported_openai_params(model="openai/gpt-4.1-nano", custom_llm_provider="edenai") + + assert supported is not None + assert {"tools", "tool_choice", "response_format", "stream_options", "max_completion_tokens"} <= set(supported) + + def test_reasoning_effort_is_supported_only_for_models_the_price_map_flags_as_reasoning(self): + reasoning = litellm.get_supported_openai_params(model="openai/gpt-5-mini", custom_llm_provider="edenai") + plain = litellm.get_supported_openai_params(model="openai/gpt-4.1-nano", custom_llm_provider="edenai") + + assert reasoning is not None and plain is not None + assert "reasoning_effort" in reasoning + assert "reasoning_effort" not in plain + + def test_validate_environment_names_the_eden_key(self, monkeypatch): + monkeypatch.delenv("EDENAI_API_KEY", raising=False) + missing = litellm.validate_environment(model="edenai/openai/gpt-4.1-nano") + monkeypatch.setenv("EDENAI_API_KEY", "eden-test-key") + present = litellm.validate_environment(model="edenai/openai/gpt-4.1-nano") + + assert (missing["keys_in_environment"], missing["missing_keys"]) == (False, ["EDENAI_API_KEY"]) + assert (present["keys_in_environment"], present["missing_keys"]) == (True, []) + + def test_a_model_registered_from_a_cost_map_still_asks_for_the_eden_key(self, monkeypatch): + """A cost map may name an Eden model without the `edenai/` prefix, leaving the provider + registry as the only way key validation can tell whose key the model needs.""" + alias = "eden-cost-map-alias" + litellm.register_model( + {alias: {"litellm_provider": "edenai", "mode": "chat", "input_cost_per_token": 1e-06}}, + persist_across_reloads=False, + ) + try: + monkeypatch.delenv("EDENAI_API_KEY", raising=False) + missing = litellm.validate_environment(model=alias) + monkeypatch.setenv("EDENAI_API_KEY", "eden-test-key") + present = litellm.validate_environment(model=alias) + finally: + litellm.edenai_models.discard(alias) + litellm.model_cost.pop(alias, None) + litellm.add_known_models(model_cost_map={}) + + assert (missing["keys_in_environment"], missing["missing_keys"]) == (False, ["EDENAI_API_KEY"]) + assert (present["keys_in_environment"], present["missing_keys"]) == (True, []) + + def test_a_cost_map_reload_reaches_wildcard_expansion(self, eden_key): + """Wildcard expansion reads the provider registry, which a cost map reload rebuilds in + place, so models added after startup have to show up without a restart.""" + alias = "edenai/openai/gpt-4.1-nano-from-cost-map" + wildcard = LiteLLM_Params(model="edenai/*", api_key="wildcard-key") + assert alias not in (get_provider_models("edenai", wildcard) or []) + + litellm.add_known_models(model_cost_map={alias: {"litellm_provider": "edenai", "mode": "chat"}}) + try: + expanded = get_provider_models("edenai", wildcard) + finally: + litellm.edenai_models.discard(alias) + litellm.add_known_models(model_cost_map={}) + + assert expanded is not None + assert alias in expanded + assert alias not in (get_provider_models("edenai", wildcard) or []) + + +class TestRequestTransformation: + def _request(self, optional_params: dict) -> dict: + return EdenAIChatConfig().transform_request( + model="openai/gpt-4.1-nano", + messages=MESSAGES, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + def test_streaming_request_asks_eden_for_the_usage_frame(self): + assert self._request({"stream": True})["stream_options"] == {"include_usage": True} + + def test_streaming_request_overrides_a_caller_opt_out(self): + body = self._request({"stream": True, "stream_options": {"include_usage": False}}) + + assert body["stream_options"] == {"include_usage": True} + + def test_non_streaming_request_carries_no_stream_options(self): + assert "stream_options" not in self._request({"max_tokens": 5}) + + +class TestCompletion: + def test_posts_to_eden_with_the_bearer_key_and_the_seller_model_id(self, eden_key, respx_mock): + respx_mock.post(EDEN_CHAT_URL).mock(return_value=httpx.Response(200, json=_eden_chat_completion())) + + response = litellm.completion(model="edenai/openai/gpt-4.1-nano", messages=MESSAGES, max_tokens=5) + + assert response.choices[0].message.content == "OK" + assert respx_mock.calls.last.request.headers["Authorization"] == f"Bearer {eden_key}" + body = _request_body(respx_mock) + assert (body["model"], body["messages"], body["max_tokens"]) == ("openai/gpt-4.1-nano", MESSAGES, 5) + + def test_reasoning_effort_reaches_eden_without_drop_params(self, eden_key, respx_mock): + respx_mock.post(EDEN_CHAT_URL).mock(return_value=httpx.Response(200, json=_eden_chat_completion())) + + litellm.completion(model="edenai/openai/gpt-5-mini", messages=MESSAGES, reasoning_effort="low") + + assert _request_body(respx_mock)["reasoning_effort"] == "low" + + def test_eden_reported_cost_beats_the_price_map(self, eden_key, respx_mock): + respx_mock.post(EDEN_CHAT_URL).mock(return_value=httpx.Response(200, json=_eden_chat_completion())) + + response = litellm.completion(model="edenai/openai/gpt-4.1-nano", messages=MESSAGES, max_tokens=5) + + assert get_response_cost_from_hidden_params(response._hidden_params) == EDEN_REPORTED_COST + assert ( + response_cost_calculator( + response_object=response, + model="openai/gpt-4.1-nano", + custom_llm_provider="edenai", + call_type="completion", + optional_params={}, + ) + == EDEN_REPORTED_COST + ) + + def test_a_body_without_cost_leaves_pricing_to_the_price_map(self, eden_key, respx_mock): + respx_mock.post(EDEN_CHAT_URL).mock(return_value=httpx.Response(200, json=_eden_chat_completion(cost=None))) + + response = litellm.completion(model="edenai/openai/gpt-4.1-nano", messages=MESSAGES, max_tokens=5) + + assert response.choices[0].message.content == "OK" + assert get_response_cost_from_hidden_params(response._hidden_params) is None + + def test_extra_body_forwards_eden_only_fields(self, eden_key, respx_mock): + respx_mock.post(EDEN_CHAT_URL).mock(return_value=httpx.Response(200, json=_eden_chat_completion())) + + litellm.completion( + model="edenai/openai/gpt-4.1-nano", + messages=MESSAGES, + extra_body={"fallbacks": ["anthropic/claude-sonnet-latest"], "routing": {"sort": "latency"}}, + ) + + body = _request_body(respx_mock) + assert body["fallbacks"] == ["anthropic/claude-sonnet-latest"] + assert body["routing"] == {"sort": "latency"} + assert "extra_body" not in body + + def test_unknown_kwargs_ride_along_as_eden_fields(self, eden_key, respx_mock): + respx_mock.post(EDEN_CHAT_URL).mock(return_value=httpx.Response(200, json=_eden_chat_completion())) + + litellm.completion(model="edenai/openai/gpt-4.1-nano", messages=MESSAGES, routing={"sort": "latency"}) + + assert _request_body(respx_mock)["routing"] == {"sort": "latency"} + + +class TestStreaming: + def test_include_usage_surfaces_eden_cost_on_the_usage_chunk(self, eden_key, respx_mock): + respx_mock.post(EDEN_CHAT_URL).mock(return_value=_sse(_eden_stream_frames())) + + chunks = list( + litellm.completion( + model="edenai/openai/gpt-4.1-nano", + messages=MESSAGES, + stream=True, + stream_options={"include_usage": True}, + ) + ) + + assert _request_body(respx_mock)["stream_options"] == {"include_usage": True} + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "OK" + usage_chunks = [chunk for chunk in chunks if getattr(chunk, "usage", None) is not None] + assert len(usage_chunks) == 1 + assert (usage_chunks[0].usage.total_tokens, usage_chunks[0].usage.cost) == (10, EDEN_REPORTED_COST) + + def test_without_include_usage_eden_cost_is_still_tracked_but_hidden(self, eden_key, respx_mock): + respx_mock.post(EDEN_CHAT_URL).mock(return_value=_sse(_eden_stream_frames())) + + chunks = list(litellm.completion(model="edenai/openai/gpt-4.1-nano", messages=MESSAGES, stream=True)) + + assert _request_body(respx_mock)["stream_options"] == {"include_usage": True} + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "OK" + assert all(getattr(chunk, "usage", None) is None for chunk in chunks) + hidden_usage = chunks[-1]._hidden_params["usage"] + assert (hidden_usage.total_tokens, hidden_usage.cost) == (10, EDEN_REPORTED_COST) + + +class TestStreamingHandler: + def _parse(self, chunk: dict): + return EdenAIChatCompletionStreamingHandler(streaming_response=None, sync_stream=True).chunk_parser(chunk) + + def test_moves_top_level_cost_onto_the_usage_object(self): + parsed = self._parse(_eden_stream_chunk({"content": None}, usage=EDEN_USAGE, cost=EDEN_REPORTED_COST)) + + assert parsed.usage is not None + assert (parsed.usage.prompt_tokens, parsed.usage.cost) == (9, EDEN_REPORTED_COST) + + def test_usage_without_cost_stays_unpriced(self): + parsed = self._parse(_eden_stream_chunk({"content": None}, usage=EDEN_USAGE)) + + assert parsed.usage is not None + assert getattr(parsed.usage, "cost", None) is None + + def test_content_chunks_are_passed_through(self): + parsed = self._parse(_eden_stream_chunk({"content": "OK"})) + + assert parsed.choices[0].delta.content == "OK" + assert getattr(parsed, "usage", None) is None + + +class TestErrors: + def test_middleware_401_detail_body_maps_to_authentication_error(self, eden_key, respx_mock): + respx_mock.post(EDEN_CHAT_URL).mock(return_value=httpx.Response(401, json={"detail": "Invalid token"})) + + with pytest.raises(litellm.AuthenticationError, match="Invalid token"): + litellm.completion(model="edenai/openai/gpt-4.1-nano", messages=MESSAGES) + + def test_unknown_model_envelope_maps_to_bad_request(self, eden_key, respx_mock): + envelope = { + "error": { + "message": "Model(s) not found or inactive: openai/does-not-exist", + "type": "invalid_request_error", + "param": None, + "code": "invalid_parameter", + } + } + respx_mock.post(EDEN_CHAT_URL).mock(return_value=httpx.Response(400, json=envelope)) + + with pytest.raises(litellm.BadRequestError, match="not found or inactive"): + litellm.completion(model="edenai/openai/does-not-exist", messages=MESSAGES) + + def test_429_maps_to_rate_limit_error(self, eden_key, respx_mock): + envelope = { + "error": {"message": "Rate limit exceeded", "type": "rate_limit_error", "code": "rate_limit_exceeded"} + } + respx_mock.post(EDEN_CHAT_URL).mock( + return_value=httpx.Response(429, json=envelope, headers={"Retry-After": "7"}) + ) + + with pytest.raises(litellm.RateLimitError, match="Rate limit exceeded"): + litellm.completion(model="edenai/openai/gpt-4.1-nano", messages=MESSAGES, num_retries=0) + + def test_error_class_is_the_eden_exception(self): + error = EdenAIChatConfig().get_error_class("boom", 503, {"Content-Type": "application/json"}) + + assert isinstance(error, EdenAIException) + assert isinstance(error, BaseLLMException) + assert (error.message, error.status_code, error.headers) == ("boom", 503, {"Content-Type": "application/json"}) + + +class TestModelListing: + CATALOG = {"data": [{"id": "openai/gpt-4.1-nano", "object": "model"}, {"id": "anthropic/claude-sonnet-latest"}]} + ROUTABLE = ["edenai/openai/gpt-4.1-nano", "edenai/anthropic/claude-sonnet-latest"] + + def test_lists_the_public_catalog_as_routable_model_names(self, eden_key, respx_mock): + respx_mock.get(f"{EDEN_BASE}/models").mock(return_value=httpx.Response(200, json=self.CATALOG)) + + assert EdenAIChatConfig().get_models() == self.ROUTABLE + + def test_lists_from_the_configured_endpoint(self, eden_key, monkeypatch, respx_mock): + monkeypatch.setenv("EDENAI_API_BASE", EDEN_EU_BASE) + respx_mock.get(f"{EDEN_EU_BASE}/models").mock(return_value=httpx.Response(200, json=self.CATALOG)) + + assert EdenAIChatConfig().get_models() == self.ROUTABLE + + def test_get_valid_models_reads_the_live_catalog(self, eden_key, respx_mock): + respx_mock.get(f"{EDEN_BASE}/models").mock(return_value=httpx.Response(200, json=self.CATALOG)) + + models = litellm.get_valid_models( + custom_llm_provider="edenai", check_provider_endpoint=True, api_key="listing-key" + ) + + assert models == self.ROUTABLE + + def test_a_rejected_catalog_request_surfaces_edens_status_and_body(self, eden_key, respx_mock): + """A bad key has to reach the caller as an Eden error, not as a parse failure on the + rejection body that never held a catalog.""" + respx_mock.get(f"{EDEN_BASE}/models").mock(return_value=httpx.Response(401, json={"detail": "Invalid token"})) + + with pytest.raises(EdenAIException) as rejected: + EdenAIChatConfig().get_models() + + assert rejected.value.status_code == 401 + assert "Invalid token" in rejected.value.message + + def test_proxy_wildcard_expands_to_the_live_catalog(self, eden_key, monkeypatch, respx_mock): + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + respx_mock.get(f"{EDEN_BASE}/models").mock(return_value=httpx.Response(200, json=self.CATALOG)) + + models = get_provider_models("edenai", LiteLLM_Params(model="edenai/*", api_key="wildcard-key")) + + assert models == self.ROUTABLE + + +class TestDashboardRegistration: + def test_add_model_form_offers_eden_with_a_required_key_and_optional_base(self): + fields_path = REPO_ROOT / "litellm" / "proxy" / "public_endpoints" / "provider_create_fields.json" + entries = [e for e in json.loads(fields_path.read_text()) if e["litellm_provider"] == "edenai"] + + assert len(entries) == 1 + entry = entries[0] + assert (entry["provider"], entry["provider_display_name"]) == ("EDENAI", "Eden AI") + assert entry["default_model_placeholder"].startswith("edenai/") + fields = {f["key"]: f for f in entry["credential_fields"]} + assert (fields["api_key"]["required"], fields["api_key"]["field_type"]) == (True, "password") + assert (fields["api_base"]["required"], fields["api_base"]["placeholder"]) == (False, EDEN_BASE) + + @pytest.mark.parametrize( + "matrix_path", + [ + REPO_ROOT / "provider_endpoints_support.json", + REPO_ROOT / "litellm" / "provider_endpoints_support_backup.json", + ], + ids=["root", "backup"], + ) + def test_endpoint_matrix_documents_every_served_surface(self, matrix_path): + entry = json.loads(matrix_path.read_text())["providers"]["edenai"] + + assert entry["url"] == "https://docs.litellm.ai/docs/providers/edenai" + served = {name for name, flag in entry["endpoints"].items() if flag} + assert served == { + "chat_completions", + "messages", + "responses", + "embeddings", + "image_generations", + "audio_transcriptions", + "audio_speech", + "video_generations", + } diff --git a/tests/test_litellm/llms/edenai/conftest.py b/tests/test_litellm/llms/edenai/conftest.py new file mode 100644 index 00000000000..5ca5728354c --- /dev/null +++ b/tests/test_litellm/llms/edenai/conftest.py @@ -0,0 +1,61 @@ +import asyncio +import uuid + +import pytest +import pytest_asyncio + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + +@pytest.fixture +def eden_key(monkeypatch) -> str: + monkeypatch.delenv("EDENAI_API_BASE", raising=False) + monkeypatch.setenv("EDENAI_API_KEY", "eden-test-key") + monkeypatch.setattr(litellm, "api_key", None) + return "eden-test-key" + + +@pytest.fixture +def no_eden_key(monkeypatch) -> None: + monkeypatch.delenv("EDENAI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "api_key", None) + + +class SpendCapture(CustomLogger): + """Records the cost the spend logs would store for one call, matched by its call id.""" + + def __init__(self, call_id: str): + super().__init__() + self.call_id = call_id + self.costs: list[object] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if kwargs.get("litellm_call_id") == self.call_id: + self.costs.append((kwargs.get("standard_logging_object") or {}).get("response_cost")) + + async def settle(self) -> None: + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + + +@pytest_asyncio.fixture +async def spend_capture(monkeypatch) -> SpendCapture: + GLOBAL_LOGGING_WORKER.start() # rebinds the worker's queue to this test's event loop + capture = SpendCapture(call_id=f"eden-{uuid.uuid4()}") + monkeypatch.setattr(litellm, "callbacks", [capture]) + return capture + + +@pytest.fixture +def httpx_transport(monkeypatch): + """respx fakes httpx, so the async client must not sit on LiteLLM's default aiohttp transport.""" + monkeypatch.setattr( # test-quality-ok: respx needs HTTPX enabled to fake the provider HTTP boundary. + litellm, + "disable_aiohttp_transport", + True, + ) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() diff --git a/tests/test_litellm/llms/edenai/embedding/test_edenai_embedding_transformation.py b/tests/test_litellm/llms/edenai/embedding/test_edenai_embedding_transformation.py new file mode 100644 index 00000000000..efdb428db32 --- /dev/null +++ b/tests/test_litellm/llms/edenai/embedding/test_edenai_embedding_transformation.py @@ -0,0 +1,113 @@ +"""Eden AI `/v3/embeddings`: OpenAI's embeddings API served by Eden's gateway, which reports the +real per-request cost at the top level of the body.""" + +import json + +import httpx +import pytest + +import litellm +from litellm.cost_calculator import get_response_cost_from_hidden_params +from litellm.llms.edenai.embedding.transformation import EdenAIEmbeddingConfig +from litellm.types.utils import EmbeddingResponse, LlmProviders +from litellm.utils import ProviderConfigManager + +EDEN_BASE = "https://api.edenai.run/v3" +EDEN_EMBEDDINGS_URL = f"{EDEN_BASE}/embeddings" +EDEN_REPORTED_COST = 0.0042 +MODEL = "edenai/openai/text-embedding-3-small" +SELLER_MODEL = "openai/text-embedding-3-small" +VECTOR = [0.016754150390625, -0.055755615234375] + + +def _eden_embedding(cost: float | None = EDEN_REPORTED_COST) -> dict: + """Live `/v3/embeddings` body: OpenAI shape plus Eden's top-level `cost`, `provider` and `status`.""" + body = { + "status": "success", + "model": "text-embedding-3-small", + "data": [{"embedding": VECTOR, "index": 0, "object": "embedding"}], + "object": "list", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + "provider": "openai", + } + return body if cost is None else {**body, "cost": cost} + + +def _request_body(respx_mock) -> dict: + return json.loads(respx_mock.calls.last.request.content) + + +class TestRegistration: + def test_eden_is_a_native_embedding_provider(self): + config = ProviderConfigManager.get_provider_embedding_config(model=SELLER_MODEL, provider=LlmProviders.EDENAI) + + assert isinstance(config, EdenAIEmbeddingConfig) + + +class TestAuthentication: + def test_missing_key_is_an_authentication_error_before_any_request(self, no_eden_key, respx_mock): + with pytest.raises(litellm.AuthenticationError, match="EDENAI_API_KEY"): + litellm.embedding(model=MODEL, input="hello") + assert not respx_mock.calls + + +class TestEmbedding: + def test_posts_to_eden_with_the_bearer_key_and_the_seller_model_id(self, eden_key, respx_mock): + respx_mock.post(EDEN_EMBEDDINGS_URL).mock(return_value=httpx.Response(200, json=_eden_embedding())) + + response = litellm.embedding(model=MODEL, input="hello", dimensions=2) + + assert isinstance(response, EmbeddingResponse) + assert response.data[0]["embedding"] == VECTOR + request = respx_mock.calls.last.request + assert request.headers["Authorization"] == f"Bearer {eden_key}" + assert request.headers["Content-Type"] == "application/json" + body = _request_body(respx_mock) + assert (body["model"], body["input"], body["dimensions"]) == (SELLER_MODEL, "hello", 2) + + def test_eden_reported_cost_beats_the_price_map(self, eden_key, respx_mock): + respx_mock.post(EDEN_EMBEDDINGS_URL).mock(return_value=httpx.Response(200, json=_eden_embedding())) + + response = litellm.embedding(model=MODEL, input="hello") + + assert get_response_cost_from_hidden_params(response._hidden_params) == EDEN_REPORTED_COST + assert response._hidden_params["response_cost"] == EDEN_REPORTED_COST + + def test_a_body_without_cost_leaves_pricing_to_the_price_map(self, eden_key, respx_mock): + respx_mock.post(EDEN_EMBEDDINGS_URL).mock(return_value=httpx.Response(200, json=_eden_embedding(cost=None))) + + response = litellm.embedding(model=MODEL, input="hello") + + assert get_response_cost_from_hidden_params(response._hidden_params) is None + + def test_extra_body_forwards_eden_only_fields(self, eden_key, respx_mock): + respx_mock.post(EDEN_EMBEDDINGS_URL).mock(return_value=httpx.Response(200, json=_eden_embedding())) + + litellm.embedding(model=MODEL, input="hello", extra_body={"metadata": {"trace": "abc"}}) + + assert _request_body(respx_mock)["metadata"] == {"trace": "abc"} + + @pytest.mark.asyncio + async def test_async_call_tracks_the_same_cost(self, eden_key, httpx_transport, respx_mock): + respx_mock.post(EDEN_EMBEDDINGS_URL).mock(return_value=httpx.Response(200, json=_eden_embedding())) + + response = await litellm.aembedding(model=MODEL, input=["hello", "world"]) + + assert _request_body(respx_mock)["input"] == ["hello", "world"] + assert response._hidden_params["response_cost"] == EDEN_REPORTED_COST + + +class TestErrors: + def test_middleware_401_maps_to_authentication_error(self, eden_key, respx_mock): + respx_mock.post(EDEN_EMBEDDINGS_URL).mock(return_value=httpx.Response(401, json={"detail": "Invalid token."})) + + with pytest.raises(litellm.AuthenticationError, match="Invalid token"): + litellm.embedding(model=MODEL, input="hello") + + def test_429_maps_to_rate_limit_error(self, eden_key, respx_mock): + respx_mock.post(EDEN_EMBEDDINGS_URL).mock( + return_value=httpx.Response(429, json={"error": {"message": "Rate limit exceeded", "type": "rate_limit"}}) + ) + + with pytest.raises(litellm.RateLimitError): + litellm.embedding(model=MODEL, input="hello") diff --git a/tests/test_litellm/llms/edenai/image_generation/test_edenai_image_generation_transformation.py b/tests/test_litellm/llms/edenai/image_generation/test_edenai_image_generation_transformation.py new file mode 100644 index 00000000000..d7b32affc70 --- /dev/null +++ b/tests/test_litellm/llms/edenai/image_generation/test_edenai_image_generation_transformation.py @@ -0,0 +1,124 @@ +"""Eden AI `/v3/images/generations`: OpenAI's image generation API served by Eden's gateway, which +reports the real per-request cost at the top level of the body.""" + +import json + +import httpx +import pytest + +import litellm +from litellm.cost_calculator import get_response_cost_from_hidden_params +from litellm.llms.edenai.image_generation.transformation import EdenAIImageGenerationConfig +from litellm.types.utils import ImageResponse, LlmProviders +from litellm.utils import ProviderConfigManager + +EDEN_BASE = "https://api.edenai.run/v3" +EDEN_IMAGES_URL = f"{EDEN_BASE}/images/generations" +EDEN_REPORTED_COST = 0.0042 +MODEL = "edenai/openai/gpt-image-1-mini" +SELLER_MODEL = "openai/gpt-image-1-mini" +PNG_B64 = "iVBORw0KGgoAAAANSUhE" + + +def _eden_image(cost: float | None = EDEN_REPORTED_COST) -> dict: + """Live `/v3/images/generations` body: OpenAI shape plus Eden's top-level `cost` and `provider`.""" + body = { + "created": 1788818607, + "background": None, + "data": [{"b64_json": PNG_B64, "revised_prompt": None, "url": None}], + "output_format": "png", + "quality": "low", + "size": "1024x1024", + "usage": { + "total_tokens": 281, + "input_tokens": 9, + "input_tokens_details": {"image_tokens": 0, "text_tokens": 9}, + "output_tokens": 272, + "output_tokens_details": {"image_tokens": 272, "text_tokens": 0}, + }, + "provider": "openai", + } + return body if cost is None else {**body, "cost": cost} + + +def _request_body(respx_mock) -> dict: + return json.loads(respx_mock.calls.last.request.content) + + +class TestRegistration: + def test_eden_is_a_native_image_generation_provider(self): + config = ProviderConfigManager.get_provider_image_generation_config( + model=SELLER_MODEL, provider=LlmProviders.EDENAI + ) + + assert isinstance(config, EdenAIImageGenerationConfig) + + +class TestAuthentication: + def test_missing_key_is_an_authentication_error_before_any_request(self, no_eden_key, respx_mock): + with pytest.raises(litellm.AuthenticationError, match="EDENAI_API_KEY"): + litellm.image_generation(model=MODEL, prompt="a red square") + assert not respx_mock.calls + + +class TestImageGeneration: + def test_a_param_outside_the_openai_image_set_is_rejected_unless_dropped(self, eden_key, respx_mock): + respx_mock.post(EDEN_IMAGES_URL).mock(return_value=httpx.Response(200, json=_eden_image())) + + with pytest.raises(litellm.UnsupportedParamsError, match="imageConfig"): + litellm.image_generation(model=MODEL, prompt="a red square", imageConfig={"aspectRatio": "16:9"}) + litellm.image_generation( + model=MODEL, prompt="a red square", imageConfig={"aspectRatio": "16:9"}, drop_params=True + ) + + assert "imageConfig" not in _request_body(respx_mock) + + def test_posts_to_eden_with_the_bearer_key_and_the_seller_model_id(self, eden_key, respx_mock): + respx_mock.post(EDEN_IMAGES_URL).mock(return_value=httpx.Response(200, json=_eden_image())) + + response = litellm.image_generation(model=MODEL, prompt="a red square", size="1024x1024", quality="low", n=1) + + assert isinstance(response, ImageResponse) + assert response.data[0].b64_json == PNG_B64 + assert respx_mock.calls.last.request.headers["Authorization"] == f"Bearer {eden_key}" + assert _request_body(respx_mock) == { + "model": SELLER_MODEL, + "prompt": "a red square", + "size": "1024x1024", + "quality": "low", + "n": 1, + } + + def test_eden_reported_cost_beats_the_price_map(self, eden_key, respx_mock): + respx_mock.post(EDEN_IMAGES_URL).mock(return_value=httpx.Response(200, json=_eden_image())) + + response = litellm.image_generation(model=MODEL, prompt="a red square") + + assert get_response_cost_from_hidden_params(response._hidden_params) == EDEN_REPORTED_COST + assert response._hidden_params["response_cost"] == EDEN_REPORTED_COST + + def test_a_body_without_cost_leaves_pricing_to_the_price_map(self, eden_key, respx_mock): + respx_mock.post(EDEN_IMAGES_URL).mock(return_value=httpx.Response(200, json=_eden_image(cost=None))) + + response = litellm.image_generation(model=MODEL, prompt="a red square") + + assert get_response_cost_from_hidden_params(response._hidden_params) is None + assert response.usage is not None + assert response.usage.output_tokens == 272 + + @pytest.mark.asyncio + async def test_async_call_tracks_the_same_cost(self, eden_key, httpx_transport, respx_mock): + respx_mock.post(EDEN_IMAGES_URL).mock(return_value=httpx.Response(200, json=_eden_image())) + + response = await litellm.aimage_generation(model=MODEL, prompt="a red square") + + assert response.data[0].b64_json == PNG_B64 + assert response._hidden_params["response_cost"] == EDEN_REPORTED_COST + + +class TestErrors: + def test_middleware_401_maps_to_authentication_error(self, eden_key, respx_mock): + respx_mock.post(EDEN_IMAGES_URL).mock(return_value=httpx.Response(401, json={"detail": "Invalid token."})) + + with pytest.raises(litellm.AuthenticationError, match="Invalid token"): + litellm.image_generation(model=MODEL, prompt="a red square") diff --git a/tests/test_litellm/llms/edenai/messages/test_edenai_anthropic_messages_transformation.py b/tests/test_litellm/llms/edenai/messages/test_edenai_anthropic_messages_transformation.py new file mode 100644 index 00000000000..e795ba70fb4 --- /dev/null +++ b/tests/test_litellm/llms/edenai/messages/test_edenai_anthropic_messages_transformation.py @@ -0,0 +1,247 @@ +"""Eden AI `/v3/v1/messages`: Anthropic's Messages API served by Eden's gateway for every model in +its catalog. The Anthropic payload is forwarded untranslated, and Eden reports the real per-request +cost at the top level of a non-streaming body.""" + +import asyncio +import json +import time +import uuid + +import httpx +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.edenai.messages.transformation import EdenAIAnthropicMessagesConfig +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +EDEN_BASE = "https://api.edenai.run/v3" +EDEN_EU_BASE = "https://api.eu.edenai.run/v3" +EDEN_MESSAGES_URL = f"{EDEN_BASE}/v1/messages" +EDEN_REPORTED_COST = 0.0042 +MODEL = "edenai/openai/gpt-4.1-nano" +SELLER_MODEL = "openai/gpt-4.1-nano" +MESSAGES = [{"role": "user", "content": "Say OK"}] +BILLING_BLOCK = {"type": "text", "text": "x-anthropic-billing-header: cc_version=2.1.0; cc_entrypoint=cli"} +SYSTEM_BLOCK = {"type": "text", "text": "Be terse", "cache_control": {"type": "ephemeral"}} + + +def _eden_message(cost: float | None = EDEN_REPORTED_COST) -> dict: + """Live body: Anthropic shape with the id sent to Eden echoed in `model` and Eden's top-level `cost`.""" + body = { + "id": "chatcmpl-eden-1", + "type": "message", + "role": "assistant", + "model": SELLER_MODEL, + "stop_sequence": None, + "stop_reason": "end_turn", + "usage": {"input_tokens": 12, "output_tokens": 1}, + "content": [{"type": "text", "text": "OK"}], + } + return body if cost is None else {**body, "cost": cost} + + +def _eden_stream() -> httpx.Response: + """Live stream: Anthropic events with token usage on `message_delta` and no cost anywhere.""" + message = { + "id": "msg_eden_1", + "type": "message", + "role": "assistant", + "content": [], + "model": SELLER_MODEL, + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 0, "output_tokens": 0}, + } + events = ( + {"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": {"input_tokens": 12, "output_tokens": 1}, + }, + {"type": "message_stop"}, + ) + body = "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events) + return httpx.Response(200, content=body.encode(), headers={"content-type": "text/event-stream"}) + + +def _request_body(respx_mock) -> dict: + return json.loads(respx_mock.calls.last.request.content) + + +def _logging_obj() -> Logging: + return Logging( + model=SELLER_MODEL, + messages=MESSAGES, + stream=False, + call_type="anthropic_messages", + start_time=time.time(), + litellm_call_id="eden-messages-unit", + function_id="eden-messages-unit", + ) + + +class TestRegistration: + @pytest.mark.parametrize("model", [SELLER_MODEL, "anthropic/claude-sonnet-latest"]) + def test_eden_serves_anthropic_messages_natively_for_every_catalog_model(self, model): + config = ProviderConfigManager.get_provider_anthropic_messages_config(model=model, provider=LlmProviders.EDENAI) + + assert isinstance(config, EdenAIAnthropicMessagesConfig) + assert config.custom_llm_provider == "edenai" + + +class TestEndpointResolution: + def _url(self, api_base: str | None) -> str: + return EdenAIAnthropicMessagesConfig().get_complete_url( + api_base=api_base, api_key=None, model=SELLER_MODEL, optional_params={}, litellm_params={} + ) + + def test_defaults_to_the_global_endpoint(self, eden_key): + assert self._url(None) == EDEN_MESSAGES_URL + + def test_env_api_base_moves_to_the_eu_endpoint(self, eden_key, monkeypatch): + monkeypatch.setenv("EDENAI_API_BASE", EDEN_EU_BASE) + + assert self._url(None) == f"{EDEN_EU_BASE}/v1/messages" + + def test_explicit_api_base_wins_over_env(self, eden_key, monkeypatch): + monkeypatch.setenv("EDENAI_API_BASE", EDEN_EU_BASE) + + assert self._url("https://eden.internal/v3/") == "https://eden.internal/v3/v1/messages" + + +class TestAuthentication: + def _headers(self, headers: dict, api_key: str | None = None) -> dict: + resolved, _ = EdenAIAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers=headers, + model=SELLER_MODEL, + messages=MESSAGES, + optional_params={}, + litellm_params={}, + api_key=api_key, + ) + return resolved + + def test_env_key_becomes_the_bearer_header_with_the_anthropic_version(self, eden_key): + headers = self._headers({}) + + assert headers == { + "authorization": f"Bearer {eden_key}", + "anthropic-version": "2023-06-01", + "content-type": "application/json", + } + + def test_explicit_key_wins_over_env(self, eden_key): + assert self._headers({}, api_key="explicit-key")["authorization"] == "Bearer explicit-key" + + def test_a_caller_supplied_authorization_header_is_kept(self, eden_key): + headers = self._headers({"Authorization": "Bearer caller-token"}) + + assert headers["Authorization"] == "Bearer caller-token" + assert "authorization" not in headers + + def test_missing_key_is_an_authentication_error(self, no_eden_key): + with pytest.raises(litellm.AuthenticationError, match="EDENAI_API_KEY"): + self._headers({}) + + +class TestResponseTransformation: + def test_eden_reported_cost_becomes_the_call_spend(self): + logging_obj = _logging_obj() + + response = EdenAIAnthropicMessagesConfig().transform_anthropic_messages_response( + model=SELLER_MODEL, raw_response=httpx.Response(200, json=_eden_message()), logging_obj=logging_obj + ) + + assert response["content"] == [{"type": "text", "text": "OK"}] + assert response["cost"] == EDEN_REPORTED_COST + assert logging_obj.model_call_details["response_cost"] == EDEN_REPORTED_COST + + def test_a_body_without_cost_leaves_pricing_to_the_price_map(self): + logging_obj = _logging_obj() + + EdenAIAnthropicMessagesConfig().transform_anthropic_messages_response( + model=SELLER_MODEL, raw_response=httpx.Response(200, json=_eden_message(cost=None)), logging_obj=logging_obj + ) + + assert "response_cost" not in logging_obj.model_call_details + + +class TestMessages: + @pytest.mark.asyncio + async def test_posts_the_anthropic_payload_untranslated_with_the_bearer_key( + self, eden_key, httpx_transport, respx_mock + ): + respx_mock.post(EDEN_MESSAGES_URL).mock(return_value=httpx.Response(200, json=_eden_message())) + + response = await litellm.anthropic.messages.acreate( + model=MODEL, + max_tokens=16, + messages=MESSAGES, + system=[SYSTEM_BLOCK], + thinking={"type": "enabled", "budget_tokens": 1024}, + ) + + assert response["content"] == [{"type": "text", "text": "OK"}] + assert response["cost"] == EDEN_REPORTED_COST + request = respx_mock.calls.last.request + assert request.headers["authorization"] == f"Bearer {eden_key}" + assert request.headers["anthropic-version"] == "2023-06-01" + body = _request_body(respx_mock) + assert (body["model"], body["messages"], body["max_tokens"]) == (SELLER_MODEL, MESSAGES, 16) + assert body["system"] == [SYSTEM_BLOCK] + assert body["thinking"] == {"type": "enabled", "budget_tokens": 1024} + + @pytest.mark.asyncio + async def test_claude_code_billing_blocks_are_stripped_from_the_system_prompt( + self, eden_key, httpx_transport, respx_mock + ): + respx_mock.post(EDEN_MESSAGES_URL).mock(return_value=httpx.Response(200, json=_eden_message())) + + await litellm.anthropic.messages.acreate( + model=MODEL, max_tokens=16, messages=MESSAGES, system=[BILLING_BLOCK, SYSTEM_BLOCK] + ) + + assert _request_body(respx_mock)["system"] == [SYSTEM_BLOCK] + + @pytest.mark.asyncio + async def test_eden_reported_cost_is_logged_as_the_call_spend( + self, eden_key, httpx_transport, respx_mock, spend_capture + ): + respx_mock.post(EDEN_MESSAGES_URL).mock(return_value=httpx.Response(200, json=_eden_message())) + await litellm.anthropic.messages.acreate( + model=MODEL, max_tokens=16, messages=MESSAGES, litellm_call_id=spend_capture.call_id + ) + await spend_capture.settle() + + assert spend_capture.costs == [EDEN_REPORTED_COST] + + +class TestStreaming: + @pytest.mark.asyncio + async def test_stream_forwards_eden_events_verbatim(self, eden_key, httpx_transport, respx_mock): + respx_mock.post(EDEN_MESSAGES_URL).mock(return_value=_eden_stream()) + + stream = await litellm.anthropic.messages.acreate(model=MODEL, max_tokens=16, messages=MESSAGES, stream=True) + body = b"".join([chunk async for chunk in stream]).decode() + + assert _request_body(respx_mock)["stream"] is True + assert "event: message_start" in body + assert '"text_delta", "text": "OK"' in body or '"text_delta","text":"OK"' in body + assert "event: message_stop" in body + + +class TestErrors: + @pytest.mark.asyncio + async def test_401_detail_body_is_an_authentication_error(self, eden_key, httpx_transport, respx_mock): + respx_mock.post(EDEN_MESSAGES_URL).mock(return_value=httpx.Response(401, json={"detail": "Invalid token"})) + + with pytest.raises(litellm.AuthenticationError, match="Invalid token"): + await litellm.anthropic.messages.acreate(model=MODEL, max_tokens=16, messages=MESSAGES) diff --git a/tests/test_litellm/llms/edenai/responses/test_edenai_responses_transformation.py b/tests/test_litellm/llms/edenai/responses/test_edenai_responses_transformation.py new file mode 100644 index 00000000000..3fe9e226da9 --- /dev/null +++ b/tests/test_litellm/llms/edenai/responses/test_edenai_responses_transformation.py @@ -0,0 +1,268 @@ +"""Eden AI `/v3/responses`: OpenAI's Responses API served by Eden's gateway. Eden reports the real +per-request cost at the top level of the body and, on streams, on the final usage frame.""" + +import json + +import httpx +import pytest + +import litellm +from litellm.cost_calculator import get_response_cost_from_hidden_params +from litellm.llms.edenai.responses.transformation import EdenAIResponsesAPIConfig +from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStreamEvents +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +EDEN_BASE = "https://api.edenai.run/v3" +EDEN_EU_BASE = "https://api.eu.edenai.run/v3" +EDEN_RESPONSES_URL = f"{EDEN_BASE}/responses" +EDEN_REPORTED_COST = 0.0042 +MODEL = "edenai/openai/gpt-4.1-nano" +SELLER_MODEL = "openai/gpt-4.1-nano" + + +def _usage(cost: float | None) -> dict: + usage = {"input_tokens": 12, "output_tokens": 2, "total_tokens": 14} + return usage if cost is None else {**usage, "cost": cost} + + +def _output(text: str = "OK") -> list[dict]: + return [ + { + "id": "msg_eden_1", + "type": "message", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ] + + +def _eden_response(cost: float | None = EDEN_REPORTED_COST) -> dict: + """Live `/v3/responses` body: OpenAI shape plus Eden's top-level `cost` and `provider`.""" + body = { + "id": "resp_eden_1", + "object": "response", + "created_at": 1788443790, + "status": "completed", + "model": "gpt-4.1-nano", + "provider": "openai", + "output": _output(), + "usage": _usage(cost), + } + return body if cost is None else {**body, "cost": cost} + + +def _eden_stream_events(cost: float | None = EDEN_REPORTED_COST) -> tuple[dict, ...]: + """Live stream: the `response.completed` frame carries Eden's cost on `usage` only.""" + in_progress = { + "id": "resp_eden_1", + "object": "response", + "created_at": 1788443790, + "status": "in_progress", + "model": SELLER_MODEL, + "output": [], + } + return ( + {"type": "response.created", "sequence_number": 0, "response": in_progress}, + { + "type": "response.output_item.added", + "sequence_number": 1, + "output_index": 0, + "item": { + "id": "msg_eden_1", + "type": "message", + "status": "in_progress", + "role": "assistant", + "content": [], + }, + }, + { + "type": "response.output_text.delta", + "sequence_number": 2, + "item_id": "msg_eden_1", + "output_index": 0, + "content_index": 0, + "delta": "OK", + }, + { + "type": "response.completed", + "sequence_number": 3, + "response": {**in_progress, "status": "completed", "output": _output(), "usage": _usage(cost)}, + }, + ) + + +def _sse(events: tuple[dict, ...]) -> httpx.Response: + body = "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events) + return httpx.Response(200, content=body.encode(), headers={"content-type": "text/event-stream"}) + + +def _request_body(respx_mock) -> dict: + return json.loads(respx_mock.calls.last.request.content) + + +class TestRegistration: + def test_eden_is_a_native_responses_provider(self): + config = ProviderConfigManager.get_provider_responses_api_config( + provider=LlmProviders.EDENAI, model=SELLER_MODEL + ) + + assert isinstance(config, EdenAIResponsesAPIConfig) + assert config.custom_llm_provider == LlmProviders.EDENAI + + def test_the_provider_string_resolves_too(self): + assert isinstance( + ProviderConfigManager.get_provider_responses_api_config(provider="edenai"), EdenAIResponsesAPIConfig + ) + + def test_websocket_callers_get_the_managed_handler(self): + """Eden serves the Responses API over HTTP only, so a websocket client has to be bridged + rather than dialled straight through to a wss:// endpoint Eden does not have.""" + assert EdenAIResponsesAPIConfig().supports_native_websocket() is False + + +class TestEndpointResolution: + def test_defaults_to_the_global_endpoint(self, eden_key): + assert EdenAIResponsesAPIConfig().get_complete_url(api_base=None, litellm_params={}) == EDEN_RESPONSES_URL + + def test_env_api_base_moves_to_the_eu_endpoint(self, eden_key, monkeypatch): + monkeypatch.setenv("EDENAI_API_BASE", EDEN_EU_BASE) + + assert ( + EdenAIResponsesAPIConfig().get_complete_url(api_base=None, litellm_params={}) == f"{EDEN_EU_BASE}/responses" + ) + + def test_explicit_api_base_wins_and_loses_its_trailing_slash(self, eden_key, monkeypatch): + monkeypatch.setenv("EDENAI_API_BASE", EDEN_EU_BASE) + + url = EdenAIResponsesAPIConfig().get_complete_url(api_base="https://eden.internal/v3/", litellm_params={}) + + assert url == "https://eden.internal/v3/responses" + + +class TestAuthentication: + def test_env_key_becomes_the_bearer_header(self, eden_key): + headers = EdenAIResponsesAPIConfig().validate_environment( + headers={"x-trace": "1"}, model=SELLER_MODEL, litellm_params=None + ) + + assert headers == {"x-trace": "1", "Authorization": f"Bearer {eden_key}"} + + def test_explicit_key_wins_over_env(self, eden_key): + headers = EdenAIResponsesAPIConfig().validate_environment( + headers={}, model=SELLER_MODEL, litellm_params=GenericLiteLLMParams(api_key="explicit-key") + ) + + assert headers["Authorization"] == "Bearer explicit-key" + + def test_missing_key_is_an_authentication_error(self, no_eden_key): + with pytest.raises(litellm.AuthenticationError, match="EDENAI_API_KEY"): + EdenAIResponsesAPIConfig().validate_environment(headers={}, model=SELLER_MODEL, litellm_params=None) + + +class TestResponses: + def test_posts_to_eden_with_the_bearer_key_and_the_seller_model_id(self, eden_key, respx_mock): + respx_mock.post(EDEN_RESPONSES_URL).mock(return_value=httpx.Response(200, json=_eden_response())) + + response = litellm.responses(model=MODEL, input="Say OK", max_output_tokens=16) + + assert isinstance(response, ResponsesAPIResponse) + assert response.output[0].content[0].text == "OK" + assert respx_mock.calls.last.request.headers["Authorization"] == f"Bearer {eden_key}" + body = _request_body(respx_mock) + assert (body["model"], body["input"], body["max_output_tokens"]) == (SELLER_MODEL, "Say OK", 16) + + def test_eden_reported_cost_beats_the_price_map(self, eden_key, respx_mock): + respx_mock.post(EDEN_RESPONSES_URL).mock(return_value=httpx.Response(200, json=_eden_response())) + + response = litellm.responses(model=MODEL, input="Say OK", max_output_tokens=16) + + assert get_response_cost_from_hidden_params(response._hidden_params) == EDEN_REPORTED_COST + assert response._hidden_params["response_cost"] == EDEN_REPORTED_COST + + def test_a_body_without_cost_leaves_pricing_to_the_price_map(self, eden_key, respx_mock): + respx_mock.post(EDEN_RESPONSES_URL).mock(return_value=httpx.Response(200, json=_eden_response(cost=None))) + + response = litellm.responses(model=MODEL, input="Say OK", max_output_tokens=16) + + assert response.output[0].content[0].text == "OK" + assert get_response_cost_from_hidden_params(response._hidden_params) is None + + def test_stateful_params_pass_through_to_eden(self, eden_key, respx_mock): + respx_mock.post(EDEN_RESPONSES_URL).mock(return_value=httpx.Response(200, json=_eden_response())) + + litellm.responses( + model=MODEL, + input="Say OK", + previous_response_id="resp_previous", + store=False, + reasoning={"effort": "low"}, + ) + + body = _request_body(respx_mock) + assert (body["previous_response_id"], body["store"], body["reasoning"]) == ( + "resp_previous", + False, + {"effort": "low"}, + ) + + def test_extra_body_forwards_eden_only_fields(self, eden_key, respx_mock): + respx_mock.post(EDEN_RESPONSES_URL).mock(return_value=httpx.Response(200, json=_eden_response())) + + litellm.responses( + model=MODEL, + input="Say OK", + extra_body={"fallbacks": ["anthropic/claude-sonnet-latest"], "routing": {"sort": "latency"}}, + ) + + body = _request_body(respx_mock) + assert body["fallbacks"] == ["anthropic/claude-sonnet-latest"] + assert body["routing"] == {"sort": "latency"} + assert "extra_body" not in body + + +class TestStreaming: + def test_stream_forwards_eden_events_and_bills_the_usage_cost(self, eden_key, respx_mock): + respx_mock.post(EDEN_RESPONSES_URL).mock(return_value=_sse(_eden_stream_events())) + + stream = litellm.responses(model=MODEL, input="Say OK", stream=True) + events = list(stream) + + assert _request_body(respx_mock)["stream"] is True + assert [event.type for event in events] == [ + ResponsesAPIStreamEvents.RESPONSE_CREATED, + ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + ] + assert events[2].delta == "OK" + assert events[-1].response.usage.cost == EDEN_REPORTED_COST + assert stream.logging_obj.model_call_details["response_cost"] == EDEN_REPORTED_COST + + +class TestErrors: + def test_401_detail_body_is_an_authentication_error(self, eden_key, respx_mock): + respx_mock.post(EDEN_RESPONSES_URL).mock(return_value=httpx.Response(401, json={"detail": "Invalid token"})) + + with pytest.raises(litellm.AuthenticationError, match="Invalid token"): + litellm.responses(model=MODEL, input="Say OK") + + def test_400_envelope_is_a_bad_request_error(self, eden_key, respx_mock): + respx_mock.post(EDEN_RESPONSES_URL).mock( + return_value=httpx.Response( + 400, + json={ + "error": { + "message": "Model(s) not found or inactive: openai/does-not-exist", + "type": "invalid_request_error", + "param": None, + "code": "invalid_parameter", + } + }, + ) + ) + + with pytest.raises(litellm.BadRequestError, match="not found or inactive"): + litellm.responses(model="edenai/openai/does-not-exist", input="Say OK") diff --git a/tests/test_litellm/llms/edenai/test_edenai_common_utils.py b/tests/test_litellm/llms/edenai/test_edenai_common_utils.py new file mode 100644 index 00000000000..01b7ef55f81 --- /dev/null +++ b/tests/test_litellm/llms/edenai/test_edenai_common_utils.py @@ -0,0 +1,63 @@ +"""Credential, endpoint and cost helpers shared by every Eden AI config.""" + +import httpx +import pytest + +import litellm +from litellm.llms.edenai.common_utils import authorized_headers, endpoint_url, json_headers, reported_cost + +EDEN_BASE = "https://api.edenai.run/v3" +EDEN_EU_BASE = "https://api.eu.edenai.run/v3" + + +class TestEndpointUrl: + def test_defaults_to_the_global_endpoint(self, eden_key): + assert endpoint_url(None, "embeddings") == f"{EDEN_BASE}/embeddings" + + def test_env_api_base_moves_to_the_eu_endpoint(self, eden_key, monkeypatch): + monkeypatch.setenv("EDENAI_API_BASE", EDEN_EU_BASE) + + assert endpoint_url(None, "audio/speech") == f"{EDEN_EU_BASE}/audio/speech" + + def test_explicit_api_base_wins_and_loses_its_trailing_slash(self, eden_key, monkeypatch): + monkeypatch.setenv("EDENAI_API_BASE", EDEN_EU_BASE) + + assert ( + endpoint_url("https://proxy.example/v3/", "images/generations") + == "https://proxy.example/v3/images/generations" + ) + + +class TestAuthorizedHeaders: + def test_env_key_becomes_the_bearer_header_and_caller_headers_are_kept(self, eden_key): + assert authorized_headers({"X-Trace": "abc"}, None, "openai/tts-1") == { + "X-Trace": "abc", + "Authorization": f"Bearer {eden_key}", + } + + def test_explicit_key_wins_over_env(self, eden_key): + assert authorized_headers({}, "explicit-key", "openai/tts-1")["Authorization"] == "Bearer explicit-key" + + def test_json_headers_add_the_content_type(self, eden_key): + assert json_headers({}, None, "openai/tts-1") == { + "Authorization": f"Bearer {eden_key}", + "Content-Type": "application/json", + } + + def test_missing_key_is_an_authentication_error(self, no_eden_key): + with pytest.raises(litellm.AuthenticationError, match="EDENAI_API_KEY"): + authorized_headers({}, None, "openai/tts-1") + + +class TestReportedCost: + def test_reads_the_top_level_cost_of_a_body(self): + assert reported_cost({"cost": 0.0042, "provider": "openai"}) == 0.0042 + assert reported_cost(b'{"cost": 0.0042, "text": "hi"}') == 0.0042 + + def test_reads_the_speech_cost_header(self): + assert reported_cost(httpx.Headers({"x-edenai-cost": "0.00015", "content-type": "audio/mpeg"})) == 0.00015 + + def test_no_cost_anywhere_is_none(self): + assert reported_cost({"provider": "openai"}) is None + assert reported_cost(httpx.Headers({"content-type": "audio/mpeg"})) is None + assert reported_cost(b"not json") is None diff --git a/tests/test_litellm/llms/edenai/text_to_speech/test_edenai_text_to_speech_transformation.py b/tests/test_litellm/llms/edenai/text_to_speech/test_edenai_text_to_speech_transformation.py new file mode 100644 index 00000000000..922713da63e --- /dev/null +++ b/tests/test_litellm/llms/edenai/text_to_speech/test_edenai_text_to_speech_transformation.py @@ -0,0 +1,139 @@ +"""Eden AI `/v3/audio/speech`: OpenAI's text-to-speech API served by Eden's gateway. The answer is +raw audio, so Eden reports the real per-request cost in the `x-edenai-cost` response header.""" + +import asyncio +import json +import uuid + +import httpx +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.edenai.common_utils import EdenAIException +from litellm.llms.edenai.text_to_speech.transformation import EdenAITextToSpeechConfig +from litellm.types.llms.openai import HttpxBinaryResponseContent +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +EDEN_BASE = "https://api.edenai.run/v3" +EDEN_SPEECH_URL = f"{EDEN_BASE}/audio/speech" +EDEN_REPORTED_COST = 0.00015 +MODEL = "edenai/openai/tts-1" +SELLER_MODEL = "openai/tts-1" +AUDIO = b"ID3\x04\x00fake-mp3-bytes" + + +def _eden_audio(cost: float | None = EDEN_REPORTED_COST) -> httpx.Response: + """Live `/v3/audio/speech` answer: audio bytes, with the cost and provider in `x-edenai-*` headers.""" + headers = {"content-type": "audio/mpeg", "x-edenai-provider": "openai"} + return httpx.Response( + 200, content=AUDIO, headers=headers if cost is None else {**headers, "x-edenai-cost": str(cost)} + ) + + +def _request_body(respx_mock) -> dict: + return json.loads(respx_mock.calls.last.request.content) + + +class TestRegistration: + def test_eden_is_a_native_text_to_speech_provider(self): + config = ProviderConfigManager.get_provider_text_to_speech_config( + model=SELLER_MODEL, provider=LlmProviders.EDENAI + ) + + assert isinstance(config, EdenAITextToSpeechConfig) + + +class TestRequestTransformation: + def test_body_is_the_openai_speech_request_without_empty_fields(self): + request = EdenAITextToSpeechConfig().transform_text_to_speech_request( + model=SELLER_MODEL, + input="hello there", + voice="alloy", + optional_params={"response_format": "wav", "speed": None}, + litellm_params={}, + headers={}, + ) + + assert request["dict_body"] == { + "model": SELLER_MODEL, + "input": "hello there", + "voice": "alloy", + "response_format": "wav", + } + + def test_a_missing_voice_is_left_for_eden_to_reject(self): + request = EdenAITextToSpeechConfig().transform_text_to_speech_request( + model=SELLER_MODEL, input="hello", voice=None, optional_params={}, litellm_params={}, headers={} + ) + + assert "voice" not in request["dict_body"] + + def test_missing_key_is_an_authentication_error_before_any_request(self, no_eden_key, respx_mock): + with pytest.raises(litellm.AuthenticationError, match="EDENAI_API_KEY"): + litellm.speech(model=MODEL, input="hello", voice="alloy") + assert not respx_mock.calls + + +class TestSpeech: + def test_posts_to_eden_with_the_bearer_key_and_returns_the_audio(self, eden_key, respx_mock): + respx_mock.post(EDEN_SPEECH_URL).mock(return_value=_eden_audio()) + + response = litellm.speech(model=MODEL, input="hello there", voice="alloy", response_format="mp3", speed=1.2) + + assert isinstance(response, HttpxBinaryResponseContent) + assert response.content == AUDIO + assert respx_mock.calls.last.request.headers["Authorization"] == f"Bearer {eden_key}" + assert _request_body(respx_mock) == { + "model": SELLER_MODEL, + "input": "hello there", + "voice": "alloy", + "response_format": "mp3", + "speed": 1.2, + } + + def test_the_cost_header_becomes_the_response_cost(self, eden_key, respx_mock): + respx_mock.post(EDEN_SPEECH_URL).mock(return_value=_eden_audio()) + + response = litellm.speech(model=MODEL, input="hello there", voice="alloy") + + assert response._hidden_params["response_cost"] == EDEN_REPORTED_COST + + def test_an_answer_without_the_cost_header_leaves_pricing_to_the_price_map(self): + response = EdenAITextToSpeechConfig().transform_text_to_speech_response( + model=SELLER_MODEL, raw_response=_eden_audio(cost=None), logging_obj=None + ) + + assert "response_cost" not in response._hidden_params + + @pytest.mark.asyncio + async def test_async_call_logs_the_header_cost_as_spend(self, eden_key, httpx_transport, spend_capture, respx_mock): + respx_mock.post(EDEN_SPEECH_URL).mock(return_value=_eden_audio()) + + response = await litellm.aspeech( + model=MODEL, input="hello there", voice="alloy", litellm_call_id=spend_capture.call_id + ) + await spend_capture.settle() + + assert response.content == AUDIO + assert spend_capture.costs == [EDEN_REPORTED_COST] + + +class TestErrors: + def test_middleware_401_surfaces_as_an_eden_error_with_the_status_code(self, eden_key, respx_mock): + """`litellm.speech` does not map provider errors onto the OpenAI exception classes the way + chat does, so the proxy relies on the status code the provider exception carries.""" + respx_mock.post(EDEN_SPEECH_URL).mock(return_value=httpx.Response(401, json={"detail": "Invalid token."})) + + with pytest.raises(EdenAIException, match="Invalid token") as excinfo: + litellm.speech(model=MODEL, input="hello", voice="alloy") + assert excinfo.value.status_code == 401 + + @pytest.mark.asyncio + async def test_async_401_maps_to_authentication_error(self, eden_key, httpx_transport, respx_mock): + respx_mock.post(EDEN_SPEECH_URL).mock(return_value=httpx.Response(401, json={"detail": "Invalid token."})) + + with pytest.raises(litellm.AuthenticationError, match="Invalid token"): + await litellm.aspeech(model=MODEL, input="hello", voice="alloy") diff --git a/tests/test_litellm/llms/edenai/videos/test_edenai_video_transformation.py b/tests/test_litellm/llms/edenai/videos/test_edenai_video_transformation.py new file mode 100644 index 00000000000..360e4d07f24 --- /dev/null +++ b/tests/test_litellm/llms/edenai/videos/test_edenai_video_transformation.py @@ -0,0 +1,308 @@ +"""Eden AI `/v3/videos`: OpenAI's video jobs API served by Eden's gateway, which reports `cost` as 0 +while a job is queued and the settled amount on the status read once it completes.""" + +import json +from io import BytesIO + +import httpx +import pytest + +import litellm +from litellm.llms.edenai.videos.transformation import EdenAIVideoConfig +from litellm.types.utils import LlmProviders +from litellm.types.videos.main import VideoObject +from litellm.types.videos.utils import decode_video_id_with_provider, encode_video_id_with_provider +from litellm.utils import ProviderConfigManager + +EDEN_BASE = "https://api.edenai.run/v3" +EDEN_VIDEOS_URL = f"{EDEN_BASE}/videos" +MODEL = "edenai/pruna/p-video" +SELLER_MODEL = "pruna/p-video" +JOB_ID = "fcd74ecd-23df-4eea-a372-478a1e842d42" +SETTLED_COST = 0.08 +FILE_URL = "https://files.example.net/60b11f54/video.mp4" +MP4_BYTES = b"\x00\x00\x00\x18ftypmp42" +PROMPT = "a red ball rolling on a wooden table" + + +def _eden_video(status: str = "queued", cost: float = 0.0, **overrides: object) -> dict: + """Live `/v3/videos` body: OpenAI's video object plus Eden's top-level `provider` and `cost`.""" + return { + "id": JOB_ID, + "object": "video", + "status": status, + "progress": 100 if status == "completed" else 0, + "created_at": 1789067483, + "completed_at": 1789067493 if status == "completed" else None, + "expires_at": None, + "model": SELLER_MODEL, + "seconds": "4", + "size": "1280x720", + "remixed_from_video_id": None, + "error": None, + "provider": "pruna", + "cost": cost, + **overrides, + } + + +def _encoded(job_id: str = JOB_ID) -> str: + return encode_video_id_with_provider(job_id, "edenai", SELLER_MODEL) + + +def _request_body(respx_mock) -> dict: + return json.loads(respx_mock.calls.last.request.content) + + +class TestRegistration: + def test_eden_is_a_native_video_provider(self): + config = ProviderConfigManager.get_provider_video_config(model=SELLER_MODEL, provider=LlmProviders.EDENAI) + + assert isinstance(config, EdenAIVideoConfig) + + +class TestAuthentication: + def test_missing_key_is_an_authentication_error_before_any_request(self, no_eden_key, respx_mock): + with pytest.raises(litellm.AuthenticationError, match="EDENAI_API_KEY"): + litellm.video_generation(model=MODEL, prompt=PROMPT) + assert not respx_mock.calls + + +class TestCreate: + def test_posts_json_to_eden_with_the_bearer_key_and_the_seller_model_id(self, eden_key, respx_mock): + respx_mock.post(EDEN_VIDEOS_URL).mock(return_value=httpx.Response(200, json=_eden_video())) + + response = litellm.video_generation(model=MODEL, prompt=PROMPT, seconds="4", size="1280x720") + + assert isinstance(response, VideoObject) + assert response.status == "queued" + request = respx_mock.calls.last.request + assert request.headers["Authorization"] == f"Bearer {eden_key}" + assert request.headers["Content-Type"] == "application/json" + assert json.loads(request.content) == { + "model": SELLER_MODEL, + "prompt": PROMPT, + "seconds": "4", + "size": "1280x720", + } + + def test_the_returned_id_routes_later_calls_back_to_eden(self, eden_key, respx_mock): + respx_mock.post(EDEN_VIDEOS_URL).mock(return_value=httpx.Response(200, json=_eden_video())) + + response = litellm.video_generation(model=MODEL, prompt=PROMPT) + + assert decode_video_id_with_provider(response.id) == { + "custom_llm_provider": "edenai", + "model_id": SELLER_MODEL, + "video_id": JOB_ID, + } + + def test_eden_extensions_go_through_as_kwargs_and_extra_body(self, eden_key, respx_mock): + respx_mock.post(EDEN_VIDEOS_URL).mock(return_value=httpx.Response(200, json=_eden_video())) + + litellm.video_generation(model=MODEL, prompt=PROMPT, seed=7, extra_body={"provider_params": {"guidance": 2}}) + + body = _request_body(respx_mock) + assert (body["seed"], body["provider_params"]) == (7, {"guidance": 2}) + + def test_a_reference_image_file_makes_the_request_multipart(self, eden_key, respx_mock): + respx_mock.post(EDEN_VIDEOS_URL).mock(return_value=httpx.Response(200, json=_eden_video())) + reference = BytesIO(b"\x89PNG\r\n\x1a\n" + b"\x00" * 16) + + litellm.video_generation(model=MODEL, prompt="animate this", input_reference=reference, seconds="4") + + request = respx_mock.calls.last.request + assert request.headers["Content-Type"].startswith("multipart/form-data") + assert b'name="input_reference"; filename="input_reference.png"' in request.content + assert b'name="model"\r\n\r\n' + SELLER_MODEL.encode() in request.content + assert b'name="seconds"\r\n\r\n4' in request.content + + def test_a_reference_image_url_stays_in_the_json_body(self, eden_key, respx_mock): + respx_mock.post(EDEN_VIDEOS_URL).mock(return_value=httpx.Response(200, json=_eden_video())) + + litellm.video_generation( + model=MODEL, prompt="animate this", input_reference={"image_url": "https://img.example.net/start.png"} + ) + + request = respx_mock.calls.last.request + assert request.headers["Content-Type"] == "application/json" + assert json.loads(request.content)["input_reference"] == {"image_url": "https://img.example.net/start.png"} + + def test_a_queued_job_reports_edens_zero_cost_and_the_requested_duration(self, eden_key, respx_mock): + respx_mock.post(EDEN_VIDEOS_URL).mock(return_value=httpx.Response(200, json=_eden_video())) + + response = litellm.video_generation(model=MODEL, prompt=PROMPT, seconds="4") + + assert response.usage == {"duration_seconds": 4.0, "provider_reported_cost_usd": 0.0} + + @pytest.mark.asyncio + async def test_a_queued_job_bills_nothing_until_it_settles( + self, eden_key, httpx_transport, respx_mock, spend_capture + ): + respx_mock.post(EDEN_VIDEOS_URL).mock(return_value=httpx.Response(200, json=_eden_video())) + + await litellm.avideo_generation(model=MODEL, prompt=PROMPT, seconds="4", litellm_call_id=spend_capture.call_id) + await spend_capture.settle() + + assert spend_capture.costs == [0.0] + + @pytest.mark.asyncio + async def test_a_cost_settled_on_the_create_response_is_billed( + self, eden_key, httpx_transport, respx_mock, spend_capture + ): + respx_mock.post(EDEN_VIDEOS_URL).mock( + return_value=httpx.Response(200, json=_eden_video(status="completed", cost=SETTLED_COST)) + ) + + await litellm.avideo_generation(model=MODEL, prompt=PROMPT, seconds="4", litellm_call_id=spend_capture.call_id) + await spend_capture.settle() + + assert spend_capture.costs == [SETTLED_COST] + + +class TestStatus: + def test_reads_the_job_with_the_bearer_key_and_surfaces_the_settled_cost(self, eden_key, respx_mock): + respx_mock.get(f"{EDEN_VIDEOS_URL}/{JOB_ID}").mock( + return_value=httpx.Response( + 200, json=_eden_video(status="completed", cost=SETTLED_COST, seconds=None, size=None) + ) + ) + + response = litellm.video_status(video_id=_encoded()) + + assert respx_mock.calls.last.request.headers["Authorization"] == f"Bearer {eden_key}" + assert (response.status, response.progress) == ("completed", 100) + assert response.usage == {"provider_reported_cost_usd": SETTLED_COST} + assert decode_video_id_with_provider(response.id)["video_id"] == JOB_ID + + @pytest.mark.asyncio + async def test_polling_a_finished_job_does_not_bill_it_again( + self, eden_key, httpx_transport, respx_mock, spend_capture + ): + respx_mock.get(f"{EDEN_VIDEOS_URL}/{JOB_ID}").mock( + return_value=httpx.Response(200, json=_eden_video(status="completed", cost=SETTLED_COST)) + ) + + await litellm.avideo_status(video_id=_encoded(), litellm_call_id=spend_capture.call_id) + await spend_capture.settle() + + assert len(spend_capture.costs) == 1 + assert not spend_capture.costs[0] + + def test_an_unknown_job_is_a_not_found_error(self, eden_key, respx_mock): + respx_mock.get(f"{EDEN_VIDEOS_URL}/{JOB_ID}").mock( + return_value=httpx.Response( + 404, + json={ + "error": { + "message": f"Video {JOB_ID} not found", + "type": "invalid_request_error", + "param": None, + "code": "model_not_found", + } + }, + ) + ) + + with pytest.raises(litellm.NotFoundError, match="not found"): + litellm.video_status(video_id=_encoded()) + + +class TestContent: + def test_follows_edens_redirect_to_the_file_without_forwarding_the_key(self, eden_key, respx_mock): + respx_mock.get(f"{EDEN_VIDEOS_URL}/{JOB_ID}/content").mock( + return_value=httpx.Response(302, headers={"location": FILE_URL}) + ) + respx_mock.get(FILE_URL).mock( + return_value=httpx.Response(200, content=MP4_BYTES, headers={"content-type": "binary/octet-stream"}) + ) + + video = litellm.video_content(video_id=_encoded()) + + assert video == MP4_BYTES + eden_request, file_request = (call.request for call in respx_mock.calls) + assert eden_request.headers["Authorization"] == f"Bearer {eden_key}" + assert "Authorization" not in file_request.headers + + @pytest.mark.asyncio + async def test_async_download_follows_the_same_redirect(self, eden_key, httpx_transport, respx_mock): + respx_mock.get(f"{EDEN_VIDEOS_URL}/{JOB_ID}/content").mock( + return_value=httpx.Response(302, headers={"location": FILE_URL}) + ) + respx_mock.get(FILE_URL).mock(return_value=httpx.Response(200, content=MP4_BYTES)) + + assert await litellm.avideo_content(video_id=_encoded()) == MP4_BYTES + + +class TestList: + def test_lists_jobs_newest_first_with_encoded_ids_and_their_costs(self, eden_key, httpx_transport, respx_mock): + """The sync entry point runs the async handler, so the client must sit on httpx for respx to see it.""" + older = "d544c281-9099-487e-b537-5f2291b603c8" + respx_mock.get(host="api.edenai.run", path="/v3/videos").mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [ + _eden_video(status="completed", cost=0.02, seconds=None, size=None), + _eden_video(status="completed", cost=0.1, id=older, seconds=None, size=None), + ], + "first_id": JOB_ID, + "last_id": older, + "has_more": True, + }, + ) + ) + + page = litellm.video_list(custom_llm_provider="edenai", limit=2) + + assert respx_mock.calls.last.request.url.params["limit"] == "2" + assert [decode_video_id_with_provider(video["id"])["video_id"] for video in page["data"]] == [JOB_ID, older] + assert [video["cost"] for video in page["data"]] == [0.02, 0.1] + assert decode_video_id_with_provider(page["last_id"]) == { + "custom_llm_provider": "edenai", + "model_id": SELLER_MODEL, + "video_id": older, + } + + +class TestErrors: + def test_middleware_401_maps_to_authentication_error(self, eden_key, respx_mock): + respx_mock.post(EDEN_VIDEOS_URL).mock(return_value=httpx.Response(401, json={"detail": "Invalid token"})) + + with pytest.raises(litellm.AuthenticationError, match="Invalid token"): + litellm.video_generation(model=MODEL, prompt=PROMPT) + + def test_a_401_on_a_read_is_an_authentication_error_too(self, eden_key, httpx_transport, respx_mock): + respx_mock.get(host="api.edenai.run", path="/v3/videos").mock( + return_value=httpx.Response(401, json={"detail": "Invalid token"}) + ) + respx_mock.get(f"{EDEN_VIDEOS_URL}/{JOB_ID}/content").mock( + return_value=httpx.Response(401, json={"detail": "Invalid token"}) + ) + + with pytest.raises(litellm.AuthenticationError, match="Invalid token"): + litellm.video_list(custom_llm_provider="edenai") + with pytest.raises(litellm.AuthenticationError, match="Invalid token"): + litellm.video_content(video_id=_encoded()) + + def test_an_openai_param_eden_does_not_accept_yet_is_forwarded_and_eden_answers(self, eden_key, respx_mock): + """OpenAI's full video param set goes through untouched, so Eden's own validation is what a caller + sees today and nothing here needs to change once Eden accepts these fields.""" + respx_mock.post(EDEN_VIDEOS_URL).mock( + return_value=httpx.Response( + 422, + json={ + "error": { + "message": "Extra inputs are not permitted", + "type": "invalid_request_error", + "param": "user", + "code": "invalid_parameter", + } + }, + ) + ) + + with pytest.raises(litellm.BadRequestError, match="Extra inputs"): + litellm.video_generation(model=MODEL, prompt=PROMPT, user="u1") + assert _request_body(respx_mock)["user"] == "u1" diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py index 5e2e4532265..86ecbf6701b 100644 --- a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py +++ b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -1,4 +1,5 @@ -from unittest.mock import Mock +from typing import Final +from unittest.mock import AsyncMock, Mock import httpx import pytest @@ -11,12 +12,15 @@ from litellm.llms.fal_ai.videos.transformation import ( FalAIVideoError, _queue_request_base_path, ) +from litellm.llms.openai.cost_calculation import video_generation_cost from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders from litellm.types.videos.utils import decode_video_id_with_provider from litellm.utils import ProviderConfigManager MODEL = "bytedance/seedance-2.5/text-to-video" +H3_TEXT_MODEL = "minimax/h3/text-to-video" +H3_REFERENCE_MODEL = "minimax/h3/reference-to-video" class TestFalAIVideoTransformation: @@ -64,6 +68,24 @@ class TestFalAIVideoTransformation: with pytest.raises(ValueError, match="public image URL"): self.config.map_openai_params({"input_reference": b"image"}, MODEL, False) + def test_map_openai_params_supports_h3_profiles(self): + url = "https://example.com/image.png" + + assert self.config.map_openai_params({"size": "2k"}, H3_TEXT_MODEL, False) == {"resolution": "2K"} + assert self.config.map_openai_params({"size": "1024x768"}, H3_TEXT_MODEL, False) == { + "resolution": "768P", + "aspect_ratio": "4:3", + } + mapped = self.config.map_openai_params( + {"seconds": 6, "input_reference": url}, + H3_REFERENCE_MODEL, + False, + ) + assert mapped["duration"] == 6 + assert isinstance(mapped["duration"], int) + assert mapped["reference_image_urls"] == [url] + assert "image_url" not in mapped + def test_transform_video_create_request(self): body, files, url = self.config.transform_video_create_request( model=MODEL, @@ -141,6 +163,20 @@ class TestFalAIVideoTransformation: assert auto_video.seconds is None assert auto_video.size is None + def test_transform_video_create_response_uses_h3_default_resolution(self): + response = Mock(spec=httpx.Response) + response.json.return_value = {"request_id": "abc"} + + video = self.config.transform_video_create_response( + model=H3_TEXT_MODEL, + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + request_data={"duration": 5}, + ) + + assert video.usage == {"duration_seconds": 5.0, "video_resolution": "2K"} + def test_status_request_uses_queue_base_path(self): response = Mock(spec=httpx.Response) response.json.return_value = {"request_id": "abc"} @@ -183,8 +219,18 @@ class TestFalAIVideoTransformation: def test_status_response_mapping(self, response_data, expected_status): status_url = "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status" response = httpx.Response(200, json=response_data, request=httpx.Request("GET", status_url)) + config = self.config + if expected_status == "completed": + result_response: Final = httpx.Response( + 200, + json={"video": {"url": "https://cdn.example.com/video.mp4"}}, + request=httpx.Request("GET", status_url.removesuffix("/status")), + ) + client: Final = Mock() + client.get.return_value = result_response + config = FalAIVideoConfig(sync_client_factory=lambda: client) - video = self.config.transform_video_status_retrieve_response( + video = config.transform_video_status_retrieve_response( raw_response=response, logging_obj=self.logging_obj, custom_llm_provider="fal_ai", @@ -210,16 +256,22 @@ class TestFalAIVideoTransformation: "status": "COMPLETED", "error": "generation failed", } + status_url = "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status" response = httpx.Response( 200, json=response_data, - request=httpx.Request( - "GET", - "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status", - ), + request=httpx.Request("GET", status_url), ) + result_response: Final = httpx.Response( + 200, + json={"video": {"url": "https://cdn.example.com/video.mp4"}}, + request=httpx.Request("GET", status_url.removesuffix("/status")), + ) + client: Final = Mock() + client.get.return_value = result_response + config = FalAIVideoConfig(sync_client_factory=lambda: client) - video = self.config.transform_video_status_retrieve_response( + video = config.transform_video_status_retrieve_response( raw_response=response, logging_obj=self.logging_obj, custom_llm_provider="fal_ai", @@ -228,8 +280,125 @@ class TestFalAIVideoTransformation: assert video.status == "failed" assert video.error == {"code": "fal_error", "message": "generation failed"} - def test_status_response_uses_namespaced_request_url(self): + def test_status_completed_result_error_surfaces_fal_message(self): + status_url = "https://queue.fal.run/minimax/h3/requests/abc/status" + auth_headers: Final = {"Authorization": "Key synthetic-fal-key", "Content-Type": "application/json"} + response: Final = httpx.Response( + 200, + json={"request_id": "abc", "status": "COMPLETED"}, + request=httpx.Request("GET", status_url, headers=auth_headers), + ) + result_url: Final = status_url.removesuffix("/status") + result_response: Final = httpx.Response( + 422, + json={ + "detail": [ + { + "loc": ["body", "input.reference_image_urls"], + "msg": "Failed to download the file. Please check if the URL is accessible and try again.", + } + ] + }, + request=httpx.Request("GET", result_url, headers=auth_headers), + ) + client: Final = Mock() + client.get.return_value = result_response + config = FalAIVideoConfig(sync_client_factory=lambda: client) + + video = config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + assert video.status == "failed" + assert "input.reference_image_urls: Failed to download the file" in video.error["message"] + client.get.assert_called_once_with(url=result_url, headers=auth_headers) + + @pytest.mark.parametrize("status_code", [429, 503]) + def test_status_completed_transient_result_error_keeps_completed(self, status_code): + status_url = "https://queue.fal.run/minimax/h3/requests/abc/status" + response: Final = httpx.Response( + 200, + json={"request_id": "abc", "status": "COMPLETED"}, + request=httpx.Request("GET", status_url), + ) + result_response: Final = httpx.Response( + status_code, + json={"detail": "temporary fal failure"}, + request=httpx.Request("GET", status_url.removesuffix("/status")), + ) + client: Final = Mock() + client.get.return_value = result_response + config = FalAIVideoConfig(sync_client_factory=lambda: client) + + video = config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + assert video.status == "completed" + assert video.error is None + + @pytest.mark.asyncio + async def test_async_status_completed_result_error_surfaces_fal_message(self): + status_url = "https://queue.fal.run/minimax/h3/requests/abc/status" + auth_headers: Final = {"Authorization": "Key synthetic-fal-key", "Content-Type": "application/json"} + response: Final = httpx.Response( + 200, + json={"request_id": "abc", "status": "COMPLETED"}, + request=httpx.Request("GET", status_url, headers=auth_headers), + ) + result_url: Final = status_url.removesuffix("/status") + result_response: Final = httpx.Response( + 422, + json={ + "detail": [ + { + "loc": ["body", "input.reference_image_urls"], + "msg": "Failed to download the file. Please check if the URL is accessible and try again.", + } + ] + }, + request=httpx.Request("GET", result_url, headers=auth_headers), + ) + client: Final = Mock() + client.get = AsyncMock(return_value=result_response) + config = FalAIVideoConfig(async_client_factory=lambda: client) + + video = await config.async_transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + assert video.status == "failed" + assert "input.reference_image_urls: Failed to download the file" in video.error["message"] + client.get.assert_awaited_once_with(url=result_url, headers=auth_headers) + + def test_status_in_progress_does_not_fetch_result(self): + status_url = "https://queue.fal.run/minimax/h3/requests/abc/status" response = httpx.Response( + 200, + json={"request_id": "abc", "status": "IN_PROGRESS"}, + request=httpx.Request("GET", status_url), + ) + + client: Final = Mock() + config = FalAIVideoConfig(sync_client_factory=lambda: client) + + video = config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + assert video.status == "in_progress" + client.get.assert_not_called() + + def test_status_response_uses_namespaced_request_url(self): + response: Final = httpx.Response( 200, json={"status": "IN_PROGRESS"}, request=httpx.Request( @@ -249,7 +418,7 @@ class TestFalAIVideoTransformation: assert decoded["video_id"] == "xyz" assert video.model == "workflows/owner/app" - def test_content_response_downloads_video_url(self, monkeypatch): + def test_content_response_downloads_video_url(self): content_response = httpx.Response( 200, content=b"video-bytes", @@ -261,11 +430,11 @@ class TestFalAIVideoTransformation: assert url == "https://cdn.example.com/video.mp4" return content_response - monkeypatch.setattr(fal_video_module, "_get_httpx_client", lambda: FakeHTTPClient()) + config = FalAIVideoConfig(sync_client_factory=FakeHTTPClient) response = Mock(spec=httpx.Response) response.json.return_value = {"video": {"url": "https://cdn.example.com/video.mp4"}} - assert self.config.transform_video_content_response(response, self.logging_obj) == b"video-bytes" + assert config.transform_video_content_response(response, self.logging_obj) == b"video-bytes" def test_content_response_rejects_missing_video(self): response = Mock(spec=httpx.Response) @@ -274,6 +443,87 @@ class TestFalAIVideoTransformation: with pytest.raises(ValueError, match="generation failed"): self.config.transform_video_content_response(response, self.logging_obj) + def test_content_response_surfaces_list_detail_error(self): + response: Final = httpx.Response( + 422, + json={ + "detail": [ + { + "loc": ["body", "input.reference_image_urls"], + "msg": "Failed to download the file. Please check if the URL is accessible and try again.", + } + ] + }, + request=httpx.Request("GET", "https://queue.fal.run/minimax/h3/requests/abc"), + ) + + with pytest.raises(FalAIVideoError) as error: + self.config.transform_video_content_response(response, self.logging_obj) + + assert error.value.status_code == 422 + assert "input.reference_image_urls: Failed to download the file" in error.value.message + assert "Failed to download the file" in error.value.response.text + + def test_content_response_surfaces_string_detail_error(self): + response: Final = httpx.Response( + 400, + json={"detail": "Request is still in progress"}, + request=httpx.Request("GET", "https://queue.fal.run/minimax/h3/requests/abc"), + ) + + with pytest.raises(FalAIVideoError) as error: + self.config.transform_video_content_response(response, self.logging_obj) + + assert error.value.status_code == 400 + assert error.value.message == "Request is still in progress" + assert "Request is still in progress" in error.value.response.text + + @pytest.mark.asyncio + async def test_async_content_response_surfaces_list_detail_error(self): + response: Final = httpx.Response( + 422, + json={ + "detail": [ + { + "loc": ["body", "input.reference_image_urls"], + "msg": "Failed to download the file. Please check if the URL is accessible and try again.", + } + ] + }, + request=httpx.Request("GET", "https://queue.fal.run/minimax/h3/requests/abc"), + ) + + with pytest.raises(FalAIVideoError) as error: + await self.config.async_transform_video_content_response(response, self.logging_obj) + + assert error.value.status_code == 422 + assert "input.reference_image_urls: Failed to download the file" in error.value.message + assert "Failed to download the file" in error.value.response.text + + @pytest.mark.asyncio + async def test_async_content_response_surfaces_string_detail_error(self): + response = httpx.Response( + 400, + json={"detail": "Request is still in progress"}, + request=httpx.Request("GET", "https://queue.fal.run/minimax/h3/requests/abc"), + ) + + with pytest.raises(FalAIVideoError) as error: + await self.config.async_transform_video_content_response(response, self.logging_obj) + + assert error.value.status_code == 400 + assert error.value.message == "Request is still in progress" + assert "Request is still in progress" in error.value.response.text + + def test_extract_video_url_surfaces_list_detail_error(self): + response: Final = Mock(spec=httpx.Response) + response.json.return_value = { + "detail": [{"loc": ["body", "input.reference_image_urls"], "msg": "Failed to download the file"}] + } + + with pytest.raises(ValueError, match=r"input\.reference_image_urls: Failed to download the file"): + self.config.transform_video_content_response(response, self.logging_obj) + def test_provider_config_and_error_class(self): provider_config = ProviderConfigManager.get_provider_video_config( model=MODEL, @@ -290,9 +540,29 @@ class TestFalAIVideoTransformation: } assert rows for model, row in rows.items(): - assert default_video_cost_calculator(model, 5, "fal_ai", video_resolution="480p") == ( - 5 * row["output_cost_per_second_480p"] - ) - assert default_video_cost_calculator(model, 5, "fal_ai", video_resolution="720p") == ( + for key, value in row.items(): + if key.startswith("output_cost_per_second_") and value is not None: + tier = key.removeprefix("output_cost_per_second_") + assert default_video_cost_calculator(model, 5, "fal_ai", video_resolution=tier) == 5 * value + assert default_video_cost_calculator(model, 5, "fal_ai", video_resolution="9999p") == ( 5 * row["output_cost_per_second"] ) + + def test_h3_video_cost_uses_model_info_tiers(self, local_model_cost_map): + row = litellm.model_cost[f"fal_ai/{H3_TEXT_MODEL}"] + model_info = litellm.get_model_info(model=H3_TEXT_MODEL, custom_llm_provider="fal_ai") + + assert video_generation_cost( + model=H3_TEXT_MODEL, + duration_seconds=5, + custom_llm_provider="fal_ai", + model_info=model_info, + video_resolution="2K", + ) == 5 * row["output_cost_per_second_2k"] + assert video_generation_cost( + model=H3_TEXT_MODEL, + duration_seconds=5, + custom_llm_provider="fal_ai", + model_info=model_info, + video_resolution="768p", + ) == 5 * row["output_cost_per_second_768p"] diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index f6da1bbcd0e..e3ae891f0d9 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -3,6 +3,8 @@ import json import os from unittest.mock import MagicMock, patch +import pytest + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( VertexAIPartnerModelsAnthropicMessagesConfig, ) @@ -67,6 +69,63 @@ def test_web_search_header_added_for_messages_endpoint(): ) +@pytest.mark.parametrize( + "client_headers", + [{"anthropic-beta": "dangerous-tool-use-2026-09-03"}, {}], + ids=["client_sends_beta", "client_omits_beta"], +) +def test_safeguards_add_dangerous_tool_use_beta_header(client_headers): + """Vertex rejects `safeguards` without the dangerous-tool-use beta, so the beta rides along with the field the way the web search and context management betas do.""" + config = VertexAIPartnerModelsAnthropicMessagesConfig() + litellm_params = { + "vertex_ai_project": "test-project", + "vertex_ai_location": "global", + "vertex_credentials": "{}", + } + optional_params = { + "safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}] + } + + with ( + patch.object(config, "_ensure_access_token", return_value=("token", "test-project")), + patch.object(config, "get_complete_vertex_url", return_value="https://mock-url"), + ): + updated_headers, _ = config.validate_anthropic_messages_environment( + headers=client_headers, + model="claude-sonnet-5", + messages=[], + optional_params=optional_params, + litellm_params=litellm_params, + api_base=None, + ) + + assert updated_headers["anthropic-beta"].split(",").count("dangerous-tool-use-2026-09-03") == 1 + + +def test_no_safeguards_leaves_dangerous_tool_use_beta_header_out(): + config = VertexAIPartnerModelsAnthropicMessagesConfig() + litellm_params = { + "vertex_ai_project": "test-project", + "vertex_ai_location": "global", + "vertex_credentials": "{}", + } + + with ( + patch.object(config, "_ensure_access_token", return_value=("token", "test-project")), + patch.object(config, "get_complete_vertex_url", return_value="https://mock-url"), + ): + updated_headers, _ = config.validate_anthropic_messages_environment( + headers={}, + model="claude-sonnet-5", + messages=[], + optional_params={"max_tokens": 64}, + litellm_params=litellm_params, + api_base=None, + ) + + assert "dangerous-tool-use-2026-09-03" not in updated_headers.get("anthropic-beta", "") + + def test_web_search_header_not_added_without_tool(): """Test that beta header is NOT added when web search tool is not present""" config = VertexAIPartnerModelsAnthropicMessagesConfig() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 7418cf67e5f..ac23831b5f4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -65,6 +65,8 @@ from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPAuth, MCPAuthType from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer from litellm.caching.caching import DualCache +from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler import litellm from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.utils import ProxyLogging @@ -10483,11 +10485,16 @@ def test_build_mcp_server_table_carries_oauth2_flow(): transport=MCPTransport.http, auth_type=MCPAuth.oauth2, oauth2_flow="client_credentials", + client_id="client-123", + client_secret="secret-xyz", + scopes=["scope:a", "scope:b"], + configured_scopes=("scope:a", "scope:b"), ) table = manager._build_mcp_server_table(server) assert table.oauth2_flow == "client_credentials" + assert table.credentials == {"scopes": ["scope:a", "scope:b"]} def test_build_mcp_server_table_carries_null_oauth2_flow(): @@ -10511,6 +10518,226 @@ def test_build_mcp_server_table_carries_null_oauth2_flow(): assert table.oauth2_flow is None +async def _mock_oauth_discovery( + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, + *, + server_url: str, + scopes: list[str], +) -> None: + resource_metadata_url: Final[str] = "https://up.example.com/.well-known/oauth-protected-resource" + authorization_server_url: Final[str] = "https://up.example.com" + authorization_metadata_url: Final[str] = f"{authorization_server_url}/.well-known/oauth-authorization-server" + respx_mock.get(server_url).respond( + status_code=401, + headers={"WWW-Authenticate": f'Bearer resource_metadata="{resource_metadata_url}"'}, + ) + respx_mock.get(resource_metadata_url).respond( + json={"authorization_servers": [authorization_server_url], "scopes_supported": scopes} + ) + respx_mock.get(authorization_metadata_url).respond( + json={ + "issuer": authorization_server_url, + "authorization_endpoint": f"{authorization_server_url}/authorize", + "token_endpoint": f"{authorization_server_url}/token", + } + ) + clients: Final[LLMClientCache] = LLMClientCache() + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", clients) + http_handler: Final[AsyncHTTPHandler] = AsyncHTTPHandler() + await http_handler.client.aclose() + http_handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respx_mock.async_handler)) + http_handler._owns_client = True + cache_key: Final[str] = f"async_httpx_clienttimeout_{MCP_METADATA_TIMEOUT}{httpxSpecialProvider.MCP.value}" + clients.set_cache(cache_key, http_handler) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("discovery_on_startup", [True, False]) +async def test_management_view_serves_configured_scopes_not_discovered_ones_from_db( + discovery_on_startup: bool, + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + row: Final[LiteLLM_MCPServerTable] = LiteLLM_MCPServerTable( + server_id="discovered-scopes-db", + alias="discovered_scopes_db", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + await _mock_oauth_discovery(respx_mock, monkeypatch, server_url=row.url or "", scopes=["discovered.read"]) + env: Final[dict[str, str]] = {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": "1"} if discovery_on_startup else {} + with patch.dict(os.environ, env, clear=True): + manager: Final[MCPServerManager] = MCPServerManager() + built: Final[MCPServer] = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + manager.registry[built.server_id] = built + resolved: Final[MCPServer] = await manager.ensure_oauth_metadata_discovered(built) + + assert resolved.scopes == ["discovered.read"] + view: Final[LiteLLM_MCPServerTable] = manager._build_mcp_server_table(resolved) + assert view.credentials is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stored_scopes", "runtime_scopes"), + [ + (None, ["openid"]), + ([], ["openid"]), + ([""], ["openid"]), + (["read", ""], ["read"]), + (["read", 7], ["read"]), + ("read", ["read"]), + ], +) +async def test_management_view_omits_invalid_or_absent_db_scopes( + stored_scopes: list[str | int] | str | None, + runtime_scopes: list[str], + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + row: Final[LiteLLM_MCPServerTable] = LiteLLM_MCPServerTable.model_construct( + server_id="empty-scopes-db", + alias="empty_scopes_db", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + credentials=json.dumps({"scopes": stored_scopes}), + created_at=datetime.now(), + updated_at=datetime.now(), + ) + await _mock_oauth_discovery(respx_mock, monkeypatch, server_url=row.url or "", scopes=["openid"]) + env: Final[dict[str, str]] = {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": "1"} + with patch.dict(os.environ, env, clear=True): + manager: Final[MCPServerManager] = MCPServerManager() + built: Final[MCPServer] = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + assert built.scopes == runtime_scopes + view: Final[LiteLLM_MCPServerTable] = manager._build_mcp_server_table(built) + assert view.credentials is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("discovery_on_startup", [True, False]) +@pytest.mark.parametrize( + ("stored_scopes", "runtime_scopes"), + [ + (["calendar.read"], ["calendar.read"]), + ([" "], ["discovered.read"]), + (["read", " "], ["read"]), + (["read", "read"], ["read", "read"]), + ], +) +async def test_management_view_serves_explicitly_configured_scopes_from_db( + stored_scopes: list[str], + runtime_scopes: list[str], + discovery_on_startup: bool, + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + row: Final[LiteLLM_MCPServerTable] = LiteLLM_MCPServerTable( + server_id="configured-scopes-db", + alias="configured_scopes_db", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + credentials={"scopes": stored_scopes}, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + await _mock_oauth_discovery(respx_mock, monkeypatch, server_url=row.url or "", scopes=["discovered.read"]) + env: Final[dict[str, str]] = {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": "1"} if discovery_on_startup else {} + with patch.dict(os.environ, env, clear=True): + manager: Final[MCPServerManager] = MCPServerManager() + built: Final[MCPServer] = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + manager.registry[built.server_id] = built + resolved: Final[MCPServer] = await manager.ensure_oauth_metadata_discovered(built) + + assert resolved.scopes == runtime_scopes + view: Final[LiteLLM_MCPServerTable] = manager._build_mcp_server_table(resolved) + assert view.credentials == {"scopes": stored_scopes} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("discovery_on_startup", [True, False]) +@pytest.mark.parametrize( + ("configured_scopes", "expected_view_scopes"), + [ + (None, None), + (["calendar.read"], ["calendar.read"]), + ([" "], None), + ([""], None), + (["calendar.read", " "], ["calendar.read"]), + ], +) +async def test_management_view_scopes_follow_yaml_config_not_discovery( + configured_scopes: list[str] | None, + expected_view_scopes: list[str] | None, + discovery_on_startup: bool, + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config: Final[dict[str, dict[str, object]]] = { + "yamlscopes": { + "url": "https://up.example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2, + "oauth2_flow": "authorization_code", + "client_id": "cid", + "client_secret": "csec", + **({"scopes": configured_scopes} if configured_scopes is not None else {}), + } + } + await _mock_oauth_discovery( + respx_mock, monkeypatch, server_url="https://up.example.com/mcp", scopes=["discovered.read"] + ) + env: Final[dict[str, str]] = {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": "1"} if discovery_on_startup else {} + with patch.dict(os.environ, env, clear=True): + manager: Final[MCPServerManager] = MCPServerManager() + await manager.load_servers_from_config(config) + server: Final[MCPServer] = next(iter(manager.config_mcp_servers.values())) + resolved: Final[MCPServer] = await manager.ensure_oauth_metadata_discovered(server) + + assert resolved.scopes == (expected_view_scopes or ["discovered.read"]) + view: Final[LiteLLM_MCPServerTable] = manager._build_mcp_server_table(resolved) + assert view.credentials == ({"scopes": expected_view_scopes} if expected_view_scopes else None) + + +@pytest.mark.asyncio +async def test_lazy_yaml_discovery_keeps_configured_scopes_out_of_the_management_view( + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config: Final[dict[str, dict[str, object]]] = { + "lazyyamlscopes": { + "url": "https://up.example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2, + "oauth2_flow": "authorization_code", + "client_id": "cid", + "client_secret": "csec", + } + } + await _mock_oauth_discovery( + respx_mock, monkeypatch, server_url="https://up.example.com/mcp", scopes=["discovered.read"] + ) + with patch.dict(os.environ, {}, clear=True): + manager: Final[MCPServerManager] = MCPServerManager() + await manager.load_servers_from_config(config) + server: Final[MCPServer] = next(iter(manager.config_mcp_servers.values())) + resolved: Final[MCPServer] = await manager.ensure_oauth_metadata_discovered(server) + + assert resolved.scopes == ["discovered.read"] + view: Final[LiteLLM_MCPServerTable] = manager._build_mcp_server_table(resolved) + assert view.credentials is None + + @pytest.mark.asyncio async def test_resolve_toolset_tool_permissions_single_db_fetch_across_checks(): """The server-level and tool-level permission primitives each resolve the diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index f10622e954b..13171a42cda 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -523,6 +523,92 @@ def test_wildcard_credential_hydration_preserves_missing_credential_name( } +def test_hydrate_credential_name_none_leaves_params_untouched(monkeypatch): + import litellm + from litellm.proxy.auth.model_checks import _hydrate_litellm_credential_name + from litellm.types.router import LiteLLM_Params + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="shared-credential", + credential_info={}, + credential_values={"api_key": "sk-shared"}, + ) + ], + ) + params = LiteLLM_Params(model="openai/gpt-4o", litellm_credential_name=None) + + result = _hydrate_litellm_credential_name(params) + + assert result is not None + assert result.api_key is None + assert result.litellm_credential_name is None + + +def test_hydrate_replaced_credential_uses_new_credential_values(monkeypatch): + import litellm + from litellm.proxy.auth.model_checks import _hydrate_litellm_credential_name + from litellm.types.router import LiteLLM_Params + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="shared-credential", + credential_info={}, + credential_values={"api_key": "sk-shared"}, + ), + CredentialItem( + credential_name="other-credential", + credential_info={}, + credential_values={"api_key": "sk-other"}, + ), + ], + ) + params = LiteLLM_Params(model="openai/gpt-4o", litellm_credential_name="other-credential") + + result = _hydrate_litellm_credential_name(params) + + assert result is not None + assert result.api_key == "sk-other" + assert result.litellm_credential_name is None + + +def test_hydrate_inline_api_key_wins_over_stored_credential(monkeypatch): + import litellm + from litellm.proxy.auth.model_checks import _hydrate_litellm_credential_name + from litellm.types.router import LiteLLM_Params + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="shared-credential", + credential_info={}, + credential_values={"api_key": "sk-shared"}, + ) + ], + ) + params = LiteLLM_Params( + model="openai/gpt-4o", + api_key="sk-inline", + litellm_credential_name="shared-credential", + ) + + result = _hydrate_litellm_credential_name(params) + + assert result is not None + assert result.api_key == "sk-inline" + + @pytest.mark.asyncio async def test_get_available_models_for_user_expands_query_team_wildcard( monkeypatch, diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 603a8686692..87bf4595af5 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -120,6 +120,33 @@ def test_user_banner_read_open_to_non_admin_roles(role): ) +@pytest.mark.parametrize( + "role", + [ + LitellmUserRoles.INTERNAL_USER.value, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + ], +) +def test_latest_release_info_read_open_to_non_admin_roles(role): # test-quality-ok: allowed path returns None, not raising is the observable + user_obj = LiteLLM_UserTable( + user_id="test_user", + user_email="test@example.com", + user_role=role, + ) + valid_token = UserAPIKeyAuth(user_id="test_user", user_role=role) + request = MagicMock(spec=Request) + request.query_params = {} + + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=role, + route="/get/latest_release_info", + request=request, + valid_token=valid_token, + request_data={}, + ) + + def test_user_banner_update_rejected_for_non_admin(): """Publishing the banner stays admin-only at the route layer.""" user_obj = LiteLLM_UserTable( diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index afadd6f3d19..80773f314d8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -6,7 +6,7 @@ import logging from contextlib import ExitStack from datetime import datetime, timedelta from types import SimpleNamespace -from typing import List, Optional +from typing import List, Optional, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -29,7 +29,7 @@ from litellm.proxy._types import ( UpdateMCPServerRequest, UserAPIKeyAuth, ) -from litellm.types.mcp import MCPAuth +from litellm.types.mcp import MCPAuth, MCPCredentials from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -834,6 +834,83 @@ class TestListMCPServers: mock_health_result.health_check_error = None mock_user_auth = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch( # test-quality-ok: endpoint test must patch module globals + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( # test-quality-ok: endpoint test must patch module globals + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=mock_server), + ), + patch( # test-quality-ok: endpoint test must patch module globals + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.health_check_server", + AsyncMock(return_value=mock_health_result), + ), + patch( # test-quality-ok: endpoint test must patch module globals + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_server, + ) + + result = await fetch_mcp_server( + request=_make_mock_request(), + server_id="server-mal", + user_api_key_dict=mock_user_auth, + ) + + assert result.credentials == expected + + @pytest.mark.parametrize( + "stored_credentials, expected", + [ + ( + { + "client_id": "cid", + "client_secret": "csecret", + "scopes": ["read", "write"], + "upstream_token_header": "esb-oauth", + }, + {"scopes": ["read", "write"], "upstream_token_header": "esb-oauth"}, + ), + ( + '{"client_id": "cid", "client_secret": "csecret", "scopes": ["read", "write"], ' + '"upstream_token_header": "esb-oauth"}', + {"scopes": ["read", "write"], "upstream_token_header": "esb-oauth"}, + ), + ( + {"client_id": "cid", "client_secret": "csecret", "scopes": []}, + None, + ), + ( + '{"client_id": "cid", "client_secret": "csecret", "scopes": []}', + None, + ), + ( + {"client_id": "cid", "client_secret": "csecret", "scopes": ["read", ""]}, + None, + ), + ( + {"client_id": "cid", "client_secret": "csecret", "scopes": "read"}, + None, + ), + ], + ) + @pytest.mark.asyncio + async def test_fetch_single_mcp_server_preserves_valid_oauth_scopes( + self, stored_credentials: object, expected: object + ): + mock_server = generate_mock_mcp_server_db_record(server_id="server-scopes", alias="Scopes") + mock_server.credentials = cast(MCPCredentials, stored_credentials) + mock_health_result = generate_mock_mcp_server_db_record(server_id="server-scopes", alias="Scopes") + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None + mock_user_auth = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + with ( patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", @@ -858,11 +935,11 @@ class TestListMCPServers: result = await fetch_mcp_server( request=_make_mock_request(), - server_id="server-mal", + server_id="server-scopes", user_api_key_dict=mock_user_auth, ) - assert result.credentials == expected + assert result.credentials == expected @pytest.mark.asyncio async def test_fetch_single_mcp_server_strips_upstream_resource_for_non_admin(self): @@ -1635,14 +1712,26 @@ class TestTemporaryMCPSessionEndpoints: return _inherit_credentials_from_existing_server(payload) - def test_admin_config_alone_does_not_suppress_credential_inheritance(self): - """The edit form round-trips upstream_resource, which is admin config rather than a credential. - Treating the blob as "credentials supplied" left the Authorize session with no declared app on - the exact path where this knob is configured.""" - updated = self._inherit_with({"upstream_resource": "api://audience"}) + @pytest.mark.parametrize( + "credentials", + [ + {"upstream_resource": "api://audience"}, + {"scopes": ["scope:a", "scope:b"]}, + {"scopes": ["scope:edited"], "upstream_resource": "api://audience"}, + {"scopes": ["scope:edited"], "upstream_token_header": "esb-oauth"}, + {"scopes": []}, + {"scopes": None}, + ], + ) + def test_admin_config_alone_does_not_suppress_credential_inheritance(self, credentials: MCPCredentials): + updated = self._inherit_with(credentials, scopes=["scope:stored"]) - assert updated.credentials["client_id"] == "client-123" - assert updated.credentials["client_secret"] == "secret-xyz" + assert updated.credentials == { + "client_id": "client-123", + "client_secret": "secret-xyz", + "scopes": ["scope:stored"], + **credentials, + } def test_upstream_token_header_is_inherited_like_other_admin_config(self): """It is admin config rather than a credential, so a session server derived from an existing @@ -1661,11 +1750,18 @@ class TestTemporaryMCPSessionEndpoints: assert updated.credentials["client_secret"] == "secret-xyz" assert updated.credentials["upstream_token_header"] == "esb-oauth" - def test_supplied_credential_still_wins_over_inheritance(self): - """A caller that supplies a real credential keeps it; inheritance must not overwrite it.""" - updated = self._inherit_with({"auth_value": "caller-token"}) + @pytest.mark.parametrize( + "credentials", + [ + {"auth_value": "caller-token"}, + {"client_id": "caller-client", "scopes": ["scope:edited"]}, + {"client_secret": "caller-secret", "scopes": ["scope:edited"]}, + ], + ) + def test_supplied_credential_still_wins_over_inheritance(self, credentials: MCPCredentials): + updated = self._inherit_with(credentials) - assert updated.credentials == {"auth_value": "caller-token"} + assert updated.credentials == credentials def test_inheritance_carries_upstream_resource_to_the_session_server(self): """Without this the temporary server omits the resource indicator and the Authorize leg it @@ -2339,7 +2435,7 @@ class TestTemporaryMCPSessionEndpoints: "client_secret": "client-secret", "scopes": ["scope1"], } - assert response.credentials is None + assert response.credentials == {"scopes": ["scope1"]} @pytest.mark.asyncio async def test_add_session_mcp_server_rejects_non_admins(self): @@ -4494,13 +4590,9 @@ class TestMCPApprovalWorkflow: assert result.total == 1 assert result.pending_review == 1 + @pytest.mark.parametrize("allowed_routes", [None, [], ["llm_api_routes"], ["mcp_routes"]]) @pytest.mark.asyncio - async def test_get_submissions_sanitizes_for_view_only_admin(self): - """PROXY_ADMIN_VIEW_ONLY reviewing the submission queue must go through - the non-admin sanitizer that fetch/list endpoints use: url, - static_headers, env, env_vars, and credentials are all dropped. A - mutation swapping the gate back to the old partial-blank pattern (which - left url/static_headers/env and env-var names intact) would fail this.""" + async def test_get_submissions_sanitizes_for_view_only_admin(self, allowed_routes: list[str] | None): from litellm.proxy._types import MCPSubmissionsSummary from litellm.proxy.management_endpoints.mcp_management_endpoints import ( get_mcp_server_submissions, @@ -4508,6 +4600,7 @@ class TestMCPApprovalWorkflow: item = _leaky_list_server() item.approval_status = "pending_review" + item.spec_path = "https://example.com/spec.json?key=private" summary = MCPSubmissionsSummary(total=1, pending_review=1, active=0, rejected=0, items=[item]) with ( @@ -4521,11 +4614,15 @@ class TestMCPApprovalWorkflow: ), ): result = await get_mcp_server_submissions( - user_api_key_dict=generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, allowed_routes=allowed_routes + ), ) + assert (result.total, result.pending_review, result.active, result.rejected) == (1, 1, 0, 0) assert len(result.items) == 1 sanitized = result.items[0] + assert sanitized.spec_path is None assert sanitized.url is None assert sanitized.static_headers is None assert sanitized.env == {} @@ -4536,11 +4633,9 @@ class TestMCPApprovalWorkflow: assert item.url == "https://leaky.example.com/mcp?api_key=sk-embedded-in-url" assert item.static_headers == {"Authorization": "Bearer sk-secret-header"} + @pytest.mark.parametrize("allowed_routes", [None, [], ["llm_api_routes"], ["mcp_routes"]]) @pytest.mark.asyncio - async def test_get_submissions_full_admin_still_sees_secrets(self): - """The view-only redaction must not over-redact for a full PROXY_ADMIN, - who needs url/static_headers/env/env_vars to review the pending - submission. Only the explicit credentials field is cleared.""" + async def test_get_submissions_full_admin_preserves_review_fields(self, allowed_routes: list[str] | None): from litellm.proxy._types import MCPSubmissionsSummary from litellm.proxy.management_endpoints.mcp_management_endpoints import ( get_mcp_server_submissions, @@ -4548,6 +4643,7 @@ class TestMCPApprovalWorkflow: item = _leaky_list_server() item.approval_status = "pending_review" + item.spec_path = "https://example.com/spec.json?key=private" summary = MCPSubmissionsSummary(total=1, pending_review=1, active=0, rejected=0, items=[item]) with ( @@ -4561,11 +4657,14 @@ class TestMCPApprovalWorkflow: ), ): result = await get_mcp_server_submissions( - user_api_key_dict=generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, allowed_routes=allowed_routes), ) + assert (result.total, result.pending_review, result.active, result.rejected) == (1, 1, 0, 0) assert len(result.items) == 1 raw = result.items[0] + assert raw.spec_path == item.spec_path + assert raw.approval_status == "pending_review" assert raw.url == "https://leaky.example.com/mcp?api_key=sk-embedded-in-url" assert raw.static_headers == {"Authorization": "Bearer sk-secret-header"} assert raw.env == {"UPSTREAM_TOKEN": "sk-secret-env"} diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 376309d8a7e..e6b5fb25c3e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -10,6 +10,7 @@ import pytest from fastapi.testclient import TestClient from litellm._uuid import uuid +from litellm.models.credentials import CredentialItem from litellm.proxy._types import ( LiteLLM_ModelTable, @@ -308,6 +309,62 @@ class TestModelManagementAuthChecks: ) assert result is True + def test_can_user_attach_credential_non_admin_explicit_null_clear_fails(self): + from litellm.proxy._types import ProxyException + from litellm.types.router import updateLiteLLMParams as litellm_params + + with pytest.raises(ProxyException) as exc_info: + ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=litellm_params(litellm_credential_name=None), + user_api_key_dict=self.team_admin_user, + existing_litellm_params=LiteLLM_Params( + model="test_model", litellm_credential_name="shared-credential" + ), + null_detaches=True, + ) + + assert exc_info.value.code == "403" + assert exc_info.value.param == "litellm_credential_name" + + def test_can_user_attach_credential_admin_explicit_null_clear_succeeds(self): + from litellm.types.router import updateLiteLLMParams as litellm_params + + result = ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=litellm_params(litellm_credential_name=None), + user_api_key_dict=self.admin_user, + existing_litellm_params=LiteLLM_Params( + model="test_model", litellm_credential_name="shared-credential" + ), + null_detaches=True, + ) + + assert result is True + + def test_can_user_attach_credential_null_without_existing_allows_any_role(self): + from litellm.types.router import updateLiteLLMParams as litellm_params + + result = ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=litellm_params(litellm_credential_name=None), + user_api_key_dict=self.team_admin_user, + existing_litellm_params=LiteLLM_Params(model="test_model"), + null_detaches=True, + ) + + assert result is True + + def test_can_user_attach_credential_null_is_noop_when_null_does_not_detach(self): + from litellm.types.router import updateLiteLLMParams as litellm_params + + result = ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=litellm_params(litellm_credential_name=None), + user_api_key_dict=self.team_admin_user, + existing_litellm_params=LiteLLM_Params( + model="test_model", litellm_credential_name="shared-credential" + ), + ) + + assert result is True + def test_can_user_attach_credential_unchanged_encrypted_existing_allows_any_role(self, monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234") encrypted_name = encrypt_value_helper(value="shared-credential") @@ -1249,6 +1306,60 @@ class TestUpdateModel: mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once() mock_clear_cache.assert_awaited_once_with() + @pytest.mark.asyncio + async def test_update_model_legacy_null_credential_name_is_not_a_detach_for_non_admin(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_model + + model_id = "legacy-null-credential" + existing = Deployment( + model_name="legacy-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini", litellm_credential_name="shared-credential"), + model_info={"id": model_id}, + ) + existing_row = MagicMock() + existing_row.litellm_params = existing.litellm_params.model_dump() + existing_row.model_dump.return_value = existing.model_dump() + updated_row = MagicMock() + updated_row.model_dump_json.return_value = "{}" + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row) + mock_router = MagicMock() + mock_router.get_model_ids.return_value = [model_id] + team_admin = UserAPIKeyAuth(user_id="team-admin", user_role=LitellmUserRoles.INTERNAL_USER) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value: value, + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + await update_model( + model_params=updateDeployment( + litellm_params=updateLiteLLMParams( + model="openai/gpt-4o-mini", litellm_credential_name=None + ), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=team_admin, + ) + + mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once() + persisted = json.loads(mock_prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"]) + assert persisted["litellm_credential_name"] == "shared-credential" + class TestUpdatePublicModelGroups: """Test that update_public_model_groups correctly sets litellm.public_model_groups @@ -4000,6 +4111,401 @@ class TestUpdateDBModelClearCacheControlInjectionPoints: assert params["tpm"] == 10 +class TestUpdateDBModelClearCredentialName: + def test_explicit_null_removes_stored_credential_name(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + api_key="sk-real", + tpm=100, + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1", team_id="team-keep", access_groups=["prod"]), + ) + update_patch: Final = updateDeployment( + litellm_params=updateLiteLLMParams(litellm_credential_name=None) + ) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + info: Final = json.loads(result["model_info"]) + assert "litellm_credential_name" not in params + assert params["model"] == "openai/gpt-4o" + assert params["api_base"] == "https://api.openai.com/v1" + assert params["api_key"] == "sk-real" + assert params["tpm"] == 100 + assert info["team_id"] == "team-keep" + assert info["access_groups"] == ["prod"] + + def test_omitted_credential_name_keeps_stored_association(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + api_key="sk-real", + tpm=100, + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1", team_id="team-keep", access_groups=["prod"]), + ) + update_patch: Final = updateDeployment(litellm_params=updateLiteLLMParams(tpm=10)) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + assert params["litellm_credential_name"] == "shared-credential" + assert params["tpm"] == 10 + + def test_null_clear_on_model_without_credential_is_noop(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params(model="openai/gpt-4o", api_base="https://api.openai.com/v1"), + model_info=ModelInfo(id="dep-cred-1", team_id="team-keep", access_groups=["prod"]), + ) + update_patch: Final = updateDeployment( + litellm_params=updateLiteLLMParams(litellm_credential_name=None) + ) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + assert "litellm_credential_name" not in params + assert params["model"] == "openai/gpt-4o" + assert params["api_base"] == "https://api.openai.com/v1" + + def test_null_credential_clear_alongside_pricing_clear(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + input_cost_per_token=0.000001, + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1", input_cost_per_token=0.000001), + ) + update_patch: Final = updateDeployment( + litellm_params=updateLiteLLMParams( + litellm_credential_name=None, + input_cost_per_token=None, + ) + ) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + info: Final = json.loads(result["model_info"]) + assert "litellm_credential_name" not in params + assert "input_cost_per_token" not in params + assert "input_cost_per_token" not in info + + def test_replace_credential_name_keeps_other_params(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + api_key="sk-real", + tpm=100, + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1", team_id="team-keep", access_groups=["prod"]), + ) + update_patch: Final = updateDeployment( + litellm_params=updateLiteLLMParams(litellm_credential_name="other-credential") + ) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + assert params["litellm_credential_name"] == "other-credential" + assert params["api_base"] == "https://api.openai.com/v1" + assert params["api_key"] == "sk-real" + assert params["tpm"] == 100 + + +class TestPatchModelCredentialName: + @staticmethod + async def _patch_model( + monkeypatch, + db_model: Deployment, + user_api_key_dict: UserAPIKeyAuth, + credential_name: str | None, + db_credential: CredentialItem | None = None, + credentials_repository: MagicMock | None = None, + ) -> list[dict[str, object]]: + import litellm + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model, update_db_model + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="shared-credential", + credential_info={}, + credential_values={"api_key": "sk-shared"}, + ), + CredentialItem( + credential_name="other-credential", + credential_info={}, + credential_values={"api_key": "sk-other"}, + ), + ], + ) + credentials_repository = credentials_repository or MagicMock() + credentials_repository.find_by_name = AsyncMock(return_value=db_credential) + persisted: Final[list[dict[str, object]]] = [] + + async def persist_model(**kwargs): + row: Final = update_db_model(db_model=kwargs["db_model"], updated_patch=kwargs["patch_data"]) + persisted.append(row) + updated_row: Final = MagicMock() + updated_row.model_dump_json.return_value = "{}" + return updated_row + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.CredentialsRepository", + return_value=credentials_repository, + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", + new=AsyncMock(return_value=db_model), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints._update_team_model_in_db", + new=AsyncMock(side_effect=persist_model), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value, **kwargs: value, + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.raise_if_reload_degraded_serving" + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log", + new=AsyncMock(), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.live_model_ids_snapshot", + return_value=frozenset(), + ), + ): + await patch_model( + model_id="dep-cred-1", + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams(litellm_credential_name=credential_name) + ), + user_api_key_dict=user_api_key_dict, + ) + + return persisted + + @pytest.mark.asyncio + async def test_patch_model_rejects_empty_string_credential_name(self, monkeypatch): + from litellm.proxy._types import ProxyException + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + with pytest.raises(ProxyException) as exc_info: + await self._patch_model( + monkeypatch, + db_model, + self._admin_user(), + "", + ) + + assert exc_info.value.code == "400" + assert exc_info.value.param == "litellm_credential_name" + assert "empty" in exc_info.value.message.lower() + + @staticmethod + def _admin_user() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + @staticmethod + def _team_admin_user() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id="team-admin", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-keep") + + @pytest.mark.asyncio + async def test_patch_model_rejects_unknown_credential_name(self, monkeypatch): + from litellm.proxy._types import ProxyException + + credentials_repository = MagicMock() + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + with pytest.raises(ProxyException) as exc_info: + await self._patch_model( + monkeypatch, + db_model, + self._admin_user(), + "ghost-credential", + credentials_repository=credentials_repository, + ) + + assert exc_info.value.code == "400" + assert "not found" in exc_info.value.message.lower() + credentials_repository.find_by_name.assert_awaited_once_with("ghost-credential") + + @pytest.mark.asyncio + async def test_patch_model_accepts_credential_known_only_in_db(self, monkeypatch): + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + persisted: Final = await self._patch_model( + monkeypatch, + db_model, + self._admin_user(), + "db-only-credential", + db_credential=CredentialItem( + credential_name="db-only-credential", + credential_info={}, + credential_values={"api_key": "sk-db"}, + ), + ) + params: Final = json.loads(persisted[0]["litellm_params"]) + assert params["litellm_credential_name"] == "db-only-credential" + + @pytest.mark.asyncio + async def test_patch_model_replaces_credential_name_and_preserves_other_params(self, monkeypatch): + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + persisted: Final = await self._patch_model(monkeypatch, db_model, self._admin_user(), "other-credential") + params: Final = json.loads(persisted[0]["litellm_params"]) + assert params["litellm_credential_name"] == "other-credential" + assert params["api_base"] == "https://api.openai.com/v1" + + @pytest.mark.asyncio + async def test_patch_model_admin_null_clear_persists_without_credential(self, monkeypatch): + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + persisted: Final = await self._patch_model(monkeypatch, db_model, self._admin_user(), None) + params: Final = json.loads(persisted[0]["litellm_params"]) + assert "litellm_credential_name" not in params + assert params["api_base"] == "https://api.openai.com/v1" + + @pytest.mark.asyncio + async def test_patch_model_rejects_non_admin_explicit_null_clear(self, monkeypatch): + from litellm.proxy._types import ProxyException + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + with pytest.raises(ProxyException) as exc_info: + await self._patch_model(monkeypatch, db_model, self._team_admin_user(), None) + + assert exc_info.value.code == "403" + assert exc_info.value.param == "litellm_credential_name" + + @pytest.mark.asyncio + async def test_patch_model_clear_then_reattach_round_trip(self, monkeypatch): + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + cleared: Final = await self._patch_model(monkeypatch, db_model, self._admin_user(), None) + cleared_model: Final = Deployment.model_validate( + { + "model_name": db_model.model_name, + "litellm_params": json.loads(cleared[0]["litellm_params"]), + "model_info": json.loads(cleared[0]["model_info"]), + } + ) + reattached: Final = await self._patch_model( + monkeypatch, + cleared_model, + self._admin_user(), + "shared-credential", + ) + params: Final = json.loads(reattached[0]["litellm_params"]) + assert params["litellm_credential_name"] == "shared-credential" + + class TestGetModelInfoWithIdBlocked: """`ProxyConfig.get_model_info_with_id` must propagate the DB-level `blocked` column into the in-memory `model_info` dict so the router filter can read it.""" diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 680dd4df0ae..0f1ff3b024d 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -1,3 +1,4 @@ +import json import re from datetime import datetime, timezone from typing import Final @@ -339,6 +340,25 @@ def test_cognition_provider_fields(): assert fields_by_key["api_base"]["required"] is False +def test_qwen_mainland_provider_fields_carry_the_qianwen_brand(): + app_instance = FastAPI() + app_instance.include_router(router) + test_client = TestClient(app_instance) + + providers = test_client.get("/public/providers/fields").json() + + mainland = next(p for p in providers if p["litellm_provider"] == "qwen_ai_platform") + international = next(p for p in providers if p["litellm_provider"] == "qwencloud") + + assert mainland["provider_display_name"] == "Qianwen AI Platform" + assert international["provider_display_name"] == "QwenCloud" + + mainland_fields = {f["key"]: f for f in mainland["credential_fields"]} + assert mainland_fields["api_key"]["label"] == "Qianwen AI Platform API Key" + assert "Qianwen AI Platform" in mainland_fields["api_base"]["tooltip"] + assert "Qwen AI Platform" not in json.dumps(mainland) + + def test_chatgpt_provider_fields(): app_instance = FastAPI() app_instance.include_router(router) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index c834ac05f0a..86bf896188f 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -1,6 +1,7 @@ import asyncio import threading -from collections.abc import Mapping +import time +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -10,6 +11,7 @@ from fastapi import HTTPException import litellm from litellm.caching.dual_cache import DualCache +from litellm.types.caching import RedisPipelineIncrementOperation from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( AgenticAnthropicStreamingIterator, @@ -2348,35 +2350,139 @@ async def test_release_non_numeric_counter_reseeds_from_db(spend_counter_state): class _ExpiringRedisCache: - def __init__(self) -> None: + """In-memory stand-in for RedisCache with real wall-clock key expiry.""" + + def __init__(self, default_ttl: float = 60.0, fail_first_refresh: bool = False) -> None: + self.default_ttl = default_ttl self.store: dict[str, float] = {} + self.expires_at: dict[str, float] = {} + self.refresh_attempts = 0 + self.refresh_count = 0 + self.fail_first_refresh = fail_first_refresh + + def _evict_expired(self, key: str) -> None: + if self.expires_at.get(key, float("inf")) <= time.monotonic(): + self.store.pop(key, None) + self.expires_at.pop(key, None) async def async_get_cache(self, key: str, *args: object, **kwargs: object) -> float | None: + self._evict_expired(key) return self.store.get(key) async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + self._evict_expired(key) self.store[key] = self.store.get(key, 0.0) + float(value) + self.expires_at[key] = time.monotonic() + self.default_ttl return self.store[key] async def async_set_max(self, key: str, value: float, **kwargs: object) -> float: + self._evict_expired(key) self.store[key] = max(self.store.get(key, float("-inf")), float(value)) + self.expires_at[key] = time.monotonic() + self.default_ttl return self.store[key] async def async_set_cache(self, key: str, value: float, *args: object, **kwargs: object) -> bool: self.store[key] = float(value) + self.expires_at[key] = time.monotonic() + self.default_ttl return True async def async_delete_cache(self, key: str, *args: object, **kwargs: object) -> None: self.store.pop(key, None) + self.expires_at.pop(key, None) - async def async_increment_pipeline(self, increment_list, **kwargs): - results = [] - for op in increment_list: - results.append(await self.async_increment(op["key"], op["increment_value"])) - return results + async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: + self.refresh_attempts += 1 + if self.fail_first_refresh and self.refresh_attempts == 1: + raise ConnectionError("Redis circuit breaker is open") + self._evict_expired(key) + if key not in self.store: + return False + self.refresh_count += 1 + self.expires_at[key] = time.monotonic() + (ttl if ttl is not None else self.default_ttl) + return True - def get_ttl(self, **kwargs) -> None: - return None + async def async_increment_pipeline( + self, increment_list: Sequence[RedisPipelineIncrementOperation], **kwargs: object + ) -> list[float]: + return [await self.async_increment(op["key"], op["increment_value"]) for op in increment_list] + + def get_ttl(self, **kwargs: object) -> int | None: + return int(self.default_ttl) + + +@pytest.mark.asyncio +async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( + spend_counter_state, +): + """A request that runs longer than the counter TTL must keep its reservation in Redis + (so a concurrent request on any worker still sees it), and renewal must stop once the + reservation is reconciled so an idle counter still expires on its own.""" + counter_cache, key_cache = spend_counter_state + redis_cache = _ExpiringRedisCache(default_ttl=0.2) + counter_cache.redis_cache = redis_cache + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease", spend=0.0, max_budget=1.0) + counter_key = "spend:key:key-lease" + + reservation = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert reservation is not None + + await asyncio.sleep(0.5) + assert await redis_cache.async_get_cache(key=counter_key) == pytest.approx(0.6) + concurrent = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert concurrent is not None + assert concurrent["reserved_cost"] == pytest.approx(0.4) + + await release_budget_reservation(reservation) + await release_budget_reservation(concurrent) + await asyncio.sleep(0.15) + refreshes_after_release = redis_cache.refresh_count + await asyncio.sleep(0.35) + assert redis_cache.refresh_count == refreshes_after_release + assert await redis_cache.async_get_cache(key=counter_key) is None + + +@pytest.mark.asyncio +async def test_reservation_lease_keeps_renewing_after_transient_redis_failure( + spend_counter_state, +): + """One failed EXPIRE (Redis blip, open circuit breaker) must not end renewal for the + rest of the request.""" + counter_cache, key_cache = spend_counter_state + redis_cache = _ExpiringRedisCache(default_ttl=0.2, fail_first_refresh=True) + counter_cache.redis_cache = redis_cache + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease-blip", spend=0.0, max_budget=1.0) + + reservation = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert reservation is not None + + await asyncio.sleep(0.5) + assert redis_cache.refresh_attempts >= 3 + await release_budget_reservation(reservation) + + +@pytest.mark.asyncio +async def test_reservation_lease_stops_when_request_task_ends_without_reconciling( + spend_counter_state, +): + """A request whose task ends without reconciling (client disconnect path that skips the + cost callbacks) must not keep renewing: the counter falls back to its plain TTL instead of + pinning the reservation until the request timeout.""" + counter_cache, key_cache = spend_counter_state + redis_cache = _ExpiringRedisCache(default_ttl=0.2) + counter_cache.redis_cache = redis_cache + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease-orphan", spend=0.0, max_budget=1.0) + counter_key = "spend:key:key-lease-orphan" + + reservation = await asyncio.create_task(_reserve(valid_token, 0.6, key_cache, proxy_logging_obj)) + assert reservation is not None + assert reservation["finalized"] is False + + await asyncio.sleep(0.5) + assert redis_cache.refresh_count == 0 + assert await redis_cache.async_get_cache(key=counter_key) is None class _TeamMembershipFloorDb: diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_latest_release_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_latest_release_endpoints.py new file mode 100644 index 00000000000..c966b8b7135 --- /dev/null +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_latest_release_endpoints.py @@ -0,0 +1,260 @@ +import asyncio +import json +import time +from typing import Final + +import httpx +import pytest +from fastapi.testclient import TestClient + +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.proxy_server import app +from litellm.proxy.ui_crud_endpoints.latest_release_endpoints import ( + LATEST_RELEASE_CACHE_KEY, + LATEST_RELEASE_CACHE_TTL_SECONDS, + LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS, + LATEST_RELEASE_URL, + LatestReleaseInfo, + LatestReleaseUnavailable, + _default_cache, + _default_client, + _default_fetch_lock, + count_release_bullets, + get_latest_release_info, +) + +SAMPLE_BODY: Final = """## What's Changed +* feat(proxy): add upgrade banner by @kerry in https://github.com/BerriAI/litellm/pull/1 +* fix(azure): retry on 429 by @a in https://github.com/BerriAI/litellm/pull/2 +* fix: handle empty body by @b in https://github.com/BerriAI/litellm/pull/3 +* Feat(ui)!: drop legacy theme by @c in https://github.com/BerriAI/litellm/pull/4 +* chore(deps): bump httpx by @d in https://github.com/BerriAI/litellm/pull/5 +* docs: fix typo by @e in https://github.com/BerriAI/litellm/pull/6 +* Litellm dev 09 08 2026 by @f in https://github.com/BerriAI/litellm/pull/7 +* refactor(router) : spaced colon does not match by @g in https://github.com/BerriAI/litellm/pull/8 + +## New Contributors +* @kerry made their first contribution in https://github.com/BerriAI/litellm/pull/1 + +**Full Changelog**: https://github.com/BerriAI/litellm/compare/v1.101.0...v1.102.0 +""" + +SAMPLE_RELEASE: Final = { + "tag_name": "v1.102.0", + "html_url": "https://github.com/BerriAI/litellm/releases/tag/v1.102.0", + "body": SAMPLE_BODY, +} +EXPECTED_INFO: Final = { + "version": "1.102.0", + "new_features": 2, + "bug_fixes": 2, + "other_updates": 4, + "release_url": SAMPLE_RELEASE["html_url"], +} + + +class _RecordingClient: + def __init__(self, outcomes: list[httpx.Response | Exception]) -> None: + self._outcomes = outcomes + self.calls: list[tuple[str, float | None]] = [] + + async def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: + self.calls.append((url, timeout)) + outcome = self._outcomes[min(len(self.calls) - 1, len(self._outcomes) - 1)] + if isinstance(outcome, Exception): + raise outcome + return outcome + + +def _github_response(status: int = 200, payload: object = SAMPLE_RELEASE) -> httpx.Response: + return httpx.Response(status, content=json.dumps(payload).encode()) + + +def _fresh_cache() -> InMemoryCache: + return InMemoryCache(max_size_in_memory=1, default_ttl=LATEST_RELEASE_CACHE_TTL_SECONDS) + + +def _override_dependencies(client: _RecordingClient, cache: InMemoryCache, role: LitellmUserRoles) -> None: + async def auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-test", user_id="test-user", user_role=role) + + app.dependency_overrides[user_api_key_auth] = auth + app.dependency_overrides[_default_client] = lambda: client + app.dependency_overrides[_default_cache] = lambda: cache + app.dependency_overrides[_default_fetch_lock] = lambda: asyncio.Lock() + + +@pytest.fixture +def http_client(): + yield TestClient(app) + app.dependency_overrides.pop(user_api_key_auth, None) + app.dependency_overrides.pop(_default_client, None) + app.dependency_overrides.pop(_default_cache, None) + app.dependency_overrides.pop(_default_fetch_lock, None) + + +class TestCountReleaseBullets: + def test_buckets_by_conventional_commit_type(self): + counts = count_release_bullets(SAMPLE_BODY) + assert counts["new_features"] == 2 + assert counts["bug_fixes"] == 2 + assert counts["other_updates"] == 4 + + def test_unprefixed_bullets_count_as_other_updates(self): + counts = count_release_bullets("* Litellm dev 09 08 2026 by @f in https://x/pull/7\n") + assert (counts["new_features"], counts["bug_fixes"], counts["other_updates"]) == (0, 0, 1) + + def test_ignores_non_bullet_lines_and_contributor_entries(self): + assert ( + sum( + count_release_bullets( + "## What's Changed\n\n* @x made their first contribution in url\n" + "\n**Full Changelog**: https://github.com/BerriAI/litellm/compare/v1...v2\n" + ).values() + ) + == 0 + ) + + def test_empty_body_yields_zero_counts(self): + counts = count_release_bullets("") + assert (counts["new_features"], counts["bug_fixes"], counts["other_updates"]) == (0, 0, 0) + + +class TestGetLatestReleaseInfo: + @pytest.mark.asyncio + async def test_fetches_and_parses_github_release(self): + client = _RecordingClient([_github_response()]) + result = await get_latest_release_info(client=client, cache=_fresh_cache(), fetch_lock=asyncio.Lock()) + assert isinstance(result, LatestReleaseInfo) + assert result.model_dump() == EXPECTED_INFO + assert client.calls == [(LATEST_RELEASE_URL, 5)] + + @pytest.mark.asyncio + async def test_second_call_within_ttl_does_not_refetch(self): + client = _RecordingClient([_github_response()]) + cache = _fresh_cache() + fetch_lock = asyncio.Lock() + first = await get_latest_release_info(client=client, cache=cache, fetch_lock=fetch_lock) + second = await get_latest_release_info(client=client, cache=cache, fetch_lock=fetch_lock) + assert first == second + assert len(client.calls) == 1 + + @pytest.mark.asyncio + async def test_success_is_cached_for_the_full_ttl(self): + cache = _fresh_cache() + await get_latest_release_info( + client=_RecordingClient([_github_response()]), cache=cache, fetch_lock=asyncio.Lock() + ) + remaining = await cache.async_get_ttl(LATEST_RELEASE_CACHE_KEY) - time.time() + assert LATEST_RELEASE_CACHE_TTL_SECONDS - 5 < remaining <= LATEST_RELEASE_CACHE_TTL_SECONDS + + @pytest.mark.asyncio + async def test_failure_is_cached_briefly_so_github_is_not_hammered(self): + client = _RecordingClient([httpx.ConnectError("boom")]) + cache = _fresh_cache() + fetch_lock = asyncio.Lock() + first = await get_latest_release_info(client=client, cache=cache, fetch_lock=fetch_lock) + second = await get_latest_release_info(client=client, cache=cache, fetch_lock=fetch_lock) + assert isinstance(first, LatestReleaseUnavailable) + assert first == second + assert len(client.calls) == 1 + remaining = await cache.async_get_ttl(LATEST_RELEASE_CACHE_KEY) - time.time() + assert ( + LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS - 5 < remaining <= LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "response", + [ + _github_response(status=403, payload={"message": "rate limited"}), + _github_response(status=500, payload={}), + _github_response(payload={"tag_name": "v1.0.0"}), + httpx.Response(200, content=b"not json"), + ], + ids=["rate_limited", "server_error", "missing_fields", "not_json"], + ) + async def test_bad_github_responses_are_unavailable(self, response: httpx.Response): + result = await get_latest_release_info( + client=_RecordingClient([response]), cache=_fresh_cache(), fetch_lock=asyncio.Lock() + ) + assert isinstance(result, LatestReleaseUnavailable) + + @pytest.mark.asyncio + async def test_concurrent_misses_share_one_fetch(self): + event = asyncio.Event() + + class _BlockingClient(_RecordingClient): + async def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: + self.calls.append((url, timeout)) + await event.wait() + return _github_response() + + client = _BlockingClient([]) + cache = _fresh_cache() + fetch_lock = asyncio.Lock() + tasks = [ + asyncio.create_task(get_latest_release_info(client=client, cache=cache, fetch_lock=fetch_lock)) + for _ in range(5) + ] + await asyncio.sleep(0) + await asyncio.sleep(0) + event.set() + results = await asyncio.gather(*tasks) + expected: Final = LatestReleaseInfo.model_validate(EXPECTED_INFO) + assert results == [expected] * 5 + assert len(client.calls) == 1 + + @pytest.mark.asyncio + async def test_failure_under_lock_is_also_coalesced(self): + event = asyncio.Event() + + class _FailingBlockingClient(_RecordingClient): + async def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: + self.calls.append((url, timeout)) + await event.wait() + raise httpx.ConnectError("boom") + + client = _FailingBlockingClient([]) + cache = _fresh_cache() + fetch_lock = asyncio.Lock() + tasks = [ + asyncio.create_task(get_latest_release_info(client=client, cache=cache, fetch_lock=fetch_lock)) + for _ in range(5) + ] + await asyncio.sleep(0) + await asyncio.sleep(0) + event.set() + results = await asyncio.gather(*tasks) + assert all(isinstance(result, LatestReleaseUnavailable) for result in results) + assert len(client.calls) == 1 + + +class TestLatestReleaseInfoEndpoint: + def test_returns_release_stats_for_authenticated_user(self, http_client): + _override_dependencies(_RecordingClient([_github_response()]), _fresh_cache(), LitellmUserRoles.INTERNAL_USER) + response = http_client.get("/get/latest_release_info") + assert response.status_code == 200 + assert response.json() == EXPECTED_INFO + + def test_returns_null_when_github_is_unreachable(self, http_client): + _override_dependencies( + _RecordingClient([httpx.ConnectError("boom")]), _fresh_cache(), LitellmUserRoles.PROXY_ADMIN + ) + response = http_client.get("/get/latest_release_info") + assert response.status_code == 200 + assert response.json() is None + + def test_repeated_requests_reuse_cache(self, http_client): + client = _RecordingClient([_github_response()]) + _override_dependencies(client, _fresh_cache(), LitellmUserRoles.PROXY_ADMIN) + assert http_client.get("/get/latest_release_info").json() == EXPECTED_INFO + assert http_client.get("/get/latest_release_info").json() == EXPECTED_INFO + assert len(client.calls) == 1 + + def test_rejects_unauthenticated_requests(self, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-1234") + response = TestClient(app).get("/get/latest_release_info") + assert response.status_code in (401, 403) diff --git a/tests/test_litellm/test_anthropic_beta_headers_filtering.py b/tests/test_litellm/test_anthropic_beta_headers_filtering.py index d404edb1281..19b26120672 100644 --- a/tests/test_litellm/test_anthropic_beta_headers_filtering.py +++ b/tests/test_litellm/test_anthropic_beta_headers_filtering.py @@ -443,6 +443,20 @@ class TestAnthropicBetaHeadersFiltering: assert filtered == ["thinking-binding-controls-2026-08-01"] + @pytest.mark.parametrize("provider", ["anthropic", "bedrock", "bedrock_mantle", "vertex_ai"]) + def test_dangerous_tool_use_forwarded(self, provider): + """Claude Code's server-side auto-mode classifier sends `safeguards` together with + dangerous-tool-use-2026-09-03. Bedrock Invoke, Bedrock Mantle, and Vertex rawPredict + all answer "safeguards: Extra inputs are not permitted" when the body field arrives + without the beta (probed 2026-09-21), so dropping the header turned every auto-mode + turn into a 400 on Vertex and silently disabled the classifier on Bedrock.""" + filtered = filter_and_transform_beta_headers( + beta_headers=["dangerous-tool-use-2026-09-03"], + provider=provider, + ) + + assert filtered == ["dangerous-tool-use-2026-09-03"] + def test_null_value_headers_filtered(self): """Test that headers with null values are always filtered out.""" for provider in [ diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 2ccb88b29db..3eba1bcfced 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -619,6 +619,8 @@ def validate_model_cost_values(model_data, exceptions=None): "output_cost_per_second", "output_cost_per_second_480p", "output_cost_per_second_720p", + "output_cost_per_second_768p", + "output_cost_per_second_2k", "output_cost_per_second_1080p", "output_cost_per_second_4k", "input_cost_per_query", @@ -838,6 +840,8 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_second": {"type": "number"}, "output_cost_per_second_480p": {"type": "number"}, "output_cost_per_second_720p": {"type": "number"}, + "output_cost_per_second_768p": {"type": "number"}, + "output_cost_per_second_2k": {"type": "number"}, "output_cost_per_second_1080p": {"type": "number"}, "output_cost_per_second_4k": {"type": "number"}, "output_cost_per_token": {"type": "number"}, @@ -2994,6 +2998,35 @@ class TestAdditionalDropParamsForNonOpenAIProviders: assert result.get("custom_param") == "value" +class TestExtraBodyCannotOverrideModel: + @pytest.mark.parametrize("custom_llm_provider", ["edenai", "openai", "azure"]) + def test_extra_body_model_is_dropped_for_openai_compatible_providers(self, custom_llm_provider: str) -> None: + from litellm.utils import add_provider_specific_params_to_optional_params + + result = add_provider_specific_params_to_optional_params( + optional_params={"extra_body": {"model": "edenai/openai/gpt-4o", "provider_flag": True}}, + passed_params={ + "model": "edenai/openai/gpt-4o-mini", + "extra_body": {"model": "edenai/anthropic/claude-3-opus", "top_k": 5}, + "custom_param": "kept", + }, + custom_llm_provider=custom_llm_provider, + openai_params=["model", "temperature"], + additional_drop_params=None, + ) + + assert result == {"extra_body": {"provider_flag": True, "top_k": 5, "custom_param": "kept"}}, result + + def test_get_optional_params_strips_extra_body_model_for_edenai(self) -> None: + result = litellm.get_optional_params( + model="openai/gpt-4o-mini", + custom_llm_provider="edenai", + extra_body={"model": "anthropic/claude-opus-4-1", "top_k": 5}, + ) + + assert result["extra_body"] == {"top_k": 5}, result + + class TestDropParamsWithPromptCacheKey: """ Test that drop_params: true correctly drops prompt_cache_key for non-OpenAI providers. diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 13f5f686464..bb022c26a3a 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -2,6 +2,7 @@ import asyncio import contextvars import gc import json +import os import threading import time import weakref @@ -20,6 +21,7 @@ import litellm from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache from litellm.caching.disk_cache import DiskCache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType from tests.test_litellm_rust.support.isolation import rebound @@ -48,6 +50,14 @@ def redis_url() -> Generator[str]: worker.join(timeout=5) +@pytest.fixture +def cluster_nodes() -> tuple[tuple[str, int], ...]: + configured: Final = os.environ.get("LITELLM_TEST_REDIS_CLUSTER_NODES") + if not configured: + pytest.skip("LITELLM_TEST_REDIS_CLUSTER_NODES is not set") + return tuple((host, int(port)) for host, _, port in (node.partition(":") for node in configured.split(","))) + + def test_existing_constructor_and_global_are_unchanged() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) assert type(facade.cache) is InMemoryCache @@ -396,8 +406,6 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: assert client.get("second") is not None await facade.cache.disconnect() client.close() - - async def test_disk_reads_python_entries_and_python_reads_native_entries(tmp_path: Path) -> None: disk_cache: Final = DiskCache(disk_cache_dir=str(tmp_path)) response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}} @@ -504,3 +512,64 @@ async def test_disk_native_batch_lookup_and_store_report_partial_hits(tmp_path: "values": [{"value": 1}, {"value": 2}, None], "missing_indices": [2], } + + +async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_natively( + cluster_nodes: tuple[tuple[str, int], ...], +) -> None: + startup_nodes: Final = [{"host": host, "port": port} for host, port in cluster_nodes] + url: Final = f"redis://{cluster_nodes[0][0]}:{cluster_nodes[0][1]}" + with rebound(litellm, "default_redis_ttl", 60): + facade: Final = Cache(type=LiteLLMCacheType.REDIS, redis_startup_nodes=startup_nodes, namespace="parity") + assert type(facade.cache) is RedisClusterCache + with pytest.raises(TypeError, match="types must match"): + _native._CacheTestHandle.redis(url, namespace="parity")._bind_facade(facade) + _native._CacheTestHandle.redis(url, namespace="parity", startup_nodes=list(cluster_nodes))._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "native" + + manager: Final = facade.cache.redis_client.nodes_manager + with rebound(manager, "connection_kwargs", {**manager.connection_kwargs, "db": 1}): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "startup_nodes": startup_nodes[:1]}): + assert resolver.resolve().kind == "python_callback" + binding: Final = resolver.resolve() + assert binding.kind == "native" + + client: Final = redis.RedisCluster(startup_nodes=[redis.cluster.ClusterNode(*node) for node in cluster_nodes]) + keys: Final = tuple(f"slot-{index}" for index in range(12)) + slots: Final = {client.keyslot(f"parity:{key}") for key in keys} + assert len(slots) > 1, slots + requests: Final = [request(key) for key in keys] + values: Final = [{"index": index} for index in range(len(keys))] + await binding.async_store_batch(requests, values) + client.set("parity:slot-3", "not a cache entry") + client.set("parity:slot-7", json.dumps({"timestamp": time.time(), "response": {"index": 7, "python": True}})) + + batch: Final = await binding.async_lookup_batch(requests) + assert batch == { + "values": [ + None if index == 3 else {"index": 7, "python": True} if index == 7 else value + for index, value in enumerate(values) + ], + "missing_indices": [3], + } + assert facade.cache.get_cache("parity:slot-0")["response"] == {"index": 0} + assert (await facade.cache.async_get_cache("parity:slot-11"))["response"] == {"index": 11} + assert facade.cache.redis_client.mget_nonatomic([f"parity:{key}" for key in keys[:2]]) == [ + client.get("parity:slot-0"), + client.get("parity:slot-1"), + ] + + await binding.async_store({**request("pinned"), "ttl_seconds": 12.0}, {"pinned": True}) + assert 0 < client.ttl("parity:pinned") <= 12 + client.set("unscoped", "stays") + + await binding.async_flush() + + remaining: Final = tuple(sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node))) + assert remaining == (), remaining + assert client.get("unscoped") == b"stays" + client.delete("unscoped") + client.close() + facade.cache.redis_client.close() diff --git a/ui/litellm-dashboard/public/assets/logos/edenai.svg b/ui/litellm-dashboard/public/assets/logos/edenai.svg new file mode 100644 index 00000000000..957bd800e00 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/edenai.svg @@ -0,0 +1,5 @@ + + + + + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo.ts new file mode 100644 index 00000000000..5186baa605c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo.ts @@ -0,0 +1,16 @@ +import { $api } from "@/lib/http/api"; +import type { components } from "@/lib/http/schema"; + +export type LatestReleaseInfo = components["schemas"]["LatestReleaseInfo"]; + +export const useLatestReleaseInfo = (accessToken: string | null | undefined) => + $api.useQuery( + "get", + "/get/latest_release_info", + {}, + { + enabled: Boolean(accessToken), + staleTime: 60 * 60 * 1000, + retry: false, + }, + ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx index 3fe34610260..d7e1bb82564 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx @@ -41,6 +41,10 @@ vi.mock("@/components/UserBanner", () => ({ UserBanner: () => null, })); +vi.mock("@/components/UpgradeBanner", () => ({ + UpgradeBanner: () => null, +})); + vi.mock("@/contexts/ThemeContext", () => ({ ThemeProvider: ({ children }: { children: React.ReactNode }) => <>{children}, })); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index fa6df7f176a..6d326f5280e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -13,6 +13,7 @@ import { NoRedisWarningBanner } from "@/components/NoRedisWarningBanner"; import { EnvCredentialLoginWarningBanner } from "@/components/EnvCredentialLoginWarningBanner"; import { LicenseExpiryBanner } from "@/components/LicenseExpiryBanner"; import { UserBanner } from "@/components/UserBanner"; +import { UpgradeBanner } from "@/components/UpgradeBanner"; import { uiHref } from "@/utils/uiHref"; import { PluginModeProvider, usePluginMode } from "@/contexts/PluginModeContext"; import { createApiClient } from "@/lib/http/client"; @@ -117,6 +118,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) { +
@@ -137,6 +139,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) { +
{children}
diff --git a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx index d07e49e4712..b4fefe1d2c3 100644 --- a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx +++ b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx @@ -102,7 +102,7 @@ export interface ModelEditFormValues { vector_store_ids?: string[]; tags?: string[]; health_check_model?: string | null; - litellm_credential_name?: string; + litellm_credential_name?: string | null; litellm_extra_params?: string; model_info?: string; team_id?: string; @@ -139,7 +139,7 @@ const modelEditShape = { vector_store_ids: z.array(z.string()).optional(), tags: z.array(z.string()).optional(), health_check_model: z.string().nullish(), - litellm_credential_name: textish, + litellm_credential_name: z.string().nullish(), litellm_extra_params: textish, model_info: textish, team_id: textish, @@ -254,7 +254,7 @@ export const toModelEditFormValues = (localModelData: any, isWildcardModel: bool tags: Array.isArray(localModelData.litellm_params?.tags) ? localModelData.litellm_params.tags : [], // antd never mounted this field for a non-wildcard model, so the key must be absent, not null. ...(isWildcardModel ? { health_check_model: localModelData.model_info?.health_check_model } : {}), - litellm_credential_name: localModelData.litellm_params?.litellm_credential_name || "", + litellm_credential_name: localModelData.litellm_params?.litellm_credential_name ?? null, litellm_extra_params: JSON.stringify( Object.fromEntries( Object.entries(localModelData.litellm_params || {}).filter( @@ -635,8 +635,8 @@ const ModelInfoEditForm: React.FC = ({ {isEditing ? ( {({ id, value, onChange, onBlur }) => { - const items = [ - { value: "", label: "None" }, + const items: { value: string | null; label: string }[] = [ + { value: null, label: "None" }, ...credentialsList.map((credential) => ({ value: credential.credential_name, label: credential.credential_name, @@ -645,15 +645,15 @@ const ModelInfoEditForm: React.FC = ({ return (