mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-27 01:22:18 +00:00
feat(rust): add config router and gateway crates (#43289)
Co-authored-by: Yujong Lee <yujong@berri.ai>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
7ae721bf79
commit
affb547525
34 changed files with 2133 additions and 3 deletions
161
litellm-rust/Cargo.lock
generated
161
litellm-rust/Cargo.lock
generated
|
|
@ -710,14 +710,20 @@ dependencies = [
|
|||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
"hyper 1.10.1",
|
||||
"hyper-util",
|
||||
"itoa",
|
||||
"matchit",
|
||||
"memchr",
|
||||
"mime",
|
||||
"multer",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"serde_core",
|
||||
"serde_json",
|
||||
"serde_path_to_error",
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tower",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
|
|
@ -1180,7 +1186,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "e75b2483e97a5a7da73ac68a05b629f9c53cff58d8ed1c77866079e18b00dba5"
|
||||
dependencies = [
|
||||
"digest 0.10.7",
|
||||
"spin",
|
||||
"spin 0.10.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -1581,6 +1587,15 @@ dependencies = [
|
|||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "encoding_rs"
|
||||
version = "0.8.35"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "75030f3c4f45dafd7586dd6780965a8c7e8e285a5ecb86713e63a79c5b2766f3"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "equivalent"
|
||||
version = "1.0.2"
|
||||
|
|
@ -3096,6 +3111,18 @@ dependencies = [
|
|||
"strum",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-config"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-auth-types",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_yaml_ng",
|
||||
"tempfile",
|
||||
"thiserror 2.0.19",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-core"
|
||||
version = "0.1.0"
|
||||
|
|
@ -3183,6 +3210,66 @@ dependencies = [
|
|||
"tokio-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-gateway"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"axum",
|
||||
"litellm-config",
|
||||
"litellm-core",
|
||||
"litellm-gateway-auth",
|
||||
"litellm-gateway-inference",
|
||||
"litellm-http",
|
||||
"litellm-llms",
|
||||
"litellm-secrets",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"tower-http 0.7.1",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-gateway-auth"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"axum",
|
||||
"futures-util",
|
||||
"litellm-auth-types",
|
||||
"litellm-config",
|
||||
"litellm-secrets",
|
||||
"rstest",
|
||||
"sha2 0.10.9",
|
||||
"subtle",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tower",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-gateway-inference"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"axum",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"litellm-auth",
|
||||
"litellm-core",
|
||||
"litellm-host",
|
||||
"litellm-http",
|
||||
"litellm-llms",
|
||||
"litellm-router",
|
||||
"litellm-secrets",
|
||||
"litellm-types",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tower",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-host"
|
||||
version = "0.1.0"
|
||||
|
|
@ -3348,6 +3435,15 @@ dependencies = [
|
|||
"thiserror 2.0.19",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-router"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-config",
|
||||
"litellm-core",
|
||||
"rstest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-secrets"
|
||||
version = "0.1.0"
|
||||
|
|
@ -3766,6 +3862,23 @@ dependencies = [
|
|||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "multer"
|
||||
version = "3.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "83e87776546dc87511aa5ee218730c92b666d7264ab6ed41f9d215af9cd5224b"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"encoding_rs",
|
||||
"futures-util",
|
||||
"http 1.4.2",
|
||||
"httparse",
|
||||
"memchr",
|
||||
"mime",
|
||||
"spin 0.9.9",
|
||||
"version_check",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nom"
|
||||
version = "7.1.3"
|
||||
|
|
@ -4767,7 +4880,7 @@ dependencies = [
|
|||
"tokio-rustls 0.26.4",
|
||||
"tokio-util",
|
||||
"tower",
|
||||
"tower-http",
|
||||
"tower-http 0.6.11",
|
||||
"tower-service",
|
||||
"url",
|
||||
"wasm-bindgen",
|
||||
|
|
@ -4809,7 +4922,7 @@ dependencies = [
|
|||
"tokio-rustls 0.26.4",
|
||||
"tokio-util",
|
||||
"tower",
|
||||
"tower-http",
|
||||
"tower-http 0.6.11",
|
||||
"tower-service",
|
||||
"url",
|
||||
"wasm-bindgen",
|
||||
|
|
@ -5331,6 +5444,19 @@ dependencies = [
|
|||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_yaml_ng"
|
||||
version = "0.10.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7b4db627b98b36d4203a7b458cf3573730f2bb591b28871d916dfa9efabfd41f"
|
||||
dependencies = [
|
||||
"indexmap 2.14.0",
|
||||
"itoa",
|
||||
"ryu",
|
||||
"serde",
|
||||
"unsafe-libyaml",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha1"
|
||||
version = "0.10.7"
|
||||
|
|
@ -5460,6 +5586,12 @@ dependencies = [
|
|||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "spin"
|
||||
version = "0.9.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3763264f6b73151db08c50ff20d7d8a0b8796e021cdea7ceedad07b80155fa0e"
|
||||
|
||||
[[package]]
|
||||
name = "spin"
|
||||
version = "0.10.1"
|
||||
|
|
@ -6031,6 +6163,23 @@ dependencies = [
|
|||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tower-http"
|
||||
version = "0.7.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08a05a66a4fdd61cbbe0a1d755ffe0ca6aba159dd4820936a0ff8a8278245b9c"
|
||||
dependencies = [
|
||||
"bitflags 2.13.1",
|
||||
"bytes",
|
||||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tower-layer"
|
||||
version = "0.3.3"
|
||||
|
|
@ -6252,6 +6401,12 @@ version = "0.1.1"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
|
||||
|
||||
[[package]]
|
||||
name = "unsafe-libyaml"
|
||||
version = "0.2.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861"
|
||||
|
||||
[[package]]
|
||||
name = "untrusted"
|
||||
version = "0.9.0"
|
||||
|
|
|
|||
|
|
@ -9,8 +9,13 @@ license = "MIT"
|
|||
repository = "https://github.com/BerriAI/litellm"
|
||||
|
||||
[workspace.dependencies]
|
||||
litellm-config = { path = "crates/config" }
|
||||
litellm-router = { path = "crates/router" }
|
||||
litellm-tracing = { path = "crates/tracing" }
|
||||
litellm-core = { path = "crates/core" }
|
||||
litellm-gateway = { path = "crates/gateway" }
|
||||
litellm-gateway-inference = { path = "crates/gateway-inference" }
|
||||
litellm-gateway-auth = { path = "crates/gateway-auth" }
|
||||
litellm-coroutine = { path = "crates/coroutine" }
|
||||
litellm-host = { path = "crates/host" }
|
||||
litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" }
|
||||
|
|
@ -50,6 +55,7 @@ litellm-host-python = { path = "crates/host-python" }
|
|||
litellm-python-compat = { path = "crates/python-compat" }
|
||||
|
||||
tracing = "0.1"
|
||||
axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] }
|
||||
bytes = "1"
|
||||
http = "1"
|
||||
google-cloud-auth = { version = "1.16.0", default-features = false }
|
||||
|
|
|
|||
16
litellm-rust/crates/config/Cargo.toml
Normal file
16
litellm-rust/crates/config/Cargo.toml
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
[package]
|
||||
name = "litellm-config"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-auth-types.workspace = true
|
||||
serde.workspace = true
|
||||
serde_yaml_ng = "0.10.0"
|
||||
thiserror.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
tempfile.workspace = true
|
||||
7
litellm-rust/crates/config/src/error.rs
Normal file
7
litellm-rust/crates/config/src/error.rs
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("could not read config")]
|
||||
Read(#[from] std::io::Error),
|
||||
#[error("invalid YAML config")]
|
||||
Parse(#[from] serde_yaml_ng::Error),
|
||||
}
|
||||
48
litellm-rust/crates/config/src/lib.rs
Normal file
48
litellm-rust/crates/config/src/lib.rs
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
mod error;
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use litellm_auth_types::SecretValue;
|
||||
use serde::Deserialize;
|
||||
|
||||
pub use error::Error;
|
||||
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct Config {
|
||||
pub model_list: Box<[Model]>,
|
||||
#[serde(default)]
|
||||
pub general_settings: GeneralSettings,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct GeneralSettings {
|
||||
pub master_key: Option<SecretValue>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct Model {
|
||||
pub model_name: String,
|
||||
pub litellm_params: LiteLlmParams,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct LiteLlmParams {
|
||||
pub model: String,
|
||||
pub api_key: Option<SecretValue>,
|
||||
pub api_base: Option<String>,
|
||||
pub custom_llm_provider: Option<String>,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
pub fn from_yaml(yaml: &str) -> Result<Self, Error> {
|
||||
Ok(serde_yaml_ng::from_str(yaml)?)
|
||||
}
|
||||
|
||||
pub fn load(path: impl AsRef<Path>) -> Result<Self, Error> {
|
||||
Self::from_yaml(&std::fs::read_to_string(path)?)
|
||||
}
|
||||
}
|
||||
119
litellm-rust/crates/config/tests/config.rs
Normal file
119
litellm-rust/crates/config/tests/config.rs
Normal file
|
|
@ -0,0 +1,119 @@
|
|||
use litellm_config::{Config, Error};
|
||||
use rstest::{fixture, rstest};
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[fixture]
|
||||
fn directory() -> TempDir {
|
||||
tempfile::tempdir().unwrap()
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn model_list_yaml() -> &'static str {
|
||||
r#"
|
||||
model_list:
|
||||
- model_name: assistant
|
||||
litellm_params:
|
||||
model: anthropic/test-model
|
||||
api_key: os.environ/ANTHROPIC_API_KEY
|
||||
- model_name: local
|
||||
litellm_params:
|
||||
model: test-model
|
||||
api_base: http://localhost:8000/v1
|
||||
custom_llm_provider: openai
|
||||
"#
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn loads_model_list_from_file(directory: TempDir, model_list_yaml: &str) {
|
||||
let path = directory.path().join("config.yaml");
|
||||
std::fs::write(&path, model_list_yaml).unwrap();
|
||||
|
||||
let config = Config::load(path).unwrap();
|
||||
assert_eq!(config.model_list.len(), 2);
|
||||
let anthropic = &config.model_list[0];
|
||||
assert_eq!(anthropic.model_name, "assistant");
|
||||
assert_eq!(anthropic.litellm_params.model, "anthropic/test-model");
|
||||
assert_eq!(
|
||||
anthropic.litellm_params.api_key.as_ref().unwrap().expose(),
|
||||
"os.environ/ANTHROPIC_API_KEY"
|
||||
);
|
||||
assert!(anthropic.litellm_params.api_base.is_none());
|
||||
assert!(anthropic.litellm_params.custom_llm_provider.is_none());
|
||||
let local = &config.model_list[1];
|
||||
assert_eq!(local.model_name, "local");
|
||||
assert_eq!(local.litellm_params.model, "test-model");
|
||||
assert!(local.litellm_params.api_key.is_none());
|
||||
assert_eq!(
|
||||
local.litellm_params.api_base.as_deref(),
|
||||
Some("http://localhost:8000/v1")
|
||||
);
|
||||
assert_eq!(
|
||||
local.litellm_params.custom_llm_provider.as_deref(),
|
||||
Some("openai")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn config_debug_redacts_api_keys() {
|
||||
let config = Config::from_yaml(
|
||||
"model_list: [{model_name: assistant, litellm_params: {model: anthropic/test-model, api_key: secret-value}}]",
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
config.model_list[0]
|
||||
.litellm_params
|
||||
.api_key
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.expose(),
|
||||
"secret-value"
|
||||
);
|
||||
assert!(!format!("{config:?}").contains("secret-value"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::malformed_yaml("model_list: [")]
|
||||
#[case::missing_model_list("{}")]
|
||||
#[case::missing_params("model_list: [{model_name: assistant}]")]
|
||||
#[case::missing_model("model_list: [{model_name: assistant, litellm_params: {api_key: key}}]")]
|
||||
#[case::unsupported_settings("model_list: []\ngeneral_settings: {unknown: true}")]
|
||||
#[case::misspelled_param(
|
||||
"model_list: [{model_name: assistant, litellm_params: {model: test, api_bsae: url}}]"
|
||||
)]
|
||||
fn rejects_malformed_incomplete_and_unsupported_config(#[case] yaml: &str) {
|
||||
assert!(matches!(Config::from_yaml(yaml), Err(Error::Parse(_))));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn distinguishes_read_errors_from_parse_errors(directory: TempDir) {
|
||||
assert!(matches!(
|
||||
Config::load(directory.path().join("missing.yaml")),
|
||||
Err(Error::Read(error)) if error.kind() == std::io::ErrorKind::NotFound
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::literal("secret-master-key")]
|
||||
#[case::reference("os.environ/LITELLM_MASTER_KEY")]
|
||||
fn loads_and_redacts_the_master_key(#[case] key: &str) {
|
||||
let config = Config::from_yaml(&format!(
|
||||
"model_list: []\ngeneral_settings:\n master_key: {key}\n"
|
||||
))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
config
|
||||
.general_settings
|
||||
.master_key
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.expose(),
|
||||
key
|
||||
);
|
||||
assert!(!format!("{config:?}").contains(key));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn missing_general_settings_has_no_master_key() {
|
||||
let config = Config::from_yaml("model_list: []").unwrap();
|
||||
assert!(config.general_settings.master_key.is_none());
|
||||
}
|
||||
21
litellm-rust/crates/gateway-auth/Cargo.toml
Normal file
21
litellm-rust/crates/gateway-auth/Cargo.toml
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
[package]
|
||||
name = "litellm-gateway-auth"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
axum.workspace = true
|
||||
litellm-auth-types.workspace = true
|
||||
litellm-config.workspace = true
|
||||
litellm-secrets.workspace = true
|
||||
sha2.workspace = true
|
||||
subtle.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
futures-util.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
tower = { version = "0.5.3", features = ["util"] }
|
||||
24
litellm-rust/crates/gateway-auth/src/error.rs
Normal file
24
litellm-rust/crates/gateway-auth/src/error.rs
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
use axum::{
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("gateway auth not configured")]
|
||||
Unconfigured,
|
||||
#[error("missing or invalid bearer token")]
|
||||
InvalidToken,
|
||||
#[error("gateway authentication unavailable")]
|
||||
Secret(#[from] litellm_secrets::Error),
|
||||
}
|
||||
|
||||
impl IntoResponse for Error {
|
||||
fn into_response(self) -> Response {
|
||||
let status = match &self {
|
||||
Self::InvalidToken => StatusCode::UNAUTHORIZED,
|
||||
Self::Unconfigured | Self::Secret(_) => StatusCode::INTERNAL_SERVER_ERROR,
|
||||
};
|
||||
(status, self.to_string()).into_response()
|
||||
}
|
||||
}
|
||||
74
litellm-rust/crates/gateway-auth/src/lib.rs
Normal file
74
litellm-rust/crates/gateway-auth/src/lib.rs
Normal file
|
|
@ -0,0 +1,74 @@
|
|||
mod error;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
extract::FromRequestParts,
|
||||
http::{header::AUTHORIZATION, request::Parts},
|
||||
};
|
||||
use litellm_auth_types::SecretValue;
|
||||
use litellm_config::Config;
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use sha2::{Digest, Sha256};
|
||||
use subtle::ConstantTimeEq;
|
||||
|
||||
pub use error::Error;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Auth {
|
||||
master_key: Option<SecretValue>,
|
||||
secrets: Arc<dyn SecretSource>,
|
||||
}
|
||||
|
||||
impl Auth {
|
||||
pub fn from_config(config: &Config, secrets: Arc<dyn SecretSource>) -> Self {
|
||||
Self {
|
||||
master_key: config.general_settings.master_key.clone(),
|
||||
secrets,
|
||||
}
|
||||
}
|
||||
|
||||
async fn master_key(&self) -> Result<SecretValue, Error> {
|
||||
let configured = self.master_key.as_ref().ok_or(Error::Unconfigured)?;
|
||||
let resolved = match configured.expose().strip_prefix("os.environ/") {
|
||||
Some(name) if !name.is_empty() => self
|
||||
.secrets
|
||||
.get_secret_str(name)
|
||||
.await?
|
||||
.ok_or(Error::Unconfigured)?,
|
||||
Some(_) => return Err(Error::Unconfigured),
|
||||
None => configured.clone(),
|
||||
};
|
||||
if resolved.expose().trim().is_empty() {
|
||||
return Err(Error::Unconfigured);
|
||||
}
|
||||
Ok(resolved)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn hash_token(token: &str) -> String {
|
||||
format!("{:x}", Sha256::digest(token.as_bytes()))
|
||||
}
|
||||
|
||||
pub struct RequireMasterKey;
|
||||
|
||||
impl FromRequestParts<Auth> for RequireMasterKey {
|
||||
type Rejection = Error;
|
||||
|
||||
async fn from_request_parts(parts: &mut Parts, state: &Auth) -> Result<Self, Error> {
|
||||
let expected = state.master_key().await?;
|
||||
let provided = parts
|
||||
.headers
|
||||
.get(AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.strip_prefix("Bearer "))
|
||||
.map(str::trim)
|
||||
.ok_or(Error::InvalidToken)?;
|
||||
let actual_hash = Sha256::digest(provided.as_bytes());
|
||||
let expected_hash = Sha256::digest(expected.expose().as_bytes());
|
||||
match bool::from(actual_hash.ct_eq(&expected_hash)) {
|
||||
true => Ok(Self),
|
||||
false => Err(Error::InvalidToken),
|
||||
}
|
||||
}
|
||||
}
|
||||
100
litellm-rust/crates/gateway-auth/tests/auth.rs
Normal file
100
litellm-rust/crates/gateway-auth/tests/auth.rs
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
Router,
|
||||
body::{Body, to_bytes},
|
||||
http::{Request, StatusCode},
|
||||
middleware::from_extractor_with_state,
|
||||
routing::get,
|
||||
};
|
||||
use futures_util::future::BoxFuture;
|
||||
use litellm_auth_types::SecretValue;
|
||||
use litellm_config::Config;
|
||||
use litellm_gateway_auth::{Auth, RequireMasterKey, hash_token};
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use rstest::{fixture, rstest};
|
||||
use tower::ServiceExt;
|
||||
|
||||
struct Secrets;
|
||||
|
||||
impl SecretSource for Secrets {
|
||||
fn get_secret_str<'a>(
|
||||
&'a self,
|
||||
name: &'a str,
|
||||
) -> BoxFuture<'a, Result<Option<SecretValue>, litellm_secrets::Error>> {
|
||||
Box::pin(async move {
|
||||
match name {
|
||||
"MASTER_KEY" => Ok(Some(SecretValue::new("resolved-key"))),
|
||||
"EMPTY" => Ok(Some(SecretValue::new(""))),
|
||||
"ERROR" => Err(litellm_secrets::Error::ExternalRead(Box::new(
|
||||
std::io::Error::other("private-backend-detail"),
|
||||
))),
|
||||
_ => Ok(None),
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn secrets() -> Arc<dyn SecretSource> {
|
||||
Arc::new(Secrets)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::literal("literal-key", Some("Bearer literal-key"), 204)]
|
||||
#[case::reference("os.environ/MASTER_KEY", Some("Bearer resolved-key"), 204)]
|
||||
#[case::reference_is_not_a_token(
|
||||
"os.environ/MASTER_KEY",
|
||||
Some("Bearer os.environ/MASTER_KEY"),
|
||||
401
|
||||
)]
|
||||
#[case::wrong("literal-key", Some("Bearer other-key"), 401)]
|
||||
#[case::missing("literal-key", None, 401)]
|
||||
#[case::wrong_scheme("literal-key", Some("Basic literal-key"), 401)]
|
||||
#[case::empty_token("literal-key", Some("Bearer "), 401)]
|
||||
#[case::missing_reference("os.environ/MISSING", Some("Bearer os.environ/MISSING"), 500)]
|
||||
#[case::empty_reference("os.environ/EMPTY", Some("Bearer "), 500)]
|
||||
#[case::empty_key("", Some("Bearer "), 500)]
|
||||
#[case::whitespace_key(" ", Some("Bearer "), 500)]
|
||||
#[case::empty_reference_name("os.environ/", Some("Bearer os.environ/"), 500)]
|
||||
#[case::secret_failure("os.environ/ERROR", Some("Bearer private-backend-detail"), 500)]
|
||||
#[tokio::test]
|
||||
async fn enforces_configured_keys_without_exposing_secrets(
|
||||
secrets: Arc<dyn SecretSource>,
|
||||
#[case] key: &str,
|
||||
#[case] authorization: Option<&str>,
|
||||
#[case] status: u16,
|
||||
) {
|
||||
let config = Config::from_yaml(&format!(
|
||||
"model_list: []\ngeneral_settings:\n master_key: '{key}'\n"
|
||||
))
|
||||
.unwrap();
|
||||
let app = Router::new()
|
||||
.route("/protected", get(|| async { StatusCode::NO_CONTENT }))
|
||||
.layer(from_extractor_with_state::<RequireMasterKey, _>(
|
||||
Auth::from_config(&config, secrets),
|
||||
));
|
||||
let request = Request::get("/protected");
|
||||
let request = match authorization {
|
||||
Some(value) => request.header("authorization", value),
|
||||
None => request,
|
||||
};
|
||||
let response = app
|
||||
.oneshot(request.body(Body::empty()).unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status().as_u16(), status);
|
||||
let body = to_bytes(response.into_body(), 4096).await.unwrap();
|
||||
let text = std::str::from_utf8(&body).unwrap();
|
||||
assert!(!text.contains("private-backend-detail"));
|
||||
assert!(!text.contains("literal-key"));
|
||||
assert!(!text.contains("resolved-key"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn hash_token_matches_python_sha256_hexdigest() {
|
||||
assert_eq!(
|
||||
hash_token("sk-1234"),
|
||||
"88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"
|
||||
);
|
||||
}
|
||||
5
litellm-rust/crates/gateway-inference/AGENTS.md
Normal file
5
litellm-rust/crates/gateway-inference/AGENTS.md
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
- Expose a mountable Axum router; listener binding, server lifecycle, and shared inbound middleware belong to `gateway`
|
||||
- Own the public inference HTTP boundary: endpoint paths, request parsing, model alias resolution, response envelopes, and SSE delivery
|
||||
- Delegate inference execution to `core` and provider transformations and authentication to `llms` and the auth crates; do not duplicate them in handlers
|
||||
- Use injected deployments, HTTP pools, settings, and secret sources; do not load process configuration or construct independent clients in handlers
|
||||
- Test HTTP contracts here, including status codes, forwarded headers, error envelopes, and streaming behavior; keep core and provider tests in their owning crates
|
||||
29
litellm-rust/crates/gateway-inference/Cargo.toml
Normal file
29
litellm-rust/crates/gateway-inference/Cargo.toml
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
[package]
|
||||
name = "litellm-gateway-inference"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
axum = { workspace = true, features = ["json", "multipart"] }
|
||||
base64.workspace = true
|
||||
bytes.workspace = true
|
||||
futures-util.workspace = true
|
||||
litellm-auth.workspace = true
|
||||
litellm-core.workspace = true
|
||||
litellm-host.workspace = true
|
||||
litellm-http.workspace = true
|
||||
litellm-llms.workspace = true
|
||||
litellm-router.workspace = true
|
||||
litellm-secrets.workspace = true
|
||||
litellm-types.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio = { workspace = true, features = ["sync"] }
|
||||
|
||||
[dev-dependencies]
|
||||
futures-util.workspace = true
|
||||
rstest.workspace = true
|
||||
tower = { version = "0.5.3", features = ["util"] }
|
||||
wiremock = "0.6.5"
|
||||
|
|
@ -0,0 +1,59 @@
|
|||
use std::{path::Path, sync::Arc};
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
extract::{Request, State},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use base64::{Engine, engine::general_purpose::STANDARD};
|
||||
use litellm_core::audio_transcription::{audio_transcription, types::AudioTranscriptionRequest};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::{Error, Gateway, request};
|
||||
|
||||
pub(crate) async fn create(State(gateway): State<Arc<Gateway>>, request: Request) -> Response {
|
||||
match handle(&gateway, request).await {
|
||||
Ok(response) => Json(response).into_response(),
|
||||
Err(error) => error.openai_response(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle(gateway: &Gateway, request: Request) -> Result<Value, Error> {
|
||||
let (body, upload) = request::parse(request).await?;
|
||||
let deployment = request::deployment(gateway, &body)?;
|
||||
let audio = match upload {
|
||||
Some(upload) => {
|
||||
let format = upload
|
||||
.file_name
|
||||
.as_deref()
|
||||
.and_then(|name| Path::new(name).extension())
|
||||
.and_then(|extension| extension.to_str())
|
||||
.ok_or_else(|| {
|
||||
Error::InvalidBody("audio file requires a filename extension".into())
|
||||
})?;
|
||||
json!({"data": STANDARD.encode(upload.bytes), "format": format.to_ascii_lowercase()})
|
||||
}
|
||||
None => body
|
||||
.get("audio")
|
||||
.cloned()
|
||||
.ok_or_else(|| Error::InvalidBody("audio is required".into()))?,
|
||||
};
|
||||
Ok(audio_transcription(
|
||||
&gateway.resources,
|
||||
&gateway.http,
|
||||
AudioTranscriptionRequest {
|
||||
model: &deployment.model,
|
||||
audio,
|
||||
api_key: deployment.api_key.as_deref(),
|
||||
api_base: deployment.api_base.as_deref(),
|
||||
custom_llm_provider: deployment.custom_llm_provider.as_deref(),
|
||||
extra_headers: None,
|
||||
optional_params: body
|
||||
.into_iter()
|
||||
.filter(|(name, _)| !matches!(name.as_str(), "model" | "audio"))
|
||||
.collect(),
|
||||
timeout: deployment.timeout,
|
||||
},
|
||||
)
|
||||
.await?)
|
||||
}
|
||||
|
|
@ -0,0 +1,83 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
body::Bytes,
|
||||
extract::{Path, State},
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use litellm_core::chat_completions::{chat_completions, types::ChatCompletionsRequest};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{Error, Gateway, request};
|
||||
|
||||
pub(crate) async fn create(State(gateway): State<Arc<Gateway>>, body: Bytes) -> Response {
|
||||
respond(&gateway, request::object(&body)).await
|
||||
}
|
||||
|
||||
pub(crate) async fn deployment(
|
||||
State(gateway): State<Arc<Gateway>>,
|
||||
Path(path): Path<String>,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
if let Some(model) = path
|
||||
.strip_suffix("/chat/completions")
|
||||
.filter(|model| !model.is_empty())
|
||||
{
|
||||
let body = request::object(&body).map(|body| {
|
||||
if body.get("model").is_some_and(|model| !model.is_null()) {
|
||||
return body;
|
||||
}
|
||||
body.into_iter()
|
||||
.chain([("model".into(), Value::String(model.into()))])
|
||||
.collect()
|
||||
});
|
||||
return respond(&gateway, body).await;
|
||||
}
|
||||
if path.ends_with("/embeddings") || path.ends_with("/completions") {
|
||||
return Error::Unsupported(path).openai_response();
|
||||
}
|
||||
StatusCode::NOT_FOUND.into_response()
|
||||
}
|
||||
|
||||
async fn respond(gateway: &Gateway, body: Result<Map<String, Value>, Error>) -> Response {
|
||||
let result = match body {
|
||||
Ok(body) => handle(gateway, body).await,
|
||||
Err(error) => Err(error),
|
||||
};
|
||||
match result {
|
||||
Ok(response) => response,
|
||||
Err(error) => error.openai_response(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle(gateway: &Gateway, body: Map<String, Value>) -> Result<Response, Error> {
|
||||
let deployment = request::deployment(gateway, &body)?;
|
||||
if body.get("stream").and_then(Value::as_bool) == Some(true) {
|
||||
return Err(Error::Unsupported("streaming chat completions".into()));
|
||||
}
|
||||
let messages = body
|
||||
.get("messages")
|
||||
.cloned()
|
||||
.ok_or_else(|| Error::InvalidBody("messages is required".into()))?;
|
||||
let response = chat_completions(
|
||||
&gateway.resources,
|
||||
&gateway.http,
|
||||
ChatCompletionsRequest {
|
||||
model: &deployment.model,
|
||||
messages,
|
||||
optional_params: body
|
||||
.into_iter()
|
||||
.filter(|(name, _)| !matches!(name.as_str(), "model" | "messages" | "stream"))
|
||||
.collect(),
|
||||
api_key: deployment.api_key.as_deref(),
|
||||
api_base: deployment.api_base.as_deref(),
|
||||
custom_llm_provider: deployment.custom_llm_provider.as_deref(),
|
||||
extra_headers: None,
|
||||
timeout: deployment.timeout,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
Ok(Json(response).into_response())
|
||||
}
|
||||
225
litellm-rust/crates/gateway-inference/src/error.rs
Normal file
225
litellm-rust/crates/gateway-inference/src/error.rs
Normal file
|
|
@ -0,0 +1,225 @@
|
|||
use axum::http::StatusCode;
|
||||
use axum::{
|
||||
Json,
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use litellm_core::RouteError;
|
||||
use litellm_http::transport::Error as TransportError;
|
||||
use litellm_llms::base_llm::ocr::error::Error as OcrError;
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("invalid request body: {0}")]
|
||||
InvalidBody(String),
|
||||
#[error(
|
||||
"/v1/messages: Invalid model name passed in model={0}. Call `/v1/models` to view available models for your key."
|
||||
)]
|
||||
UnknownModel(String),
|
||||
#[error(transparent)]
|
||||
Route(#[from] RouteError),
|
||||
#[error(transparent)]
|
||||
Ocr(#[from] OcrError),
|
||||
#[error("{0} is not implemented by the Rust gateway")]
|
||||
Unsupported(String),
|
||||
#[error("request body exceeds the size limit")]
|
||||
BodyTooLarge,
|
||||
#[error("{0}")]
|
||||
Internal(String),
|
||||
}
|
||||
|
||||
impl Error {
|
||||
pub fn status(&self) -> StatusCode {
|
||||
match self {
|
||||
Self::Unsupported(_)
|
||||
| Self::Route(RouteError::Unsupported(_))
|
||||
| Self::Ocr(OcrError::Unsupported(_)) => StatusCode::NOT_IMPLEMENTED,
|
||||
Self::BodyTooLarge => StatusCode::PAYLOAD_TOO_LARGE,
|
||||
Self::Ocr(
|
||||
OcrError::Auth(litellm_auth::Error::MissingApiKey { .. })
|
||||
| OcrError::MissingAzureAiCredentials
|
||||
| OcrError::MissingAzureDocumentIntelligenceCredentials
|
||||
| OcrError::MissingReductoApiKey,
|
||||
) => StatusCode::UNAUTHORIZED,
|
||||
Self::Ocr(error) => error
|
||||
.http_status_code()
|
||||
.and_then(|status| StatusCode::from_u16(status).ok())
|
||||
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR),
|
||||
Self::InvalidBody(_) | Self::UnknownModel(_) => StatusCode::BAD_REQUEST,
|
||||
Self::Route(RouteError::Transport(TransportError::Http { status, .. })) => {
|
||||
StatusCode::from_u16(*status).unwrap_or(StatusCode::BAD_GATEWAY)
|
||||
}
|
||||
Self::Route(RouteError::Auth(litellm_auth::Error::MissingApiKey { .. })) => {
|
||||
StatusCode::UNAUTHORIZED
|
||||
}
|
||||
Self::Route(error) if error.is_request() => StatusCode::BAD_REQUEST,
|
||||
Self::Route(_) | Self::Internal(_) => StatusCode::INTERNAL_SERVER_ERROR,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn openai_response(self) -> Response {
|
||||
let status = self.status();
|
||||
let message = match &self {
|
||||
Self::UnknownModel(model) => format!("Invalid model name passed in model={model}"),
|
||||
_ => self.to_string(),
|
||||
};
|
||||
(
|
||||
status,
|
||||
Json(json!({"error": {
|
||||
"message": message,
|
||||
"type": error_type(status),
|
||||
"param": null,
|
||||
"code": status.as_u16(),
|
||||
}})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
/// The Anthropic error envelope Python's `AnthropicExceptionMapping` builds: an upstream
|
||||
/// body already in that shape passes through, any other has its message extracted.
|
||||
pub fn body(&self, request_id: Option<&str>) -> Value {
|
||||
let raw = match self {
|
||||
Self::Route(RouteError::Transport(TransportError::Http { body, .. })) => body.clone(),
|
||||
other => other.to_string(),
|
||||
};
|
||||
let parsed = serde_json::from_str::<Value>(&raw).ok();
|
||||
let envelope = match parsed {
|
||||
Some(Value::Object(object)) if is_anthropic_error(&object) => object,
|
||||
Some(Value::Object(object)) => {
|
||||
envelope(self.status(), provider_message(&object).unwrap_or(&raw))
|
||||
}
|
||||
_ => envelope(self.status(), &raw),
|
||||
};
|
||||
Value::Object(with_request_id(envelope, request_id))
|
||||
}
|
||||
|
||||
/// An `event: error` frame, for a stream that fails after its headers went out.
|
||||
pub fn sse_frame(&self) -> String {
|
||||
format!("event: error\ndata: {}\n\n", self.body(None))
|
||||
}
|
||||
}
|
||||
|
||||
fn error_type(status: StatusCode) -> &'static str {
|
||||
match status.as_u16() {
|
||||
400 => "invalid_request_error",
|
||||
401 => "authentication_error",
|
||||
403 => "permission_error",
|
||||
404 => "not_found_error",
|
||||
413 => "request_too_large",
|
||||
429 => "rate_limit_error",
|
||||
529 => "overloaded_error",
|
||||
_ => "api_error",
|
||||
}
|
||||
}
|
||||
|
||||
fn envelope(status: StatusCode, message: &str) -> Map<String, Value> {
|
||||
let Value::Object(envelope) = json!({
|
||||
"type": "error",
|
||||
"error": {"type": error_type(status), "message": message},
|
||||
}) else {
|
||||
unreachable!("a json object literal is an object")
|
||||
};
|
||||
envelope
|
||||
}
|
||||
|
||||
fn is_anthropic_error(object: &Map<String, Value>) -> bool {
|
||||
object.get("type").and_then(Value::as_str) == Some("error")
|
||||
&& object
|
||||
.get("error")
|
||||
.and_then(Value::as_object)
|
||||
.is_some_and(|error| error.contains_key("type") && error.contains_key("message"))
|
||||
}
|
||||
|
||||
fn provider_message(object: &Map<String, Value>) -> Option<&str> {
|
||||
if let Some(detail) = object.get("detail").and_then(Value::as_object) {
|
||||
return detail.get("message").and_then(Value::as_str);
|
||||
}
|
||||
["Message", "message"]
|
||||
.into_iter()
|
||||
.filter_map(|key| object.get(key).and_then(Value::as_str))
|
||||
.find(|message| !message.is_empty())
|
||||
}
|
||||
|
||||
fn with_request_id(envelope: Map<String, Value>, request_id: Option<&str>) -> Map<String, Value> {
|
||||
match request_id {
|
||||
Some(id) if !id.is_empty() && !envelope.contains_key("request_id") => envelope
|
||||
.into_iter()
|
||||
.chain([("request_id".to_string(), Value::from(id))])
|
||||
.collect(),
|
||||
_ => envelope,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn upstream(status: u16, body: &str) -> Error {
|
||||
Error::Route(RouteError::Transport(TransportError::Http {
|
||||
status,
|
||||
body: body.into(),
|
||||
}))
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::anthropic_body_passes_through(
|
||||
upstream(529, r#"{"type":"error","error":{"type":"overloaded_error","message":"busy","extra":1}}"#),
|
||||
Some("req_1"),
|
||||
json!({"type": "error", "error": {"type": "overloaded_error", "message": "busy", "extra": 1}, "request_id": "req_1"}),
|
||||
)]
|
||||
#[case::upstream_request_id_wins(
|
||||
upstream(400, r#"{"type":"error","error":{"type":"x","message":"m"},"request_id":"upstream"}"#),
|
||||
Some("caller"),
|
||||
json!({"type": "error", "error": {"type": "x", "message": "m"}, "request_id": "upstream"}),
|
||||
)]
|
||||
#[case::bedrock_detail(
|
||||
upstream(403, r#"{"detail":{"message":"denied"}}"#),
|
||||
None,
|
||||
json!({"type": "error", "error": {"type": "permission_error", "message": "denied"}}),
|
||||
)]
|
||||
#[case::aws_message(
|
||||
upstream(429, r#"{"Message":"slow down"}"#),
|
||||
None,
|
||||
json!({"type": "error", "error": {"type": "rate_limit_error", "message": "slow down"}}),
|
||||
)]
|
||||
#[case::plain_text_with_unmapped_status(
|
||||
upstream(502, "bad gateway"),
|
||||
None,
|
||||
json!({"type": "error", "error": {"type": "api_error", "message": "bad gateway"}}),
|
||||
)]
|
||||
#[case::unknown_model(
|
||||
Error::UnknownModel("nope".into()),
|
||||
None,
|
||||
json!({"type": "error", "error": {
|
||||
"type": "invalid_request_error",
|
||||
"message": "/v1/messages: Invalid model name passed in model=nope. Call `/v1/models` to view available models for your key.",
|
||||
}}),
|
||||
)]
|
||||
fn body_follows_the_anthropic_exception_mapping(
|
||||
#[case] error: Error,
|
||||
#[case] request_id: Option<&str>,
|
||||
#[case] expected: Value,
|
||||
) {
|
||||
assert_eq!(error.body(request_id), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::upstream_status(upstream(429, ""), StatusCode::TOO_MANY_REQUESTS)]
|
||||
#[case::rejected_request(Error::Route(RouteError::InvalidRequest("top_k".into())), StatusCode::BAD_REQUEST)]
|
||||
#[case::missing_key(
|
||||
Error::Route(RouteError::Auth(litellm_auth::Error::MissingApiKey {
|
||||
provider: "Anthropic",
|
||||
environment_variable: "ANTHROPIC_API_KEY",
|
||||
})),
|
||||
StatusCode::UNAUTHORIZED,
|
||||
)]
|
||||
#[case::lost_connection(
|
||||
Error::Route(RouteError::Transport(TransportError::Network("reset".into()))),
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
)]
|
||||
fn status_follows_who_is_at_fault(#[case] error: Error, #[case] status: StatusCode) {
|
||||
assert_eq!(error.status(), status);
|
||||
}
|
||||
}
|
||||
59
litellm-rust/crates/gateway-inference/src/lib.rs
Normal file
59
litellm-rust/crates/gateway-inference/src/lib.rs
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
//! The proxy's inference endpoints as an axum [`Router`] a server mounts.
|
||||
//!
|
||||
//! Authentication, rate limiting and logging are the mounting server's layers; this crate
|
||||
//! maps a public model name to its deployment and runs the core route.
|
||||
|
||||
mod audio_transcription;
|
||||
mod chat_completions;
|
||||
mod error;
|
||||
pub mod messages;
|
||||
mod ocr;
|
||||
mod request;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{Router, routing::post};
|
||||
use litellm_core::resources::CoreResources;
|
||||
use litellm_http::HttpClientConfig;
|
||||
use litellm_llms::base_llm::ocr::handler::OcrClient;
|
||||
use litellm_secrets::source::SecretSource;
|
||||
|
||||
pub use error::Error;
|
||||
pub use litellm_router::{Deployment, Router as ModelList};
|
||||
|
||||
pub struct Gateway {
|
||||
pub resources: CoreResources,
|
||||
pub http: HttpClientConfig,
|
||||
pub secrets: Arc<dyn SecretSource>,
|
||||
pub models: ModelList,
|
||||
pub ocr: OcrClient,
|
||||
}
|
||||
|
||||
pub fn router(gateway: Arc<Gateway>) -> Router {
|
||||
Router::new()
|
||||
.route("/v1/messages", post(messages::create))
|
||||
.route("/ocr", post(ocr::create))
|
||||
.route("/v1/ocr", post(ocr::create))
|
||||
.route("/chat/completions", post(chat_completions::create))
|
||||
.route("/v1/chat/completions", post(chat_completions::create))
|
||||
.route("/engines/{*path}", post(chat_completions::deployment))
|
||||
.route(
|
||||
"/openai/deployments/{*path}",
|
||||
post(chat_completions::deployment),
|
||||
)
|
||||
.route("/audio/transcriptions", post(audio_transcription::create))
|
||||
.route(
|
||||
"/v1/audio/transcriptions",
|
||||
post(audio_transcription::create),
|
||||
)
|
||||
.route("/responses", post(request::unsupported))
|
||||
.route("/v1/responses", post(request::unsupported))
|
||||
.route("/embeddings", post(request::unsupported))
|
||||
.route("/v1/embeddings", post(request::unsupported))
|
||||
.route("/completions", post(request::unsupported))
|
||||
.route("/v1/completions", post(request::unsupported))
|
||||
.layer(axum::extract::DefaultBodyLimit::max(
|
||||
request::MAX_BODY_BYTES,
|
||||
))
|
||||
.with_state(gateway)
|
||||
}
|
||||
69
litellm-rust/crates/gateway-inference/src/messages/host.rs
Normal file
69
litellm-rust/crates/gateway-inference/src/messages/host.rs
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
use std::{convert::Infallible, sync::Mutex};
|
||||
|
||||
use bytes::Bytes;
|
||||
use litellm_core::messages::{
|
||||
Error,
|
||||
route::{LocalMessagesHost, Messages, MessagesCall, MessagesStreamHead},
|
||||
};
|
||||
use litellm_host::host::{Demand, Host};
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
|
||||
/// Hands a streamed response to the HTTP body: the head once, then each chunk. A dropped
|
||||
/// receiver means the client went away, which detaches the call.
|
||||
pub(super) struct ChannelHost {
|
||||
local: LocalMessagesHost,
|
||||
head: Mutex<Option<oneshot::Sender<MessagesStreamHead>>>,
|
||||
pub(super) chunks: mpsc::Sender<Bytes>,
|
||||
}
|
||||
|
||||
impl ChannelHost {
|
||||
pub(super) fn new(
|
||||
call: MessagesCall,
|
||||
head: oneshot::Sender<MessagesStreamHead>,
|
||||
chunks: mpsc::Sender<Bytes>,
|
||||
) -> Self {
|
||||
Self {
|
||||
local: LocalMessagesHost::new(call),
|
||||
head: Mutex::new(Some(head)),
|
||||
chunks,
|
||||
}
|
||||
}
|
||||
|
||||
fn take_head(&self) -> Option<oneshot::Sender<MessagesStreamHead>> {
|
||||
self.head
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner())
|
||||
.take()
|
||||
}
|
||||
|
||||
pub(super) fn opened(&self) -> bool {
|
||||
self.head
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner())
|
||||
.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
impl Host<Messages> for ChannelHost {
|
||||
async fn project(&self) -> Result<MessagesCall, Error> {
|
||||
self.local.project().await
|
||||
}
|
||||
|
||||
async fn custom_op(&self, op: Infallible) -> Result<(), Error> {
|
||||
match op {}
|
||||
}
|
||||
|
||||
async fn open(&self, head: MessagesStreamHead) -> Result<Demand, Error> {
|
||||
Ok(match self.take_head().map(|sender| sender.send(head)) {
|
||||
Some(Ok(())) => Demand::More,
|
||||
Some(Err(_)) | None => Demand::Detached,
|
||||
})
|
||||
}
|
||||
|
||||
async fn deliver(&self, chunk: Bytes) -> Result<Demand, Error> {
|
||||
Ok(match self.chunks.send(chunk).await {
|
||||
Ok(()) => Demand::More,
|
||||
Err(_) => Demand::Detached,
|
||||
})
|
||||
}
|
||||
}
|
||||
137
litellm-rust/crates/gateway-inference/src/messages/mod.rs
Normal file
137
litellm-rust/crates/gateway-inference/src/messages/mod.rs
Normal file
|
|
@ -0,0 +1,137 @@
|
|||
//! `POST /v1/messages`, as the Python proxy's `anthropic_response` serves it.
|
||||
|
||||
mod host;
|
||||
|
||||
use std::{convert::Infallible, sync::Arc};
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
body::{Body, Bytes},
|
||||
extract::State,
|
||||
http::{HeaderMap, StatusCode, header},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use host::ChannelHost;
|
||||
use litellm_core::messages::route::{
|
||||
MessagesCall, MessagesOutput, messages_body, messages_machine,
|
||||
};
|
||||
use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders};
|
||||
use serde_json::{Map, Value};
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
|
||||
use crate::{Deployment, Error, Gateway};
|
||||
|
||||
/// Client headers Python forwards to Anthropic-speaking providers on every call.
|
||||
const ANTHROPIC_API_HEADERS: [&str; 2] = ["anthropic-version", "anthropic-beta"];
|
||||
const ANTHROPIC_API_HEADER_PROVIDERS: &str = "anthropic,bedrock,bedrock_mantle,vertex_ai";
|
||||
|
||||
pub async fn create(
|
||||
State(gateway): State<Arc<Gateway>>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
let request_id = headers
|
||||
.get("x-request-id")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_owned);
|
||||
match handle(&gateway, &headers, &body).await {
|
||||
Ok(response) => response,
|
||||
Err(error) => (error.status(), Json(error.body(request_id.as_deref()))).into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle(gateway: &Gateway, headers: &HeaderMap, body: &[u8]) -> Result<Response, Error> {
|
||||
let body = match serde_json::from_slice(body) {
|
||||
Ok(Value::Object(body)) => body,
|
||||
Ok(_) => return Err(Error::InvalidBody("expected a JSON object".into())),
|
||||
Err(error) => return Err(Error::InvalidBody(error.to_string())),
|
||||
};
|
||||
let model_name = body
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or_else(|| Error::InvalidBody("model is required".into()))?;
|
||||
let deployment = gateway
|
||||
.models
|
||||
.get(model_name)
|
||||
.ok_or_else(|| Error::UnknownModel(model_name.to_owned()))?;
|
||||
let call = project(deployment, body, headers)?;
|
||||
let machine = messages_machine(&gateway.resources, &gateway.http, gateway.secrets.clone())
|
||||
.map_err(|error| Error::Route(error.into()))?;
|
||||
|
||||
let (head_sender, head) = oneshot::channel();
|
||||
let (chunk_sender, chunks) = mpsc::channel(1);
|
||||
let host = ChannelHost::new(call, head_sender, chunk_sender);
|
||||
let call = tokio::spawn(async move {
|
||||
let outcome = litellm_host::run::run(machine, &host).await;
|
||||
if let Err(error) = &outcome
|
||||
&& host.opened()
|
||||
{
|
||||
let _ = host
|
||||
.chunks
|
||||
.send(Bytes::from(Error::Route(error.clone()).sse_frame()))
|
||||
.await;
|
||||
}
|
||||
outcome
|
||||
});
|
||||
tokio::select! {
|
||||
biased;
|
||||
Ok(_) = head => Ok(stream(chunks)),
|
||||
joined = call => match joined.map_err(|error| Error::Internal(error.to_string()))?? {
|
||||
MessagesOutput::Message(message) => Ok(Json(message).into_response()),
|
||||
MessagesOutput::Streamed => Err(Error::Internal("the stream ended before it opened".into())),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn project(
|
||||
deployment: &Deployment,
|
||||
body: Map<String, Value>,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<MessagesCall, Error> {
|
||||
let body = body
|
||||
.into_iter()
|
||||
.map(|(name, value)| match name.as_str() {
|
||||
"model" => (name, Value::from(deployment.model.as_str())),
|
||||
_ => (name, value),
|
||||
})
|
||||
.collect();
|
||||
Ok(MessagesCall {
|
||||
body: messages_body(body)?,
|
||||
api_key: deployment.api_key.clone(),
|
||||
api_base: deployment.api_base.clone(),
|
||||
custom_llm_provider: deployment.custom_llm_provider.clone(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: anthropic_api_headers(headers),
|
||||
timeout: deployment.timeout,
|
||||
shaping: deployment.shaping.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn anthropic_api_headers(headers: &HeaderMap) -> Option<ProviderSpecificHeaders> {
|
||||
let extra_headers: Map<String, Value> = ANTHROPIC_API_HEADERS
|
||||
.into_iter()
|
||||
.filter_map(|name| {
|
||||
let value = headers.get(name)?.to_str().ok()?;
|
||||
Some((name.to_owned(), Value::from(value)))
|
||||
})
|
||||
.collect();
|
||||
(!extra_headers.is_empty()).then(|| {
|
||||
ProviderSpecificHeaders::One(ProviderSpecificHeader {
|
||||
custom_llm_provider: ANTHROPIC_API_HEADER_PROVIDERS.into(),
|
||||
extra_headers,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn stream(chunks: mpsc::Receiver<Bytes>) -> Response {
|
||||
let body = futures_util::stream::unfold(chunks, |mut chunks| async move {
|
||||
let chunk = chunks.recv().await?;
|
||||
Some((Ok::<_, Infallible>(chunk), chunks))
|
||||
});
|
||||
(
|
||||
StatusCode::OK,
|
||||
[(header::CONTENT_TYPE, "text/event-stream")],
|
||||
Body::from_stream(body),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
77
litellm-rust/crates/gateway-inference/src/ocr.rs
Normal file
77
litellm-rust/crates/gateway-inference/src/ocr.rs
Normal file
|
|
@ -0,0 +1,77 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
extract::{Request, State},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use litellm_auth::SecretValue;
|
||||
use litellm_core::ocr::{
|
||||
client::perform,
|
||||
types::{LiteLLMOcrRequest, OcrConnectionInputs, OcrDocumentInput},
|
||||
};
|
||||
use litellm_llms::base_llm::ocr::transformation::OcrDocument;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{Error, Gateway, request};
|
||||
|
||||
pub(crate) async fn create(State(gateway): State<Arc<Gateway>>, request: Request) -> Response {
|
||||
match handle(&gateway, request).await {
|
||||
Ok(response) => Json(response).into_response(),
|
||||
Err(error) => error.openai_response(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle(gateway: &Gateway, request: Request) -> Result<Value, Error> {
|
||||
let header_format = request
|
||||
.headers()
|
||||
.get("x-req-format")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_owned);
|
||||
let (body, upload) = request::parse(request).await?;
|
||||
let deployment = request::deployment(gateway, &body)?;
|
||||
let document = match upload {
|
||||
Some(upload) => OcrDocumentInput::Bytes {
|
||||
bytes: upload.bytes,
|
||||
file_name: upload.file_name,
|
||||
mime_type: upload.mime_type,
|
||||
},
|
||||
None => OcrDocument::try_from(
|
||||
body.get("document")
|
||||
.cloned()
|
||||
.ok_or_else(|| Error::InvalidBody("document is required".into()))?,
|
||||
)?
|
||||
.into(),
|
||||
};
|
||||
let format = body
|
||||
.get("req_format")
|
||||
.filter(|value| !value.is_null())
|
||||
.cloned()
|
||||
.or_else(|| header_format.map(Value::String));
|
||||
let format = format.map(|value| match value {
|
||||
Value::String(value) => Value::String(value.trim().to_ascii_lowercase()),
|
||||
value => value,
|
||||
});
|
||||
let options = body
|
||||
.into_iter()
|
||||
.filter(|(name, _)| !matches!(name.as_str(), "model" | "document" | "req_format"))
|
||||
.chain(format.map(|value| ("req_format".into(), value)))
|
||||
.collect();
|
||||
let call = LiteLLMOcrRequest::from_inputs(
|
||||
deployment.model.clone(),
|
||||
document,
|
||||
deployment.custom_llm_provider.as_deref(),
|
||||
options,
|
||||
OcrConnectionInputs {
|
||||
api_key: deployment.api_key.clone().map(SecretValue::new),
|
||||
api_base: deployment.api_base.clone(),
|
||||
timeout: deployment.timeout,
|
||||
..Default::default()
|
||||
},
|
||||
)?;
|
||||
let response = perform(&gateway.ocr, call).await?;
|
||||
match response.provider_native_response {
|
||||
Some(native) => Ok(Value::Object(native)),
|
||||
None => Ok(response.into_json()),
|
||||
}
|
||||
}
|
||||
108
litellm-rust/crates/gateway-inference/src/request.rs
Normal file
108
litellm-rust/crates/gateway-inference/src/request.rs
Normal file
|
|
@ -0,0 +1,108 @@
|
|||
use axum::{
|
||||
body::{Bytes, to_bytes},
|
||||
extract::{FromRequest, Multipart, Request},
|
||||
http::Uri,
|
||||
response::Response,
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{Deployment, Error, Gateway};
|
||||
|
||||
pub(crate) const MAX_FILE_BYTES: usize = 50 * 1024 * 1024;
|
||||
pub(crate) const MAX_BODY_BYTES: usize = MAX_FILE_BYTES + 1024 * 1024;
|
||||
|
||||
pub(crate) struct Upload {
|
||||
pub bytes: Bytes,
|
||||
pub file_name: Option<String>,
|
||||
pub mime_type: Option<String>,
|
||||
}
|
||||
|
||||
pub(crate) fn object(body: &[u8]) -> Result<Map<String, Value>, Error> {
|
||||
match serde_json::from_slice(body) {
|
||||
Ok(Value::Object(body)) => Ok(body),
|
||||
Ok(_) => Err(Error::InvalidBody("expected a JSON object".into())),
|
||||
Err(error) => Err(Error::InvalidBody(error.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn deployment<'a>(
|
||||
gateway: &'a Gateway,
|
||||
body: &Map<String, Value>,
|
||||
) -> Result<&'a Deployment, Error> {
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or_else(|| Error::InvalidBody("model is required".into()))?;
|
||||
gateway
|
||||
.models
|
||||
.get(model)
|
||||
.ok_or_else(|| Error::UnknownModel(model.to_owned()))
|
||||
}
|
||||
|
||||
pub(crate) async fn parse(request: Request) -> Result<(Map<String, Value>, Option<Upload>), Error> {
|
||||
let multipart = request
|
||||
.headers()
|
||||
.get("content-type")
|
||||
.and_then(|header| header.to_str().ok())
|
||||
.is_some_and(|value| {
|
||||
value
|
||||
.to_ascii_lowercase()
|
||||
.starts_with("multipart/form-data")
|
||||
});
|
||||
if !multipart {
|
||||
let body = to_bytes(request.into_body(), MAX_BODY_BYTES)
|
||||
.await
|
||||
.map_err(|_| Error::BodyTooLarge)?;
|
||||
return Ok((object(&body)?, None));
|
||||
}
|
||||
let mut multipart = Multipart::from_request(request, &())
|
||||
.await
|
||||
.map_err(|error| Error::InvalidBody(error.to_string()))?;
|
||||
let mut fields = Map::new();
|
||||
let mut upload = None;
|
||||
while let Some(field) = multipart.next_field().await.map_err(multipart_error)? {
|
||||
let name = field.name().unwrap_or_default().to_owned();
|
||||
if name == "file" {
|
||||
let file_name = field.file_name().map(str::to_owned);
|
||||
let mime_type = field
|
||||
.content_type()
|
||||
.and_then(|value| value.split(';').next())
|
||||
.map(str::trim)
|
||||
.filter(|value| *value != "application/octet-stream")
|
||||
.map(str::to_owned);
|
||||
let bytes = field.bytes().await.map_err(multipart_error)?;
|
||||
if bytes.len() > MAX_FILE_BYTES {
|
||||
return Err(Error::BodyTooLarge);
|
||||
}
|
||||
if bytes.is_empty() {
|
||||
return Err(Error::InvalidBody("uploaded file is empty".into()));
|
||||
}
|
||||
upload = Some(Upload {
|
||||
bytes,
|
||||
file_name,
|
||||
mime_type,
|
||||
});
|
||||
} else if name != "document" {
|
||||
let text = field.text().await.map_err(multipart_error)?;
|
||||
let value = serde_json::from_str(&text).unwrap_or(Value::String(text));
|
||||
fields.insert(name, value);
|
||||
}
|
||||
}
|
||||
if upload.is_none() {
|
||||
return Err(Error::InvalidBody(
|
||||
"multipart request requires a file field".into(),
|
||||
));
|
||||
}
|
||||
Ok((fields, upload))
|
||||
}
|
||||
|
||||
fn multipart_error(error: axum::extract::multipart::MultipartError) -> Error {
|
||||
if error.status() == axum::http::StatusCode::PAYLOAD_TOO_LARGE {
|
||||
return Error::BodyTooLarge;
|
||||
}
|
||||
Error::InvalidBody(error.to_string())
|
||||
}
|
||||
|
||||
pub(crate) async fn unsupported(uri: Uri) -> Response {
|
||||
Error::Unsupported(uri.path().to_owned()).openai_response()
|
||||
}
|
||||
66
litellm-rust/crates/gateway-inference/tests/messages.rs
Normal file
66
litellm-rust/crates/gateway-inference/tests/messages.rs
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
mod support;
|
||||
|
||||
use axum::{
|
||||
body::{Body, to_bytes},
|
||||
http::Request,
|
||||
};
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
use tower::ServiceExt;
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{body_json, header, method, path},
|
||||
};
|
||||
|
||||
#[rstest]
|
||||
#[case(false)]
|
||||
#[case(true)]
|
||||
#[tokio::test]
|
||||
async fn messages_reaches_the_provider_and_preserves_json_or_sse(#[case] streaming: bool) {
|
||||
let upstream = MockServer::start().await;
|
||||
let message = json!({"id": "msg_test", "type": "message", "role": "assistant",
|
||||
"model": "test-model", "content": [{"type": "text", "text": "hello"}],
|
||||
"stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1}});
|
||||
let sse = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
|
||||
let template = if streaming {
|
||||
ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")
|
||||
} else {
|
||||
ResponseTemplate::new(200).set_body_json(message.clone())
|
||||
};
|
||||
let messages = json!([{"role": "user", "content": "hi"}]);
|
||||
Mock::given(method("POST")).and(path("/v1/messages"))
|
||||
.and(header("x-api-key", "test-key"))
|
||||
.and(header("anthropic-beta", "test-feature"))
|
||||
.and(body_json(json!({"model": "test-model", "messages": messages, "max_tokens": 16, "stream": streaming})))
|
||||
.respond_with(template).expect(1).mount(&upstream).await;
|
||||
let request = Request::post("/v1/messages")
|
||||
.header("content-type", "application/json").header("anthropic-beta", "test-feature")
|
||||
.body(Body::from(json!({"model": "public/model", "messages": messages, "max_tokens": 16, "stream": streaming}).to_string())).unwrap();
|
||||
let response = support::app("anthropic/test-model", &upstream.uri())
|
||||
.oneshot(request)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), 200);
|
||||
if streaming {
|
||||
assert_eq!(response.headers()["content-type"], "text/event-stream");
|
||||
assert_eq!(to_bytes(response.into_body(), 4096).await.unwrap(), sse);
|
||||
} else {
|
||||
let body = support::json(response).await;
|
||||
assert_eq!(body["content"], message["content"]);
|
||||
assert_eq!(body["usage"], message["usage"]);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_messages_stays_an_anthropic_error() {
|
||||
let response = support::post(
|
||||
support::app("anthropic/test-model", "http://127.0.0.1:1"),
|
||||
"/v1/messages",
|
||||
json!({"model": "public/model", "messages": "invalid", "max_tokens": 16}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), 400);
|
||||
let body = support::json(response).await;
|
||||
assert_eq!(body["type"], "error");
|
||||
assert_eq!(body["error"]["type"], "invalid_request_error");
|
||||
}
|
||||
118
litellm-rust/crates/gateway-inference/tests/ocr.rs
Normal file
118
litellm-rust/crates/gateway-inference/tests/ocr.rs
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
mod support;
|
||||
|
||||
use axum::{body::Body, http::Request};
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
use tower::ServiceExt;
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{body_json, header, method, path},
|
||||
};
|
||||
|
||||
const DOCUMENT: &str = "data:application/pdf;base64,YWJj";
|
||||
|
||||
#[rstest]
|
||||
#[case("/ocr", false)]
|
||||
#[case("/v1/ocr", true)]
|
||||
#[tokio::test]
|
||||
async fn json_and_multipart_reach_ocr_with_the_deployment(
|
||||
#[case] route: &str,
|
||||
#[case] multipart: bool,
|
||||
) {
|
||||
let upstream = MockServer::start().await;
|
||||
Mock::given(method("POST")).and(path("/v1/ocr"))
|
||||
.and(header("authorization", "Bearer test-key"))
|
||||
.and(body_json(json!({"model": "test-ocr", "document": {"type": "document_url", "document_url": DOCUMENT}, "pages": [0]})))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"pages": [{"index": 0, "markdown": "recognized text"}]})))
|
||||
.expect(1).mount(&upstream).await;
|
||||
let app = support::app("mistral/test-ocr", &upstream.uri());
|
||||
let response = if multipart {
|
||||
let body = "--boundary\r\nContent-Disposition: form-data; name=\"model\"\r\n\r\npublic/model\r\n--boundary\r\nContent-Disposition: form-data; name=\"pages\"\r\n\r\n[0]\r\n--boundary\r\nContent-Disposition: form-data; name=\"file\"; filename=\"test.pdf\"\r\nContent-Type: application/pdf\r\n\r\nabc\r\n--boundary--\r\n";
|
||||
app.oneshot(
|
||||
Request::post(route)
|
||||
.header("content-type", "multipart/form-data; boundary=boundary")
|
||||
.body(Body::from(body))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
} else {
|
||||
support::post(app, route, json!({"model": "public/model", "document": {"type": "document_url", "document_url": DOCUMENT}, "pages": [0]})).await
|
||||
};
|
||||
assert_eq!(response.status(), 200);
|
||||
let body = support::json(response).await;
|
||||
assert_eq!(body["pages"][0]["markdown"], "recognized text");
|
||||
assert_eq!(body["model"], "test-ocr");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(None, true)]
|
||||
#[case(Some("litellm"), false)]
|
||||
#[tokio::test]
|
||||
async fn native_format_header_is_used_unless_the_body_overrides_it(
|
||||
#[case] format: Option<&str>,
|
||||
#[case] native: bool,
|
||||
) {
|
||||
let upstream = MockServer::start().await;
|
||||
let payload = json!({"pages": [{"index": 0, "markdown": "text"}], "provider_only": true});
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(payload.clone()))
|
||||
.expect(1)
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
let body = json!({"model": "public/model", "document": {"type": "document_url", "document_url": DOCUMENT}, "req_format": format});
|
||||
let response = support::app("mistral/test-ocr", &upstream.uri())
|
||||
.oneshot(
|
||||
Request::post("/ocr")
|
||||
.header("x-req-format", " Native ")
|
||||
.body(Body::from(body.to_string()))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), 200);
|
||||
let body = support::json(response).await;
|
||||
if native {
|
||||
assert_eq!(body, payload);
|
||||
} else {
|
||||
assert_eq!(body["object"], "ocr");
|
||||
assert_eq!(
|
||||
body["pages"][0]["markdown"],
|
||||
payload["pages"][0]["markdown"]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ocr_keeps_upstream_status_in_an_openai_error_envelope() {
|
||||
let upstream = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(ResponseTemplate::new(429).set_body_json(json!({"message": "busy"})))
|
||||
.expect(1)
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
let response = support::post(support::app("mistral/test-ocr", &upstream.uri()), "/ocr",
|
||||
json!({"model": "public/model", "document": {"type": "document_url", "document_url": DOCUMENT}})).await;
|
||||
assert_eq!(response.status(), 429);
|
||||
let body = support::json(response).await;
|
||||
assert_eq!(body["error"]["code"], 429);
|
||||
assert!(body["error"]["message"].as_str().unwrap().contains("busy"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(json!({"model": "public/model"}))]
|
||||
#[case(json!({"model": "public/model", "document": "/etc/passwd"}))]
|
||||
#[case(json!({"model": "missing", "document": {"type": "document_url", "document_url": DOCUMENT}}))]
|
||||
#[tokio::test]
|
||||
async fn invalid_ocr_requests_do_not_call_the_provider(#[case] body: Value) {
|
||||
let upstream = MockServer::start().await;
|
||||
let response = support::post(
|
||||
support::app("mistral/test-ocr", &upstream.uri()),
|
||||
"/ocr",
|
||||
body,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), 400);
|
||||
assert!(support::json(response).await["error"]["message"].is_string());
|
||||
assert!(upstream.received_requests().await.unwrap().is_empty());
|
||||
}
|
||||
92
litellm-rust/crates/gateway-inference/tests/routes.rs
Normal file
92
litellm-rust/crates/gateway-inference/tests/routes.rs
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
mod support;
|
||||
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{body_partial_json, method},
|
||||
};
|
||||
|
||||
#[rstest]
|
||||
#[case("/chat/completions", Some("public/model"))]
|
||||
#[case("/v1/chat/completions", Some("public/model"))]
|
||||
#[case("/engines/public/model/chat/completions", None)]
|
||||
#[case("/openai/deployments/public/model/chat/completions", None)]
|
||||
#[case("/openai/deployments/unused/chat/completions", Some("public/model"))]
|
||||
#[tokio::test]
|
||||
async fn chat_aliases_call_core_and_use_the_body_model_before_the_path(
|
||||
#[case] route: &str,
|
||||
#[case] model: Option<&str>,
|
||||
) {
|
||||
let upstream = MockServer::start().await;
|
||||
let messages = json!([{"role": "user", "content": "hi"}]);
|
||||
Mock::given(method("POST"))
|
||||
.and(body_partial_json(
|
||||
json!({"model": "test-model", "max_tokens": 16}),
|
||||
))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"id": "msg_test", "model": "test-model", "content": [{"type": "text", "text": "hello"}],
|
||||
"stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1}
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
let response = support::post(
|
||||
support::app("anthropic/test-model", &upstream.uri()),
|
||||
route,
|
||||
json!({"model": model, "messages": messages, "max_tokens": 16}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), 200);
|
||||
assert_eq!(
|
||||
support::json(response).await["choices"][0]["message"]["content"],
|
||||
"hello"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("/responses")]
|
||||
#[case("/v1/responses")]
|
||||
#[case("/embeddings")]
|
||||
#[case("/v1/embeddings")]
|
||||
#[case("/completions")]
|
||||
#[case("/v1/completions")]
|
||||
#[case("/engines/public/model/embeddings")]
|
||||
#[case("/openai/deployments/public/model/completions")]
|
||||
#[tokio::test]
|
||||
async fn unimplemented_routes_return_an_explicit_error(#[case] path: &str) {
|
||||
let response = support::post(
|
||||
support::app("anthropic/test-model", "http://127.0.0.1:1"),
|
||||
path,
|
||||
json!({}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), 501);
|
||||
assert!(
|
||||
support::json(response).await["error"]["message"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("not implemented")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("/audio/transcriptions")]
|
||||
#[case("/v1/audio/transcriptions")]
|
||||
#[tokio::test]
|
||||
async fn transcription_aliases_reach_core_validation(#[case] path: &str) {
|
||||
let response = support::post(
|
||||
support::app("bedrock/test-model", "http://127.0.0.1:1"),
|
||||
path,
|
||||
json!({"model": "public/model", "audio": {"data": "YWJj", "format": "invalid"}}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), 400);
|
||||
let body: Value = support::json(response).await;
|
||||
assert!(
|
||||
body["error"]["message"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("audio.format")
|
||||
);
|
||||
}
|
||||
75
litellm-rust/crates/gateway-inference/tests/support/mod.rs
Normal file
75
litellm-rust/crates/gateway-inference/tests/support/mod.rs
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use axum::{
|
||||
Router,
|
||||
body::{Body, to_bytes},
|
||||
http::Request,
|
||||
response::Response,
|
||||
};
|
||||
use futures_util::future::BoxFuture;
|
||||
use litellm_core::resources::CoreResources;
|
||||
use litellm_gateway_inference::{Deployment, Gateway, router};
|
||||
use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver};
|
||||
use litellm_llms::base_llm::ocr::settings::OcrSettings;
|
||||
use litellm_secrets::{SecretValue, source::SecretSource};
|
||||
use serde_json::Value;
|
||||
use tower::ServiceExt;
|
||||
|
||||
struct NoSecrets;
|
||||
|
||||
impl SecretSource for NoSecrets {
|
||||
fn get_secret_str<'a>(
|
||||
&'a self,
|
||||
_: &'a str,
|
||||
) -> BoxFuture<'a, Result<Option<SecretValue>, litellm_secrets::Error>> {
|
||||
Box::pin(async { Ok(None) })
|
||||
}
|
||||
}
|
||||
|
||||
pub fn app(model: &str, api_base: &str) -> Router {
|
||||
let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver)));
|
||||
let http = Resolution::from(&HttpSettings::default()).config;
|
||||
let secrets = Arc::new(NoSecrets);
|
||||
let resources = CoreResources::new(pool);
|
||||
let ocr = resources
|
||||
.ocr_client(
|
||||
&http,
|
||||
Default::default(),
|
||||
OcrSettings::default(),
|
||||
secrets.clone(),
|
||||
)
|
||||
.unwrap();
|
||||
router(Arc::new(Gateway {
|
||||
resources,
|
||||
http,
|
||||
secrets,
|
||||
ocr,
|
||||
models: [(
|
||||
"public/model".into(),
|
||||
Deployment {
|
||||
model: model.into(),
|
||||
api_base: Some(api_base.into()),
|
||||
api_key: Some("test-key".into()),
|
||||
timeout: Some(Duration::from_secs(5)),
|
||||
..Default::default()
|
||||
},
|
||||
)]
|
||||
.into_iter()
|
||||
.collect(),
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn post(app: Router, path: &str, body: Value) -> Response {
|
||||
app.oneshot(
|
||||
Request::post(path)
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(body.to_string()))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
pub async fn json(response: Response) -> Value {
|
||||
serde_json::from_slice(&to_bytes(response.into_body(), 1024 * 1024).await.unwrap()).unwrap()
|
||||
}
|
||||
5
litellm-rust/crates/gateway/AGENTS.md
Normal file
5
litellm-rust/crates/gateway/AGENTS.md
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
- Keep this crate a thin composition layer: mount endpoint routers and serve the supplied listener
|
||||
- Server lifecycle and shared inbound middleware belong here, including client authentication, rate limiting, and request logging
|
||||
- Endpoint paths, request handling, model resolution, and response encoding belong to the mounted crates; provider execution belongs to `core` and `llms`
|
||||
- Inject shared state and infrastructure; avoid global runtimes, duplicate client pools, and abstractions for hypothetical endpoint groups
|
||||
- Test mounting and server lifecycle through public HTTP behavior; test endpoint semantics in the owning crate
|
||||
24
litellm-rust/crates/gateway/Cargo.toml
Normal file
24
litellm-rust/crates/gateway/Cargo.toml
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
[package]
|
||||
name = "litellm-gateway"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
axum.workspace = true
|
||||
litellm-core.workspace = true
|
||||
litellm-gateway-inference.workspace = true
|
||||
litellm-gateway-auth.workspace = true
|
||||
litellm-config.workspace = true
|
||||
litellm-http.workspace = true
|
||||
litellm-llms.workspace = true
|
||||
litellm-secrets.workspace = true
|
||||
tower-http = { version = "0.7.1", default-features = false, features = ["trace"] }
|
||||
tracing.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
serde_json.workspace = true
|
||||
tokio = { workspace = true, features = ["sync"] }
|
||||
52
litellm-rust/crates/gateway/src/lib.rs
Normal file
52
litellm-rust/crates/gateway/src/lib.rs
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use axum::{Router, extract::Request};
|
||||
use tower_http::trace::{DefaultOnResponse, TraceLayer};
|
||||
|
||||
use litellm_config::Config;
|
||||
use litellm_core::resources::CoreResources;
|
||||
use litellm_gateway_auth::{Auth, RequireMasterKey};
|
||||
use litellm_gateway_inference::{Gateway, ModelList};
|
||||
use litellm_http::{
|
||||
ClientVariant, HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver,
|
||||
};
|
||||
use litellm_llms::base_llm::ocr::settings::OcrSettings;
|
||||
use litellm_secrets::source::EnvironmentSecrets;
|
||||
|
||||
pub fn build_inference(config: &Config) -> Result<Arc<Gateway>, litellm_http::Error> {
|
||||
let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver)));
|
||||
let http = Resolution::from(&HttpSettings::default()).config;
|
||||
let client = pool.client(&http, ClientVariant::Provider)?;
|
||||
let secrets = Arc::new(EnvironmentSecrets::python_compatible(client));
|
||||
let resources = CoreResources::new(pool);
|
||||
let ocr = resources.ocr_client(
|
||||
&http,
|
||||
Default::default(),
|
||||
OcrSettings::default(),
|
||||
secrets.clone(),
|
||||
)?;
|
||||
|
||||
Ok(Arc::new(Gateway {
|
||||
resources,
|
||||
http,
|
||||
secrets,
|
||||
models: ModelList::from_model_list(&config.model_list),
|
||||
ocr,
|
||||
}))
|
||||
}
|
||||
|
||||
pub fn router(inference: Arc<Gateway>, config: &Config) -> Router {
|
||||
let auth = Auth::from_config(config, inference.secrets.clone());
|
||||
litellm_gateway_inference::router(inference)
|
||||
.route_layer(axum::middleware::from_extractor_with_state::<
|
||||
RequireMasterKey,
|
||||
_,
|
||||
>(auth))
|
||||
.layer(
|
||||
TraceLayer::new_for_http()
|
||||
.make_span_with(|request: &Request| {
|
||||
tracing::info_span!("request", method = %request.method(), path = request.uri().path())
|
||||
})
|
||||
.on_response(DefaultOnResponse::new().level(tracing::Level::INFO)),
|
||||
)
|
||||
}
|
||||
18
litellm-rust/crates/gateway/src/main.rs
Normal file
18
litellm-rust/crates/gateway/src/main.rs
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
use std::error::Error;
|
||||
|
||||
use litellm_config::Config;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn Error>> {
|
||||
let config_path = std::env::var("LITELLM_CONFIG").unwrap_or_else(|_| "config.yaml".into());
|
||||
let config = Config::load(config_path)?;
|
||||
let inference = litellm_gateway::build_inference(&config)?;
|
||||
let host = std::env::var("HOST").unwrap_or_else(|_| "0.0.0.0".into());
|
||||
let port = std::env::var("PORT")
|
||||
.unwrap_or_else(|_| "4000".into())
|
||||
.parse::<u16>()?;
|
||||
let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?;
|
||||
|
||||
axum::serve(listener, litellm_gateway::router(inference, &config)).await?;
|
||||
Ok(())
|
||||
}
|
||||
90
litellm-rust/crates/gateway/tests/server.rs
Normal file
90
litellm-rust/crates/gateway/tests/server.rs
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_config::Config;
|
||||
use litellm_gateway_inference::{Error, Gateway};
|
||||
use litellm_http::ClientVariant;
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
use tokio::{net::TcpListener, sync::oneshot, time::timeout};
|
||||
|
||||
#[fixture]
|
||||
fn inference() -> Arc<Gateway> {
|
||||
litellm_gateway::build_inference(&Config::from_yaml("model_list: []").unwrap()).unwrap()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::authorized("/v1/messages", Some("Bearer gateway-key"), Some("gateway-key"), 400)]
|
||||
#[case::missing_token("/v1/messages", None, Some("gateway-key"), 401)]
|
||||
#[case::wrong_token("/v1/messages", Some("Bearer wrong"), Some("gateway-key"), 401)]
|
||||
#[case::ocr("/ocr", None, Some("gateway-key"), 401)]
|
||||
#[case::chat("/v1/chat/completions", None, Some("gateway-key"), 401)]
|
||||
#[case::deployment(
|
||||
"/openai/deployments/model/chat/completions",
|
||||
None,
|
||||
Some("gateway-key"),
|
||||
401
|
||||
)]
|
||||
#[case::transcription("/audio/transcriptions", None, Some("gateway-key"), 401)]
|
||||
#[case::unsupported_route("/responses", None, Some("gateway-key"), 401)]
|
||||
#[case::unknown_path("/unknown", None, Some("gateway-key"), 404)]
|
||||
#[case::unknown_path_unconfigured("/unknown", None, None, 404)]
|
||||
#[case::unconfigured("/v1/messages", Some("Bearer gateway-key"), None, 500)]
|
||||
#[tokio::test]
|
||||
async fn authenticates_before_serving_mounted_inference_routes(
|
||||
inference: Arc<Gateway>,
|
||||
#[case] path: &str,
|
||||
#[case] authorization: Option<&str>,
|
||||
#[case] master_key: Option<&str>,
|
||||
#[case] status: u16,
|
||||
) {
|
||||
let config = Config::from_yaml(&format!(
|
||||
"model_list: []\ngeneral_settings:\n master_key: {}\n",
|
||||
master_key.unwrap_or("null")
|
||||
))
|
||||
.unwrap();
|
||||
let client = inference
|
||||
.resources
|
||||
.pool
|
||||
.client(&inference.http, ClientVariant::Provider)
|
||||
.unwrap();
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let (shutdown, stopped) = oneshot::channel();
|
||||
let server = tokio::spawn(async move {
|
||||
axum::serve(listener, litellm_gateway::router(inference, &config))
|
||||
.with_graceful_shutdown(async move {
|
||||
let _ = stopped.await;
|
||||
})
|
||||
.await
|
||||
});
|
||||
|
||||
let request = client
|
||||
.post(format!("http://{address}{path}"))
|
||||
.timeout(Duration::from_secs(5))
|
||||
.header("x-request-id", "gateway-test")
|
||||
.json(&json!({"model": "unconfigured-model"}));
|
||||
let request = match authorization {
|
||||
Some(value) => request.header("authorization", value),
|
||||
None => request,
|
||||
};
|
||||
let response = request.send().await.unwrap();
|
||||
assert_eq!(response.status().as_u16(), status);
|
||||
if status == 400 {
|
||||
let expected = Error::UnknownModel("unconfigured-model".into());
|
||||
assert_eq!(
|
||||
response.json::<Value>().await.unwrap(),
|
||||
expected.body(Some("gateway-test"))
|
||||
);
|
||||
} else {
|
||||
let text = response.text().await.unwrap();
|
||||
assert!(!text.contains("gateway-key"));
|
||||
assert!(!text.contains("unconfigured-model"));
|
||||
}
|
||||
|
||||
shutdown.send(()).unwrap();
|
||||
timeout(Duration::from_secs(5), server)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
}
|
||||
13
litellm-rust/crates/router/Cargo.toml
Normal file
13
litellm-rust/crates/router/Cargo.toml
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
[package]
|
||||
name = "litellm-router"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-config.workspace = true
|
||||
litellm-core.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
5
litellm-rust/crates/router/README.md
Normal file
5
litellm-rust/crates/router/README.md
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
`litellm-router` scaffolds the model-list setup and deployment lookup portion of Python's `litellm.Router`. `Router::from_model_list(&config.model_list)` maps configured public names to provider deployments. Programmatic callers can collect `(String, Deployment)` entries into a `Router`
|
||||
|
||||
Lookup is exact and returns `None` for an unknown name. This extraction preserves the gateway's existing behavior: the last entry wins when public names repeat. Multiple deployments per model group, routing strategies, retries, cooldowns, and fallbacks are not implemented yet
|
||||
|
||||
The router owns deployment configuration and selection. The gateway handles HTTP errors and responses, while `core` executes provider calls and resolves credentials
|
||||
13
litellm-rust/crates/router/src/deployment.rs
Normal file
13
litellm-rust/crates/router/src/deployment.rs
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_core::messages::types::MessagesShaping;
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct Deployment {
|
||||
pub model: String,
|
||||
pub api_key: Option<String>,
|
||||
pub api_base: Option<String>,
|
||||
pub custom_llm_provider: Option<String>,
|
||||
pub timeout: Option<Duration>,
|
||||
pub shaping: MessagesShaping,
|
||||
}
|
||||
44
litellm-rust/crates/router/src/lib.rs
Normal file
44
litellm-rust/crates/router/src/lib.rs
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
mod deployment;
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use litellm_config::Model;
|
||||
|
||||
pub use deployment::Deployment;
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct Router(HashMap<String, Deployment>);
|
||||
|
||||
impl Router {
|
||||
pub fn from_model_list(model_list: &[Model]) -> Self {
|
||||
model_list
|
||||
.iter()
|
||||
.map(|model| {
|
||||
(
|
||||
model.model_name.clone(),
|
||||
Deployment {
|
||||
model: model.litellm_params.model.clone(),
|
||||
api_key: model
|
||||
.litellm_params
|
||||
.api_key
|
||||
.as_ref()
|
||||
.map(|value| value.expose().to_string()),
|
||||
api_base: model.litellm_params.api_base.clone(),
|
||||
custom_llm_provider: model.litellm_params.custom_llm_provider.clone(),
|
||||
..Deployment::default()
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn get(&self, model_name: &str) -> Option<&Deployment> {
|
||||
self.0.get(model_name)
|
||||
}
|
||||
}
|
||||
|
||||
impl FromIterator<(String, Deployment)> for Router {
|
||||
fn from_iter<I: IntoIterator<Item = (String, Deployment)>>(entries: I) -> Self {
|
||||
Self(entries.into_iter().collect())
|
||||
}
|
||||
}
|
||||
94
litellm-rust/crates/router/tests/router.rs
Normal file
94
litellm-rust/crates/router/tests/router.rs
Normal file
|
|
@ -0,0 +1,94 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_config::Config;
|
||||
use litellm_core::messages::types::MessagesShaping;
|
||||
use litellm_router::{Deployment, Router};
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
#[case::minimal("")]
|
||||
#[case::configured(
|
||||
"api_key: test-key\n api_base: https://provider.example/v1\n custom_llm_provider: test-provider"
|
||||
)]
|
||||
#[case::secret_reference("api_key: os.environ/ROUTER_TEST_API_KEY")]
|
||||
fn configuration_preserves_deployment_parameters(#[case] parameters: &str) {
|
||||
let config = Config::from_yaml(&format!(
|
||||
"model_list:\n - model_name: public-model\n litellm_params:\n model: provider/model\n {parameters}"
|
||||
))
|
||||
.unwrap();
|
||||
let router = Router::from_model_list(&config.model_list);
|
||||
let deployment = router.get(&config.model_list[0].model_name).unwrap();
|
||||
let params = &config.model_list[0].litellm_params;
|
||||
|
||||
assert_eq!(deployment.model, params.model);
|
||||
assert_eq!(
|
||||
deployment.api_key.as_deref(),
|
||||
params.api_key.as_ref().map(|key| key.expose())
|
||||
);
|
||||
assert_eq!(deployment.api_base, params.api_base);
|
||||
assert_eq!(deployment.custom_llm_provider, params.custom_llm_provider);
|
||||
assert_eq!(deployment.timeout, Deployment::default().timeout);
|
||||
assert_eq!(deployment.shaping, Deployment::default().shaping);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::first("public-a", Some("provider/a"))]
|
||||
#[case::second("public-b", Some("provider/b"))]
|
||||
#[case::unknown("missing", None)]
|
||||
#[case::provider_name_is_not_an_alias("provider/a", None)]
|
||||
#[case::case_sensitive("PUBLIC-A", None)]
|
||||
fn lookup_uses_public_names(#[case] name: &str, #[case] expected: Option<&str>) {
|
||||
let config = Config::from_yaml(
|
||||
"model_list:
|
||||
- model_name: public-a
|
||||
litellm_params:
|
||||
model: provider/a
|
||||
- model_name: public-b
|
||||
litellm_params:
|
||||
model: provider/b",
|
||||
)
|
||||
.unwrap();
|
||||
let router = Router::from_model_list(&config.model_list);
|
||||
|
||||
assert_eq!(router.get(name).map(|entry| entry.model.as_str()), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn empty_configuration_has_no_deployment() {
|
||||
let config = Config::from_yaml("model_list: []").unwrap();
|
||||
|
||||
assert!(
|
||||
Router::from_model_list(&config.model_list)
|
||||
.get("")
|
||||
.is_none()
|
||||
);
|
||||
assert!(Router::default().get("unknown").is_none());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn programmatic_deployments_preserve_overrides_and_last_entry_wins() {
|
||||
let deployment = Deployment {
|
||||
model: "provider/selected".into(),
|
||||
api_key: Some("test-key".into()),
|
||||
api_base: Some("https://provider.example/v1".into()),
|
||||
custom_llm_provider: Some("test-provider".into()),
|
||||
timeout: Some(Duration::from_secs(7)),
|
||||
shaping: MessagesShaping {
|
||||
drop_params: true,
|
||||
additional_drop_params: vec!["metadata.test".into()],
|
||||
..Default::default()
|
||||
},
|
||||
};
|
||||
let router = Router::from_iter([
|
||||
("public-model".into(), Deployment::default()),
|
||||
("public-model".into(), deployment.clone()),
|
||||
]);
|
||||
let selected = router.get("public-model").unwrap();
|
||||
|
||||
assert_eq!(selected.model, deployment.model);
|
||||
assert_eq!(selected.api_key, deployment.api_key);
|
||||
assert_eq!(selected.api_base, deployment.api_base);
|
||||
assert_eq!(selected.custom_llm_provider, deployment.custom_llm_provider);
|
||||
assert_eq!(selected.timeout, deployment.timeout);
|
||||
assert_eq!(selected.shaping, deployment.shaping);
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue