merge(main): integrate upstream secret manager changes

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee 2026-09-21 21:48:18 +00:00
commit e0e4850ec1
178 changed files with 19421 additions and 2479 deletions

View file

@ -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")

View file

@ -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 = 40_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),
(f"Native extension does not exceed {native_size_limit / 1_000_000:.0f} MB", native_size_within_limit),
("Wheel contents are valid", not unexpected_members),
)
@ -267,7 +267,8 @@ 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 {native_size_limit / 1_000_000:.0f} MB: "
f"{native_member.file_size / 1_000_000:.2f} MB",
),
(bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"),
)

View file

@ -130,7 +130,7 @@ jobs:
- name: Test secret manager feature combinations
run: |
cargo test -p litellm-auth-gcp --locked --no-default-features
for features in '' aws google hashicorp aws,google,hashicorp; do
for features in '' aws google hashicorp cyberark aws,google,hashicorp aws,google,cyberark aws,google,hashicorp,cyberark; do
cargo test -p litellm-secrets --locked --no-default-features --features "$features"
done

View file

@ -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) | ✅ | ✅ | ✅ | | | | | | | |

View file

@ -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"
@ -2682,6 +2688,7 @@ dependencies = [
"rstest",
"serde",
"serde_json",
"serde_with",
"tokio",
"tokio-tungstenite",
]
@ -2697,6 +2704,7 @@ dependencies = [
"jsonwebtoken",
"litellm-core-utils",
"litellm-secrets-aws",
"litellm-secrets-cyberark",
"litellm-secrets-google",
"litellm-secrets-hashicorp",
"litellm-secrets-types",
@ -2732,6 +2740,26 @@ dependencies = [
"wiremock",
]
[[package]]
name = "litellm-secrets-cyberark"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"litellm-core-utils",
"litellm-secrets-types",
"moka",
"percent-encoding",
"reqwest 0.12.28",
"rstest",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
"tracing",
"veil",
"wiremock",
]
[[package]]
name = "litellm-secrets-google"
version = "0.1.0"
@ -3730,9 +3758,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",

View file

@ -23,6 +23,7 @@ litellm-secrets-types = { path = "crates/secrets-types" }
litellm-secrets-aws = { path = "crates/secrets-aws" }
litellm-secrets-google = { path = "crates/secrets-google" }
litellm-secrets-hashicorp = { path = "crates/secrets-hashicorp" }
litellm-secrets-cyberark = { path = "crates/secrets-cyberark" }
litellm-http = { path = "crates/http" }
litellm-llms = { path = "crates/llms" }
litellm-types = { path = "crates/types" }

View file

@ -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

View file

@ -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<PooledConnection, redis::RedisError> {
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::<String>(&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<C> {
Pool(r2d2::Pool<ConnectionManager>),
Cluster(r2d2::Pool<ClusterConnectionManager>),
Fixed(Mutex<C>),
}
struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike);
impl redis::ConnectionLike for ConnectionRef<'_> {
fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
self.0.req_packed_command(cmd)
}
fn req_packed_commands(
&mut self,
cmd: &[u8],
offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
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<C> Connections<C>
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<S, C = redis::Connection> {
default_ttl: Duration,
codec: S,
namespace: Option<String>,
topology: RedisTopology,
}
impl<S: CacheCodec> RedisCache<S> {
pub fn new(url: &str, default_ttl: Option<Duration>, codec: S) -> Result<Self, Error> {
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<Duration>,
codec: S,
) -> Result<Self, Error> {
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<M: r2d2::ManageConnection>(manager: M) -> Result<r2d2::Pool<M>, 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<S, C> RedisCache<S, C>
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<String>) = 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<Option<S::Value>, Error> {
@ -350,19 +308,19 @@ where
})
.collect::<Result<Vec<_>, _>>()?;
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<CacheConnectionResult, Error> {
match Self::run_blocking(Arc::clone(&self.connections), |connection| {
Ok(match redis::cmd("PING").query::<String>(connection) {
Ok(match connection.ping() {
Ok(_) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Redis cache connection test successful".into(),

View file

@ -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<C> {
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<Self, Error> {
redis::Client::open(url)
.map(Self)
.map_err(|_| Error::Unavailable)
}
}
impl r2d2::ManageConnection for ConnectionManager {
type Connection = PooledConnection<redis::Connection>;
type Error = redis::RedisError;
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
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::<String>(&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<Self, Error> {
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::<Result<Vec<_>, _>>()?;
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<ConnectionInfo, Error> {
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<ClusterConnection>;
type Error = redis::RedisError;
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
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::<String>(&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<redis::Value> {
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<Vec<redis::Value>> {
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<redis::Cmd>,
) -> Result<Vec<redis::Value>, Error> {
match self {
Self::Node(connection) => {
let mut pipeline = redis::pipe();
for command in &commands {
pipeline.add_command(command.clone());
}
pipeline
.query::<Vec<redis::Value>>(*connection)
.map_err(|_| Error::Unavailable)
}
Self::Cluster(connection) => {
let mut replies: Vec<Option<redis::Value>> = 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::<Option<Vec<_>>>()
.ok_or(Error::Unavailable)
}
}
}
pub(crate) fn scan(
&mut self,
pattern: &str,
count: usize,
mut visit: impl FnMut(&mut Self, Vec<String>) -> Result<bool, Error>,
) -> Result<(), Error> {
let pages = match self {
Self::Node(connection) => {
let page = scan_command(0, pattern, count)
.query::<ScanPage>(*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<bool, redis::RedisError> {
let command = redis::cmd("PING");
match self {
Self::Node(connection) => command
.query::<String>(*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<String, Error> {
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::<String>(node)
.map_err(|_| Error::Unavailable)?,
redis::from_redis_value::<String>(reply)
.map_err(|_| Error::Unavailable)?,
))
})
.collect::<Result<Vec<(String, String)>, Error>>()?;
replies.sort();
Ok(replies
.into_iter()
.map(|(_, reply)| reply)
.collect::<Vec<_>>()
.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<ScanPage, Error> {
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<String>);
fn primary_pages(value: redis::Value) -> Result<Vec<(NodeAddress, ScanPage)>, Error> {
let redis::Value::Map(entries) = value else {
return Err(Error::Unavailable);
};
entries
.into_iter()
.map(|(node, page)| {
let node = redis::from_redis_value::<String>(node).map_err(|_| Error::Unavailable)?;
let node = NodeAddress::try_from(node.as_str()).map_err(|_| Error::Unavailable)?;
let page = redis::from_redis_value::<ScanPage>(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<Slot, Vec<usize>> {
let mut groups: HashMap<Slot, Vec<usize>> = 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
}

View file

@ -183,20 +183,13 @@ where
}
pub fn sync_ping(&self) -> Result<bool, Error> {
self.connections.execute(|connection| {
redis::cmd("PING")
.query::<String>(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<bool, Error> {
Self::run_blocking(Arc::clone(&self.connections), |connection| {
redis::cmd("PING")
.query::<String>(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<Vec<String>, 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<String>) = 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::<usize>)
.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::<Vec<_>>();
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::<Vec<redis::Value>>(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<String, Error> {
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<String, Error> {
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
}

View file

@ -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<JsonCodec<serde_json::Value>>;
fn topology() -> Option<RedisTopology> {
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<Cache> {
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<RedisCache<JsonCodec<f64>>> {
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<String> {
let keys: Vec<String> = (0..count).map(|index| format!("key-{index}")).collect();
let slots: std::collections::HashSet<Slot> = 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<String> = 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::<Vec<_>>());
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<Vec<u8>> = (0..=index)
.take(2)
.map(|value| value.to_string().into_bytes())
.collect();
assert_eq!(values, expected);
}
other => panic!("queue {index}: {other:?}"),
}
}
let counters: Vec<String> = 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<f64> = (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<String> = 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);
}

View file

@ -1,33 +1,91 @@
use serde::{Deserialize, Deserializer, de::Error};
use serde_json::Value;
use serde::{
Deserializer,
de::{Error, Visitor},
};
use serde_with::DeserializeAs;
pub struct LaxI64;
pub struct FiniteF64;
pub fn parse_str_bool(value: &str) -> Option<bool> {
let token = value.trim_matches(|character: char| {
character.is_whitespace() || matches!(character, '\u{1c}'..='\u{1f}')
});
if token.eq_ignore_ascii_case("true") {
return Some(true);
}
token.eq_ignore_ascii_case("false").then_some(false)
}
impl<'de> DeserializeAs<'de, i64> for LaxI64 {
fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<i64, D::Error> {
match Value::deserialize(deserializer)? {
Value::Number(number) if number.is_f64() => number.as_f64().and_then(integral_float),
Value::Number(number) => number.as_i64(),
Value::String(value) => integer_string(value.trim()),
Value::Bool(value) => Some(i64::from(value)),
_ => None,
}
.ok_or_else(|| D::Error::custom("expected an integer in the i64 range"))
deserializer.deserialize_any(Self)
}
}
impl<'de> Visitor<'de> for LaxI64 {
type Value = i64;
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("an integer in the i64 range")
}
fn visit_i64<E: Error>(self, value: i64) -> Result<i64, E> {
Ok(value)
}
fn visit_u64<E: Error>(self, value: u64) -> Result<i64, E> {
i64::try_from(value).map_err(E::custom)
}
fn visit_f64<E: Error>(self, value: f64) -> Result<i64, E> {
integral_float(value).ok_or_else(|| E::custom("expected an integer in the i64 range"))
}
fn visit_str<E: Error>(self, value: &str) -> Result<i64, E> {
integer_string(value.trim())
.ok_or_else(|| E::custom("expected an integer in the i64 range"))
}
fn visit_bool<E: Error>(self, value: bool) -> Result<i64, E> {
Ok(i64::from(value))
}
}
impl<'de> DeserializeAs<'de, f64> for FiniteF64 {
fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<f64, D::Error> {
match Value::deserialize(deserializer)? {
Value::Number(number) => number.as_f64(),
Value::String(value) => value.trim().parse::<f64>().ok(),
Value::Bool(value) => Some(f64::from(value)),
_ => None,
}
.filter(|value| value.is_finite())
.ok_or_else(|| D::Error::custom("expected a finite number"))
deserializer.deserialize_any(Self)
}
}
impl<'de> Visitor<'de> for FiniteF64 {
type Value = f64;
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("a finite number")
}
fn visit_i64<E: Error>(self, value: i64) -> Result<f64, E> {
Ok(value as f64)
}
fn visit_u64<E: Error>(self, value: u64) -> Result<f64, E> {
Ok(value as f64)
}
fn visit_f64<E: Error>(self, value: f64) -> Result<f64, E> {
value
.is_finite()
.then_some(value)
.ok_or_else(|| E::custom("expected a finite number"))
}
fn visit_str<E: Error>(self, value: &str) -> Result<f64, E> {
self.visit_f64(value.trim().parse::<f64>().map_err(E::custom)?)
}
fn visit_bool<E: Error>(self, value: bool) -> Result<f64, E> {
Ok(f64::from(value))
}
}
@ -66,7 +124,7 @@ fn integral_float(value: f64) -> Option<i64> {
#[cfg(test)]
mod tests {
use serde::Serialize;
use serde::{Deserialize, Serialize};
use serde_json::json;
use serde_with::serde_as;
@ -81,6 +139,22 @@ mod tests {
float: Option<f64>,
}
#[test]
fn boolean_tokens_follow_python_string_trimming_without_redis_tokens() {
for (input, expected) in [
(" True ", Some(true)),
("\u{1c}TRUE\u{1f}", Some(true)),
("\u{a0}False\u{2003}", Some(false)),
("true\u{200b}", None),
("yes", None),
("1", None),
("", None),
("unknown", None),
] {
assert_eq!(parse_str_bool(input), expected, "{input:?}");
}
}
#[test]
fn adapters_compose_and_serialize_as_numbers() {
let numbers: Numbers = serde_json::from_value(json!({

View file

@ -1,5 +1,7 @@
use std::str::FromStr;
use crate::serde_compat::parse_str_bool;
pub trait Lookup {
fn get(&self, name: &str) -> Option<String>;
@ -9,7 +11,7 @@ pub trait Lookup {
fn enabled(&self, name: &str) -> Option<bool> {
self.get(name)
.is_some_and(|value| value.trim().eq_ignore_ascii_case("true"))
.is_some_and(|value| parse_str_bool(&value) == Some(true))
.then_some(true)
}

View file

@ -45,7 +45,7 @@ where
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
catch_unwind(AssertUnwindSafe(|| pythonize::pythonize(py, &self.0)))
.map_err(panic_to_pyerr)?
.map_err(|error| PyValueError::new_err(error.to_string()))
.map_err(PyErr::from)
}
}
@ -87,6 +87,19 @@ mod tests {
});
}
#[test]
fn pythonized_preserves_python_serialization_error_types() {
crate::initialize_python();
Python::attach(|py| {
let value = std::collections::BTreeMap::from([(vec![1], "value")]);
let direct = to_py(py, &value).unwrap_err();
let wrapped = Pythonized(value).into_pyobject(py).unwrap_err();
assert!(direct.is_instance_of::<pyo3::exceptions::PyTypeError>(py));
assert!(wrapped.is_instance_of::<pyo3::exceptions::PyTypeError>(py));
assert_eq!(wrapped.to_string(), direct.to_string());
});
}
#[test]
fn pythonized_maps_serializer_panics_to_a_base_exception() {
crate::initialize_python();

View file

@ -129,6 +129,7 @@ mod tests {
use rstest::rstest;
use super::*;
use crate::TlsSource;
fn settings(ssl_verify: Option<SslVerify>, ssl_cert_file: Option<&str>) -> HttpSettings {
HttpSettings {
@ -298,7 +299,11 @@ mod tests {
};
assert!(matches!(
reqwest::ClientBuilder::try_from(&config),
Err(Error::Read { path: reported, .. }) if reported == path
Err(Error::Read {
path: reported,
tls_source: TlsSource::CaBundle,
..
}) if reported == path
));
}
@ -315,7 +320,11 @@ mod tests {
std::fs::remove_file(&path).unwrap();
assert!(matches!(
result,
Err(Error::InvalidPem { path: reported, .. }) if reported == path
Err(Error::InvalidPem {
path: reported,
tls_source: TlsSource::CaBundle,
..
}) if reported == path
));
}
}

View file

@ -1,11 +1,25 @@
use std::path::PathBuf;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TlsSource {
CaBundle,
ClientIdentity,
}
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
pub enum Error {
#[error("could not read {}: {message}", path.display())]
Read { path: PathBuf, message: String },
Read {
path: PathBuf,
message: String,
tls_source: TlsSource,
},
#[error("{} is not a PEM file: {message}", path.display())]
InvalidPem { path: PathBuf, message: String },
InvalidPem {
path: PathBuf,
message: String,
tls_source: TlsSource,
},
#[error("could not build the HTTP client: {0}")]
Client(String),
#[error("request body could not be serialized: {0}")]

View file

@ -10,7 +10,7 @@ mod tls;
pub mod transport;
pub use config::{HttpClientConfig, Resolution, Verify};
pub use error::Error;
pub use error::{Error, TlsSource};
pub use pool::{ClientVariant, HttpClientPool};
pub use proxy::EnvironmentProxies;
pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive};

View file

@ -54,16 +54,39 @@ impl Default for UrlPolicy {
impl UrlPolicy {
fn allows(&self, host: &str, port: u16) -> bool {
let host = normalize_host(host);
let with_port = format!("{host}:{port}");
self.allowed_hosts
.iter()
.map(|entry| normalize_host(entry))
.any(|entry| entry == host || entry == with_port)
.filter_map(|entry| parse_allowed_host(entry))
.any(|(entry_host, entry_port)| {
entry_host == host && entry_port.is_none_or(|entry_port| entry_port == port)
})
}
}
fn normalize_host(host: &str) -> String {
host.to_ascii_lowercase().trim_end_matches('.').to_owned()
pub fn normalize_host(host: &str) -> String {
let host = host.trim().trim_end_matches('.');
let host = host
.strip_prefix('[')
.and_then(|host| host.strip_suffix(']'))
.unwrap_or(host);
host.to_ascii_lowercase()
}
fn parse_allowed_host(entry: &str) -> Option<(String, Option<u16>)> {
let entry = entry.trim();
if let Some(entry) = entry.strip_prefix('[') {
let (host, suffix) = entry.split_once(']')?;
let port = match suffix {
"" => None,
suffix => Some(suffix.strip_prefix(':')?.parse().ok()?),
};
return Some((normalize_host(host), port));
}
let (host, port) = match entry.rsplit_once(':') {
Some((host, port)) if !host.contains(':') => (host, Some(port.parse().ok()?)),
_ => (entry, None),
};
Some((normalize_host(host), port))
}
type ProxyMatch = Arc<dyn Fn(&Url) -> bool + Send + Sync>;
@ -670,6 +693,21 @@ mod tests {
assert!(matches!(result, Err(Error::BlockedUrl)));
}
#[test]
fn allowlist_matches_bracketed_ipv6_hosts_and_ports() {
let policy = UrlPolicy {
validate: true,
allowed_hosts: vec!["[2001:db8::1]".into(), "[2001:db8::1]:8443".into()],
};
assert!(policy.allows("2001:db8::1", 443));
assert!(policy.allows("2001:db8::1", 8443));
let port_specific = UrlPolicy {
validate: true,
allowed_hosts: vec!["[2001:db8::1]:8443".into()],
};
assert!(!port_specific.allows("2001:db8::1", 9443));
}
#[tokio::test]
async fn validation_off_fetches_private_hosts_and_follows_redirects() {
let (url, server, _) = serve_named(

View file

@ -3,7 +3,10 @@ use std::{
time::Duration,
};
use litellm_core_utils::settings::{Layer, Lookup, merge};
use litellm_core_utils::{
serde_compat::parse_str_bool,
settings::{Layer, Lookup, merge},
};
use crate::proxy::EnvironmentProxies;
@ -16,9 +19,9 @@ pub enum SslVerify {
impl SslVerify {
pub fn parse(value: &str) -> Self {
match value.trim().to_ascii_lowercase().as_str() {
"true" => Self::Enabled,
"false" => Self::Disabled,
match parse_str_bool(value) {
Some(true) => Self::Enabled,
Some(false) => Self::Disabled,
_ => Self::CaBundle(PathBuf::from(value)),
}
}
@ -152,9 +155,7 @@ impl HttpSettings {
Self {
ssl_verify: merged.ssl_verify,
ssl_cert_file: merged.ssl_cert_file,
ssl_certificate: merged
.ssl_certificate
.filter(|path| !path.as_os_str().is_empty()),
ssl_certificate: merged.ssl_certificate,
ssl_security_level: merged.ssl_security_level.filter(|level| !level.is_empty()),
ssl_ecdh_curve: merged.ssl_ecdh_curve.filter(|curve| !curve.is_empty()),
force_ipv4: merged.force_ipv4.unwrap_or(defaults.force_ipv4),
@ -287,7 +288,7 @@ mod tests {
}
#[test]
fn empty_environment_values_clear_the_setting_like_python_truthiness() {
fn empty_certificate_is_retained_for_validation_while_empty_tuning_is_absent() {
let configured = HttpSettingsLayer {
ssl_certificate: Some("/configured/client.pem".into()),
ssl_security_level: Some("configured".into()),
@ -300,7 +301,7 @@ mod tests {
("SSL_ECDH_CURVE", ""),
]));
let settings = HttpSettings::from_layers([environment, configured]);
assert_eq!(settings.ssl_certificate, None);
assert_eq!(settings.ssl_certificate, Some(PathBuf::new()));
assert_eq!(settings.ssl_security_level, None);
assert_eq!(settings.ssl_ecdh_curve, None);
}

View file

@ -9,7 +9,7 @@ use rustls::{
use crate::{
config::{HttpClientConfig, Verify},
error::Error,
error::{Error, TlsSource},
};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
@ -197,15 +197,17 @@ impl TryFrom<&HttpClientConfig> for ClientConfig {
Verify::BuiltInRoots => builder.with_root_certificates(RootCertStore {
roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(),
}),
Verify::CaBundle(path) => builder.with_root_certificates(bundle_roots(path)?),
Verify::CaBundle(path) => {
builder.with_root_certificates(bundle_roots(path, TlsSource::CaBundle)?)
}
};
let mut tls = match &config.client_certificate {
None => verified.with_no_client_auth(),
Some(path) => {
let (chain, key) = identity(path)?;
let (chain, key) = identity(path, TlsSource::ClientIdentity)?;
verified
.with_client_auth_cert(chain, key)
.map_err(|error| invalid_pem(path, error))?
.map_err(|error| invalid_pem(path, TlsSource::ClientIdentity, error))?
}
};
tls.alpn_protocols = if config.http2 {
@ -217,47 +219,52 @@ impl TryFrom<&HttpClientConfig> for ClientConfig {
}
}
fn bundle_roots(path: &Path) -> Result<RootCertStore, Error> {
let certificates = certificates(path)?;
fn bundle_roots(path: &Path, source: TlsSource) -> Result<RootCertStore, Error> {
let certificates = certificates(path, source)?;
if certificates.is_empty() {
return Err(invalid_pem(path, "no certificates found"));
return Err(invalid_pem(path, source, "no certificates found"));
}
let mut store = RootCertStore::empty();
for certificate in certificates {
store
.add(certificate)
.map_err(|error| invalid_pem(path, error))?;
.map_err(|error| invalid_pem(path, source, error))?;
}
Ok(store)
}
fn identity(path: &Path) -> Result<(Vec<CertificateDer<'static>>, PrivateKeyDer<'static>), Error> {
let chain = certificates(path)?;
fn identity(
path: &Path,
source: TlsSource,
) -> Result<(Vec<CertificateDer<'static>>, PrivateKeyDer<'static>), Error> {
let chain = certificates(path, source)?;
if chain.is_empty() {
return Err(invalid_pem(path, "no certificates found"));
return Err(invalid_pem(path, source, "no certificates found"));
}
let key =
PrivateKeyDer::from_pem_slice(&read(path)?).map_err(|error| invalid_pem(path, error))?;
let key = PrivateKeyDer::from_pem_slice(&read(path, source)?)
.map_err(|error| invalid_pem(path, source, error))?;
Ok((chain, key))
}
fn certificates(path: &Path) -> Result<Vec<CertificateDer<'static>>, Error> {
CertificateDer::pem_slice_iter(&read(path)?)
fn certificates(path: &Path, source: TlsSource) -> Result<Vec<CertificateDer<'static>>, Error> {
CertificateDer::pem_slice_iter(&read(path, source)?)
.collect::<Result<_, _>>()
.map_err(|error| invalid_pem(path, error))
.map_err(|error| invalid_pem(path, source, error))
}
fn read(path: &Path) -> Result<Vec<u8>, Error> {
fn read(path: &Path, source: TlsSource) -> Result<Vec<u8>, Error> {
std::fs::read(path).map_err(|error| Error::Read {
path: path.to_path_buf(),
message: error.to_string(),
tls_source: source,
})
}
fn invalid_pem(path: &Path, message: impl fmt::Display) -> Error {
fn invalid_pem(path: &Path, source: TlsSource, message: impl fmt::Display) -> Error {
Error::InvalidPem {
path: path.to_path_buf(),
message: message.to_string(),
tls_source: source,
}
}
@ -405,7 +412,11 @@ mod tests {
std::fs::remove_file(&path).unwrap();
assert!(matches!(
result,
Err(Error::InvalidPem { path: reported, .. }) if reported == path
Err(Error::InvalidPem {
path: reported,
tls_source: TlsSource::ClientIdentity,
..
}) if reported == path
));
}
}

View file

@ -41,6 +41,8 @@ serde_json.workspace = true
tokio = { workspace = true, features = ["sync"] }
[dev-dependencies]
serde.workspace = true
serde_with.workspace = true
criterion.workspace = true
futures-util.workspace = true
rstest.workspace = true

View file

@ -1,26 +1,154 @@
{
"http_settings": [
"ssl_verify",
"ssl_certificate",
"ssl_security_level",
"ssl_ecdh_curve",
"force_ipv4",
"http2",
"aiohttp_trust_env",
"disable_aiohttp_trust_env",
"disable_aiohttp_transport",
"user_agent"
],
"url_policy": [
"user_url_validation",
"user_url_allowed_hosts"
],
"provider_defaults": [
"vertex_project",
"vertex_location",
"enable_azure_ad_token_refresh"
],
"secret_manager": [
"readable"
]
"http_settings": {
"version": 1,
"fields": {
"ssl_verify": {
"adapter": "SslVerifyInput",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [
"none",
"bool",
"str"
],
"unsupported_live": "configuration_error"
},
"ssl_certificate": {
"adapter": "OptionalStrictString",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"ssl_security_level": {
"adapter": "TuningString",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"ssl_ecdh_curve": {
"adapter": "TuningString",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"force_ipv4": {
"adapter": "Truthy",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"http2": {
"adapter": "ExactTrue",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"aiohttp_trust_env": {
"adapter": "Truthy",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"disable_aiohttp_trust_env": {
"adapter": "Truthy",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"disable_aiohttp_transport": {
"adapter": "ExactTrue",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"user_agent": {
"adapter": "StrictString",
"required": true,
"precedence": "accessor",
"sensitive": false,
"shapes": [],
"unsupported_live": null
}
}
},
"url_policy": {
"version": 1,
"fields": {
"user_url_validation": {
"adapter": "Truthy",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"user_url_allowed_hosts": {
"adapter": "HostCollection",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
}
}
},
"provider_defaults": {
"version": 1,
"fields": {
"vertex_project": {
"adapter": "FalsyOptionalString",
"required": true,
"precedence": "module_global",
"sensitive": true,
"shapes": [],
"unsupported_live": null
},
"vertex_location": {
"adapter": "FalsyOptionalString",
"required": true,
"precedence": "module_global",
"sensitive": true,
"shapes": [],
"unsupported_live": null
},
"enable_azure_ad_token_refresh": {
"adapter": "ExactTrue",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
}
}
},
"secret_manager": {
"version": 1,
"fields": {
"readable": {
"adapter": "StrictBool",
"required": true,
"precedence": "accessor",
"sensitive": false,
"shapes": [],
"unsupported_live": null
}
}
}
}

View file

@ -1,10 +1,11 @@
use std::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};
@ -70,9 +71,21 @@ pub(super) struct RedisCacheConfig {
pub(super) default_ttl: Duration,
pub(super) namespace: Option<String>,
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<RedisTlsConfig>,
}
const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31;
pub(super) enum CacheBackendConfig {
Memory(MemoryCacheConfig),
Redis(Box<RedisCacheConfig>),
@ -182,6 +195,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"),
@ -206,9 +222,6 @@ fn project_redis(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<RedisCacheConfig, UnsupportedCacheConfig>> {
let source = backend.getattr("redis_kwargs")?.cast_into::<PyDict>()?;
if has_value(&source, "startup_nodes")? {
return Ok(Err(UnsupportedCacheConfig::RedisTopology));
}
if has_value(&source, "sentinel_nodes")? {
return Ok(Err(UnsupportedCacheConfig::RedisTopology));
}
@ -248,26 +261,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::<PyDict>()?;
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,
@ -280,15 +292,15 @@ fn project_redis(
default_ttl: duration(backend.getattr("default_ttl")?.extract::<f64>()?)?,
namespace: optional_attribute_string(backend, "namespace")?,
flush_size: backend.getattr("redis_flush_size")?.extract::<usize>()?,
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::<usize>()?,
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")?,
@ -299,6 +311,128 @@ fn project_redis(
}))
}
#[inline(never)]
fn project_standalone_client<'py>(
client: &Bound<'py, PyAny>,
) -> PyResult<Result<RedisClientProjection<'py>, 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::<PyDict>()?;
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::<usize>()?,
resolved,
tls,
}))
}
#[inline(never)]
fn project_cluster_client<'py>(
source: &Bound<'py, PyDict>,
client: &Bound<'py, PyAny>,
) -> PyResult<Result<RedisClientProjection<'py>, 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::<PyDict>()?;
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<Option<Vec<RedisNode>>> {
let Some(nodes) = source.get_item("startup_nodes")? else {
return Ok(None);
};
let Ok(nodes) = nodes.cast_into::<PyList>() 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::<PyDict>() 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> {
u16::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis port"))
}
#[inline(never)]
fn project_tls(values: &Bound<'_, PyDict>) -> PyResult<RedisTlsConfig> {
Ok(RedisTlsConfig {
@ -467,12 +601,27 @@ mod tests {
use pyo3::{prelude::*, types::PyDict};
use litellm_cache_redis::{RedisNode, RedisTopology};
use super::{
CacheBackendConfig, CacheConfigProjection, CertificateRequirement, 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(
@ -591,4 +740,87 @@ mod tests {
assert_eq!(reason.message(), "native Redis credentials require Python");
});
}
#[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}");
}
});
}
}

View file

@ -1,3 +1,4 @@
use litellm_cache_redis::RedisTopology;
use litellm_host_python::from_py;
use pyo3::{
PyTraverseError, PyVisit,
@ -29,9 +30,28 @@ struct RedisPoolGuard {
reference: Py<PyAny>,
connection_class: Py<PyAny>,
connection_kwargs: Py<PyAny>,
max_connections: usize,
max_connections: Option<usize>,
attributes: RedisPoolAttributes,
}
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,
@ -138,31 +158,40 @@ impl ObjectGuard {
}
impl RedisPoolGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let pool = backend
.getattr("redis_client")?
.getattr("connection_pool")?;
fn capture(backend: &Bound<'_, PyAny>, attributes: RedisPoolAttributes) -> PyResult<Self> {
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::<usize>()?,
max_connections: Self::max_connections(&pool, &attributes)?,
attributes,
})
}
fn max_connections(
pool: &Bound<'_, PyAny>,
attributes: &RedisPoolAttributes,
) -> PyResult<Option<usize>> {
attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
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::<usize>()?
.is(&pool.getattr(self.attributes.connection_class)?)
&& self.max_connections == Self::max_connections(&pool, &self.attributes)?
&& self
.connection_kwargs
.bind(py)
@ -189,9 +218,15 @@ 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"),
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",
),
_ => unreachable!(),
};
let backend = facade.getattr("cache")?;
@ -237,9 +272,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,
},
})
}

View file

@ -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<String>,
startup_nodes: Option<Vec<(String, u16)>>,
) -> PyResult<Self> {
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,

View file

@ -2,7 +2,7 @@ use std::{sync::Arc, time::Duration};
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
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,
};
@ -34,10 +34,12 @@ impl NativeResponseCache {
pub fn redis(
url: &str,
topology: &RedisTopology,
ttl: Option<Duration>,
namespace: Option<String>,
) -> Result<Self, Error> {
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,
@ -67,6 +69,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<usize> {
match self {
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),

View file

@ -0,0 +1,231 @@
use std::collections::BTreeSet;
use litellm_core_utils::serde_compat::parse_str_bool;
use litellm_http::SslVerify;
use pyo3::{
exceptions::{PyAttributeError, PyRuntimeError, PyValueError},
prelude::*,
types::{PyBool, PyString},
};
#[derive(Debug)]
pub(crate) enum ProjectionError {
Python(PyErr),
InvalidConfiguration(String),
UnsupportedLiveObject(String),
InternalSchemaFailure(String),
}
impl From<PyErr> for ProjectionError {
fn from(error: PyErr) -> Self {
Self::Python(error)
}
}
impl From<ProjectionError> for PyErr {
fn from(error: ProjectionError) -> Self {
match error {
ProjectionError::Python(error) => error,
ProjectionError::InvalidConfiguration(message)
| ProjectionError::UnsupportedLiveObject(message) => PyValueError::new_err(message),
ProjectionError::InternalSchemaFailure(message) => PyRuntimeError::new_err(message),
}
}
}
pub(crate) struct Truthy(pub bool);
pub(crate) struct ExactTrue(pub bool);
pub(crate) struct StrBool(pub Option<bool>);
pub(crate) struct OptionalStrictString(pub Option<String>);
pub(crate) struct FalsyOptionalString(pub Option<String>);
pub(crate) struct TuningString(pub Option<String>);
pub(crate) struct StringCollection(pub Vec<String>);
pub(crate) struct SslVerifyInput(pub Option<SslVerify>);
pub(crate) struct Field<'py> {
path: &'static str,
value: Bound<'py, PyAny>,
}
impl<'py> Field<'py> {
pub(crate) fn new(path: &'static str, value: Bound<'py, PyAny>) -> Self {
Self { path, value }
}
pub(crate) fn read(
snapshot: &Bound<'py, PyAny>,
path: &'static str,
) -> Result<Self, ProjectionError> {
let name = path.rsplit('.').next().unwrap_or(path);
match snapshot.getattr(name) {
Ok(value) => Ok(Self::new(path, value)),
Err(error) if error.is_instance_of::<PyAttributeError>(snapshot.py()) => {
match Self::missing_field(snapshot, name) {
Ok(true) => Err(ProjectionError::InternalSchemaFailure(format!(
"{path}: missing snapshot field"
))),
_ => Err(error.into()),
}
}
Err(error) => Err(error.into()),
}
}
fn missing_field(snapshot: &Bound<'_, PyAny>, name: &str) -> PyResult<bool> {
let py = snapshot.py();
let object = py.import("builtins")?.getattr("object")?;
let missing = object.call0()?;
let lookup = py.import("inspect")?.getattr("getattr_static")?;
let declared = lookup.call1((snapshot, name, &missing))?;
let fallback = lookup.call1((snapshot.get_type(), "__getattr__", &missing))?;
let getter = lookup.call1((snapshot.get_type(), "__getattribute__"))?;
Ok(declared.is(&missing)
&& fallback.is(&missing)
&& getter.is(object.getattr("__getattribute__")?))
}
fn expected(&self, expected: &'static str) -> Result<String, ProjectionError> {
Ok(format!(
"{}: expected {expected}, got {}",
self.path,
self.value.get_type().name()?
))
}
fn invalid(&self, expected: &'static str) -> ProjectionError {
match self.expected(expected) {
Ok(message) => ProjectionError::InvalidConfiguration(message),
Err(error) => error,
}
}
pub(crate) fn truthy(&self) -> Result<Truthy, ProjectionError> {
Ok(Truthy(self.value.is_truthy()?))
}
pub(crate) fn exact_true(&self) -> ExactTrue {
ExactTrue(self.value.is(PyBool::new(self.value.py(), true)))
}
pub(crate) fn strict_string(&self) -> Result<String, ProjectionError> {
let value = self
.value
.cast::<PyString>()
.map_err(|_| self.invalid("a string"))?;
Ok(value.to_str()?.to_owned())
}
pub(crate) fn schema_string(&self) -> Result<String, ProjectionError> {
if !self.value.is_instance_of::<PyString>() {
return Err(ProjectionError::InternalSchemaFailure(
self.expected("a string")?,
));
}
self.strict_string()
}
pub(crate) fn schema_bool(&self) -> Result<bool, ProjectionError> {
if !self.value.is_instance_of::<PyBool>() {
return Err(ProjectionError::InternalSchemaFailure(
self.expected("a Boolean")?,
));
}
Ok(self.exact_true().0)
}
pub(crate) fn str_bool(&self) -> Result<StrBool, ProjectionError> {
if self.value.is_none() {
return Ok(StrBool(None));
}
Ok(StrBool(parse_str_bool(&self.strict_string()?)))
}
pub(crate) fn optional_strict_string(&self) -> Result<OptionalStrictString, ProjectionError> {
if self.value.is_none() {
return Ok(OptionalStrictString(None));
}
self.strict_string().map(Some).map(OptionalStrictString)
}
pub(crate) fn falsy_optional_string(&self) -> Result<FalsyOptionalString, ProjectionError> {
if !self.truthy()?.0 {
return Ok(FalsyOptionalString(None));
}
self.strict_string().map(Some).map(FalsyOptionalString)
}
pub(crate) fn tuning_string(&self) -> Result<TuningString, ProjectionError> {
if !self.truthy()?.0 || !self.value.is_instance_of::<PyString>() {
return Ok(TuningString(None));
}
self.strict_string().map(Some).map(TuningString)
}
pub(crate) fn string_collection(&self) -> Result<StringCollection, ProjectionError> {
if !self.truthy()?.0 {
return Ok(StringCollection(Vec::new()));
}
if self.value.is_instance_of::<PyString>() {
return self
.strict_string()
.map(|value| StringCollection(vec![value]));
}
let values = self
.value
.try_iter()?
.filter_map(|item| {
let member = match item {
Ok(value) => Self::new(self.path, value),
Err(error) => return Some(Err(error.into())),
};
match member.truthy() {
Ok(Truthy(false)) => None,
Ok(Truthy(true)) => Some(member.strict_string()),
Err(error) => Some(Err(error)),
}
})
.collect::<Result<Vec<_>, ProjectionError>>()?;
Ok(StringCollection(values))
}
pub(crate) fn host_collection(&self) -> Result<StringCollection, ProjectionError> {
let values = self
.string_collection()?
.0
.into_iter()
.map(|host| litellm_http::media::normalize_host(&host))
.collect::<BTreeSet<_>>();
Ok(StringCollection(values.into_iter().collect()))
}
pub(crate) fn ssl_verify(&self) -> Result<SslVerifyInput, ProjectionError> {
if self.value.is_none() {
return Ok(SslVerifyInput(None));
}
if self.value.is_instance_of::<PyBool>() {
return Ok(SslVerifyInput(Some(if self.exact_true().0 {
SslVerify::Enabled
} else {
SslVerify::Disabled
})));
}
if self.value.is_instance_of::<PyString>() {
let parsed = match self.str_bool()?.0 {
Some(true) => SslVerify::Enabled,
Some(false) => SslVerify::Disabled,
None => SslVerify::CaBundle(self.strict_string()?.into()),
};
return Ok(SslVerifyInput(Some(parsed)));
}
let context = self.value.py().import("ssl")?.getattr("SSLContext")?;
if self.value.is_instance(&context)? {
return Err(ProjectionError::UnsupportedLiveObject(self.expected(
"a Boolean, Boolean string, CA path, or None; live SSLContext is unsupported",
)?));
}
Err(self.invalid("a Boolean, Boolean string, CA path, or None"))
}
}
#[cfg(test)]
mod tests;

View file

@ -0,0 +1,372 @@
use std::ffi::CString;
use pyo3::{
exceptions::{PyLookupError, PyRuntimeError, PyValueError},
types::PyDict,
};
use rstest::rstest;
use super::*;
fn evaluate<'py>(py: Python<'py>, source: &str) -> Bound<'py, PyAny> {
py.eval(&CString::new(source).unwrap(), None, None).unwrap()
}
#[rstest]
#[case("None", false, false)]
#[case("False", false, false)]
#[case("True", true, true)]
#[case("0", false, false)]
#[case("1", true, false)]
#[case("''", false, false)]
#[case("'false'", true, false)]
#[case("[]", false, false)]
#[case("[0]", true, false)]
#[case("{}", false, false)]
#[case("object()", true, false)]
fn boolean_operations_have_distinct_python_semantics(
#[case] source: &str,
#[case] truth: bool,
#[case] exact: bool,
) {
Python::initialize();
Python::attach(|py| {
let value = evaluate(py, source);
let field = Field::new("test.flag", value.clone());
assert_eq!(field.truthy().unwrap().0, truth);
assert_eq!(field.exact_true().0, exact);
assert_eq!(
field.truthy().unwrap().0,
py.import("builtins")
.unwrap()
.getattr("bool")
.unwrap()
.call1((value,))
.unwrap()
.extract::<bool>()
.unwrap()
);
});
}
#[rstest]
#[case("None", Ok(None), Ok(None), Ok(None))]
#[case("''", Ok(Some("")), Ok(None), Ok(None))]
#[case(
"' value '",
Ok(Some(" value ")),
Ok(Some(" value ")),
Ok(Some(" value "))
)]
#[case("[]", Err(()), Ok(None), Ok(None))]
#[case("0", Err(()), Ok(None), Ok(None))]
#[case("1", Err(()), Err(()), Ok(None))]
#[case("object()", Err(()), Err(()), Ok(None))]
fn string_operations_do_not_conflate_absence_and_type_checks(
#[case] source: &str,
#[case] strict: Result<Option<&str>, ()>,
#[case] fallback: Result<Option<&str>, ()>,
#[case] tuning: Result<Option<&str>, ()>,
) {
Python::initialize();
Python::attach(|py| {
let field = Field::new("test.string", evaluate(py, source));
let owned =
|expected: Result<Option<&str>, ()>| expected.map(|value| value.map(str::to_owned));
assert_eq!(
field
.optional_strict_string()
.map(|value| value.0)
.map_err(|_| ()),
owned(strict)
);
assert_eq!(
field
.falsy_optional_string()
.map(|value| value.0)
.map_err(|_| ()),
owned(fallback)
);
assert_eq!(
field.tuning_string().map(|value| value.0).map_err(|_| ()),
owned(tuning)
);
});
}
#[rstest]
#[case("None", None)]
#[case("' True '", Some(true))]
#[case("' fAlSe '", Some(false))]
#[case("'yes'", None)]
#[case("'1'", None)]
#[case("'unknown'", None)]
fn string_boolean_tokens_remain_separate_from_truthiness(
#[case] source: &str,
#[case] expected: Option<bool>,
) {
Python::initialize();
Python::attach(|py| {
assert_eq!(
Field::new("test.flag", evaluate(py, source))
.str_bool()
.unwrap()
.0,
expected
);
});
}
#[rstest]
#[case("'EXAMPLE.TEST.'", vec!["example.test"])]
#[case("['B.test', '', None, 0, [], 'A.test.', 'b.test']", vec!["a.test", "b.test"])]
#[case("('B.test', 'a.test')", vec!["a.test", "b.test"])]
#[case("{'B.test', 'a.test'}", vec!["a.test", "b.test"])]
#[case("(host for host in ['B.test', 'a.test'])", vec!["a.test", "b.test"])]
#[case("None", vec![])]
#[case("False", vec![])]
fn host_collection_is_owned_normalized_and_deterministic(
#[case] source: &str,
#[case] expected: Vec<&str>,
) {
Python::initialize();
Python::attach(|py| {
assert_eq!(
Field::new("url_policy.user_url_allowed_hosts", evaluate(py, source))
.host_collection()
.unwrap()
.0,
expected
);
});
}
#[test]
fn protocol_errors_preserve_exception_identity_traceback_cause_and_context() {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
py.run(
c"
failure = LookupError('protocol failed')
cause = ValueError('cause')
context = RuntimeError('context')
def fail():
try:
raise context
except RuntimeError:
raise failure from cause
class Bool:
def __bool__(self): return fail()
class Length:
def __len__(self): return fail()
class Iter:
def __iter__(self): return fail()
class Next:
def __iter__(self): return self
def __next__(self): return fail()
class Descriptor:
@property
def flag(self): return fail()
values = (Bool(), Length(), Iter(), Next(), [Bool()])
descriptor = Descriptor()
",
Some(&locals),
Some(&locals),
)
.unwrap();
let values = locals.get_item("values").unwrap().unwrap();
for value in values.try_iter().unwrap() {
let error = Field::new("test.flag", value.unwrap())
.host_collection()
.err()
.unwrap();
let error = PyErr::from(error);
assert!(
error
.value(py)
.is(locals.get_item("failure").unwrap().unwrap())
);
assert!(error.is_instance_of::<PyLookupError>(py));
assert!(error.traceback(py).is_some());
assert!(
error
.value(py)
.getattr("__cause__")
.unwrap()
.is(locals.get_item("cause").unwrap().unwrap())
);
assert!(
error
.value(py)
.getattr("__context__")
.unwrap()
.is(locals.get_item("context").unwrap().unwrap())
);
}
let error = Field::read(
&locals.get_item("descriptor").unwrap().unwrap(),
"test.flag",
)
.err()
.unwrap();
assert!(
PyErr::from(error)
.value(py)
.is(locals.get_item("failure").unwrap().unwrap())
);
});
}
#[test]
fn identity_and_string_contents_do_not_invoke_unrelated_protocols() {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
py.run(
c"
class Hostile:
def __bool__(self): raise AssertionError('bool called')
def __eq__(self, other): raise AssertionError('eq called')
def __str__(self): raise AssertionError('str called')
class Text(str):
def __str__(self): raise AssertionError('str called')
def strip(self): raise AssertionError('strip called')
def lower(self): raise AssertionError('lower called')
hostile = Hostile()
text = Text(' False ')
",
Some(&locals),
Some(&locals),
)
.unwrap();
let hostile = Field::new("test.flag", locals.get_item("hostile").unwrap().unwrap());
assert!(!hostile.exact_true().0);
assert!(matches!(
hostile.strict_string(),
Err(ProjectionError::InvalidConfiguration(_))
));
let text = Field::new("test.flag", locals.get_item("text").unwrap().unwrap());
assert_eq!(text.strict_string().unwrap(), " False ");
assert_eq!(text.str_bool().unwrap().0, Some(false));
});
}
#[test]
fn missing_snapshot_fields_and_descriptor_attribute_errors_are_distinct() {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
py.run(
c"
failure = AttributeError('descriptor failed')
class Snapshot:
@property
def flag(self): raise failure
snapshot = Snapshot()
class Dynamic:
def __getattr__(self, name): raise failure
class Intercepted:
def __getattribute__(self, name): raise failure
dynamic = Dynamic()
intercepted = Intercepted()
",
Some(&locals),
Some(&locals),
)
.unwrap();
let snapshot = locals.get_item("snapshot").unwrap().unwrap();
let descriptor = PyErr::from(Field::read(&snapshot, "test.flag").err().unwrap());
assert!(
descriptor
.value(py)
.is(locals.get_item("failure").unwrap().unwrap())
);
for name in ["dynamic", "intercepted"] {
let value = locals.get_item(name).unwrap().unwrap();
let error = PyErr::from(Field::read(&value, "test.flag").err().unwrap());
assert!(
error
.value(py)
.is(locals.get_item("failure").unwrap().unwrap())
);
}
let missing = PyErr::from(Field::read(&snapshot, "test.missing").err().unwrap());
assert!(missing.is_instance_of::<PyRuntimeError>(py));
assert!(missing.to_string().contains("test.missing"));
});
}
#[test]
fn configuration_errors_name_fields_without_exposing_values() {
Python::initialize();
Python::attach(|py| {
for source in [
"{'secret': 'do-not-print'}",
"['host.test', {'secret': 'do-not-print'}]",
] {
let field = Field::new("test.setting", evaluate(py, source));
let error = PyErr::from(field.falsy_optional_string().err().unwrap());
assert!(error.is_instance_of::<PyValueError>(py));
assert!(error.to_string().contains("test.setting"));
assert!(!error.to_string().contains("do-not-print"));
}
let hosts = Field::new(
"url_policy.user_url_allowed_hosts",
evaluate(py, "['host.test', 1]"),
);
assert!(matches!(
hosts.host_collection(),
Err(ProjectionError::InvalidConfiguration(_))
));
assert!(matches!(
Field::new("test.flag", evaluate(py, "1")).str_bool(),
Err(ProjectionError::InvalidConfiguration(_))
));
});
}
#[test]
fn projection_releases_the_source_collection() {
Python::initialize();
Python::attach(|py| {
let source = evaluate(py, "['A.test']");
let projected = Field::new("test.hosts", source.clone())
.host_collection()
.unwrap()
.0;
source.call_method1("append", ("b.test",)).unwrap();
assert_eq!(projected, ["a.test"]);
assert_eq!(
Field::new("test.hosts", source)
.host_collection()
.unwrap()
.0,
["a.test", "b.test"]
);
});
}
#[rstest]
#[case("True", Some(true))]
#[case("False", Some(false))]
#[case("1", None)]
#[case("None", None)]
#[case("[]", None)]
fn accessor_booleans_are_strict_schema_values(
#[case] source: &str,
#[case] expected: Option<bool>,
) {
Python::initialize();
Python::attach(|py| {
let result = Field::new("secret_manager.readable", evaluate(py, source)).schema_bool();
match expected {
Some(expected) => assert_eq!(result.unwrap(), expected),
None => {
let error = PyErr::from(result.unwrap_err());
assert!(error.is_instance_of::<PyRuntimeError>(py));
assert!(error.to_string().contains("secret_manager.readable"));
}
}
});
}

View file

@ -7,12 +7,12 @@ use std::{
use litellm_core_utils::settings::ProcessEnvironment;
use litellm_http::{
HttpClientConfig, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify,
Unsupported,
TlsSource, Unsupported,
media::{PublicDnsResolver, UrlPolicy},
};
use pyo3::{prelude::*, types::PyDict};
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
use crate::{errors::RustBridgeDeclined, python_settings::PythonSettings};
use crate::{coercion::Field, python_settings::PythonSettings};
static POOL: LazyLock<HttpClientPool> =
LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver)));
@ -41,6 +41,30 @@ pub(crate) fn call_config(
Ok(resolution.config)
}
pub(crate) fn client_error(error: litellm_http::Error) -> PyErr {
match error {
litellm_http::Error::Read {
tls_source: TlsSource::ClientIdentity,
..
}
| litellm_http::Error::InvalidPem {
tls_source: TlsSource::ClientIdentity,
..
} => PyValueError::new_err(
"http_settings.ssl_certificate: expected a readable PEM certificate and private key",
),
litellm_http::Error::Read {
tls_source: TlsSource::CaBundle,
..
}
| litellm_http::Error::InvalidPem {
tls_source: TlsSource::CaBundle,
..
} => PyValueError::new_err("http_settings.ssl_verify: expected a readable PEM CA bundle"),
_ => PyValueError::new_err("http_settings: native HTTP client configuration is invalid"),
}
}
fn unreported(
reported: &Mutex<HashSet<Unsupported>>,
unsupported: Vec<Unsupported>,
@ -53,25 +77,25 @@ fn unreported(
}
pub(crate) fn url_policy(py: Python<'_>) -> PyResult<UrlPolicy> {
let policy: PythonUrlPolicy =
PythonSettings::UrlPolicy
.read(py)?
.extract()
.map_err(|error: PyErr| {
RustBridgeDeclined::new_err(format!(
"litellm URL policy cannot be used by the Rust route: {error}"
))
})?;
project_url_policy(&PythonSettings::UrlPolicy.read(py)?)
}
fn project_url_policy(value: &Bound<'_, PyAny>) -> PyResult<UrlPolicy> {
Ok(UrlPolicy {
validate: policy.user_url_validation,
allowed_hosts: policy.user_url_allowed_hosts,
validate: Field::read(value, "url_policy.user_url_validation")?
.truthy()?
.0,
allowed_hosts: Field::read(value, "url_policy.user_url_allowed_hosts")?
.host_collection()?
.0,
})
}
fn call_ssl_verify(kwargs: &Bound<'_, PyDict>) -> PyResult<Option<SslVerify>> {
Ok(kwargs
.get_item("ssl_verify")?
.and_then(|value| ssl_verify(&value)))
match kwargs.get_item("ssl_verify")? {
Some(value) => Ok(Field::new("request.ssl_verify", value).ssl_verify()?.0),
None => Ok(None),
}
}
fn for_call(call_ssl_verify: Option<SslVerify>, asynchronous: bool) -> HttpSettingsLayer {
@ -82,64 +106,47 @@ fn for_call(call_ssl_verify: Option<SslVerify>, asynchronous: bool) -> HttpSetti
}
}
#[derive(FromPyObject)]
struct PythonUrlPolicy {
user_url_validation: bool,
user_url_allowed_hosts: Vec<String>,
}
#[derive(FromPyObject)]
struct PythonHttpSettings<'py> {
ssl_verify: Bound<'py, PyAny>,
ssl_certificate: Option<String>,
ssl_security_level: Option<String>,
ssl_ecdh_curve: Option<String>,
force_ipv4: bool,
http2: bool,
aiohttp_trust_env: bool,
disable_aiohttp_trust_env: bool,
disable_aiohttp_transport: bool,
user_agent: String,
}
fn configured(value: &Bound<'_, PyAny>) -> PyResult<HttpSettingsLayer> {
let python: PythonHttpSettings = value.extract().map_err(|error: PyErr| {
RustBridgeDeclined::new_err(format!(
"litellm HTTP settings cannot be used by the Rust route: {error}"
))
})?;
Ok(HttpSettingsLayer {
ssl_verify: ssl_verify(&python.ssl_verify),
ssl_certificate: python.ssl_certificate.map(PathBuf::from),
ssl_security_level: python.ssl_security_level,
ssl_ecdh_curve: python.ssl_ecdh_curve,
force_ipv4: Some(python.force_ipv4),
http2: Some(python.http2),
aiohttp_trust_env: Some(python.aiohttp_trust_env),
disable_aiohttp_trust_env: Some(python.disable_aiohttp_trust_env),
disable_aiohttp_transport: Some(python.disable_aiohttp_transport),
user_agent: Some(python.user_agent),
ssl_verify: Field::read(value, "http_settings.ssl_verify")?
.ssl_verify()?
.0,
ssl_certificate: Field::read(value, "http_settings.ssl_certificate")?
.optional_strict_string()?
.0
.map(PathBuf::from),
ssl_security_level: Field::read(value, "http_settings.ssl_security_level")?
.tuning_string()?
.0,
ssl_ecdh_curve: Field::read(value, "http_settings.ssl_ecdh_curve")?
.tuning_string()?
.0,
force_ipv4: Some(Field::read(value, "http_settings.force_ipv4")?.truthy()?.0),
http2: Some(Field::read(value, "http_settings.http2")?.exact_true().0),
aiohttp_trust_env: Some(
Field::read(value, "http_settings.aiohttp_trust_env")?
.truthy()?
.0,
),
disable_aiohttp_trust_env: Some(
Field::read(value, "http_settings.disable_aiohttp_trust_env")?
.truthy()?
.0,
),
disable_aiohttp_transport: Some(
Field::read(value, "http_settings.disable_aiohttp_transport")?
.exact_true()
.0,
),
user_agent: Some(Field::read(value, "http_settings.user_agent")?.schema_string()?),
..HttpSettingsLayer::default()
})
}
fn ssl_verify(value: &Bound<'_, PyAny>) -> Option<SslVerify> {
if let Ok(enabled) = value.extract::<bool>() {
return Some(if enabled {
SslVerify::Enabled
} else {
SslVerify::Disabled
});
}
value
.extract::<String>()
.ok()
.map(|path| SslVerify::parse(&path))
}
#[cfg(test)]
mod tests {
use litellm_http::Verify;
use pyo3::exceptions::PyRuntimeError;
use rstest::rstest;
use super::*;
@ -163,7 +170,7 @@ defaults = dict(
user_agent='litellm/test',
)
defaults.update(dict({overrides}))
settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads(contract)['http_settings']}})
settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads(contract)['http_settings']['fields']}})
"
);
let locals = PyDict::new(py);
@ -189,6 +196,33 @@ settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads
});
}
#[test]
fn client_error_uses_tls_source_when_paths_match() {
Python::initialize();
Python::attach(|py| {
let path = PathBuf::from("/shared.pem");
let ca_error = client_error(litellm_http::Error::InvalidPem {
path: path.clone(),
message: "invalid".into(),
tls_source: TlsSource::CaBundle,
});
assert_eq!(
ca_error.to_string(),
"ValueError: http_settings.ssl_verify: expected a readable PEM CA bundle"
);
let client_error = client_error(litellm_http::Error::InvalidPem {
path,
message: "invalid".into(),
tls_source: TlsSource::ClientIdentity,
});
assert!(client_error.is_instance_of::<PyValueError>(py));
assert_eq!(
client_error.to_string(),
"ValueError: http_settings.ssl_certificate: expected a readable PEM certificate and private key"
);
});
}
#[test]
fn python_settings_flow_into_the_configured_layer() {
Python::initialize();
@ -259,12 +293,16 @@ user_agent='litellm/9.9.9',
});
}
#[test]
fn ssl_context_global_is_ignored_so_environment_and_defaults_apply() {
#[rstest]
#[case("ssl_verify=object()")]
#[case("ssl_verify=__import__('ssl').SSLContext(__import__('ssl').PROTOCOL_TLS_CLIENT)")]
#[case("ssl_certificate=1")]
fn invalid_http_configuration_is_terminal(#[case] overrides: &str) {
Python::initialize();
Python::attach(|py| {
let layer = configured(&python_settings(py, "ssl_verify=object()")).unwrap();
assert_eq!(layer.ssl_verify, None);
let error = configured(&python_settings(py, overrides)).unwrap_err();
assert!(error.is_instance_of::<PyValueError>(py));
assert!(error.to_string().contains("http_settings.ssl_"));
});
}
@ -281,11 +319,21 @@ user_agent='litellm/9.9.9',
}
#[test]
fn mistyped_python_settings_decline_instead_of_raising() {
fn mutable_globals_use_their_consumer_operations() {
Python::initialize();
Python::attach(|py| {
let error = configured(&python_settings(py, "force_ipv4='yes'")).unwrap_err();
assert!(error.is_instance_of::<RustBridgeDeclined>(py));
let layer = configured(&python_settings(py,
"force_ipv4='yes', http2=1, disable_aiohttp_transport=1, aiohttp_trust_env=[1], disable_aiohttp_trust_env=[], ssl_security_level=1, ssl_ecdh_curve=[]"
)).unwrap();
assert_eq!(layer.force_ipv4, Some(true));
assert_eq!(layer.http2, Some(false));
assert_eq!(layer.disable_aiohttp_transport, Some(false));
assert_eq!(layer.aiohttp_trust_env, Some(true));
assert_eq!(layer.disable_aiohttp_trust_env, Some(false));
assert_eq!(layer.ssl_security_level, None);
assert_eq!(layer.ssl_ecdh_curve, None);
let error = configured(&python_settings(py, "user_agent=1")).unwrap_err();
assert!(error.is_instance_of::<PyRuntimeError>(py));
});
}
@ -323,17 +371,36 @@ user_agent='litellm/9.9.9',
}
#[test]
fn live_ssl_context_argument_is_ignored_so_the_configured_value_applies() {
fn live_ssl_context_argument_raises_instead_of_using_another_layer() {
Python::initialize();
Python::attach(|py| {
let kwargs = PyDict::new(py);
kwargs
.set_item("ssl_verify", py.eval(c"object()", None, None).unwrap())
let ssl = py.import("ssl").unwrap();
let context = ssl
.getattr("SSLContext")
.unwrap()
.call1((ssl.getattr("PROTOCOL_TLS_CLIENT").unwrap(),))
.unwrap();
let call = for_call(call_ssl_verify(&kwargs).unwrap(), true);
let settings =
HttpSettings::from_layers([call, configured_ssl_verify(SslVerify::Disabled)]);
assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled));
kwargs.set_item("ssl_verify", context).unwrap();
let error = call_ssl_verify(&kwargs).unwrap_err();
assert!(error.is_instance_of::<PyValueError>(py));
assert!(error.to_string().contains("request.ssl_verify"));
assert!(error.to_string().contains("SSLContext"));
});
}
#[test]
fn url_policy_uses_truthiness_and_normalized_owned_hosts() {
Python::initialize();
Python::attach(|py| {
let value = py.eval(c"__import__('types').SimpleNamespace(user_url_validation=[], user_url_allowed_hosts=['B.test', 'a.test.', 'b.test'])", None, None).unwrap();
assert_eq!(
project_url_policy(&value).unwrap(),
UrlPolicy {
validate: false,
allowed_hosts: vec!["a.test".into(), "b.test".into()],
}
);
});
}

View file

@ -1,4 +1,5 @@
mod cache;
mod coercion;
mod credentials;
mod diagnostics;
mod errors;

View file

@ -172,6 +172,60 @@ mod tests {
request_input_sources(&kwargs, names.iter().copied())
}
#[serde_with::serde_as]
#[derive(Debug, serde::Deserialize, serde::Serialize, PartialEq)]
struct Numbers {
#[serde_as(deserialize_as = "Option<Vec<litellm_core_utils::serde_compat::LaxI64>>")]
integers: Option<Vec<i64>>,
#[serde_as(deserialize_as = "Option<litellm_core_utils::serde_compat::FiniteF64>")]
float: Option<f64>,
}
#[test]
fn numeric_adapters_agree_across_json_and_python_boundaries() {
Python::initialize();
Python::attach(|py| {
for input in [
json!({}),
json!({"integers": null, "float": null}),
json!({"integers": [i64::MIN, i64::MAX, "9007199254740993.0", " +1_000.00 ", true, 3.0], "float": " 1.25 "}),
json!({"integers": [u64::MAX]}),
json!({"integers": ["1.0000000000000001"]}),
json!({"integers": [2.5]}),
json!({"float": "NaN"}),
json!({"float": "inf"}),
json!({"float": "1e999"}),
json!({"float": true}),
json!({"float": u64::MAX}),
] {
let expected = serde_json::from_value::<Numbers>(input.clone());
let python = litellm_host_python::to_py(py, &input).unwrap();
let actual = from_py::<Numbers>(python.bind(py));
match (expected, actual) {
(Ok(expected), Ok(actual)) => {
assert_eq!(actual, expected);
let serialized = litellm_host_python::to_py(py, &actual).unwrap();
assert_eq!(
from_py::<Value>(serialized.bind(py)).unwrap(),
serde_json::to_value(expected).unwrap()
);
}
(Err(_), Err(_)) => {}
mismatch => panic!("boundary mismatch for {input}: {mismatch:?}"),
}
}
for source in [
c"{'float': float('nan')}",
c"{'float': float('inf')}",
c"{'integers': [float('inf')]}",
c"{'integers': [2 ** 100]}",
] {
let value = py.eval(source, None, None).unwrap();
assert!(from_py::<Numbers>(&value).is_err());
}
});
}
#[test]
fn argument_converters_keep_nested_values_and_accept_explicit_none() {
Python::initialize();

View file

@ -43,32 +43,204 @@ pub(crate) const CONTRACT: &str = include_str!("../python_settings.json");
#[cfg(test)]
mod tests {
use std::{collections::BTreeSet, ffi::CString};
use pyo3::{prelude::*, types::PyDict};
use super::{CONTRACT, PythonSettings};
use pyo3::prelude::*;
use serde_json::{Value, json};
struct SettingSpec {
group: &'static str,
name: &'static str,
adapter: &'static str,
precedence: &'static str,
sensitive: bool,
shapes: &'static [&'static str],
unsupported_live: Option<&'static str>,
}
const SETTINGS: &[SettingSpec] = &[
SettingSpec {
group: "http_settings",
name: "ssl_verify",
adapter: "SslVerifyInput",
precedence: "module_global",
sensitive: false,
shapes: &["none", "bool", "str"],
unsupported_live: Some("configuration_error"),
},
SettingSpec {
group: "http_settings",
name: "ssl_certificate",
adapter: "OptionalStrictString",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "ssl_security_level",
adapter: "TuningString",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "ssl_ecdh_curve",
adapter: "TuningString",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "force_ipv4",
adapter: "Truthy",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "http2",
adapter: "ExactTrue",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "aiohttp_trust_env",
adapter: "Truthy",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "disable_aiohttp_trust_env",
adapter: "Truthy",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "disable_aiohttp_transport",
adapter: "ExactTrue",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "user_agent",
adapter: "StrictString",
precedence: "accessor",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "url_policy",
name: "user_url_validation",
adapter: "Truthy",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "url_policy",
name: "user_url_allowed_hosts",
adapter: "HostCollection",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "provider_defaults",
name: "vertex_project",
adapter: "FalsyOptionalString",
precedence: "module_global",
sensitive: true,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "provider_defaults",
name: "vertex_location",
adapter: "FalsyOptionalString",
precedence: "module_global",
sensitive: true,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "provider_defaults",
name: "enable_azure_ad_token_refresh",
adapter: "ExactTrue",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "secret_manager",
name: "readable",
adapter: "StrictBool",
precedence: "accessor",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
];
#[test]
fn every_settings_group_is_in_the_python_contract() {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
locals.set_item("contract", CONTRACT).unwrap();
let source = CString::new("import json\nkeys = list(json.loads(contract))").unwrap();
py.run(&source, Some(&locals), Some(&locals)).unwrap();
let declared: BTreeSet<String> = locals
.get_item("keys")
fn settings_manifest_matches_the_semantic_contract() {
pyo3::Python::initialize();
let manifest: Value = pyo3::Python::attach(|py| {
let value = py
.import("json")
.unwrap()
.unwrap()
.extract::<Vec<String>>()
.unwrap()
.into_iter()
.collect();
let read: BTreeSet<String> = PythonSettings::ALL
.map(|group| group.name().to_owned())
.into();
assert_eq!(read, declared);
.call_method1("loads", (CONTRACT,))
.unwrap();
litellm_host_python::from_py(&value).unwrap()
});
let expected: serde_json::Map<String, Value> = PythonSettings::ALL
.into_iter()
.map(|group| {
let fields: serde_json::Map<String, Value> = SETTINGS
.iter()
.filter(|spec| spec.group == group.name())
.map(|spec| {
(
spec.name.to_owned(),
json!({
"adapter": spec.adapter,
"required": true,
"precedence": spec.precedence,
"sensitive": spec.sensitive,
"shapes": spec.shapes,
"unsupported_live": spec.unsupported_live,
}),
)
})
.collect();
(
group.name().to_owned(),
json!({"version": 1, "fields": fields}),
)
})
.collect();
assert_eq!(manifest, Value::Object(expected));
}
}

View file

@ -19,7 +19,7 @@ use pyo3::{
types::{PyDict, PyTuple},
};
use crate::{errors::RustBridgeDeclined, http, python_settings::PythonSettings};
use crate::{coercion::Field, errors::RustBridgeDeclined, http, python_settings::PythonSettings};
const SURFACE: LegacySurface = LegacySurface {
call_type: "ocr",
@ -51,7 +51,7 @@ fn run_ocr(
ocr_settings(py)?,
secrets,
)
.map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?;
.map_err(http::client_error)?;
run_legacy_call(
py,
if asynchronous { ASYNC_SURFACE } else { SURFACE },
@ -62,14 +62,8 @@ fn run_ocr(
)
}
#[derive(FromPyObject)]
struct PythonSecretManager {
readable: bool,
}
fn process_environment_secrets(secret_manager: &Bound<'_, PyAny>) -> PyResult<Secrets> {
let manager: PythonSecretManager = secret_manager.extract()?;
if manager.readable {
if Field::read(secret_manager, "secret_manager.readable")?.schema_bool()? {
return Err(RustBridgeDeclined::new_err(
"a readable secret manager is configured and the Rust route only reads the process environment",
));
@ -77,26 +71,24 @@ fn process_environment_secrets(secret_manager: &Bound<'_, PyAny>) -> PyResult<Se
Ok(Arc::new(ProcessEnvironment))
}
#[derive(FromPyObject)]
struct PythonProviderDefaults {
vertex_project: Option<String>,
vertex_location: Option<String>,
enable_azure_ad_token_refresh: Option<bool>,
fn ocr_settings(py: Python<'_>) -> PyResult<OcrSettings> {
project_provider_defaults(&PythonSettings::ProviderDefaults.read(py)?)
}
fn ocr_settings(py: Python<'_>) -> PyResult<OcrSettings> {
let defaults: PythonProviderDefaults = PythonSettings::ProviderDefaults
.read(py)?
.extract()
.map_err(|error: PyErr| {
RustBridgeDeclined::new_err(format!(
"litellm provider defaults cannot be used by the Rust route: {error}"
))
})?;
fn project_provider_defaults(value: &Bound<'_, PyAny>) -> PyResult<OcrSettings> {
Ok(OcrSettings {
vertex_project: defaults.vertex_project,
vertex_location: defaults.vertex_location,
enable_azure_ad_token_refresh: defaults.enable_azure_ad_token_refresh == Some(true),
vertex_project: Field::read(value, "provider_defaults.vertex_project")?
.falsy_optional_string()?
.0,
vertex_location: Field::read(value, "provider_defaults.vertex_location")?
.falsy_optional_string()?
.0,
enable_azure_ad_token_refresh: Field::read(
value,
"provider_defaults.enable_azure_ad_token_refresh",
)?
.exact_true()
.0,
..OcrSettings::from_environment(&ProcessEnvironment)
})
}
@ -140,6 +132,35 @@ mod tests {
locals.get_item("manager").unwrap().unwrap()
}
#[test]
fn provider_defaults_distinguish_falsey_values_and_exact_true() {
Python::initialize();
Python::attach(|py| {
let value = py.eval(c"__import__('types').SimpleNamespace(vertex_project=[], vertex_location=0, enable_azure_ad_token_refresh=1)", None, None).unwrap();
let projected = super::project_provider_defaults(&value).unwrap();
assert_eq!(projected.vertex_project, None);
assert_eq!(projected.vertex_location, None);
assert!(!projected.enable_azure_ad_token_refresh);
value.setattr("vertex_project", "project").unwrap();
value.setattr("vertex_location", "region").unwrap();
value
.setattr("enable_azure_ad_token_refresh", true)
.unwrap();
let next = super::project_provider_defaults(&value).unwrap();
assert_eq!(next.vertex_project.as_deref(), Some("project"));
assert_eq!(next.vertex_location.as_deref(), Some("region"));
assert!(next.enable_azure_ad_token_refresh);
value.setattr("vertex_project", 1).unwrap();
let error = super::project_provider_defaults(&value).err().unwrap();
assert!(error.is_instance_of::<pyo3::exceptions::PyValueError>(py));
assert!(
error
.to_string()
.contains("provider_defaults.vertex_project")
);
});
}
#[test]
fn a_readable_secret_manager_sends_the_call_back_to_python() {
Python::initialize();

View file

@ -0,0 +1,26 @@
[package]
name = "litellm-secrets-cyberark"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-secrets-types.workspace = true
litellm-core-utils.workspace = true
base64.workspace = true
moka.workspace = true
reqwest.workspace = true
serde_json.workspace = true
thiserror.workspace = true
veil.workspace = true
tracing = "0.1"
percent-encoding = "2.3"
tokio = { workspace = true, features = ["sync"] }
[dev-dependencies]
rstest.workspace = true
tokio.workspace = true
wiremock = "0.6.5"
serde.workspace = true
serde_json.workspace = true

View file

@ -0,0 +1,27 @@
#[derive(thiserror::Error, veil::Redact)]
pub enum Error {
#[error("CyberArk Conjur HTTP request failed")]
Http(
#[from]
#[redact]
reqwest::Error,
),
#[error("CyberArk Conjur authentication returned HTTP {0}")]
AuthStatus(u16),
#[error("CyberArk Conjur returned HTTP {0}")]
Status(u16),
#[error(
"CyberArk credentials are missing: set CYBERARK_API_KEY or both CYBERARK_CLIENT_CERT and CYBERARK_CLIENT_KEY"
)]
MissingCredentials,
#[error("CyberArk client certificate could not be loaded")]
ClientCertificate,
#[error("invalid refresh interval")]
RefreshInterval,
#[error("invalid CyberArk Conjur endpoint")]
Endpoint,
#[error("CyberArk secret manager requires an enterprise license")]
EnterpriseRequired,
#[error(transparent)]
Operation(#[from] litellm_secrets_types::Error),
}

View file

@ -0,0 +1,7 @@
#![forbid(unsafe_code)]
mod error;
mod secret_manager;
pub use error::Error;
pub use secret_manager::{CyberArkSecretManager, DeleteOutcome};

View file

@ -0,0 +1,317 @@
use std::{fs, sync::Arc, time::Duration};
use base64::{Engine, engine::general_purpose::STANDARD};
use litellm_core_utils::settings::Lookup;
use litellm_secrets_types::{BaseSecretManager, SecretValue, validate_secret_name};
use moka::future::Cache;
use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, utf8_percent_encode};
use crate::Error;
const CYBERARK_API_BASE: &str = "CYBERARK_API_BASE";
const CYBERARK_ACCOUNT: &str = "CYBERARK_ACCOUNT";
const CYBERARK_USERNAME: &str = "CYBERARK_USERNAME";
const CYBERARK_API_KEY: &str = "CYBERARK_API_KEY";
const CYBERARK_CLIENT_CERT: &str = "CYBERARK_CLIENT_CERT";
const CYBERARK_CLIENT_KEY: &str = "CYBERARK_CLIENT_KEY";
const CYBERARK_SSL_VERIFY: &str = "CYBERARK_SSL_VERIFY";
const CYBERARK_REFRESH_INTERVAL: &str = "CYBERARK_REFRESH_INTERVAL";
const DEFAULT_API_BASE: &str = "http://127.0.0.1:8080";
const DEFAULT_ACCOUNT: &str = "default";
const DEFAULT_USERNAME: &str = "admin";
const DEFAULT_REFRESH_INTERVAL: Duration = Duration::from_secs(300);
const SECRET_NAME_SAFE: &AsciiSet = &NON_ALPHANUMERIC
.remove(b'-')
.remove(b'_')
.remove(b'.')
.remove(b'~');
#[derive(Clone)]
pub struct CyberArkSecretManager {
client: reqwest::Client,
endpoint: reqwest::Url,
account: String,
username: String,
api_key: SecretValue,
token: Cache<(), SecretValue>,
secrets: Cache<String, SecretValue>,
authentication_lock: Arc<tokio::sync::Mutex<()>>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DeleteOutcome {
NotSupported,
}
impl CyberArkSecretManager {
pub fn with_client(
client: reqwest::Client,
endpoint: reqwest::Url,
account: String,
username: String,
api_key: SecretValue,
refresh_interval: Option<Duration>,
) -> Self {
let endpoint = normalize_endpoint(endpoint);
let ttl = refresh_interval
.filter(|interval| !interval.is_zero())
.unwrap_or(DEFAULT_REFRESH_INTERVAL);
let token = Cache::builder().time_to_live(ttl).build();
let secrets = Cache::builder().time_to_live(ttl).build();
Self {
client,
endpoint,
account,
username,
api_key,
token,
secrets,
authentication_lock: Arc::new(tokio::sync::Mutex::new(())),
}
}
pub fn new(
environment: Arc<dyn Lookup + Send + Sync>,
enterprise_enabled: bool,
) -> Result<Self, Error> {
let api_key = environment.get(CYBERARK_API_KEY).unwrap_or_default();
let cert = environment.get(CYBERARK_CLIENT_CERT).unwrap_or_default();
let key = environment.get(CYBERARK_CLIENT_KEY).unwrap_or_default();
if api_key.is_empty() && (cert.is_empty() || key.is_empty()) {
return Err(Error::MissingCredentials);
}
if !enterprise_enabled {
return Err(Error::EnterpriseRequired);
}
let verify = environment
.get(CYBERARK_SSL_VERIFY)
.map(|value| !value.trim().eq_ignore_ascii_case("false"))
.unwrap_or(true);
let mut builder = reqwest::Client::builder();
if !verify {
tracing::warn!(
"CyberArk SSL verification is disabled. This is insecure and should only be used for testing with self-signed certificates."
);
builder = builder.danger_accept_invalid_certs(true);
}
if !cert.is_empty() && !key.is_empty() {
let certificate = fs::read(cert).map_err(|_| Error::ClientCertificate)?;
let private_key = fs::read(key).map_err(|_| Error::ClientCertificate)?;
let identity = reqwest::Identity::from_pem(&[certificate, private_key].concat())
.map_err(|_| Error::ClientCertificate)?;
builder = builder.identity(identity);
}
let client = builder.build()?;
let endpoint = reqwest::Url::parse(
&environment
.get(CYBERARK_API_BASE)
.unwrap_or_else(|| DEFAULT_API_BASE.to_owned()),
)
.map_err(|_| Error::Endpoint)?;
let account = environment
.get(CYBERARK_ACCOUNT)
.unwrap_or_else(|| DEFAULT_ACCOUNT.to_owned());
let username = environment
.get(CYBERARK_USERNAME)
.unwrap_or_else(|| DEFAULT_USERNAME.to_owned());
let refresh_interval = environment
.get(CYBERARK_REFRESH_INTERVAL)
.map(|value| {
value
.parse::<u64>()
.map(Duration::from_secs)
.map_err(|_| Error::RefreshInterval)
})
.transpose()?;
Ok(Self::with_client(
client,
endpoint,
account,
username,
SecretValue::new(api_key),
refresh_interval,
))
}
fn secret_url(&self, name: &str) -> Result<reqwest::Url, Error> {
let encoded = utf8_percent_encode(name, SECRET_NAME_SAFE);
self.endpoint
.join(&format!("secrets/{}/variable/{}", self.account, encoded))
.map_err(|_| Error::Endpoint)
}
async fn authenticate(&self) -> Result<SecretValue, Error> {
if let Some(token) = self.token.get(&()).await {
return Ok(token);
}
let _guard = self.authentication_lock.lock().await;
if let Some(token) = self.token.get(&()).await {
return Ok(token);
}
let url = self
.endpoint
.join(&format!(
"authn/{}/{}/authenticate",
self.account, self.username
))
.map_err(|_| Error::Endpoint)?;
let response = self
.client
.post(url)
.body(self.api_key.expose().to_owned())
.send()
.await?;
if !response.status().is_success() {
return Err(Error::AuthStatus(response.status().as_u16()));
}
let token = SecretValue::new(STANDARD.encode(response.text().await?));
self.token.insert((), token.clone()).await;
Ok(token)
}
async fn authorization_header(&self) -> Result<String, Error> {
Ok(format!(
"Token token=\"{}\"",
self.authenticate().await?.expose()
))
}
pub async fn async_read_secret(&self, name: &str) -> Result<Option<SecretValue>, Error> {
if let Some(value) = self.secrets.get(name).await {
return Ok(Some(value));
}
let response = self
.client
.get(self.secret_url(name)?)
.header("Authorization", self.authorization_header().await?)
.send()
.await?;
if response.status() == reqwest::StatusCode::NOT_FOUND {
return Ok(None);
}
if !response.status().is_success() {
return Err(Error::Status(response.status().as_u16()));
}
let value = SecretValue::new(response.text().await?);
self.secrets.insert(name.to_owned(), value.clone()).await;
Ok(Some(value))
}
pub async fn async_write_secret(
&self,
name: &str,
value: &SecretValue,
_description: Option<&str>,
) -> Result<(), Error> {
validate_secret_name(name)?;
self.ensure_variable_exists(name).await;
let response = self
.client
.post(self.secret_url(name)?)
.header("Authorization", self.authorization_header().await?)
.body(value.expose().to_owned())
.send()
.await?;
if !response.status().is_success() {
return Err(Error::Status(response.status().as_u16()));
}
self.secrets.insert(name.to_owned(), value.clone()).await;
Ok(())
}
async fn ensure_variable_exists(&self, name: &str) {
let policy_url = self
.endpoint
.join(&format!("policies/{}/policy/root", self.account));
let Ok(policy_url) = policy_url else {
tracing::warn!("Could not build CyberArk policy endpoint");
return;
};
let Ok(authorization) = self.authorization_header().await else {
tracing::warn!("Could not authenticate while ensuring CyberArk variable exists");
return;
};
let body = format!(
"- !variable {}\n",
serde_json::to_string(name).expect("serializing a string cannot fail")
);
let response = self
.client
.post(policy_url)
.header("Authorization", authorization)
.header("Content-Type", "application/x-yaml")
.body(body)
.send()
.await;
match response {
Ok(response) if response.status().is_success() => {}
Ok(response)
if matches!(
response.status(),
reqwest::StatusCode::CONFLICT | reqwest::StatusCode::UNPROCESSABLE_ENTITY
) =>
{
tracing::debug!(
"CyberArk variable policy already exists or conflicts: {}",
response.status()
);
}
Ok(response) => {
tracing::warn!(
"Could not ensure CyberArk variable exists: {}",
response.status()
);
}
Err(error) => {
tracing::warn!("Error ensuring CyberArk variable exists: {error}");
}
}
}
pub async fn async_delete_secret(
&self,
name: &str,
_recovery_window_in_days: i64,
) -> Result<DeleteOutcome, Error> {
tracing::warn!(
"CyberArk Conjur does not support direct secret deletion. Secrets must be removed through policy updates."
);
self.secrets.invalidate(name).await;
Ok(DeleteOutcome::NotSupported)
}
}
impl BaseSecretManager for CyberArkSecretManager {
type Error = Error;
type WriteResponse = ();
type DeleteResponse = DeleteOutcome;
async fn async_read_secret(&self, name: &str) -> Result<Option<SecretValue>, Error> {
self.async_read_secret(name).await
}
async fn async_write_secret(
&self,
name: &str,
value: &SecretValue,
description: Option<&str>,
) -> Result<(), Error> {
self.async_write_secret(name, value, description).await
}
async fn async_delete_secret(
&self,
name: &str,
recovery_window_in_days: i64,
) -> Result<DeleteOutcome, Error> {
self.async_delete_secret(name, recovery_window_in_days)
.await
}
}
fn normalize_endpoint(mut endpoint: reqwest::Url) -> reqwest::Url {
if !endpoint.path().ends_with('/') {
endpoint.set_path(&format!("{}/", endpoint.path()));
}
endpoint
}

View file

@ -0,0 +1,32 @@
{
"endpoint": "http://conjur.test:8080",
"account": "acct",
"username": "admin",
"api_key": "k3y",
"authenticate_path": "/authn/acct/admin/authenticate",
"token_json": "{\"protected\":\"p\",\"payload\":\"q\",\"signature\":\"s\"}",
"authorization_header": "Token token=\"eyJwcm90ZWN0ZWQiOiJwIiwicGF5bG9hZCI6InEiLCJzaWduYXR1cmUiOiJzIn0=\"",
"policy_path": "/policies/acct/policy/root",
"secrets": [
{
"name": "OPENAI_API_KEY",
"path": "/secrets/acct/variable/OPENAI_API_KEY",
"policy_body": "- !variable \"OPENAI_API_KEY\"\n"
},
{
"name": "team/app/key",
"path": "/secrets/acct/variable/team%2Fapp%2Fkey",
"policy_body": "- !variable \"team/app/key\"\n"
},
{
"name": "a b+c.d-e_f~g",
"path": "/secrets/acct/variable/a%20b%2Bc.d-e_f~g",
"policy_body": "- !variable \"a b+c.d-e_f~g\"\n"
},
{
"name": "needs \"quote\"",
"path": "/secrets/acct/variable/needs%20%22quote%22",
"policy_body": "- !variable \"needs \\\"quote\\\"\"\n"
}
]
}

View file

@ -0,0 +1,516 @@
use std::{sync::Arc, time::Duration};
use base64::{Engine, engine::general_purpose::STANDARD};
use litellm_secrets_cyberark::{CyberArkSecretManager, DeleteOutcome, Error};
use litellm_secrets_types::SecretValue;
use serde::Deserialize;
use wiremock::{
Match, Mock, MockServer, Request, ResponseTemplate,
matchers::{body_string, header, method, path},
};
const TOKEN_JSON: &str = r#"{"protected":"p","payload":"q","signature":"s"}"#;
#[derive(Deserialize)]
struct ParityFixture {
endpoint: String,
account: String,
username: String,
api_key: String,
authenticate_path: String,
token_json: String,
authorization_header: String,
policy_path: String,
secrets: Vec<ParitySecret>,
}
#[derive(Deserialize)]
struct ParitySecret {
name: String,
path: String,
policy_body: String,
}
#[derive(Debug)]
struct RawPath(String);
impl Match for RawPath {
fn matches(&self, request: &Request) -> bool {
request.url.path() == self.0
}
}
fn fixture() -> ParityFixture {
serde_json::from_str(include_str!("fixtures/parity.json")).unwrap()
}
fn manager(server: &MockServer, ttl: Duration) -> CyberArkSecretManager {
CyberArkSecretManager::with_client(
reqwest::Client::new(),
server.uri().parse().unwrap(),
"acct".into(),
"admin".into(),
SecretValue::new("k3y"),
Some(ttl),
)
}
async fn mount_auth(server: &MockServer, expected: u64) {
Mock::given(method("POST"))
.and(path("/authn/acct/admin/authenticate"))
.and(body_string("k3y"))
.respond_with(ResponseTemplate::new(200).set_body_string(TOKEN_JSON))
.expect(expected)
.mount(server)
.await;
}
#[tokio::test]
async fn successful_reads_cache_auth_secret_and_redact_values() {
let server = MockServer::start().await;
mount_auth(&server, 1).await;
let token = STANDARD.encode(TOKEN_JSON);
Mock::given(path("/secrets/acct/variable/OPENAI_API_KEY"))
.and(header("authorization", format!("Token token=\"{token}\"")))
.respond_with(ResponseTemplate::new(200).set_body_string("sk-live"))
.expect(1)
.mount(&server)
.await;
let manager = manager(&server, Duration::from_secs(60));
for _ in 0..2 {
let value = manager
.async_read_secret("OPENAI_API_KEY")
.await
.unwrap()
.unwrap();
assert_eq!(value.expose(), "sk-live");
assert!(!format!("{value:?}").contains("sk-live"));
}
}
#[tokio::test]
async fn concurrent_reads_share_authentication_request() {
let server = MockServer::start().await;
Mock::given(path("/authn/acct/admin/authenticate"))
.and(body_string("k3y"))
.respond_with(
ResponseTemplate::new(200)
.set_body_string(TOKEN_JSON)
.set_delay(Duration::from_millis(20)),
)
.expect(1)
.mount(&server)
.await;
Mock::given(path("/secrets/acct/variable/key"))
.and(header(
"authorization",
format!("Token token=\"{}\"", STANDARD.encode(TOKEN_JSON)),
))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.expect(2)
.mount(&server)
.await;
let manager = manager(&server, Duration::from_secs(60));
let (first, second) = tokio::join!(
manager.async_read_secret("key"),
manager.async_read_secret("key")
);
assert_eq!(first.unwrap().unwrap().expose(), "value");
assert_eq!(second.unwrap().unwrap().expose(), "value");
}
#[rstest::rstest]
#[case::not_found(404)]
#[case::unauthorized(401)]
#[case::forbidden(403)]
#[case::server_error(500)]
#[tokio::test]
async fn failed_reads_are_not_cached(#[case] status: u16) {
let server = MockServer::start().await;
mount_auth(&server, 1).await;
let failing = Mock::given(path("/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(status))
.expect(1)
.mount_as_scoped(&server)
.await;
let manager = manager(&server, Duration::from_secs(60));
let result = manager.async_read_secret("key").await;
if status == 404 {
assert_eq!(result.unwrap(), None);
} else {
assert!(matches!(result, Err(Error::Status(actual)) if actual == status));
}
drop(failing);
Mock::given(path("/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("recovered"))
.expect(1)
.mount(&server)
.await;
for _ in 0..2 {
assert_eq!(
manager
.async_read_secret("key")
.await
.unwrap()
.unwrap()
.expose(),
"recovered"
);
}
}
#[tokio::test]
async fn failed_authentication_is_not_cached_and_does_not_read_secret() {
let server = MockServer::start().await;
let failing = Mock::given(path("/authn/acct/admin/authenticate"))
.respond_with(ResponseTemplate::new(401))
.expect(1)
.mount_as_scoped(&server)
.await;
let unused_secret = Mock::given(path("/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.expect(0)
.mount_as_scoped(&server)
.await;
let manager = manager(&server, Duration::from_secs(60));
assert!(matches!(
manager.async_read_secret("key").await,
Err(Error::AuthStatus(401))
));
drop(unused_secret);
drop(failing);
mount_auth(&server, 1).await;
Mock::given(path("/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.expect(1)
.mount(&server)
.await;
assert_eq!(
manager
.async_read_secret("key")
.await
.unwrap()
.unwrap()
.expose(),
"value"
);
}
#[tokio::test]
async fn expired_tokens_and_secrets_are_fetched_again() {
let server = MockServer::start().await;
mount_auth(&server, 2).await;
Mock::given(path("/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.expect(2)
.mount(&server)
.await;
let manager = manager(&server, Duration::from_millis(1));
for _ in 0..2 {
assert!(manager.async_read_secret("key").await.unwrap().is_some());
tokio::time::sleep(Duration::from_millis(5)).await;
}
}
#[rstest::rstest]
#[tokio::test]
async fn secret_names_use_python_quote_encoding(
#[values("OPENAI_API_KEY", "team/app/key", "a b+c.d-e_f~g", "needs \"quote\"")] name: &str,
) {
let fixture = fixture();
let secret = fixture
.secrets
.iter()
.find(|secret| secret.name == name)
.unwrap();
let server = MockServer::start().await;
mount_auth(&server, 1).await;
Mock::given(RawPath(secret.path.clone()))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.expect(1)
.mount(&server)
.await;
assert_eq!(
manager(&server, Duration::from_secs(60))
.async_read_secret(name)
.await
.unwrap()
.unwrap()
.expose(),
"value"
);
}
#[rstest::rstest]
#[case(201)]
#[case(409)]
#[case(422)]
#[case(500)]
#[tokio::test]
async fn writes_tolerate_policy_status_and_cache_value(#[case] policy_status: u16) {
let server = MockServer::start().await;
mount_auth(&server, 1).await;
Mock::given(path("/policies/acct/policy/root"))
.and(header("content-type", "application/x-yaml"))
.and(body_string("- !variable \"team/app\"\n"))
.respond_with(ResponseTemplate::new(policy_status))
.expect(1)
.mount(&server)
.await;
Mock::given(path("/secrets/acct/variable/team%2Fapp"))
.and(body_string("v"))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&server)
.await;
let manager = manager(&server, Duration::from_secs(60));
manager
.async_write_secret("team/app", &SecretValue::new("v"), None)
.await
.unwrap();
assert_eq!(
manager
.async_read_secret("team/app")
.await
.unwrap()
.unwrap()
.expose(),
"v"
);
}
#[tokio::test]
async fn failed_value_write_is_not_cached() {
let server = MockServer::start().await;
mount_auth(&server, 1).await;
Mock::given(path("/policies/acct/policy/root"))
.respond_with(ResponseTemplate::new(409))
.mount(&server)
.await;
Mock::given(path("/secrets/acct/variable/key"))
.and(body_string("v"))
.respond_with(ResponseTemplate::new(403))
.expect(1)
.mount(&server)
.await;
Mock::given(path("/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("recovered"))
.expect(1)
.mount(&server)
.await;
let manager = manager(&server, Duration::from_secs(60));
assert!(matches!(
manager
.async_write_secret("key", &SecretValue::new("v"), None)
.await,
Err(Error::Status(403))
));
assert_eq!(
manager
.async_read_secret("key")
.await
.unwrap()
.unwrap()
.expose(),
"recovered"
);
}
#[tokio::test]
async fn unsafe_names_fail_before_http_calls() {
let server = MockServer::start().await;
let manager = manager(&server, Duration::from_secs(60));
assert!(matches!(
manager
.async_write_secret("../etc", &SecretValue::new("v"), None)
.await,
Err(Error::Operation(
litellm_secrets_types::Error::UnsafeSecretName
))
));
}
#[tokio::test]
async fn delete_invalidates_cache_and_reports_not_supported() {
let server = MockServer::start().await;
mount_auth(&server, 1).await;
Mock::given(path("/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("v"))
.expect(2)
.mount(&server)
.await;
let manager = manager(&server, Duration::from_secs(60));
assert_eq!(
manager
.async_read_secret("key")
.await
.unwrap()
.unwrap()
.expose(),
"v"
);
assert_eq!(
manager.async_delete_secret("key", 7).await.unwrap(),
DeleteOutcome::NotSupported
);
assert_eq!(
manager
.async_read_secret("key")
.await
.unwrap()
.unwrap()
.expose(),
"v"
);
}
#[test]
fn new_validates_credentials_before_license_and_configuration() {
let empty: Arc<dyn litellm_core_utils::settings::Lookup + Send + Sync> =
Arc::new(|_: &str| None);
assert!(matches!(
CyberArkSecretManager::new(empty, true),
Err(Error::MissingCredentials)
));
assert!(matches!(
CyberArkSecretManager::new(
Arc::new(|name: &str| (name == "CYBERARK_API_KEY").then(|| "k3y".into())),
false
),
Err(Error::EnterpriseRequired)
));
assert!(matches!(
CyberArkSecretManager::new(
Arc::new(|name: &str| (name == "CYBERARK_CLIENT_CERT").then(|| "cert".into())),
true
),
Err(Error::MissingCredentials)
));
assert!(matches!(
CyberArkSecretManager::new(
Arc::new(|name: &str| match name {
"CYBERARK_API_KEY" => Some("k3y".into()),
"CYBERARK_REFRESH_INTERVAL" => Some("abc".into()),
_ => None,
}),
true
),
Err(Error::RefreshInterval)
));
assert!(matches!(
CyberArkSecretManager::new(
Arc::new(|name: &str| match name {
"CYBERARK_API_KEY" => Some("k3y".into()),
"CYBERARK_API_BASE" => Some("not a url".into()),
_ => None,
}),
true
),
Err(Error::Endpoint)
));
}
#[tokio::test]
async fn new_reads_environment_defaults_end_to_end() {
let server = MockServer::start().await;
Mock::given(path("/authn/default/admin/authenticate"))
.and(body_string("k3y"))
.respond_with(ResponseTemplate::new(200).set_body_string(TOKEN_JSON))
.mount(&server)
.await;
Mock::given(path("/secrets/default/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.mount(&server)
.await;
let endpoint = server.uri();
let manager = CyberArkSecretManager::new(
Arc::new(move |name: &str| match name {
"CYBERARK_API_BASE" => Some(endpoint.clone()),
"CYBERARK_API_KEY" => Some("k3y".into()),
_ => None,
}),
true,
)
.unwrap();
assert_eq!(
manager
.async_read_secret("key")
.await
.unwrap()
.unwrap()
.expose(),
"value"
);
}
#[test]
fn new_reports_missing_client_certificate_files() {
assert!(matches!(
CyberArkSecretManager::new(
Arc::new(|name: &str| match name {
"CYBERARK_CLIENT_CERT" => Some("/missing/cert".into()),
"CYBERARK_CLIENT_KEY" => Some("/missing/key".into()),
_ => None,
}),
true
),
Err(Error::ClientCertificate)
));
}
#[tokio::test]
async fn trailing_slash_endpoint_preserves_base_path() {
let server = MockServer::start().await;
Mock::given(path("/prefix/authn/acct/admin/authenticate"))
.and(body_string("k3y"))
.respond_with(ResponseTemplate::new(200).set_body_string(TOKEN_JSON))
.expect(1)
.mount(&server)
.await;
Mock::given(path("/prefix/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.mount(&server)
.await;
let endpoint = format!("{}/prefix/", server.uri()).parse().unwrap();
let manager = CyberArkSecretManager::with_client(
reqwest::Client::new(),
endpoint,
"acct".into(),
"admin".into(),
SecretValue::new("k3y"),
Some(Duration::from_secs(60)),
);
assert_eq!(
manager
.async_read_secret("key")
.await
.unwrap()
.unwrap()
.expose(),
"value"
);
}
#[test]
fn parity_fixture_matches_authentication_contract() {
let fixture = fixture();
assert_eq!(fixture.endpoint, "http://conjur.test:8080");
assert_eq!(fixture.account, "acct");
assert_eq!(fixture.username, "admin");
assert_eq!(fixture.api_key, "k3y");
assert_eq!(fixture.authenticate_path, "/authn/acct/admin/authenticate");
assert_eq!(fixture.token_json, TOKEN_JSON);
assert_eq!(
fixture.authorization_header,
format!("Token token=\"{}\"", STANDARD.encode(TOKEN_JSON))
);
assert_eq!(fixture.policy_path, "/policies/acct/policy/root");
assert_eq!(fixture.secrets.len(), 4);
assert_eq!(
fixture.secrets[1].policy_body,
"- !variable \"team/app/key\"\n"
);
}

View file

@ -10,12 +10,14 @@ default = []
aws = ["dep:litellm-secrets-aws"]
google = ["dep:litellm-secrets-google"]
hashicorp = ["dep:litellm-secrets-hashicorp"]
cyberark = ["dep:litellm-secrets-cyberark"]
[dependencies]
litellm-secrets-types.workspace = true
litellm-secrets-aws = { workspace = true, optional = true }
litellm-secrets-google = { workspace = true, optional = true }
litellm-secrets-hashicorp = { workspace = true, optional = true }
litellm-secrets-cyberark = { workspace = true, optional = true }
litellm-core-utils.workspace = true
base64.workspace = true
serde.workspace = true

View file

@ -33,4 +33,7 @@ pub enum Error {
#[cfg(feature = "hashicorp")]
#[error(transparent)]
Hashicorp(#[from] litellm_secrets_hashicorp::Error),
#[cfg(feature = "cyberark")]
#[error(transparent)]
Cyberark(#[from] litellm_secrets_cyberark::Error),
}

View file

@ -15,6 +15,8 @@ pub enum SecretManager {
GoogleSecretManager(crate::google::GoogleSecretManager),
#[cfg(feature = "hashicorp")]
HashicorpVault(crate::hashicorp::HashicorpVault),
#[cfg(feature = "cyberark")]
Cyberark(crate::cyberark::CyberArkSecretManager),
}
impl SecretManager {
@ -31,6 +33,8 @@ impl SecretManager {
Self::GoogleSecretManager(_) => KeyManagementSystem::GoogleSecretManager,
#[cfg(feature = "hashicorp")]
Self::HashicorpVault(_) => KeyManagementSystem::HashicorpVault,
#[cfg(feature = "cyberark")]
Self::Cyberark(_) => KeyManagementSystem::Cyberark,
}
}
}
@ -88,6 +92,12 @@ pub async fn get_secret_from_manager(
.await
.map(|value| value.map(Secret::String))
.map_err(Error::from),
#[cfg(feature = "cyberark")]
SecretManager::Cyberark(client) => client
.async_read_secret(secret_name)
.await
.map(|value| value.map(Secret::String))
.map_err(Error::from),
}
}

View file

@ -17,6 +17,8 @@ pub use state::{SecretManagerState, secret_manager_would_be_consulted};
#[cfg(feature = "aws")]
pub use litellm_secrets_aws as aws;
#[cfg(feature = "cyberark")]
pub use litellm_secrets_cyberark as cyberark;
#[cfg(feature = "google")]
pub use litellm_secrets_google as google;
#[cfg(feature = "hashicorp")]

View file

@ -105,7 +105,6 @@ async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whites
Err(Error::MissingCiphertext)
));
}
#[cfg(feature = "hashicorp")]
#[tokio::test]
async fn hashicorp_handler_resolves_found_missing_and_failed_values() {
@ -248,3 +247,56 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() {
))
));
}
#[cfg(feature = "cyberark")]
#[tokio::test]
async fn cyberark_handler_reads_values_and_surfaces_errors() {
use std::time::Duration;
use litellm_secrets::{
Error, KeyManagementSettings, SecretManager, SecretValue, cyberark::CyberArkSecretManager,
get_secret_from_manager,
};
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_string, path},
};
let server = MockServer::start().await;
Mock::given(path("/authn/acct/admin/authenticate"))
.and(body_string("k3y"))
.respond_with(ResponseTemplate::new(200).set_body_string("token"))
.mount(&server)
.await;
Mock::given(path("/secrets/acct/variable/KEY"))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.mount(&server)
.await;
let manager = SecretManager::Cyberark(CyberArkSecretManager::with_client(
reqwest::Client::new(),
server.uri().parse().unwrap(),
"acct".into(),
"admin".into(),
SecretValue::new("k3y"),
Some(Duration::from_secs(60)),
));
assert_eq!(
manager.system(),
litellm_secrets::KeyManagementSystem::Cyberark
);
let settings = KeyManagementSettings::default();
let value = get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| None)
.await
.unwrap()
.unwrap();
assert_eq!(value.as_str(), Some("value"));
Mock::given(path("/secrets/acct/variable/ERROR"))
.respond_with(ResponseTemplate::new(500))
.mount(&server)
.await;
assert!(matches!(
get_secret_from_manager(&manager, "ERROR", &settings, &|_: &str| None).await,
Err(Error::Cyberark(_))
));
}

View file

@ -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,
)

View file

@ -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",

View file

@ -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",

View file

@ -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,

View file

@ -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",

View file

@ -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}")

View file

@ -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

View file

@ -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

View file

@ -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,

View file

@ -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
)

View file

@ -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,

View file

@ -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,

View file

@ -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,

View file

@ -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."

View file

@ -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

View file

@ -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)

View file

@ -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

View file

@ -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})

View file

@ -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)

View file

@ -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)

View file

@ -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)

View file

@ -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

View file

@ -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)

View file

@ -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)

View file

@ -1,7 +1,7 @@
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
@ -163,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",
@ -345,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]:
@ -401,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
)
@ -419,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
)

View file

@ -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)

View file

@ -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
):

View file

@ -43011,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,
@ -76889,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
}
}

View file

@ -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",

View file

@ -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=[],

View file

@ -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

View file

@ -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,

View file

@ -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",

View file

@ -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:
@ -6878,10 +6891,27 @@ class ProxyConfig:
router_model_ids: Final = llm_router.get_model_ids()
# Check for model IDs in llm_router not present in combined_id_list and delete them
kept_config_ids: Final[frozenset[str]] = (
frozenset(
model_id
for model_id in router_model_ids
if (deployment := llm_router.get_deployment(model_id=model_id)) is not None
and deployment.model_info.db_model is False
)
if model_list is None
else frozenset()
)
if kept_config_ids:
verbose_proxy_logger.warning(
"Config read in _delete_deployment returned no model_list. "
"Keeping %d config-defined deployments to avoid removing valid models.",
len(kept_config_ids),
)
for model_id in router_model_ids:
if model_id not in combined_id_list:
if model_id not in combined_id_list and model_id not in kept_config_ids:
llm_router.delete_deployment(id=model_id)
return frozenset(combined_id_list)
return frozenset(combined_id_list) | kept_config_ids
def _resolve_db_litellm_param(self, key: str, value: object) -> object:
if not isinstance(value, str):
@ -19267,6 +19297,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)

View file

@ -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",

View file

@ -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(

View file

@ -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

View file

@ -1,34 +1,33 @@
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
@dataclass(frozen=True, slots=True)
class HttpSettings:
ssl_verify: bool | str
ssl_certificate: str | None
ssl_security_level: str | None
ssl_ecdh_curve: str | None
force_ipv4: bool
http2: bool
aiohttp_trust_env: bool
disable_aiohttp_trust_env: bool
disable_aiohttp_transport: bool
ssl_verify: object
ssl_certificate: object
ssl_security_level: object
ssl_ecdh_curve: object
force_ipv4: object
http2: object
aiohttp_trust_env: object
disable_aiohttp_trust_env: object
disable_aiohttp_transport: object
user_agent: str
@dataclass(frozen=True, slots=True)
class UrlPolicy:
user_url_validation: bool
user_url_allowed_hosts: Sequence[str]
user_url_validation: object
user_url_allowed_hosts: object
@dataclass(frozen=True, slots=True)
class ProviderDefaults:
vertex_project: str | None
vertex_location: str | None
enable_azure_ad_token_refresh: bool | None
vertex_project: object
vertex_location: object
enable_azure_ad_token_refresh: object
@dataclass(frozen=True, slots=True)

View file

@ -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)

View file

@ -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`

View file

@ -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".

View file

@ -4161,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"

View file

@ -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)
@ -6574,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
@ -6824,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:
@ -8324,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),
@ -8626,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
@ -8746,6 +8764,8 @@ class ProviderConfigManager:
)
return GithubCopilotAnthropicMessagesConfig()
elif litellm.LlmProviders.EDENAI == provider:
return litellm.EdenAIAnthropicMessagesConfig()
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
@ -8854,6 +8874,8 @@ class ProviderConfigManager:
)
return GeminiAudioTranscriptionConfig()
elif litellm.LlmProviders.EDENAI == provider:
return litellm.EdenAIAudioTranscriptionConfig()
return None
@staticmethod
@ -8956,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
@ -9028,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:
@ -9065,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
@ -9413,6 +9439,8 @@ class ProviderConfigManager:
)
return get_modelscope_image_generation_config(model)
elif LlmProviders.EDENAI == provider:
return litellm.EdenAIImageGenerationConfig()
return None
@staticmethod
@ -9448,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
@ -9768,6 +9798,8 @@ class ProviderConfigManager:
)
return AWSPollyTextToSpeechConfig()
elif litellm.LlmProviders.EDENAI == provider:
return litellm.EdenAITextToSpeechConfig()
return None
@staticmethod

View file

@ -43011,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,
@ -76889,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
}
}

View file

@ -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",

View file

@ -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",
]

View file

@ -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

View file

@ -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 <your-e2e-config>.yml --port 4000
uv run pytest tests/e2e/<your_suite>/ -v
uv run --group e2e-dev pytest tests/e2e/<your_suite>/ -v
```
4. Capture screenshots of the test run and attach them to the PR as proof

View file

@ -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()

View file

@ -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)

View file

@ -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)

View file

@ -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")

View file

@ -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):

View file

@ -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 "

View file

@ -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

View file

@ -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}"

Some files were not shown because too many files have changed in this diff Show more