mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
Merge 8e6354e56d into d56ff281de
This commit is contained in:
Alex Fournier
GitHub
commit
6bb1493bb3
15 changed files with 1419 additions and 71 deletions
237
litellm-rust/Cargo.lock
generated
237
litellm-rust/Cargo.lock
generated
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -101,11 +101,15 @@ impl litellm_host::interceptors::Interceptors<litellm_inference::RouteError> 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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<Arc<Gateway>>,
|
||||
identity: AuthenticatedRequest,
|
||||
#[cfg(feature = "switchyard")] headers: HeaderMap,
|
||||
JsonObject(body): JsonObject,
|
||||
) -> Result<impl IntoResponse, Error> {
|
||||
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<Arc<Gateway>>,
|
||||
identity: AuthenticatedRequest,
|
||||
#[cfg(feature = "switchyard")] headers: HeaderMap,
|
||||
Path(model): Path<String>,
|
||||
JsonObject(body): JsonObject,
|
||||
) -> Result<impl IntoResponse, Error> {
|
||||
|
|
@ -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<String, Value>,
|
||||
) -> Result<Response, Error> {
|
||||
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<String, Value>,
|
||||
) -> 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,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<Arc<dyn litellm_cache_response::ResponseCacheService>>,
|
||||
#[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, SwitchyardConfigError> {
|
||||
self.switchyard_routes = switchyard::Routes::from_models(models, &self.models)?;
|
||||
Ok(self)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn router(gateway: Arc<Gateway>) -> Router {
|
||||
|
|
|
|||
|
|
@ -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<JsonObject, Error>,
|
||||
) -> 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<String, Value>,
|
||||
) -> Result<Response, Error> {
|
||||
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::<Messages, _, _>::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<String, Value>,
|
||||
mut body: Map<String, Value>,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<MessagesCall, Error> {
|
||||
let body = body
|
||||
.into_iter()
|
||||
.map(|(name, value)| match name.as_str() {
|
||||
"model" => (name, Value::from(deployment.model.as_str())),
|
||||
_ => (name, value),
|
||||
})
|
||||
.collect();
|
||||
) -> Result<MessagesCall, litellm_inference::RouteError> {
|
||||
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(),
|
||||
|
|
|
|||
|
|
@ -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<Arc<Gateway>>,
|
||||
identity: AuthenticatedRequest,
|
||||
#[cfg(feature = "switchyard")] headers: HeaderMap,
|
||||
JsonObject(body): JsonObject,
|
||||
) -> Result<Response, Error> {
|
||||
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::<Responses, _, _>::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<String, Value>,
|
||||
) -> 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,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
545
litellm-rust/crates/gateway-inference/src/switchyard.rs
Normal file
545
litellm-rust/crates/gateway-inference/src/switchyard.rs
Normal file
|
|
@ -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<String, Route>);
|
||||
|
||||
impl Routes {
|
||||
pub(crate) fn from_models(
|
||||
model_list: &[Model],
|
||||
deployments: &ModelRouter,
|
||||
) -> Result<Self, SwitchyardConfigError> {
|
||||
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<dyn Algorithm>,
|
||||
runtime_models: Arc<RuntimeModels>,
|
||||
targets: Arc<HashMap<ModelId, Target>>,
|
||||
}
|
||||
|
||||
impl Route {
|
||||
pub(crate) async fn execute(
|
||||
&self,
|
||||
gateway: &Gateway,
|
||||
cache_options: CacheOptions,
|
||||
wire_format: WireFormat,
|
||||
body: Map<String, Value>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<HttpResponse, Error> {
|
||||
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<HashMap<ModelId, Target>>,
|
||||
cache_policy: CachePolicy,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RoutedLlmClient for LiteLlmClient {
|
||||
async fn call(&self, request: Request) -> Result<Response, LlmClientError> {
|
||||
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<LlmResponse, LlmClientError> {
|
||||
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<T, H>(
|
||||
response: CallOutput<T, H, bytes::Bytes, RouteError>,
|
||||
wire_format: WireFormat,
|
||||
) -> Result<LlmResponse, LlmClientError>
|
||||
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<String, ModelId>,
|
||||
by_model: HashMap<ModelId, Target>,
|
||||
}
|
||||
|
||||
fn targets(
|
||||
route: &Model,
|
||||
config: &AlgorithmSpec,
|
||||
model_list: &[Model],
|
||||
deployments: &ModelRouter,
|
||||
) -> Result<Targets, SwitchyardConfigError> {
|
||||
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<WireFormat, SwitchyardConfigError> {
|
||||
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<String>) -> 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::<RouteError>() {
|
||||
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]"));
|
||||
}
|
||||
}
|
||||
|
|
@ -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<litellm_gateway_auth::Principal>,
|
||||
) -> Router {
|
||||
router(Arc::new(gateway)).layer(axum::middleware::from_fn_with_state(
|
||||
(permissions, principal),
|
||||
test_identity,
|
||||
|
|
|
|||
405
litellm-rust/crates/gateway-inference/tests/switchyard.rs
Normal file
405
litellm-rust/crates/gateway-inference/tests/switchyard.rs
Normal file
|
|
@ -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)
|
||||
}
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -62,12 +62,15 @@ pub fn build_inference(config: &Config) -> Result<Arc<Gateway>, 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(
|
||||
|
|
|
|||
|
|
@ -33,6 +33,32 @@ fn inference() -> Arc<Gateway> {
|
|||
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)]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue