From 5bbd2810d49280f0fd3057035d1993e3cdd6335b Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Sun, 4 Oct 2026 20:40:03 -0700 Subject: [PATCH 1/5] feat(gateway): route models through Switchyard Signed-off-by: Alex Fournier --- litellm-rust/Cargo.lock | 249 +++++++++++++- .../crates/gateway-inference/Cargo.toml | 18 + .../crates/gateway-inference/src/caching.rs | 17 +- .../gateway-inference/src/chat_completions.rs | 74 ++++- .../crates/gateway-inference/src/error.rs | 9 + .../crates/gateway-inference/src/lib.rs | 19 ++ .../gateway-inference/src/switchyard.rs | 311 ++++++++++++++++++ .../gateway-inference/tests/support/mod.rs | 8 + .../gateway-inference/tests/switchyard.rs | 173 ++++++++++ litellm-rust/crates/gateway/Cargo.toml | 3 + litellm-rust/crates/gateway/src/error.rs | 3 + litellm-rust/crates/gateway/src/lib.rs | 7 +- litellm-rust/crates/gateway/tests/server.rs | 26 ++ 13 files changed, 880 insertions(+), 37 deletions(-) create mode 100644 litellm-rust/crates/gateway-inference/src/switchyard.rs create mode 100644 litellm-rust/crates/gateway-inference/tests/switchyard.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 2f77b09221a..fd4e57c0418 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" @@ -2657,9 +2667,9 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] name = "hermit-abi" -version = "0.5.3" +version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e17592d60ebacc7d5e169f4663c5f84f9161cc90328abcfe8456f41e4dfcb284" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" [[package]] name = "hex" @@ -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", @@ -3783,6 +3855,7 @@ dependencies = [ name = "litellm-gateway-inference" version = "0.1.0" dependencies = [ + "async-trait", "axum", "base64 0.22.1", "bytes", @@ -3790,6 +3863,7 @@ dependencies = [ "litellm-auth", "litellm-cache-memory", "litellm-cache-response", + "litellm-config", "litellm-gateway-auth", "litellm-host", "litellm-host-http", @@ -3807,6 +3881,11 @@ dependencies = [ "rstest", "serde", "serde_json", + "switchyard-libsy", + "switchyard-llm-client", + "switchyard-protocol", + "switchyard-runner", + "switchyard-translation", "thiserror 2.0.19", "tokio", "tower", @@ -4189,7 +4268,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", @@ -4502,7 +4581,7 @@ dependencies = [ "tempfile", "thiserror 2.0.19", "tokio", - "toml", + "toml 0.9.12+spec-1.1.0", "zip", ] @@ -5260,9 +5339,9 @@ checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" [[package]] name = "portable-atomic-util" -version = "0.2.8" +version = "0.2.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" dependencies = [ "portable-atomic", ] @@ -5805,6 +5884,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" @@ -7026,6 +7122,103 @@ 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 = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0646b0891f703d91ba97492d1d0527a7fdeb7b0269c53757e7b9388953550c16" +dependencies = [ + "async-trait", + "futures", + "http 1.4.2", + "jsonptr", + "jsonschema 0.49.9", + "opentelemetry 0.32.0", + "parking_lot", + "rand 0.10.2", + "regex", + "serde", + "serde_json", + "switchyard-protocol", + "thiserror 2.0.19", + "tokio", + "tokio-stream", + "tracing", + "tracing-opentelemetry", + "uuid", +] + +[[package]] +name = "switchyard-llm-client" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c07bffb8f677431f9530fc8335140134a743d3e1618c595cfd6953ed3b009360" +dependencies = [ + "async-trait", + "futures", + "futures-util", + "http 1.4.2", + "httpdate", + "opentelemetry 0.32.0", + "parking_lot", + "reqwest 0.13.5", + "serde_json", + "switchyard-libsy", + "switchyard-protocol", + "switchyard-translation", + "tokio", + "tracing", + "tracing-opentelemetry", +] + +[[package]] +name = "switchyard-protocol" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c49a9e7884e7f50669250463a58a21bc79084adb87ef03647cbe7ba5c3b22ee" +dependencies = [ + "async-trait", + "futures", + "http 1.4.2", + "serde", + "serde_json", + "thiserror 2.0.19", +] + +[[package]] +name = "switchyard-runner" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cbaeff25e505e0b5f06beec4dd6a6db58d844b0c08c80bc9e5c4f0f8e6c7079" +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 = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "629f9c47cb198da7e92569384fec4da8486852c38a7af922e0c9d6b62ed74fa1" +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" @@ -7339,6 +7532,7 @@ dependencies = [ "bytes", "libc", "mio", + "parking_lot", "pin-project-lite", "signal-hook-registry", "socket2 0.6.5", @@ -7433,6 +7627,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" @@ -7702,6 +7911,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" @@ -7709,9 +7929,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adbc64cba7137545b8044cb1fe9814f7aacf3c6b5f9b45be8bb5db538befdb26" dependencies = [ "js-sys", - "opentelemetry", + "opentelemetry 0.32.0", "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..3697e46336e 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -5,7 +5,19 @@ edition.workspace = true license.workspace = true repository.workspace = true +[features] +switchyard = [ + "dep:async-trait", + "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 @@ -18,6 +30,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 +40,11 @@ litellm-secrets.workspace = true litellm-llms-types.workspace = true serde.workspace = true serde_json.workspace = true +switchyard-libsy = { version = "0.3.0", optional = true } +switchyard-llm-client = { version = "0.3.0", optional = true } +switchyard-protocol = { version = "0.3.0", optional = true } +switchyard-runner = { version = "0.3.0", optional = true } +switchyard-translation = { version = "0.3.0", 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..f7bd6764fc2 100644 --- a/litellm-rust/crates/gateway-inference/src/caching.rs +++ b/litellm-rust/crates/gateway-inference/src/caching.rs @@ -101,11 +101,22 @@ impl litellm_host::interceptors::Interceptors for impl CacheHeaders { pub(crate) fn apply(&self, mut response: axum::response::Response) -> axum::response::Response { - if let Some(key) = self.0.get() - && let Ok(value) = axum::http::HeaderValue::from_str(key) - { + if let Some(value) = self.value() { response.headers_mut().insert("x-litellm-cache-key", value); } response } + + #[cfg(feature = "switchyard")] + pub(crate) fn headers(&self) -> axum::http::HeaderMap { + let mut headers = axum::http::HeaderMap::new(); + if let Some(value) = self.value() { + headers.insert("x-litellm-cache-key", value); + } + headers + } + + fn value(&self) -> Option { + axum::http::HeaderValue::from_str(self.0.get()?).ok() + } } diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index 742a373b286..bd33e31bdc9 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -1,27 +1,38 @@ 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}; -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,16 +43,29 @@ 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)?; let route = gateway.chat_completions.clone(); let route = match &gateway.cache { @@ -52,23 +76,19 @@ async fn handle( None => route, }; - let messages = body.get("messages").cloned().unwrap_or_default(); + #[cfg(feature = "switchyard")] + if let Some(switchyard_route) = switchyard_route { + let (body, upstream_headers) = switchyard_route + .execute(route, cache_options.policy, body, inbound_headers) + .await?; + let mut response = Json(body).into_response(); + response.headers_mut().extend(upstream_headers); + return Ok(response); + } 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 +99,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..66132df386d 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)] diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index 4663151f1de..2b0c0a5429f 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,17 @@ 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, Arc::new(self.models.clone()))?; + Ok(self) + } } pub fn router(gateway: Arc) -> Router { 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..5607e6159ab --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/switchyard.rs @@ -0,0 +1,311 @@ +use std::{ + collections::{BTreeMap, HashMap, HashSet}, + sync::Arc, +}; + +use async_trait::async_trait; +use axum::http::HeaderMap; +use litellm_cache_response::CachePolicy; +use litellm_config::Model; +use litellm_inference::RouteError; +use litellm_inference_chat::ChatCompletionsRoute; +use litellm_http::transport::Error as TransportError; +use serde::Deserialize; +use serde_json::{Map, Value}; +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::{ + decode_aggregated_response, decode_request, encode_aggregated_response_with_extensions, + encode_request, +}; + +use crate::{ + Error, ModelRouter, SwitchyardConfigError, caching::CacheHeaders, + chat_completions::provider_request, +}; + +const SWITCHYARD_MODEL: &str = "switchyard/algorithm"; + +#[derive(Default)] +pub(crate) struct Routes(HashMap); + +impl Routes { + pub(crate) fn from_models( + model_list: &[Model], + deployments: Arc, + ) -> 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, &deployments)?; + let (algorithm, runtime_models) = config + .build_with_runtime_models(&model.model_name, &targets) + .map_err(|error| invalid(model, error.to_string()))?; + routes.insert( + model.model_name.clone(), + Route { + algorithm, + runtime_models: Arc::new(runtime_models), + deployments: Arc::clone(&deployments), + }, + ); + } + 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, + deployments: Arc, +} + +impl Route { + pub(crate) async fn execute( + &self, + route: ChatCompletionsRoute, + cache_policy: CachePolicy, + body: Map, + headers: HeaderMap, + ) -> Result<(Value, HeaderMap), Error> { + let llm_request = decode_request(WireFormat::OpenAiChat, &Value::Object(body)) + .map_err(|error| Error::InvalidBody(error.to_string()))?; + // Reject unsupported streaming before a classifier or judge can make a model call. + if llm_request.stream { + return Err(Error::Route(RouteError::Unsupported("streaming"))); + } + let extensions = llm_request.extensions.clone(); + let mut metadata = Metadata::from_headers(&headers); + metadata.http_headers = Some(headers); + metadata.wire_format = Some(WireFormat::OpenAiChat); + let request = Request { + llm_request, + raw_request: None, + metadata: Some(metadata), + }; + let client = Arc::new(LiteLlmClient { + route, + deployments: Arc::clone(&self.deployments), + cache_policy, + }); + let (_, 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 upstream_headers = response.upstream_headers; + let response = response + .llm_response + .into_agg() + .await + .map_err(map_client_error)?; + let body = encode_aggregated_response_with_extensions( + &response, + WireFormat::OpenAiChat, + None, + &extensions, + ) + .map_err(|error| Error::Internal(error.to_string()))?; + Ok((body, upstream_headers)) + } +} + +struct LiteLlmClient { + route: ChatCompletionsRoute, + deployments: 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 deployment = + self.deployments + .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, WireFormat::OpenAiChat) + .map_err(|error| LlmClientError::RequestEncoding(error.to_string()))? + else { + return Err(LlmClientError::RequestEncoding( + "OpenAI Chat request is not an object".to_string(), + )); + }; + let cache_headers = CacheHeaders::default(); + let response = self + .route + .execute( + provider_request(deployment, body), + &cache_headers, + self.cache_policy, + ) + .await + .map_err(map_route_error)?; + let body = serde_json::to_value(response) + .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?; + let response = decode_aggregated_response(&body, WireFormat::OpenAiChat) + .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?; + Ok(Response { + llm_response: LlmResponse::Agg(response), + metadata, + upstream_headers: cache_headers.headers(), + }) + } +} + +fn targets( + route: &Model, + config: &AlgorithmSpec, + deployments: &ModelRouter, +) -> Result, SwitchyardConfigError> { + config + .callable_target_names() + .into_iter() + .map(|target| { + 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"), + )); + } + Ok((target.to_string(), ModelId::from(target))) + }) + .collect() +} + +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 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(_)) + )); + } +} 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..4985f549bd6 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/switchyard.rs @@ -0,0 +1,173 @@ +#![cfg(feature = "switchyard")] + +mod support; + +use std::sync::Arc; + +use litellm_config::Config; +use litellm_inference::resources::CoreResources; +use litellm_gateway_auth::{AccessRequest, Permissions}; +use litellm_gateway_inference::{Gateway, ModelRouter}; +use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver}; +use rstest::rstest; +use serde_json::json; +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 + - model_name: strong + litellm_params: + model: strong-model + api_base: {base}/v1 + api_key: test-key + custom_llm_provider: openai_like + - model_name: judge + litellm_params: + model: judge-model + api_base: {base}/v1 + api_key: test-key + custom_llm_provider: openai_like + - 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 pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver))); + let resources = CoreResources::new(pool); + 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"); + let app = support::authenticated_app( + gateway, + Permissions::Only(Arc::from([ + AccessRequest::Model { + name: "coding-router".into(), + deployment: "switchyard/algorithm".into(), + }, + AccessRequest::Model { + name: "noop-router".into(), + deployment: "switchyard/algorithm".into(), + }, + ])), + None, + ); + + let response = support::post( + app.clone(), + "/v1/chat/completions", + json!({ + "model": "coding-router", + "messages": [{"role": "user", "content": "fix this"}], + "stream": true + }), + ) + .await; + assert_eq!(response.status(), 501); + + 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-model"); + 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"); +} 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)] From 997c8f07bffc8b584c6d455586a5bcce078f7767 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Wed, 7 Oct 2026 12:11:16 -0700 Subject: [PATCH 2/5] feat(gateway): route buffered APIs through Switchyard Signed-off-by: Alex Fournier --- .../gateway-inference/src/chat_completions.rs | 24 +- .../crates/gateway-inference/src/lib.rs | 3 +- .../crates/gateway-inference/src/messages.rs | 41 ++- .../crates/gateway-inference/src/responses.rs | 58 +++- .../gateway-inference/src/switchyard.rs | 283 ++++++++++++++---- .../gateway-inference/tests/switchyard.rs | 155 +++++++++- 6 files changed, 443 insertions(+), 121 deletions(-) diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index bd33e31bdc9..c920b874786 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -10,6 +10,8 @@ use axum::{ }; use litellm_inference_chat::types::{ChatCompletionsCall, ChatCompletionsRequest}; use serde_json::{Map, Value}; +#[cfg(feature = "switchyard")] +use switchyard_protocol::WireFormat; use crate::{Deployment, Error, Gateway, JsonObject, request}; @@ -67,6 +69,18 @@ async fn handle( .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( @@ -75,16 +89,6 @@ async fn handle( )), None => route, }; - - #[cfg(feature = "switchyard")] - if let Some(switchyard_route) = switchyard_route { - let (body, upstream_headers) = switchyard_route - .execute(route, cache_options.policy, body, inbound_headers) - .await?; - let mut response = Json(body).into_response(); - response.headers_mut().extend(upstream_headers); - return Ok(response); - } let headers = crate::caching::CacheHeaders::default(); let response = litellm_host_http::serve_unary( route.machine( diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index 2b0c0a5429f..06d8aa5949d 100644 --- a/litellm-rust/crates/gateway-inference/src/lib.rs +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -101,8 +101,7 @@ impl Gateway { mut self, models: &[litellm_config::Model], ) -> Result { - self.switchyard_routes = - switchyard::Routes::from_models(models, Arc::new(self.models.clone()))?; + self.switchyard_routes = switchyard::Routes::from_models(models, &self.models)?; Ok(self) } } diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index c7a8e3e3294..1a6dedbde24 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 index 5607e6159ab..8725b196f3d 100644 --- a/litellm-rust/crates/gateway-inference/src/switchyard.rs +++ b/litellm-rust/crates/gateway-inference/src/switchyard.rs @@ -4,11 +4,18 @@ use std::{ }; use async_trait::async_trait; -use axum::http::HeaderMap; -use litellm_cache_response::CachePolicy; +use axum::{ + Json, + http::HeaderMap, + response::{IntoResponse, Response as HttpResponse}, +}; +use litellm_cache_response::{CacheOptions, CachePolicy, ScopedCache}; use litellm_config::Model; +use litellm_host::call::CallOutput; use litellm_inference::RouteError; use litellm_inference_chat::ChatCompletionsRoute; +use litellm_inference_messages::MessagesRoute; +use litellm_inference_responses::ResponsesRoute; use litellm_http::transport::Error as TransportError; use serde::Deserialize; use serde_json::{Map, Value}; @@ -25,8 +32,8 @@ use switchyard_translation::{ }; use crate::{ - Error, ModelRouter, SwitchyardConfigError, caching::CacheHeaders, - chat_completions::provider_request, + Deployment, Error, Gateway, ModelRouter, SwitchyardConfigError, caching::CacheHeaders, + chat_completions, messages, responses, }; const SWITCHYARD_MODEL: &str = "switchyard/algorithm"; @@ -37,7 +44,7 @@ pub(crate) struct Routes(HashMap); impl Routes { pub(crate) fn from_models( model_list: &[Model], - deployments: Arc, + deployments: &ModelRouter, ) -> Result { let mut routes = HashMap::new(); let mut seen = HashSet::new(); @@ -61,16 +68,16 @@ impl Routes { .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, &deployments)?; + let targets = targets(model, &config, model_list, deployments)?; let (algorithm, runtime_models) = config - .build_with_runtime_models(&model.model_name, &targets) + .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), - deployments: Arc::clone(&deployments), + targets: Arc::new(targets.by_model), }, ); } @@ -85,38 +92,39 @@ impl Routes { pub(crate) struct Route { algorithm: Arc, runtime_models: Arc, - deployments: Arc, + targets: Arc>, } impl Route { pub(crate) async fn execute( &self, - route: ChatCompletionsRoute, - cache_policy: CachePolicy, + gateway: &Gateway, + cache_options: CacheOptions, + wire_format: WireFormat, body: Map, headers: HeaderMap, - ) -> Result<(Value, HeaderMap), Error> { - let llm_request = decode_request(WireFormat::OpenAiChat, &Value::Object(body)) + ) -> Result { + let llm_request = decode_request(wire_format, &Value::Object(body)) .map_err(|error| Error::InvalidBody(error.to_string()))?; - // Reject unsupported streaming before a classifier or judge can make a model call. + // Switchyard routes currently support buffered calls only. if llm_request.stream { return Err(Error::Route(RouteError::Unsupported("streaming"))); } let extensions = llm_request.extensions.clone(); let mut metadata = Metadata::from_headers(&headers); metadata.http_headers = Some(headers); - metadata.wire_format = Some(WireFormat::OpenAiChat); + metadata.wire_format = None; let request = Request { llm_request, raw_request: None, metadata: Some(metadata), }; let client = Arc::new(LiteLlmClient { - route, - deployments: Arc::clone(&self.deployments), - cache_policy, + routes: ProviderRoutes::new(gateway, &cache_options), + targets: Arc::clone(&self.targets), + cache_policy: cache_options.policy, }); - let (_, response) = switchyard_llm_client::run( + let (selected_model, response) = switchyard_llm_client::run( Arc::clone(&self.algorithm), ClientRouter::single(client), request, @@ -125,6 +133,7 @@ impl Route { ) .await .map_err(map_libsy_error)?; + let served_model = response.served_model().cloned().unwrap_or(selected_model); let upstream_headers = response.upstream_headers; let response = response .llm_response @@ -133,18 +142,48 @@ impl Route { .map_err(map_client_error)?; let body = encode_aggregated_response_with_extensions( &response, - WireFormat::OpenAiChat, - None, + wire_format, + Some(served_model.as_str()), &extensions, ) .map_err(|error| Error::Internal(error.to_string()))?; - Ok((body, upstream_headers)) + let mut response = Json(body).into_response(); + 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 { - route: ChatCompletionsRoute, - deployments: Arc, + routes: ProviderRoutes, + targets: Arc>, cache_policy: CachePolicy, } @@ -156,67 +195,169 @@ impl RoutedLlmClient for LiteLlmClient { message: "routed request has no model".to_string(), }); }; - let deployment = - self.deployments - .get(model) - .ok_or_else(|| LlmClientError::Configuration { - message: format!("unknown LiteLLM deployment {model}"), - })?; + 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, WireFormat::OpenAiChat) + let Value::Object(body) = encode_request(&llm_request, target.wire_format) .map_err(|error| LlmClientError::RequestEncoding(error.to_string()))? else { - return Err(LlmClientError::RequestEncoding( - "OpenAI Chat request is not an object".to_string(), - )); + return Err(LlmClientError::RequestEncoding(format!( + "{} request is not an object", + target.wire_format + ))); }; let cache_headers = CacheHeaders::default(); - let response = self - .route - .execute( - provider_request(deployment, body), - &cache_headers, - self.cache_policy, - ) - .await - .map_err(map_route_error)?; - let body = serde_json::to_value(response) - .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?; - let response = decode_aggregated_response(&body, WireFormat::OpenAiChat) - .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?; + 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)? + } + }; Ok(Response { - llm_response: LlmResponse::Agg(response), + llm_response: response, metadata, upstream_headers: cache_headers.headers(), }) } } +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 { .. } => Err(LlmClientError::ResponseTranslation( + "LiteLLM streamed a response for a buffered Switchyard call".to_string(), + )), + } +} + +struct Targets { + models: BTreeMap, + by_model: HashMap, +} + fn targets( route: &Model, config: &AlgorithmSpec, + model_list: &[Model], deployments: &ModelRouter, -) -> Result, SwitchyardConfigError> { - config - .callable_target_names() - .into_iter() - .map(|target| { - 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"), - )); - } - Ok((target.to_string(), ModelId::from(target))) - }) - .collect() +) -> 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 { @@ -308,4 +449,16 @@ mod tests { 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")); + } } diff --git a/litellm-rust/crates/gateway-inference/tests/switchyard.rs b/litellm-rust/crates/gateway-inference/tests/switchyard.rs index 4985f549bd6..10f87b680d7 100644 --- a/litellm-rust/crates/gateway-inference/tests/switchyard.rs +++ b/litellm-rust/crates/gateway-inference/tests/switchyard.rs @@ -73,18 +73,21 @@ model_list: 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 @@ -104,20 +107,8 @@ model_list: base = upstream.uri() )) .expect("valid config"); - let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver))); - let resources = CoreResources::new(pool); - 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"); - let app = support::authenticated_app( - gateway, + let app = switchyard_app( + &config, Permissions::Only(Arc::from([ AccessRequest::Model { name: "coding-router".into(), @@ -128,7 +119,6 @@ model_list: deployment: "switchyard/algorithm".into(), }, ])), - None, ); let response = support::post( @@ -155,7 +145,7 @@ model_list: .await; assert_eq!(response.status(), 200); let body = support::json(response).await; - assert_eq!(body["model"], "strong-model"); + assert_eq!(body["model"], "strong"); assert_eq!(body["choices"][0]["message"]["content"], "fallback answer"); let response = support::post( @@ -171,3 +161,136 @@ model_list: 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"); +} + +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) +} From cc13482bd77abef60288afc56d79926ff326533d Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Tue, 13 Oct 2026 22:09:54 -0700 Subject: [PATCH 3/5] build(gateway): pin Switchyard to official main Signed-off-by: Alex Fournier --- litellm-rust/Cargo.lock | 32 +++++++++---------- .../crates/gateway-inference/Cargo.toml | 10 +++--- 2 files changed, 20 insertions(+), 22 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index fd4e57c0418..f76bd0f5ff3 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2667,9 +2667,9 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] name = "hermit-abi" -version = "0.5.2" +version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" +checksum = "e17592d60ebacc7d5e169f4663c5f84f9161cc90328abcfe8456f41e4dfcb284" [[package]] name = "hex" @@ -5339,9 +5339,9 @@ checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" [[package]] name = "portable-atomic-util" -version = "0.2.7" +version = "0.2.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715" dependencies = [ "portable-atomic", ] @@ -7125,20 +7125,20 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" [[package]] name = "switchyard-libsy" version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0646b0891f703d91ba97492d1d0527a7fdeb7b0269c53757e7b9388953550c16" +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 0.32.0", + "opentelemetry", "parking_lot", "rand 0.10.2", "regex", "serde", "serde_json", + "shlex", "switchyard-protocol", "thiserror 2.0.19", "tokio", @@ -7151,17 +7151,17 @@ dependencies = [ [[package]] name = "switchyard-llm-client" version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c07bffb8f677431f9530fc8335140134a743d3e1618c595cfd6953ed3b009360" +source = "git+https://github.com/NVIDIA-NeMo/Switchyard.git?rev=4bdf5719ce1c43902df92ef64407b858d1f99394#4bdf5719ce1c43902df92ef64407b858d1f99394" dependencies = [ "async-trait", "futures", "futures-util", "http 1.4.2", "httpdate", - "opentelemetry 0.32.0", + "opentelemetry", "parking_lot", "reqwest 0.13.5", + "serde", "serde_json", "switchyard-libsy", "switchyard-protocol", @@ -7174,8 +7174,7 @@ dependencies = [ [[package]] name = "switchyard-protocol" version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c49a9e7884e7f50669250463a58a21bc79084adb87ef03647cbe7ba5c3b22ee" +source = "git+https://github.com/NVIDIA-NeMo/Switchyard.git?rev=4bdf5719ce1c43902df92ef64407b858d1f99394#4bdf5719ce1c43902df92ef64407b858d1f99394" dependencies = [ "async-trait", "futures", @@ -7188,8 +7187,7 @@ dependencies = [ [[package]] name = "switchyard-runner" version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7cbaeff25e505e0b5f06beec4dd6a6db58d844b0c08c80bc9e5c4f0f8e6c7079" +source = "git+https://github.com/NVIDIA-NeMo/Switchyard.git?rev=4bdf5719ce1c43902df92ef64407b858d1f99394#4bdf5719ce1c43902df92ef64407b858d1f99394" dependencies = [ "reqwest 0.13.5", "serde", @@ -7206,8 +7204,7 @@ dependencies = [ [[package]] name = "switchyard-translation" version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "629f9c47cb198da7e92569384fec4da8486852c38a7af922e0c9d6b62ed74fa1" +source = "git+https://github.com/NVIDIA-NeMo/Switchyard.git?rev=4bdf5719ce1c43902df92ef64407b858d1f99394#4bdf5719ce1c43902df92ef64407b858d1f99394" dependencies = [ "async-stream", "base64 0.22.1", @@ -7929,7 +7926,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adbc64cba7137545b8044cb1fe9814f7aacf3c6b5f9b45be8bb5db538befdb26" dependencies = [ "js-sys", - "opentelemetry 0.32.0", + "opentelemetry", + "smallvec", "tracing", "tracing-core", "tracing-log", diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index 3697e46336e..c909c6816c8 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -40,11 +40,11 @@ litellm-secrets.workspace = true litellm-llms-types.workspace = true serde.workspace = true serde_json.workspace = true -switchyard-libsy = { version = "0.3.0", optional = true } -switchyard-llm-client = { version = "0.3.0", optional = true } -switchyard-protocol = { version = "0.3.0", optional = true } -switchyard-runner = { version = "0.3.0", optional = true } -switchyard-translation = { version = "0.3.0", optional = 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] From b3dc0bc5d1197e6bb1a1f5d4d9de07d8c75073f6 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Tue, 13 Oct 2026 22:10:00 -0700 Subject: [PATCH 4/5] feat(gateway): translate existing streams through Switchyard Signed-off-by: Alex Fournier --- .../crates/gateway-inference/Cargo.toml | 2 + .../crates/gateway-inference/src/caching.rs | 17 +-- .../crates/gateway-inference/src/error.rs | 22 +-- .../gateway-inference/src/switchyard.rs | 139 ++++++++++++++---- .../gateway-inference/tests/switchyard.rs | 137 +++++++++++++++-- 5 files changed, 251 insertions(+), 66 deletions(-) diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index c909c6816c8..8c6cb9efccf 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -8,6 +8,7 @@ repository.workspace = true [features] switchyard = [ "dep:async-trait", + "dep:futures-util", "dep:litellm-config", "dep:switchyard-libsy", "dep:switchyard-llm-client", @@ -22,6 +23,7 @@ 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 diff --git a/litellm-rust/crates/gateway-inference/src/caching.rs b/litellm-rust/crates/gateway-inference/src/caching.rs index f7bd6764fc2..ddf0a6ca79e 100644 --- a/litellm-rust/crates/gateway-inference/src/caching.rs +++ b/litellm-rust/crates/gateway-inference/src/caching.rs @@ -101,22 +101,15 @@ impl litellm_host::interceptors::Interceptors for impl CacheHeaders { pub(crate) fn apply(&self, mut response: axum::response::Response) -> axum::response::Response { - if let Some(value) = self.value() { - response.headers_mut().insert("x-litellm-cache-key", value); - } + self.apply_to(response.headers_mut()); response } - #[cfg(feature = "switchyard")] - pub(crate) fn headers(&self) -> axum::http::HeaderMap { - let mut headers = axum::http::HeaderMap::new(); - if let Some(value) = self.value() { + 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) + { headers.insert("x-litellm-cache-key", value); } - headers - } - - fn value(&self) -> Option { - axum::http::HeaderValue::from_str(self.0.get()?).ok() } } diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index 66132df386d..46110612fff 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -85,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/switchyard.rs b/litellm-rust/crates/gateway-inference/src/switchyard.rs index 8725b196f3d..11fd8729e02 100644 --- a/litellm-rust/crates/gateway-inference/src/switchyard.rs +++ b/litellm-rust/crates/gateway-inference/src/switchyard.rs @@ -1,5 +1,6 @@ use std::{ collections::{BTreeMap, HashMap, HashSet}, + convert::Infallible, sync::Arc, }; @@ -7,18 +8,22 @@ use async_trait::async_trait; use axum::{ Json, http::HeaderMap, - response::{IntoResponse, Response as HttpResponse}, + 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 litellm_http::transport::Error as TransportError; use serde::Deserialize; -use serde_json::{Map, Value}; +use serde_json::{Map, Value, json}; use switchyard_libsy::{Algorithm, LibsyError, RuntimeModels}; use switchyard_llm_client::ClientRouter; use switchyard_protocol::{ @@ -27,8 +32,8 @@ use switchyard_protocol::{ }; use switchyard_runner::AlgorithmSpec; use switchyard_translation::{ - decode_aggregated_response, decode_request, encode_aggregated_response_with_extensions, - encode_request, + LlmStreamError, RawEventStream, decode_aggregated_response, decode_request, decode_stream, + encode_aggregated_response_with_extensions, encode_request, encode_stream_with_extensions, }; use crate::{ @@ -106,14 +111,9 @@ impl Route { ) -> Result { let llm_request = decode_request(wire_format, &Value::Object(body)) .map_err(|error| Error::InvalidBody(error.to_string()))?; - // Switchyard routes currently support buffered calls only. - if llm_request.stream { - return Err(Error::Route(RouteError::Unsupported("streaming"))); - } let extensions = llm_request.extensions.clone(); let mut metadata = Metadata::from_headers(&headers); metadata.http_headers = Some(headers); - metadata.wire_format = None; let request = Request { llm_request, raw_request: None, @@ -135,19 +135,28 @@ impl Route { .map_err(map_libsy_error)?; let served_model = response.served_model().cloned().unwrap_or(selected_model); let upstream_headers = response.upstream_headers; - let response = response - .llm_response - .into_agg() - .await - .map_err(map_client_error)?; - let body = encode_aggregated_response_with_extensions( - &response, - wire_format, - Some(served_model.as_str()), - &extensions, - ) - .map_err(|error| Error::Internal(error.to_string()))?; - let mut response = Json(body).into_response(); + 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) } @@ -259,11 +268,13 @@ impl RoutedLlmClient for LiteLlmClient { decode_output(response, target.wire_format)? } }; - Ok(Response { + let mut response = Response { llm_response: response, metadata, - upstream_headers: cache_headers.headers(), - }) + upstream_headers: HeaderMap::new(), + }; + cache_headers.apply_to(&mut response.upstream_headers); + Ok(response) } } @@ -287,9 +298,55 @@ where { match response { CallOutput::Complete(response) => decode_complete(response, wire_format), - CallOutput::Stream { .. } => Err(LlmClientError::ResponseTranslation( - "LiteLLM streamed a response for a buffered Switchyard call".to_string(), - )), + 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) + } } } @@ -433,6 +490,7 @@ fn map_source(source: BoxError) -> Error { #[cfg(test)] mod tests { + use axum::body::to_bytes; use rstest::rstest; use super::*; @@ -461,4 +519,27 @@ mod tests { .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/switchyard.rs b/litellm-rust/crates/gateway-inference/tests/switchyard.rs index 10f87b680d7..902473f7773 100644 --- a/litellm-rust/crates/gateway-inference/tests/switchyard.rs +++ b/litellm-rust/crates/gateway-inference/tests/switchyard.rs @@ -2,15 +2,17 @@ mod support; -use std::sync::Arc; +use std::{sync::Arc, time::Duration}; +use futures_util::StreamExt; use litellm_config::Config; -use litellm_inference::resources::CoreResources; 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}, @@ -121,18 +123,6 @@ model_list: ])), ); - let response = support::post( - app.clone(), - "/v1/chat/completions", - json!({ - "model": "coding-router", - "messages": [{"role": "user", "content": "fix this"}], - "stream": true - }), - ) - .await; - assert_eq!(response.status(), 501); - let response = support::post( app.clone(), "/v1/chat/completions", @@ -280,6 +270,125 @@ model_list: 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; From 8e6354e56d346ae5d075cc7f4700a906f013d753 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Wed, 14 Oct 2026 08:04:46 -0700 Subject: [PATCH 5/5] ci(gateway): refresh checks after upstream migration fix Signed-off-by: Alex Fournier