From affb5475256af1c8035a10a77e7f66bd222f66de Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 23:12:48 -0700 Subject: [PATCH] feat(rust): add config router and gateway crates (#43289) Co-authored-by: Yujong Lee --- litellm-rust/Cargo.lock | 161 ++++++++++++- litellm-rust/Cargo.toml | 6 + litellm-rust/crates/config/Cargo.toml | 16 ++ litellm-rust/crates/config/src/error.rs | 7 + litellm-rust/crates/config/src/lib.rs | 48 ++++ litellm-rust/crates/config/tests/config.rs | 119 +++++++++ litellm-rust/crates/gateway-auth/Cargo.toml | 21 ++ litellm-rust/crates/gateway-auth/src/error.rs | 24 ++ litellm-rust/crates/gateway-auth/src/lib.rs | 74 ++++++ .../crates/gateway-auth/tests/auth.rs | 100 ++++++++ .../crates/gateway-inference/AGENTS.md | 5 + .../crates/gateway-inference/Cargo.toml | 29 +++ .../src/audio_transcription.rs | 59 +++++ .../gateway-inference/src/chat_completions.rs | 83 +++++++ .../crates/gateway-inference/src/error.rs | 225 ++++++++++++++++++ .../crates/gateway-inference/src/lib.rs | 59 +++++ .../gateway-inference/src/messages/host.rs | 69 ++++++ .../gateway-inference/src/messages/mod.rs | 137 +++++++++++ .../crates/gateway-inference/src/ocr.rs | 77 ++++++ .../crates/gateway-inference/src/request.rs | 108 +++++++++ .../gateway-inference/tests/messages.rs | 66 +++++ .../crates/gateway-inference/tests/ocr.rs | 118 +++++++++ .../crates/gateway-inference/tests/routes.rs | 92 +++++++ .../gateway-inference/tests/support/mod.rs | 75 ++++++ litellm-rust/crates/gateway/AGENTS.md | 5 + litellm-rust/crates/gateway/Cargo.toml | 24 ++ litellm-rust/crates/gateway/src/lib.rs | 52 ++++ litellm-rust/crates/gateway/src/main.rs | 18 ++ litellm-rust/crates/gateway/tests/server.rs | 90 +++++++ litellm-rust/crates/router/Cargo.toml | 13 + litellm-rust/crates/router/README.md | 5 + litellm-rust/crates/router/src/deployment.rs | 13 + litellm-rust/crates/router/src/lib.rs | 44 ++++ litellm-rust/crates/router/tests/router.rs | 94 ++++++++ 34 files changed, 2133 insertions(+), 3 deletions(-) create mode 100644 litellm-rust/crates/config/Cargo.toml create mode 100644 litellm-rust/crates/config/src/error.rs create mode 100644 litellm-rust/crates/config/src/lib.rs create mode 100644 litellm-rust/crates/config/tests/config.rs create mode 100644 litellm-rust/crates/gateway-auth/Cargo.toml create mode 100644 litellm-rust/crates/gateway-auth/src/error.rs create mode 100644 litellm-rust/crates/gateway-auth/src/lib.rs create mode 100644 litellm-rust/crates/gateway-auth/tests/auth.rs create mode 100644 litellm-rust/crates/gateway-inference/AGENTS.md create mode 100644 litellm-rust/crates/gateway-inference/Cargo.toml create mode 100644 litellm-rust/crates/gateway-inference/src/audio_transcription.rs create mode 100644 litellm-rust/crates/gateway-inference/src/chat_completions.rs create mode 100644 litellm-rust/crates/gateway-inference/src/error.rs create mode 100644 litellm-rust/crates/gateway-inference/src/lib.rs create mode 100644 litellm-rust/crates/gateway-inference/src/messages/host.rs create mode 100644 litellm-rust/crates/gateway-inference/src/messages/mod.rs create mode 100644 litellm-rust/crates/gateway-inference/src/ocr.rs create mode 100644 litellm-rust/crates/gateway-inference/src/request.rs create mode 100644 litellm-rust/crates/gateway-inference/tests/messages.rs create mode 100644 litellm-rust/crates/gateway-inference/tests/ocr.rs create mode 100644 litellm-rust/crates/gateway-inference/tests/routes.rs create mode 100644 litellm-rust/crates/gateway-inference/tests/support/mod.rs create mode 100644 litellm-rust/crates/gateway/AGENTS.md create mode 100644 litellm-rust/crates/gateway/Cargo.toml create mode 100644 litellm-rust/crates/gateway/src/lib.rs create mode 100644 litellm-rust/crates/gateway/src/main.rs create mode 100644 litellm-rust/crates/gateway/tests/server.rs create mode 100644 litellm-rust/crates/router/Cargo.toml create mode 100644 litellm-rust/crates/router/README.md create mode 100644 litellm-rust/crates/router/src/deployment.rs create mode 100644 litellm-rust/crates/router/src/lib.rs create mode 100644 litellm-rust/crates/router/tests/router.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 580b427a97e..b9d363e72ca 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -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" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 442bd620e05..ed703396c22 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -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 } diff --git a/litellm-rust/crates/config/Cargo.toml b/litellm-rust/crates/config/Cargo.toml new file mode 100644 index 00000000000..36bd68fe2a0 --- /dev/null +++ b/litellm-rust/crates/config/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/config/src/error.rs b/litellm-rust/crates/config/src/error.rs new file mode 100644 index 00000000000..61d19491abc --- /dev/null +++ b/litellm-rust/crates/config/src/error.rs @@ -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), +} diff --git a/litellm-rust/crates/config/src/lib.rs b/litellm-rust/crates/config/src/lib.rs new file mode 100644 index 00000000000..8e86e345025 --- /dev/null +++ b/litellm-rust/crates/config/src/lib.rs @@ -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, +} + +#[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, + pub api_base: Option, + pub custom_llm_provider: Option, +} + +impl Config { + pub fn from_yaml(yaml: &str) -> Result { + Ok(serde_yaml_ng::from_str(yaml)?) + } + + pub fn load(path: impl AsRef) -> Result { + Self::from_yaml(&std::fs::read_to_string(path)?) + } +} diff --git a/litellm-rust/crates/config/tests/config.rs b/litellm-rust/crates/config/tests/config.rs new file mode 100644 index 00000000000..ce6d684ec72 --- /dev/null +++ b/litellm-rust/crates/config/tests/config.rs @@ -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()); +} diff --git a/litellm-rust/crates/gateway-auth/Cargo.toml b/litellm-rust/crates/gateway-auth/Cargo.toml new file mode 100644 index 00000000000..340f7224618 --- /dev/null +++ b/litellm-rust/crates/gateway-auth/Cargo.toml @@ -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"] } diff --git a/litellm-rust/crates/gateway-auth/src/error.rs b/litellm-rust/crates/gateway-auth/src/error.rs new file mode 100644 index 00000000000..735eb7741ad --- /dev/null +++ b/litellm-rust/crates/gateway-auth/src/error.rs @@ -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() + } +} diff --git a/litellm-rust/crates/gateway-auth/src/lib.rs b/litellm-rust/crates/gateway-auth/src/lib.rs new file mode 100644 index 00000000000..7a1fb244654 --- /dev/null +++ b/litellm-rust/crates/gateway-auth/src/lib.rs @@ -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, + secrets: Arc, +} + +impl Auth { + pub fn from_config(config: &Config, secrets: Arc) -> Self { + Self { + master_key: config.general_settings.master_key.clone(), + secrets, + } + } + + async fn master_key(&self) -> Result { + 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 for RequireMasterKey { + type Rejection = Error; + + async fn from_request_parts(parts: &mut Parts, state: &Auth) -> Result { + 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), + } + } +} diff --git a/litellm-rust/crates/gateway-auth/tests/auth.rs b/litellm-rust/crates/gateway-auth/tests/auth.rs new file mode 100644 index 00000000000..58625296664 --- /dev/null +++ b/litellm-rust/crates/gateway-auth/tests/auth.rs @@ -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, 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 { + 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, + #[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::( + 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" + ); +} diff --git a/litellm-rust/crates/gateway-inference/AGENTS.md b/litellm-rust/crates/gateway-inference/AGENTS.md new file mode 100644 index 00000000000..7dc57380083 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/AGENTS.md @@ -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 diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml new file mode 100644 index 00000000000..e3f40ec5354 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -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" diff --git a/litellm-rust/crates/gateway-inference/src/audio_transcription.rs b/litellm-rust/crates/gateway-inference/src/audio_transcription.rs new file mode 100644 index 00000000000..d5fd6603e20 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/audio_transcription.rs @@ -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>, 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 { + 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?) +} diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs new file mode 100644 index 00000000000..c386f38fe06 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -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>, body: Bytes) -> Response { + respond(&gateway, request::object(&body)).await +} + +pub(crate) async fn deployment( + State(gateway): State>, + Path(path): Path, + 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, 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) -> Result { + 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()) +} diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs new file mode 100644 index 00000000000..1d40857d962 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -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::(&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 { + 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) -> 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) -> 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, request_id: Option<&str>) -> Map { + 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); + } +} diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs new file mode 100644 index 00000000000..eebe3f34a09 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -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, + pub models: ModelList, + pub ocr: OcrClient, +} + +pub fn router(gateway: Arc) -> 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) +} diff --git a/litellm-rust/crates/gateway-inference/src/messages/host.rs b/litellm-rust/crates/gateway-inference/src/messages/host.rs new file mode 100644 index 00000000000..57486d9d157 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/messages/host.rs @@ -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>>, + pub(super) chunks: mpsc::Sender, +} + +impl ChannelHost { + pub(super) fn new( + call: MessagesCall, + head: oneshot::Sender, + chunks: mpsc::Sender, + ) -> Self { + Self { + local: LocalMessagesHost::new(call), + head: Mutex::new(Some(head)), + chunks, + } + } + + fn take_head(&self) -> Option> { + 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 for ChannelHost { + async fn project(&self) -> Result { + self.local.project().await + } + + async fn custom_op(&self, op: Infallible) -> Result<(), Error> { + match op {} + } + + async fn open(&self, head: MessagesStreamHead) -> Result { + 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 { + Ok(match self.chunks.send(chunk).await { + Ok(()) => Demand::More, + Err(_) => Demand::Detached, + }) + } +} diff --git a/litellm-rust/crates/gateway-inference/src/messages/mod.rs b/litellm-rust/crates/gateway-inference/src/messages/mod.rs new file mode 100644 index 00000000000..e96cae8ba53 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/messages/mod.rs @@ -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>, + 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 { + 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, + headers: &HeaderMap, +) -> Result { + 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 { + let extra_headers: Map = 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) -> 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() +} diff --git a/litellm-rust/crates/gateway-inference/src/ocr.rs b/litellm-rust/crates/gateway-inference/src/ocr.rs new file mode 100644 index 00000000000..d666223e037 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/ocr.rs @@ -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>, 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 { + 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()), + } +} diff --git a/litellm-rust/crates/gateway-inference/src/request.rs b/litellm-rust/crates/gateway-inference/src/request.rs new file mode 100644 index 00000000000..f58c7b3ed79 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/request.rs @@ -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, + pub mime_type: Option, +} + +pub(crate) fn object(body: &[u8]) -> Result, 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, +) -> 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, Option), 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() +} diff --git a/litellm-rust/crates/gateway-inference/tests/messages.rs b/litellm-rust/crates/gateway-inference/tests/messages.rs new file mode 100644 index 00000000000..ef7e7c66681 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/messages.rs @@ -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"); +} diff --git a/litellm-rust/crates/gateway-inference/tests/ocr.rs b/litellm-rust/crates/gateway-inference/tests/ocr.rs new file mode 100644 index 00000000000..3d5a2d22b09 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/ocr.rs @@ -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()); +} diff --git a/litellm-rust/crates/gateway-inference/tests/routes.rs b/litellm-rust/crates/gateway-inference/tests/routes.rs new file mode 100644 index 00000000000..5d3b06ccadb --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/routes.rs @@ -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") + ); +} diff --git a/litellm-rust/crates/gateway-inference/tests/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs new file mode 100644 index 00000000000..d56489d28cd --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -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, 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() +} diff --git a/litellm-rust/crates/gateway/AGENTS.md b/litellm-rust/crates/gateway/AGENTS.md new file mode 100644 index 00000000000..10a42b9ec3b --- /dev/null +++ b/litellm-rust/crates/gateway/AGENTS.md @@ -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 diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml new file mode 100644 index 00000000000..554186955b4 --- /dev/null +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -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"] } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs new file mode 100644 index 00000000000..3f67f923a5c --- /dev/null +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -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, 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, 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)), + ) +} diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs new file mode 100644 index 00000000000..bae711e16c7 --- /dev/null +++ b/litellm-rust/crates/gateway/src/main.rs @@ -0,0 +1,18 @@ +use std::error::Error; + +use litellm_config::Config; + +#[tokio::main] +async fn main() -> Result<(), Box> { + 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::()?; + let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?; + + axum::serve(listener, litellm_gateway::router(inference, &config)).await?; + Ok(()) +} diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs new file mode 100644 index 00000000000..a19d8c9c4fe --- /dev/null +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -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 { + 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, + #[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::().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(); +} diff --git a/litellm-rust/crates/router/Cargo.toml b/litellm-rust/crates/router/Cargo.toml new file mode 100644 index 00000000000..cc6972a066e --- /dev/null +++ b/litellm-rust/crates/router/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/router/README.md b/litellm-rust/crates/router/README.md new file mode 100644 index 00000000000..ed5b2966fe8 --- /dev/null +++ b/litellm-rust/crates/router/README.md @@ -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 diff --git a/litellm-rust/crates/router/src/deployment.rs b/litellm-rust/crates/router/src/deployment.rs new file mode 100644 index 00000000000..4904bf4ffdd --- /dev/null +++ b/litellm-rust/crates/router/src/deployment.rs @@ -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, + pub api_base: Option, + pub custom_llm_provider: Option, + pub timeout: Option, + pub shaping: MessagesShaping, +} diff --git a/litellm-rust/crates/router/src/lib.rs b/litellm-rust/crates/router/src/lib.rs new file mode 100644 index 00000000000..da33bfb04bd --- /dev/null +++ b/litellm-rust/crates/router/src/lib.rs @@ -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); + +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>(entries: I) -> Self { + Self(entries.into_iter().collect()) + } +} diff --git a/litellm-rust/crates/router/tests/router.rs b/litellm-rust/crates/router/tests/router.rs new file mode 100644 index 00000000000..3cbd316cab0 --- /dev/null +++ b/litellm-rust/crates/router/tests/router.rs @@ -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); +}