diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 98f36607fff..4f394567c52 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -209,9 +209,9 @@ dependencies = [ [[package]] name = "async-trait" -version = "0.1.91" +version = "0.1.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae36dc4177970ef04fde5178d3e2429882def40e57a451f919c098f72baa6cec" +checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667" dependencies = [ "proc-macro2", "quote", @@ -2109,6 +2109,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fraction" +version = "0.15.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e076045bb43dac435333ed5f04caf35c7463631d0dae2deb2638d94dd0a5b872" +dependencies = [ + "lazy_static", + "num", +] + [[package]] name = "fraction" version = "0.17.0" @@ -3236,6 +3246,45 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonptr" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72c22d99362c7aac7d77efab6b49c2181674b165d7e6e7c755b3432ded64bf51" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "jsonschema" +version = "0.49.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ec8a241beed129f06114aa68007e905ca350e7baeb6e17a7631bb7978d91b2" +dependencies = [ + "ahash", + "bytecount", + "data-encoding", + "email_address", + "fancy-regex 0.19.2", + "fraction 0.15.4", + "getrandom 0.3.4", + "idna", + "itoa", + "jsonschema-regex 0.49.9", + "jsonschema-value 0.49.9", + "num-cmp", + "num-traits", + "percent-encoding", + "referencing 0.49.9", + "regex", + "serde", + "serde_json", + "strum", + "unicode-general-category", + "uuid-simd", +] + [[package]] name = "jsonschema" version = "0.55.1" @@ -3247,15 +3296,15 @@ dependencies = [ "data-encoding", "email_address", "fancy-regex 0.19.2", - "fraction", + "fraction 0.17.0", "getrandom 0.3.4", "itoa", - "jsonschema-regex", - "jsonschema-value", + "jsonschema-regex 0.55.1", + "jsonschema-value 0.55.1", "num-cmp", "num-traits", "percent-encoding", - "referencing", + "referencing 0.55.1", "regex", "serde", "serde_json", @@ -3264,6 +3313,15 @@ dependencies = [ "uuid-simd", ] +[[package]] +name = "jsonschema-regex" +version = "0.49.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91994f45017ed5e66aa8e59b8415f4cb033a6380d7200387b7cf117595fbdf85" +dependencies = [ + "regex-syntax", +] + [[package]] name = "jsonschema-regex" version = "0.55.1" @@ -3273,6 +3331,20 @@ dependencies = [ "regex-syntax", ] +[[package]] +name = "jsonschema-value" +version = "0.49.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ec7637f83e510868ae6ed625f7ebfbbde4554ee8ce49854caa5126a8b9b9ecb" +dependencies = [ + "ahash", + "bytecount", + "fraction 0.15.4", + "num-cmp", + "num-traits", + "serde_json", +] + [[package]] name = "jsonschema-value" version = "0.55.1" @@ -3281,7 +3353,7 @@ checksum = "0230ac05e09c6111e96c147b75c390579f5cbd45654b980c68ac60fe17b3f129" dependencies = [ "ahash", "bytecount", - "fraction", + "fraction 0.17.0", "getrandom 0.3.4", "num-cmp", "num-traits", @@ -3784,6 +3856,7 @@ dependencies = [ name = "litellm-gateway-inference" version = "0.1.0" dependencies = [ + "async-trait", "axum", "base64 0.22.1", "bytes", @@ -3791,6 +3864,7 @@ dependencies = [ "litellm-auth", "litellm-cache-memory", "litellm-cache-response", + "litellm-config", "litellm-gateway-auth", "litellm-host", "litellm-host-http", @@ -3808,6 +3882,11 @@ dependencies = [ "rstest", "serde", "serde_json", + "switchyard-libsy", + "switchyard-llm-client", + "switchyard-protocol", + "switchyard-runner", + "switchyard-translation", "thiserror 2.0.19", "tokio", "tower", @@ -4195,7 +4274,7 @@ name = "litellm-model-catalog" version = "0.1.0" dependencies = [ "indexmap 2.14.0", - "jsonschema", + "jsonschema 0.55.1", "litellm-llms-types", "rstest", "schemars 1.2.2", @@ -4508,7 +4587,7 @@ dependencies = [ "tempfile", "thiserror 2.0.19", "tokio", - "toml", + "toml 0.9.12+spec-1.1.0", "zip", ] @@ -5811,6 +5890,23 @@ dependencies = [ "syn 3.0.6", ] +[[package]] +name = "referencing" +version = "0.49.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6efa2154ea6f5ce0fdecdd2a8d18f2fa1a39a8fbba91564f555a592e4dce8278" +dependencies = [ + "ahash", + "fluent-uri", + "getrandom 0.3.4", + "hashbrown 0.17.1", + "itoa", + "micromap", + "parking_lot", + "percent-encoding", + "serde_json", +] + [[package]] name = "referencing" version = "0.55.1" @@ -7032,6 +7128,100 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" +[[package]] +name = "switchyard-libsy" +version = "0.3.0" +source = "git+https://github.com/NVIDIA-NeMo/Switchyard.git?rev=4bdf5719ce1c43902df92ef64407b858d1f99394#4bdf5719ce1c43902df92ef64407b858d1f99394" +dependencies = [ + "async-trait", + "futures", + "http 1.4.2", + "jsonptr", + "jsonschema 0.49.9", + "opentelemetry", + "parking_lot", + "rand 0.10.2", + "regex", + "serde", + "serde_json", + "shlex", + "switchyard-protocol", + "thiserror 2.0.19", + "tokio", + "tokio-stream", + "tracing", + "tracing-opentelemetry", + "uuid", +] + +[[package]] +name = "switchyard-llm-client" +version = "0.3.0" +source = "git+https://github.com/NVIDIA-NeMo/Switchyard.git?rev=4bdf5719ce1c43902df92ef64407b858d1f99394#4bdf5719ce1c43902df92ef64407b858d1f99394" +dependencies = [ + "async-trait", + "futures", + "futures-util", + "http 1.4.2", + "httpdate", + "opentelemetry", + "parking_lot", + "reqwest 0.13.5", + "serde", + "serde_json", + "switchyard-libsy", + "switchyard-protocol", + "switchyard-translation", + "tokio", + "tracing", + "tracing-opentelemetry", +] + +[[package]] +name = "switchyard-protocol" +version = "0.3.0" +source = "git+https://github.com/NVIDIA-NeMo/Switchyard.git?rev=4bdf5719ce1c43902df92ef64407b858d1f99394#4bdf5719ce1c43902df92ef64407b858d1f99394" +dependencies = [ + "async-trait", + "futures", + "http 1.4.2", + "serde", + "serde_json", + "thiserror 2.0.19", +] + +[[package]] +name = "switchyard-runner" +version = "0.3.0" +source = "git+https://github.com/NVIDIA-NeMo/Switchyard.git?rev=4bdf5719ce1c43902df92ef64407b858d1f99394#4bdf5719ce1c43902df92ef64407b858d1f99394" +dependencies = [ + "reqwest 0.13.5", + "serde", + "serde_json", + "strum_macros", + "switchyard-libsy", + "switchyard-llm-client", + "switchyard-protocol", + "thiserror 2.0.19", + "toml 1.1.6+spec-1.1.0", + "tracing", +] + +[[package]] +name = "switchyard-translation" +version = "0.3.0" +source = "git+https://github.com/NVIDIA-NeMo/Switchyard.git?rev=4bdf5719ce1c43902df92ef64407b858d1f99394#4bdf5719ce1c43902df92ef64407b858d1f99394" +dependencies = [ + "async-stream", + "base64 0.22.1", + "futures", + "serde", + "serde_json", + "switchyard-protocol", + "thiserror 2.0.19", + "tracing", +] + [[package]] name = "syn" version = "1.0.109" @@ -7345,6 +7535,7 @@ dependencies = [ "bytes", "libc", "mio", + "parking_lot", "pin-project-lite", "signal-hook-registry", "socket2 0.6.5", @@ -7439,6 +7630,21 @@ dependencies = [ "winnow 0.7.15", ] +[[package]] +name = "toml" +version = "1.1.6+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "920602543f0911ab71da12c50d59701da54c196d1a2bf5cb4b75667f137a406a" +dependencies = [ + "indexmap 2.14.0", + "serde_core", + "serde_spanned", + "toml_datetime 1.1.1+spec-1.1.0", + "toml_parser", + "toml_writer", + "winnow 1.0.4", +] + [[package]] name = "toml_datetime" version = "0.7.5+spec-1.1.0" @@ -7708,6 +7914,17 @@ dependencies = [ "tracing", ] +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + [[package]] name = "tracing-opentelemetry" version = "0.33.0" @@ -7716,8 +7933,10 @@ checksum = "adbc64cba7137545b8044cb1fe9814f7aacf3c6b5f9b45be8bb5db538befdb26" dependencies = [ "js-sys", "opentelemetry", + "smallvec", "tracing", "tracing-core", + "tracing-log", "tracing-subscriber", "web-time", ] diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index 312d44b9a4b..8c6cb9efccf 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -5,11 +5,25 @@ edition.workspace = true license.workspace = true repository.workspace = true +[features] +switchyard = [ + "dep:async-trait", + "dep:futures-util", + "dep:litellm-config", + "dep:switchyard-libsy", + "dep:switchyard-llm-client", + "dep:switchyard-protocol", + "dep:switchyard-runner", + "dep:switchyard-translation", +] + [dependencies] +async-trait = { version = "0.1.92", optional = true } litellm-cache-response.workspace = true axum = { workspace = true, features = ["json", "multipart", "original-uri"] } base64.workspace = true bytes.workspace = true +futures-util = { workspace = true, optional = true } litellm-auth.workspace = true litellm-gateway-auth.workspace = true litellm-inference.workspace = true @@ -18,6 +32,7 @@ litellm-inference-responses.workspace = true litellm-inference-messages.workspace = true litellm-inference-chat.workspace = true litellm-inference-ocr.workspace = true +litellm-config = { workspace = true, optional = true } litellm-host-http.workspace = true litellm-host.workspace = true litellm-http.workspace = true @@ -27,6 +42,11 @@ litellm-secrets.workspace = true litellm-llms-types.workspace = true serde.workspace = true serde_json.workspace = true +switchyard-libsy = { git = "https://github.com/NVIDIA-NeMo/Switchyard.git", rev = "4bdf5719ce1c43902df92ef64407b858d1f99394", optional = true } +switchyard-llm-client = { git = "https://github.com/NVIDIA-NeMo/Switchyard.git", rev = "4bdf5719ce1c43902df92ef64407b858d1f99394", optional = true } +switchyard-protocol = { git = "https://github.com/NVIDIA-NeMo/Switchyard.git", rev = "4bdf5719ce1c43902df92ef64407b858d1f99394", optional = true } +switchyard-runner = { git = "https://github.com/NVIDIA-NeMo/Switchyard.git", rev = "4bdf5719ce1c43902df92ef64407b858d1f99394", optional = true } +switchyard-translation = { git = "https://github.com/NVIDIA-NeMo/Switchyard.git", rev = "4bdf5719ce1c43902df92ef64407b858d1f99394", optional = true } thiserror.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/gateway-inference/src/caching.rs b/litellm-rust/crates/gateway-inference/src/caching.rs index 1abea61b164..ddf0a6ca79e 100644 --- a/litellm-rust/crates/gateway-inference/src/caching.rs +++ b/litellm-rust/crates/gateway-inference/src/caching.rs @@ -101,11 +101,15 @@ impl litellm_host::interceptors::Interceptors for impl CacheHeaders { pub(crate) fn apply(&self, mut response: axum::response::Response) -> axum::response::Response { + self.apply_to(response.headers_mut()); + response + } + + pub(crate) fn apply_to(&self, headers: &mut axum::http::HeaderMap) { if let Some(key) = self.0.get() && let Ok(value) = axum::http::HeaderValue::from_str(key) { - response.headers_mut().insert("x-litellm-cache-key", value); + headers.insert("x-litellm-cache-key", value); } - response } } diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index 742a373b286..c920b874786 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -1,27 +1,40 @@ use litellm_gateway_auth::AuthenticatedRequest; use std::sync::Arc; +#[cfg(feature = "switchyard")] +use axum::http::HeaderMap; use axum::{ Json, extract::{Path, State}, response::{IntoResponse, Response}, }; -use litellm_inference_chat::types::ChatCompletionsCall; +use litellm_inference_chat::types::{ChatCompletionsCall, ChatCompletionsRequest}; use serde_json::{Map, Value}; +#[cfg(feature = "switchyard")] +use switchyard_protocol::WireFormat; -use crate::{Error, Gateway, JsonObject, request}; +use crate::{Deployment, Error, Gateway, JsonObject, request}; pub(crate) async fn create( State(gateway): State>, identity: AuthenticatedRequest, + #[cfg(feature = "switchyard")] headers: HeaderMap, JsonObject(body): JsonObject, ) -> Result { - handle(&gateway, &identity, body).await + handle( + &gateway, + &identity, + #[cfg(feature = "switchyard")] + headers, + body, + ) + .await } pub(crate) async fn create_from_model_path( State(gateway): State>, identity: AuthenticatedRequest, + #[cfg(feature = "switchyard")] headers: HeaderMap, Path(model): Path, JsonObject(body): JsonObject, ) -> Result { @@ -32,17 +45,42 @@ pub(crate) async fn create_from_model_path( .collect(), Some(_) => body, }; - handle(&gateway, &identity, body).await + handle( + &gateway, + &identity, + #[cfg(feature = "switchyard")] + headers, + body, + ) + .await } async fn handle( gateway: &Gateway, identity: &AuthenticatedRequest, + #[cfg(feature = "switchyard")] inbound_headers: HeaderMap, body: Map, ) -> Result { let deployment = request::resolve_deployment(gateway, &body)?; request::authorize_model(identity, deployment, &body).await?; + #[cfg(feature = "switchyard")] + let switchyard_route = body + .get("model") + .and_then(Value::as_str) + .and_then(|model| gateway.switchyard_routes.get(model)); let (body, cache_options) = crate::caching::prepare(identity, body)?; + #[cfg(feature = "switchyard")] + if let Some(switchyard_route) = switchyard_route { + return switchyard_route + .execute( + gateway, + cache_options, + WireFormat::OpenAiChat, + body, + inbound_headers, + ) + .await; + } let route = gateway.chat_completions.clone(); let route = match &gateway.cache { Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new( @@ -51,24 +89,10 @@ async fn handle( )), None => route, }; - - let messages = body.get("messages").cloned().unwrap_or_default(); let headers = crate::caching::CacheHeaders::default(); let response = litellm_host_http::serve_unary( route.machine( - ChatCompletionsCall { - model: deployment.model.clone(), - messages, - optional_params: body - .into_iter() - .filter(|(name, _)| !matches!(name.as_str(), "model" | "messages")) - .collect(), - api_key: deployment.api_key.clone(), - api_base: deployment.api_base.clone(), - custom_llm_provider: deployment.custom_llm_provider.clone(), - extra_headers: None, - timeout: deployment.timeout, - }, + ChatCompletionsCall::from(provider_request(deployment, body)), cache_options.policy, ), (), @@ -79,3 +103,21 @@ async fn handle( .await?; Ok(headers.apply(response)) } + +pub(crate) fn provider_request( + deployment: &Deployment, + mut body: Map, +) -> ChatCompletionsRequest<'_> { + let messages = body.remove("messages").unwrap_or_default(); + body.remove("model"); + ChatCompletionsRequest { + model: &deployment.model, + messages, + optional_params: body, + 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, + } +} diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index 39656b933eb..46110612fff 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -8,6 +8,15 @@ use litellm_inference::RouteError; use litellm_llms::base_llm::ocr::error::Error as OcrError; use serde_json::{Map, Value, json}; +/// Invalid configuration for a Switchyard virtual model. +#[cfg(feature = "switchyard")] +#[derive(Debug, thiserror::Error)] +#[error("invalid Switchyard model {model}: {reason}")] +pub struct SwitchyardConfigError { + pub(crate) model: String, + pub(crate) reason: String, +} + #[derive(Debug, thiserror::Error)] pub enum Error { #[error(transparent)] @@ -76,21 +85,21 @@ impl Error { } pub fn openai_response(self) -> Response { + (self.status(), Json(self.openai_body())).into_response() + } + + pub(crate) fn openai_body(&self) -> Value { let status = self.status(); - let message = match &self { + 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() + json!({"error": { + "message": message, + "type": error_type(status), + "param": null, + "code": status.as_u16(), + }}) } /// The Anthropic error envelope Python's `AnthropicExceptionMapping` builds: an upstream diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index 4663151f1de..06d8aa5949d 100644 --- a/litellm-rust/crates/gateway-inference/src/lib.rs +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -11,6 +11,8 @@ pub mod messages; mod ocr; mod request; mod responses; +#[cfg(feature = "switchyard")] +mod switchyard; use std::sync::Arc; @@ -26,11 +28,15 @@ use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings}; use litellm_secrets::source::SecretSource; pub use error::Error; +#[cfg(feature = "switchyard")] +pub use error::SwitchyardConfigError; pub use litellm_router::{Deployment, Router as ModelRouter}; pub use request::{JsonObject, RequestId}; pub struct Gateway { cache: Option>, + #[cfg(feature = "switchyard")] + switchyard_routes: switchyard::Routes, pub audio_transcription: AudioTranscriptionRoute, pub chat_completions: ChatCompletionsRoute, pub messages: MessagesRoute, @@ -60,6 +66,8 @@ impl Gateway { let auth = resources.auth.clone(); Ok(Self { cache: None, + #[cfg(feature = "switchyard")] + switchyard_routes: switchyard::Routes::default(), audio_transcription: AudioTranscriptionRoute::new( provider.clone(), auth.clone(), @@ -86,6 +94,16 @@ impl Gateway { http, }) } + + /// Configure Switchyard virtual models from LiteLLM's model list. + #[cfg(feature = "switchyard")] + pub fn with_switchyard_models( + mut self, + models: &[litellm_config::Model], + ) -> Result { + self.switchyard_routes = switchyard::Routes::from_models(models, &self.models)?; + Ok(self) + } } pub fn router(gateway: Arc) -> Router { diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index d612a042e03..8e9850d769a 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -14,6 +14,8 @@ use litellm_host_http::Sse; use litellm_inference_messages::{MessagesCall, messages_body, route::Messages}; use litellm_llms_types::headers::{ProviderSpecificHeader, ProviderSpecificHeaders}; use serde_json::{Map, Value}; +#[cfg(feature = "switchyard")] +use switchyard_protocol::WireFormat; use crate::{Deployment, Error, Gateway, JsonObject, RequestId, request}; @@ -34,7 +36,7 @@ pub async fn create( body: Result, ) -> impl IntoResponse { let result = match body { - Ok(JsonObject(body)) => handle(&gateway, &identity, &headers, body).await, + Ok(JsonObject(body)) => handle(&gateway, &identity, headers, body).await, Err(error) => Err(error), }; result.map_err(|error| (error.status(), Json(error.body(request_id.as_deref())))) @@ -43,12 +45,29 @@ pub async fn create( async fn handle( gateway: &Gateway, identity: &AuthenticatedRequest, - headers: &HeaderMap, + headers: HeaderMap, body: Map, ) -> Result { let deployment = request::resolve_deployment(gateway, &body)?; request::authorize_model(identity, deployment, &body).await?; + #[cfg(feature = "switchyard")] + let switchyard_route = body + .get("model") + .and_then(Value::as_str) + .and_then(|model| gateway.switchyard_routes.get(model)); let (body, cache_options) = crate::caching::prepare(identity, body)?; + #[cfg(feature = "switchyard")] + if let Some(switchyard_route) = switchyard_route { + return switchyard_route + .execute( + gateway, + cache_options, + WireFormat::AnthropicMessages, + body, + headers, + ) + .await; + } let route = gateway.messages.clone(); let route = match &gateway.cache { Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new( @@ -58,7 +77,7 @@ async fn handle( None => route, }; - let call = project(deployment, body, headers)?; + let call = provider_request(deployment, body, &headers)?; let machine = route.machine(call, cache_options.policy); let stream = Sse::::new(Json, |error| Bytes::from(Error::from(error).sse_frame())); @@ -67,18 +86,14 @@ async fn handle( Ok(headers.apply(response)) } -fn project( +pub(crate) fn provider_request( deployment: &Deployment, - body: Map, + mut 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(); +) -> Result { + if let Some(model) = body.get_mut("model") { + *model = Value::from(deployment.model.as_str()); + } Ok(MessagesCall { body: messages_body(body)?, api_key: deployment.api_key.clone(), diff --git a/litellm-rust/crates/gateway-inference/src/responses.rs b/litellm-rust/crates/gateway-inference/src/responses.rs index 77bebf610e1..0379d4fe584 100644 --- a/litellm-rust/crates/gateway-inference/src/responses.rs +++ b/litellm-rust/crates/gateway-inference/src/responses.rs @@ -1,21 +1,43 @@ use std::sync::Arc; +#[cfg(feature = "switchyard")] +use axum::http::HeaderMap; use axum::{Json, body::Bytes, extract::State, response::Response}; use litellm_gateway_auth::AuthenticatedRequest; use litellm_host_http::Sse; use litellm_inference_responses::{route::Responses, types::ResponsesCall}; -use serde_json::json; +use serde_json::{Map, Value, json}; +#[cfg(feature = "switchyard")] +use switchyard_protocol::WireFormat; -use crate::{Error, Gateway, JsonObject, request}; +use crate::{Deployment, Error, Gateway, JsonObject, request}; pub(crate) async fn create( State(gateway): State>, identity: AuthenticatedRequest, + #[cfg(feature = "switchyard")] headers: HeaderMap, JsonObject(body): JsonObject, ) -> Result { let deployment = request::resolve_deployment(&gateway, &body)?; request::authorize_model(&identity, deployment, &body).await?; + #[cfg(feature = "switchyard")] + let switchyard_route = body + .get("model") + .and_then(Value::as_str) + .and_then(|model| gateway.switchyard_routes.get(model)); let (body, cache_options) = crate::caching::prepare(&identity, body)?; + #[cfg(feature = "switchyard")] + if let Some(switchyard_route) = switchyard_route { + return switchyard_route + .execute( + &gateway, + cache_options, + WireFormat::OpenAiResponses, + body, + headers, + ) + .await; + } let route = gateway.responses.clone(); let route = match &gateway.cache { Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new( @@ -25,19 +47,7 @@ pub(crate) async fn create( None => route, }; - let call = ResponsesCall { - model: deployment.model.clone(), - input: body.get("input").cloned().unwrap_or_default(), - optional_params: body - .into_iter() - .filter(|(name, _)| !matches!(name.as_str(), "model" | "input")) - .collect(), - api_key: deployment.api_key.clone(), - api_base: deployment.api_base.clone(), - custom_llm_provider: deployment.custom_llm_provider.clone(), - extra_headers: None, - timeout: deployment.timeout, - }; + let call = provider_request(deployment, body); let machine = route.machine(call, cache_options.policy); let stream = Sse::::new(Json, |error| { let error = Error::from(error); @@ -50,3 +60,21 @@ pub(crate) async fn create( let response = litellm_host_http::serve(machine, (), headers.clone(), stream, None).await?; Ok(headers.apply(response)) } + +pub(crate) fn provider_request( + deployment: &Deployment, + mut body: Map, +) -> ResponsesCall { + let input = body.remove("input").unwrap_or_default(); + body.remove("model"); + ResponsesCall { + model: deployment.model.clone(), + input, + optional_params: body, + api_key: deployment.api_key.clone(), + api_base: deployment.api_base.clone(), + custom_llm_provider: deployment.custom_llm_provider.clone(), + extra_headers: None, + timeout: deployment.timeout, + } +} diff --git a/litellm-rust/crates/gateway-inference/src/switchyard.rs b/litellm-rust/crates/gateway-inference/src/switchyard.rs new file mode 100644 index 00000000000..11fd8729e02 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/switchyard.rs @@ -0,0 +1,545 @@ +use std::{ + collections::{BTreeMap, HashMap, HashSet}, + convert::Infallible, + sync::Arc, +}; + +use async_trait::async_trait; +use axum::{ + Json, + http::HeaderMap, + response::{ + IntoResponse, Response as HttpResponse, + sse::{Event, Sse}, + }, +}; +use futures_util::{StreamExt, stream}; +use litellm_cache_response::{CacheOptions, CachePolicy, ScopedCache}; +use litellm_config::Model; +use litellm_host::call::CallOutput; +use litellm_http::transport::Error as TransportError; +use litellm_inference::RouteError; +use litellm_inference_chat::ChatCompletionsRoute; +use litellm_inference_messages::MessagesRoute; +use litellm_inference_responses::ResponsesRoute; +use serde::Deserialize; +use serde_json::{Map, Value, json}; +use switchyard_libsy::{Algorithm, LibsyError, RuntimeModels}; +use switchyard_llm_client::ClientRouter; +use switchyard_protocol::{ + BoxError, LlmClientError, LlmResponse, Metadata, ModelId, Request, Response, RoutedLlmClient, + WireFormat, +}; +use switchyard_runner::AlgorithmSpec; +use switchyard_translation::{ + LlmStreamError, RawEventStream, decode_aggregated_response, decode_request, decode_stream, + encode_aggregated_response_with_extensions, encode_request, encode_stream_with_extensions, +}; + +use crate::{ + Deployment, Error, Gateway, ModelRouter, SwitchyardConfigError, caching::CacheHeaders, + chat_completions, messages, responses, +}; + +const SWITCHYARD_MODEL: &str = "switchyard/algorithm"; + +#[derive(Default)] +pub(crate) struct Routes(HashMap); + +impl Routes { + pub(crate) fn from_models( + model_list: &[Model], + deployments: &ModelRouter, + ) -> Result { + let mut routes = HashMap::new(); + let mut seen = HashSet::new(); + // ModelRouter uses the final entry when public model names repeat. + for model in model_list.iter().rev() { + if !seen.insert(model.model_name.as_str()) + || model.litellm_params.model != SWITCHYARD_MODEL + { + continue; + } + if deployments + .get(&model.model_name) + .is_none_or(|deployment| deployment.model != SWITCHYARD_MODEL) + { + return Err(invalid(model, "model is not a Switchyard route")); + } + let config = model + .litellm_params + .extra + .get("switchyard_config") + .ok_or_else(|| invalid(model, "switchyard_config is required"))?; + let config = AlgorithmSpec::deserialize(config) + .map_err(|error| invalid(model, error.to_string()))?; + let targets = targets(model, &config, model_list, deployments)?; + let (algorithm, runtime_models) = config + .build_with_runtime_models(&model.model_name, &targets.models) + .map_err(|error| invalid(model, error.to_string()))?; + routes.insert( + model.model_name.clone(), + Route { + algorithm, + runtime_models: Arc::new(runtime_models), + targets: Arc::new(targets.by_model), + }, + ); + } + Ok(Self(routes)) + } + + pub(crate) fn get(&self, model: &str) -> Option<&Route> { + self.0.get(model) + } +} + +pub(crate) struct Route { + algorithm: Arc, + runtime_models: Arc, + targets: Arc>, +} + +impl Route { + pub(crate) async fn execute( + &self, + gateway: &Gateway, + cache_options: CacheOptions, + wire_format: WireFormat, + body: Map, + headers: HeaderMap, + ) -> Result { + let llm_request = decode_request(wire_format, &Value::Object(body)) + .map_err(|error| Error::InvalidBody(error.to_string()))?; + let extensions = llm_request.extensions.clone(); + let mut metadata = Metadata::from_headers(&headers); + metadata.http_headers = Some(headers); + let request = Request { + llm_request, + raw_request: None, + metadata: Some(metadata), + }; + let client = Arc::new(LiteLlmClient { + routes: ProviderRoutes::new(gateway, &cache_options), + targets: Arc::clone(&self.targets), + cache_policy: cache_options.policy, + }); + let (selected_model, response) = switchyard_llm_client::run( + Arc::clone(&self.algorithm), + ClientRouter::single(client), + request, + Arc::clone(&self.runtime_models), + None, + ) + .await + .map_err(map_libsy_error)?; + let served_model = response.served_model().cloned().unwrap_or(selected_model); + let upstream_headers = response.upstream_headers; + let mut response = match response.llm_response { + LlmResponse::Agg(response) => { + let body = encode_aggregated_response_with_extensions( + &response, + wire_format, + Some(served_model.as_str()), + &extensions, + ) + .map_err(|error| Error::Internal(error.to_string()))?; + Json(body).into_response() + } + LlmResponse::Stream(response) => { + let events = encode_stream_with_extensions( + response, + wire_format, + Some(served_model.into_string()), + &extensions, + ) + .map_err(|error| Error::Internal(error.to_string()))?; + stream_response(events, wire_format) + } + }; + response.headers_mut().extend(upstream_headers); + Ok(response) + } +} + +struct Target { + deployment: Deployment, + wire_format: WireFormat, +} + +struct ProviderRoutes { + chat: ChatCompletionsRoute, + responses: ResponsesRoute, + messages: MessagesRoute, +} + +impl ProviderRoutes { + fn new(gateway: &Gateway, options: &CacheOptions) -> Self { + let mut routes = Self { + chat: gateway.chat_completions.clone(), + responses: gateway.responses.clone(), + messages: gateway.messages.clone(), + }; + if let Some(cache) = &gateway.cache { + let scoped = || ScopedCache::new(Arc::clone(cache), options.scope.clone()); + routes.chat = routes.chat.with_cache(scoped()); + routes.responses = routes.responses.with_cache(scoped()); + routes.messages = routes.messages.with_cache(scoped()); + } + routes + } +} + +struct LiteLlmClient { + routes: ProviderRoutes, + targets: Arc>, + cache_policy: CachePolicy, +} + +#[async_trait] +impl RoutedLlmClient for LiteLlmClient { + async fn call(&self, request: Request) -> Result { + let Some(model) = request.llm_request.model.as_deref() else { + return Err(LlmClientError::InvalidRequest { + message: "routed request has no model".to_string(), + }); + }; + let target = self + .targets + .get(model) + .ok_or_else(|| LlmClientError::Configuration { + message: format!("unknown LiteLLM deployment {model}"), + })?; + let Request { + llm_request, + metadata, + .. + } = request; + let Value::Object(body) = encode_request(&llm_request, target.wire_format) + .map_err(|error| LlmClientError::RequestEncoding(error.to_string()))? + else { + return Err(LlmClientError::RequestEncoding(format!( + "{} request is not an object", + target.wire_format + ))); + }; + let cache_headers = CacheHeaders::default(); + let response = match target.wire_format { + WireFormat::OpenAiChat => { + let response = self + .routes + .chat + .execute( + chat_completions::provider_request(&target.deployment, body), + &cache_headers, + self.cache_policy, + ) + .await + .map_err(map_route_error)?; + decode_complete(response, target.wire_format)? + } + WireFormat::OpenAiResponses => { + let response = self + .routes + .responses + .execute( + responses::provider_request(&target.deployment, body), + &cache_headers, + self.cache_policy, + ) + .await + .map_err(map_route_error)?; + decode_output(response, target.wire_format)? + } + WireFormat::AnthropicMessages => { + let empty_headers = HeaderMap::new(); + let headers = metadata + .as_ref() + .and_then(|metadata| metadata.http_headers.as_ref()) + .unwrap_or(&empty_headers); + let call = messages::provider_request(&target.deployment, body, headers) + .map_err(map_route_error)?; + let response = self + .routes + .messages + .execute(call, &cache_headers, self.cache_policy) + .await + .map_err(map_route_error)?; + decode_output(response, target.wire_format)? + } + }; + let mut response = Response { + llm_response: response, + metadata, + upstream_headers: HeaderMap::new(), + }; + cache_headers.apply_to(&mut response.upstream_headers); + Ok(response) + } +} + +fn decode_complete( + response: impl serde::Serialize, + wire_format: WireFormat, +) -> Result { + let body = serde_json::to_value(response) + .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?; + decode_aggregated_response(&body, wire_format) + .map(LlmResponse::Agg) + .map_err(|error| LlmClientError::ResponseTranslation(error.to_string())) +} + +fn decode_output( + response: CallOutput, + wire_format: WireFormat, +) -> Result +where + T: serde::Serialize, +{ + match response { + CallOutput::Complete(response) => decode_complete(response, wire_format), + CallOutput::Stream { chunks, .. } => { + let bytes = + chunks.map(|chunk| chunk.map(|bytes| bytes.to_vec()).map_err(map_route_error)); + decode_stream(bytes, wire_format).map(LlmResponse::Stream) + } + } +} + +// A stream error ends the response; Chat's success sentinel must not follow it. +fn stream_response(events: RawEventStream, wire_format: WireFormat) -> HttpResponse { + let frames = stream::unfold(Some(events), move |events| async move { + let mut events = events?; + let (event, next) = match events.next().await { + Some(Ok(value)) => (frame_event(&value, wire_format), Some(events)), + Some(Err(LlmStreamError::Upstream(value))) => (frame_event(&value, wire_format), None), + Some(Err(LlmStreamError::Client(error))) => { + let error = map_client_error(error); + let value = match wire_format { + WireFormat::OpenAiChat => error.openai_body(), + WireFormat::AnthropicMessages => error.body(None), + WireFormat::OpenAiResponses => json!({ + "type": "error", "code": error.status().as_u16().to_string(), + "message": error.to_string(), "param": null + }), + }; + (frame_event(&value, wire_format), None) + } + None if wire_format == WireFormat::OpenAiChat => { + (Event::default().data("[DONE]"), None) + } + None => return None, + }; + Some((Ok::<_, Infallible>(event), next)) + }); + Sse::new(frames).into_response() +} + +fn frame_event(value: &Value, wire_format: WireFormat) -> Event { + let event = Event::default().data(value.to_string()); + match wire_format { + WireFormat::OpenAiChat => event, + WireFormat::OpenAiResponses | WireFormat::AnthropicMessages => { + let event_type = value + .get("type") + .and_then(Value::as_str) + .filter(|name| !name.contains(['\r', '\n'])) + .unwrap_or("message"); + event.event(event_type) + } + } +} + +struct Targets { + models: BTreeMap, + by_model: HashMap, +} + +fn targets( + route: &Model, + config: &AlgorithmSpec, + model_list: &[Model], + deployments: &ModelRouter, +) -> Result { + let mut models = BTreeMap::new(); + let mut by_model = HashMap::new(); + for target in config.callable_target_names() { + let configured = model_list + .iter() + .rev() + .find(|model| model.model_name == target) + .ok_or_else(|| invalid(route, format!("target {target:?} is not configured")))?; + let deployment = deployments + .get(target) + .ok_or_else(|| invalid(route, format!("target {target:?} is not configured")))?; + if deployment.model == SWITCHYARD_MODEL { + return Err(invalid( + route, + format!("target {target:?} is another Switchyard route"), + )); + } + let model = ModelId::from(target); + models.insert(target.to_string(), model.clone()); + by_model.insert( + model, + Target { + deployment: deployment.clone(), + wire_format: target_wire_format(route, configured)?, + }, + ); + } + Ok(Targets { models, by_model }) +} + +// Caller format and target execution format are independent, so each target declares its own. +fn target_wire_format(route: &Model, target: &Model) -> Result { + let value = target + .litellm_params + .extra + .get("switchyard_wire_format") + .ok_or_else(|| { + invalid( + route, + format!( + "target {:?} requires switchyard_wire_format", + target.model_name + ), + ) + })?; + WireFormat::deserialize(value).map_err(|error| { + invalid( + route, + format!( + "target {:?} has invalid switchyard_wire_format: {error}", + target.model_name, + ), + ) + }) +} + +fn invalid(model: &Model, message: impl Into) -> SwitchyardConfigError { + SwitchyardConfigError { + model: model.model_name.clone(), + reason: message.into(), + } +} + +fn map_route_error(error: RouteError) -> LlmClientError { + match error { + RouteError::Transport(TransportError::Http { status, body }) => { + match axum::http::StatusCode::from_u16(status) { + Ok(status) => LlmClientError::UpstreamHttp { status, body }, + Err(source) => LlmClientError::InvalidResponse { + source: Box::new(source), + }, + } + } + // Only Connect is known to precede dispatch, so only it is safe to fall back from. + error @ RouteError::Transport(TransportError::Connect(_)) => LlmClientError::Transport { + source: Box::new(error), + }, + error @ RouteError::InvalidResponse(_) => LlmClientError::InvalidResponse { + source: Box::new(error), + }, + error => LlmClientError::Host { + source: Box::new(error), + }, + } +} + +fn map_libsy_error(error: LibsyError) -> Error { + match error { + LibsyError::ClientCall { source, .. } => map_client_error(source), + error => Error::Internal(error.to_string()), + } +} + +fn map_client_error(error: LlmClientError) -> Error { + match error { + LlmClientError::InvalidRequest { message } + | LlmClientError::RequestTranslation(message) => { + Error::Route(RouteError::InvalidRequest(message.into())) + } + LlmClientError::UpstreamHttp { status, body } => { + Error::Route(RouteError::Transport(TransportError::Http { + status: status.as_u16(), + body, + })) + } + LlmClientError::Transport { source } + | LlmClientError::Timeout { source } + | LlmClientError::InvalidResponse { source } + | LlmClientError::Host { source } + | LlmClientError::Ffi { source } => map_source(source), + LlmClientError::ContextWindowExceeded { message, .. } => { + Error::Route(RouteError::InvalidRequest(message.into())) + } + LlmClientError::ResponseTranslation(message) => { + Error::Route(RouteError::InvalidResponse(message.into())) + } + error => Error::Internal(error.to_string()), + } +} + +fn map_source(source: BoxError) -> Error { + match source.downcast::() { + Ok(error) => Error::Route(*error), + Err(source) => Error::Internal(source.to_string()), + } +} + +#[cfg(test)] +mod tests { + use axum::body::to_bytes; + use rstest::rstest; + + use super::*; + + #[rstest] + fn host_errors_keep_their_gateway_classification() { + let error = RouteError::Auth(litellm_auth::Error::MissingApiKey { + provider: "OpenAI", + environment_variable: "OPENAI_API_KEY", + }); + + assert!(matches!( + map_client_error(map_route_error(error)), + Error::Route(RouteError::Auth(_)) + )); + } + + #[test] + fn targets_require_an_explicit_wire_format() { + let config = litellm_config::Config::from_yaml( + "model_list:\n - model_name: target\n litellm_params: { model: upstream }\n - model_name: route\n litellm_params: { model: switchyard/algorithm }\n", + ) + .expect("config"); + + let error = target_wire_format(&config.model_list[1], &config.model_list[0]) + .expect_err("ambiguous target format"); + assert!(error.reason.contains("requires switchyard_wire_format")); + } + + #[rstest] + #[tokio::test] + async fn a_stream_error_is_terminal() { + let events: RawEventStream = Box::pin(stream::iter([ + Ok(json!({"chunk": "before"})), + Err(LlmStreamError::Upstream(json!({"error": "failed"}))), + Ok(json!({"chunk": "after"})), + ])); + + let body = to_bytes( + stream_response(events, WireFormat::OpenAiChat).into_body(), + 4096, + ) + .await + .expect("read stream"); + let body = std::str::from_utf8(&body).expect("stream is UTF-8"); + + assert!(body.contains("before")); + assert!(body.contains("failed")); + assert!(!body.contains("after")); + assert!(!body.contains("[DONE]")); + } +} diff --git a/litellm-rust/crates/gateway-inference/tests/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs index 0e37bb0aff9..c7195c13472 100644 --- a/litellm-rust/crates/gateway-inference/tests/support/mod.rs +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -102,6 +102,14 @@ fn configured_app( Some(cache) => gateway.with_cache(cache), None => gateway, }; + authenticated_app(gateway, permissions, principal) +} + +pub fn authenticated_app( + gateway: Gateway, + permissions: litellm_gateway_auth::Permissions, + principal: Option, +) -> Router { router(Arc::new(gateway)).layer(axum::middleware::from_fn_with_state( (permissions, principal), test_identity, diff --git a/litellm-rust/crates/gateway-inference/tests/switchyard.rs b/litellm-rust/crates/gateway-inference/tests/switchyard.rs new file mode 100644 index 00000000000..902473f7773 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/switchyard.rs @@ -0,0 +1,405 @@ +#![cfg(feature = "switchyard")] + +mod support; + +use std::{sync::Arc, time::Duration}; + +use futures_util::StreamExt; +use litellm_config::Config; +use litellm_gateway_auth::{AccessRequest, Permissions}; +use litellm_gateway_inference::{Gateway, ModelRouter}; +use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver}; +use litellm_inference::resources::CoreResources; +use rstest::rstest; +use serde_json::json; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_partial_json, method, path}, +}; + +#[rstest] +#[tokio::test] +async fn configured_switchyard_routes_answer_and_fall_back() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(body_partial_json(json!({"model": "judge-model"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "model": "judge-model", + "choices": [{ + "message": { + "content": r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"# + } + }] + }))) + .expect(1) + .mount(&upstream) + .await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(body_partial_json(json!({ + "model": "fast-model", + "messages": [{"role": "user", "content": "fix this"}], + "service_tier": "flex" + }))) + .respond_with(ResponseTemplate::new(503).set_body_json(json!({ + "error": {"message": "unavailable"} + }))) + .expect(1) + .mount(&upstream) + .await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(body_partial_json(json!({ + "model": "strong-model", + "messages": [{"role": "user", "content": "fix this"}], + "service_tier": "flex" + }))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "model": "strong-model", + "choices": [{ + "message": {"content": "fallback answer"} + }] + }))) + .expect(1) + .mount(&upstream) + .await; + + let config = Config::from_yaml(&format!( + r#" +model_list: + - model_name: fast + litellm_params: + model: fast-model + api_base: {base}/v1 + api_key: test-key + custom_llm_provider: openai_like + switchyard_wire_format: openai_chat + - model_name: strong + litellm_params: + model: strong-model + api_base: {base}/v1 + api_key: test-key + custom_llm_provider: openai_like + switchyard_wire_format: openai_chat + - model_name: judge + litellm_params: + model: judge-model + api_base: {base}/v1 + api_key: test-key + custom_llm_provider: openai_like + switchyard_wire_format: openai_chat + - model_name: coding-router + litellm_params: + model: switchyard/algorithm + switchyard_config: + type: llm_classifier + mode: capability + classifier_target: judge + weak_target: fast + strong_target: strong + base_threshold: 0.5 + - model_name: noop-router + litellm_params: + model: switchyard/algorithm + switchyard_config: + type: noop +"#, + base = upstream.uri() + )) + .expect("valid config"); + let app = switchyard_app( + &config, + Permissions::Only(Arc::from([ + AccessRequest::Model { + name: "coding-router".into(), + deployment: "switchyard/algorithm".into(), + }, + AccessRequest::Model { + name: "noop-router".into(), + deployment: "switchyard/algorithm".into(), + }, + ])), + ); + + let response = support::post( + app.clone(), + "/v1/chat/completions", + json!({ + "model": "coding-router", + "messages": [{"role": "user", "content": "fix this"}], + "service_tier": "flex" + }), + ) + .await; + assert_eq!(response.status(), 200); + let body = support::json(response).await; + assert_eq!(body["model"], "strong"); + assert_eq!(body["choices"][0]["message"]["content"], "fallback answer"); + + let response = support::post( + app, + "/v1/chat/completions", + json!({ + "model": "noop-router", + "messages": [{"role": "user", "content": "health check"}] + }), + ) + .await; + assert_eq!(response.status(), 200); + let body = support::json(response).await; + assert_eq!(body["choices"][0]["message"]["content"], "OK"); +} + +// The caller API and selected target API are independent and use the same host bridge. +#[rstest] +#[tokio::test] +async fn configured_switchyard_routes_translate_responses_and_messages() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/responses")) + .and(body_partial_json(json!({"model": "responses-model"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "resp_1", + "model": "responses-model", + "status": "completed", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "responses answer"}] + }], + "provider_extra": true, + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2} + }))) + .expect(2) + .mount(&upstream) + .await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .and(body_partial_json(json!({"model": "messages-model"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "messages-model", + "content": [{"type": "text", "text": "messages answer"}], + "provider_extra": true, + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1} + }))) + .expect(1) + .mount(&upstream) + .await; + + let config = Config::from_yaml(&format!( + r#" +model_list: + - model_name: responses-target + litellm_params: + model: openai/responses-model + api_base: {base} + api_key: test-key + custom_llm_provider: openai + switchyard_wire_format: openai_responses + - model_name: messages-target + litellm_params: + model: anthropic/messages-model + api_base: {base} + api_key: test-key + custom_llm_provider: anthropic + switchyard_wire_format: anthropic_messages + - model_name: responses-router + litellm_params: + model: switchyard/algorithm + switchyard_config: + type: passthrough + target: responses-target + - model_name: messages-router + litellm_params: + model: switchyard/algorithm + switchyard_config: + type: passthrough + target: messages-target +"#, + base = upstream.uri() + )) + .expect("valid config"); + let app = switchyard_app(&config, Permissions::All); + + let response = support::post( + app.clone(), + "/v1/responses", + json!({"model": "responses-router", "input": "hello"}), + ) + .await; + assert_eq!(response.status(), 200); + let body = support::json(response).await; + assert_eq!(body["model"], "responses-target"); + assert_eq!(body["output"][0]["content"][0]["text"], "responses answer"); + assert_eq!(body["provider_extra"], true); + + let response = support::post( + app.clone(), + "/v1/messages", + json!({ + "model": "messages-router", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 16 + }), + ) + .await; + assert_eq!(response.status(), 200); + let body = support::json(response).await; + assert_eq!(body["model"], "messages-target"); + assert_eq!(body["content"][0]["text"], "messages answer"); + assert_eq!(body["provider_extra"], true); + + let response = support::post( + app, + "/v1/chat/completions", + json!({ + "model": "responses-router", + "messages": [{"role": "user", "content": "hello"}] + }), + ) + .await; + assert_eq!(response.status(), 200); + let body = support::json(response).await; + assert_eq!(body["model"], "responses-target"); + assert_eq!(body["choices"][0]["message"]["content"], "responses answer"); +} + +#[rstest] +#[tokio::test] +async fn streaming_chat_target_keeps_litellm_core_rejection() { + let upstream = MockServer::start().await; + let config = Config::from_yaml(&format!( + r#" +model_list: + - model_name: target + litellm_params: {{ model: test-model, api_base: {base}/v1, + custom_llm_provider: openai_like, switchyard_wire_format: openai_chat }} + - model_name: route + litellm_params: {{ model: switchyard/algorithm, + switchyard_config: {{ type: passthrough, target: target }} }} +"#, + base = upstream.uri(), + )) + .expect("valid config"); + let response = support::post( + switchyard_app(&config, Permissions::All), + "/v1/chat/completions", + json!({ + "model": "route", + "messages": [{"role": "user", "content": "hello"}], + "stream": true + }), + ) + .await; + + assert_eq!(response.status(), 501); + let body = support::json(response).await; + assert!( + body["error"]["message"] + .as_str() + .unwrap() + .contains("streaming") + ); + assert!(upstream.received_requests().await.unwrap().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn switchyard_streams_across_formats_and_cancels_on_drop() { + let (base, upstream) = streaming_responses_upstream().await; + let config = Config::from_yaml(&format!( + r#" +model_list: + - model_name: responses-target + litellm_params: {{ model: openai/responses-model, api_base: {base}, api_key: test-key, + custom_llm_provider: openai, switchyard_wire_format: openai_responses }} + - model_name: responses-router + litellm_params: {{ model: switchyard/algorithm, + switchyard_config: {{ type: passthrough, target: responses-target }} }} +"# + )) + .expect("valid config"); + let response = tokio::time::timeout( + Duration::from_secs(2), + support::post( + switchyard_app(&config, Permissions::All), + "/v1/messages", + json!({ + "model": "responses-router", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 16, + "stream": true + }), + ), + ) + .await + .expect("response starts while upstream remains open"); + + assert_eq!(response.status(), 200); + assert_eq!(response.headers()["content-type"], "text/event-stream"); + let mut body = response.into_body().into_data_stream(); + let first = tokio::time::timeout(Duration::from_secs(2), body.next()) + .await + .expect("first event arrives while upstream remains open") + .expect("stream has a first event") + .expect("first event is readable"); + assert!( + std::str::from_utf8(&first) + .expect("event is UTF-8") + .contains("event: message_start") + ); + + drop(body); + tokio::time::timeout(Duration::from_secs(2), upstream) + .await + .expect("dropping the caller closes the upstream stream") + .expect("upstream task succeeds"); +} + +async fn streaming_responses_upstream() -> (String, tokio::task::JoinHandle<()>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind upstream"); + let base = format!( + "http://{}", + listener.local_addr().expect("upstream address") + ); + let upstream = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accept request"); + let response = concat!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n", + "content-length: 4096\r\n\r\n", + "event: response.created\n", + "data: {\"type\":\"response.created\",\"response\":", + "{\"id\":\"resp_1\",\"model\":\"responses-model\"}}\n\n" + ); + socket + .write_all(response.as_bytes()) + .await + .expect("write first event"); + let mut pending = [0; 1]; + while matches!(socket.read(&mut pending).await, Ok(read) if read > 0) {} + }); + (base, upstream) +} + +fn switchyard_app(config: &Config, permissions: Permissions) -> axum::Router { + let resources = CoreResources::new(Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver)))); + let http = Resolution::from(&HttpSettings::default()).config; + let gateway = Gateway::new( + resources, + http, + Arc::new(support::NoSecrets), + ModelRouter::from_model_list(&config.model_list), + ) + .expect("gateway") + .with_switchyard_models(&config.model_list) + .expect("Switchyard routes"); + support::authenticated_app(gateway, permissions, None) +} diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 387bcd43231..b738af8f4a7 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -5,6 +5,9 @@ edition.workspace = true license.workspace = true repository.workspace = true +[features] +switchyard = ["litellm-gateway-inference/switchyard"] + [dependencies] axum.workspace = true envy = "0.4.2" diff --git a/litellm-rust/crates/gateway/src/error.rs b/litellm-rust/crates/gateway/src/error.rs index 75dc04d179d..579a55e3c1c 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -10,4 +10,7 @@ pub enum Error { Mcp(#[from] litellm_gateway_mcp::ConnectError), #[error("MCP requires a configured master key")] Auth(#[from] litellm_gateway_auth::Error), + #[cfg(feature = "switchyard")] + #[error(transparent)] + Switchyard(#[from] litellm_gateway_inference::SwitchyardConfigError), } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 0b3aba81181..1ff1b9de966 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -62,12 +62,15 @@ pub fn build_inference(config: &Config) -> Result, Error> { Arc::new(EnvironmentSecrets::python_compatible(client)), )); let resources = CoreResources::new(pool); - Ok(Arc::new(Gateway::new( + let gateway = Gateway::new( resources, http, secrets, ModelRouter::from_model_list(&config.model_list), - )?)) + )?; + #[cfg(feature = "switchyard")] + let gateway = gateway.with_switchyard_models(&config.model_list)?; + Ok(Arc::new(gateway)) } pub async fn build_mcp( diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index 71edd3c3650..31c6c27e972 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -33,6 +33,32 @@ fn inference() -> Arc { litellm_gateway::build_inference(&Config::from_yaml("model_list: []").unwrap()).unwrap() } +#[cfg(feature = "switchyard")] +#[rstest] +fn validates_only_winning_switchyard_models() { + let config = Config::from_yaml( + r#" +model_list: + - model_name: coding-router + litellm_params: + model: switchyard/algorithm + switchyard_config: { type: passthrough, target: missing } +"#, + ) + .unwrap(); + assert!(litellm_gateway::build_inference(&config).is_err()); + + let config = Config::from_yaml( + r#" +model_list: + - { model_name: coding-router, litellm_params: { model: switchyard/algorithm, switchyard_config: { type: passthrough, target: missing } } } + - { model_name: coding-router, litellm_params: { model: openai/model } } +"#, + ) + .unwrap(); + assert!(litellm_gateway::build_inference(&config).is_ok()); +} + #[rstest] #[case::authorized("/v1/messages", Some("Bearer gateway-key"), Some("gateway-key"), 400)] #[case::missing_token("/v1/messages", None, Some("gateway-key"), 401)]