diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 4259ca1965b..8d30a5e0bd0 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3732,11 +3732,14 @@ version = "0.1.0" dependencies = [ "axum", "base64 0.22.1", + "bytes", "envy", "futures-util", "http-body-util", + "hyper-util", "litellm-auth-types", "litellm-config", + "litellm-core-utils", "litellm-gateway-auth", "litellm-gateway-inference", "litellm-gateway-mcp", @@ -3789,7 +3792,6 @@ dependencies = [ "base64 0.22.1", "bytes", "futures-util", - "hyper-util", "litellm-auth", "litellm-cache-memory", "litellm-cache-response", @@ -3814,7 +3816,6 @@ dependencies = [ "thiserror 2.0.19", "tokio", "tower", - "tracing", "wiremock", ] diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index a316ff58b50..ed8cdfd60d9 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -10,8 +10,6 @@ litellm-cache-response.workspace = true axum = { workspace = true, features = ["json", "multipart", "original-uri"] } base64.workspace = true bytes.workspace = true -futures-util.workspace = true -hyper-util = { workspace = true, features = ["client-legacy", "http1", "tokio"] } litellm-auth.workspace = true litellm-core-utils.workspace = true litellm-gateway-auth.workspace = true @@ -31,10 +29,10 @@ litellm-llms-types.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true -tracing.workspace = true [dev-dependencies] litellm-cache-memory.workspace = true +futures-util.workspace = true tokio = { workspace = true, features = ["io-util"] } rstest.workspace = true tower = { version = "0.5.3", features = ["util"] } diff --git a/litellm-rust/crates/gateway-inference/src/dispatch/mod.rs b/litellm-rust/crates/gateway-inference/src/dispatch/mod.rs deleted file mode 100644 index 845394da8f9..00000000000 --- a/litellm-rust/crates/gateway-inference/src/dispatch/mod.rs +++ /dev/null @@ -1,24 +0,0 @@ -//! What the gateway does with a request Rust core does not serve. -//! -//! A core decline travels up as an error and rides on the error response as a [`Decline`] -//! extension. [`replay`] is the only place that acts on it: it hands the request, as the -//! client sent it, to the [`Fallback`]. - -mod forward; -mod replay; - -pub use forward::{Forward, InvalidUpstream}; -pub(crate) use replay::replay; - -use litellm_core_utils::decline::Decline; - -/// The reason a request with no native route carries to Python. -pub const NOT_NATIVE: Decline = Decline::new("route not native"); - -#[derive(Clone)] -pub enum Fallback { - /// Standalone: an unserved path is a 404, a decline is the gateway's own 501. - Reject, - /// Next to the Python proxy: both go to Python unchanged. - Forward(Forward), -} diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index b980ea2d186..16f4d492133 100644 --- a/litellm-rust/crates/gateway-inference/src/lib.rs +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -6,7 +6,6 @@ mod audio_transcription; mod caching; mod chat_completions; -pub mod dispatch; mod error; pub mod messages; mod ocr; @@ -28,7 +27,7 @@ use litellm_secrets::source::SecretSource; pub use error::Error; pub use litellm_router::{Deployment, Router as ModelRouter}; -pub use request::{JsonObject, RequestId}; +pub use request::{JsonObject, MAX_BODY_BYTES, RequestId}; pub struct Gateway { cache: Option>, @@ -89,7 +88,7 @@ impl Gateway { } } -pub fn router(gateway: Arc, fallback: dispatch::Fallback) -> Router { +pub fn router(gateway: Arc) -> Router { Router::new() .route("/v1/messages", post(messages::create)) .route("/ocr", post(ocr::create)) @@ -105,10 +104,6 @@ pub fn router(gateway: Arc, fallback: dispatch::Fallback) -> Router { ) .route("/responses", post(responses::create)) .route("/v1/responses", post(responses::create)) - .route_layer(axum::middleware::from_fn_with_state( - fallback, - dispatch::replay, - )) .layer(axum::extract::DefaultBodyLimit::max( request::MAX_BODY_BYTES, )) diff --git a/litellm-rust/crates/gateway-inference/src/request.rs b/litellm-rust/crates/gateway-inference/src/request.rs index 293c37ec551..41f1aab3ebe 100644 --- a/litellm-rust/crates/gateway-inference/src/request.rs +++ b/litellm-rust/crates/gateway-inference/src/request.rs @@ -8,7 +8,7 @@ use serde_json::{Map, Value}; use crate::{Deployment, Error, Gateway}; pub(crate) const MAX_FILE_BYTES: usize = 50 * 1024 * 1024; -pub(crate) const MAX_BODY_BYTES: usize = MAX_FILE_BYTES + 1024 * 1024; +pub const MAX_BODY_BYTES: usize = MAX_FILE_BYTES + 1024 * 1024; pub(crate) struct Upload { pub bytes: Bytes, diff --git a/litellm-rust/crates/gateway-inference/tests/decline.rs b/litellm-rust/crates/gateway-inference/tests/decline.rs new file mode 100644 index 00000000000..65b0b694484 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/decline.rs @@ -0,0 +1,72 @@ +mod support; + +use litellm_core_utils::decline::Decline; +use rstest::rstest; +use serde_json::{Value, json}; +use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any}; + +async fn provider() -> MockServer { + let provider = MockServer::start().await; + Mock::given(any()) + .respond_with(ResponseTemplate::new(500).set_body_string("provider failed")) + .mount(&provider) + .await; + provider +} + +#[rstest] +#[case::chat_streaming( + "anthropic/test-model", + "/v1/chat/completions", + json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], "stream": true}), + "streaming" +)] +#[case::messages_bedrock_invoke( + "bedrock/invoke/anthropic.claude-3-5-sonnet-20240620-v1:0", + "/v1/messages", + json!({"model": "public/model", "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]}), + "Bedrock invoke messages request shaping" +)] +#[case::responses_provider( + "gemini/test-model", + "/v1/responses", + json!({"model": "public/model", "input": "hi"}), + "native HTTP responses provider" +)] +#[tokio::test] +async fn a_core_decline_is_a_501_that_carries_the_decline( + #[case] model: &str, + #[case] path: &str, + #[case] body: Value, + #[case] reason: &'static str, +) { + let provider = provider().await; + let response = support::post(support::app(model, &provider.uri()), path, body).await; + + assert_eq!(response.status(), 501); + assert_eq!( + response.extensions().get::(), + Some(&Decline::new(reason)) + ); + let body = support::json(response).await; + assert!(body.to_string().contains(reason)); + assert!(provider.received_requests().await.unwrap().is_empty()); +} + +#[rstest] +#[case::invalid_request(json!({"model": "public/model", "messages": "hi"}), 400)] +#[case::unknown_model(json!({"model": "other", "messages": [{"role": "user", "content": "hi"}]}), 400)] +#[case::provider_failure(json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), 500)] +#[tokio::test] +async fn a_terminal_error_carries_no_decline(#[case] body: Value, #[case] status: u16) { + let provider = provider().await; + let response = support::post( + support::app("anthropic/test-model", &provider.uri()), + "/v1/chat/completions", + body, + ) + .await; + + assert_eq!(response.status(), status); + assert!(response.extensions().get::().is_none()); +} diff --git a/litellm-rust/crates/gateway-inference/tests/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs index 5366cbd4118..0e37bb0aff9 100644 --- a/litellm-rust/crates/gateway-inference/tests/support/mod.rs +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -10,11 +10,7 @@ use axum::{ response::Response, }; use futures_util::future::BoxFuture; -use litellm_gateway_inference::{ - Deployment, Gateway, - dispatch::{Fallback, Forward}, - router, -}; +use litellm_gateway_inference::{Deployment, Gateway, router}; use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver}; use litellm_inference::resources::CoreResources; use litellm_secrets::{SecretValue, source::SecretSource}; @@ -41,18 +37,7 @@ pub fn app_with_permissions( api_base: &str, permissions: litellm_gateway_auth::Permissions, ) -> Router { - configured_app(model, api_base, permissions, None, None, Fallback::Reject) -} - -pub fn app_forwarding_to(python: &str, model: &str, api_base: &str) -> Router { - configured_app( - model, - api_base, - litellm_gateway_auth::Permissions::All, - None, - None, - Fallback::Forward(Forward::new(python).unwrap()), - ) + configured_app(model, api_base, permissions, None, None) } pub fn app_with_cache( @@ -66,7 +51,6 @@ pub fn app_with_cache( litellm_gateway_auth::Permissions::All, Some(cache), None, - Fallback::Reject, ) } @@ -82,7 +66,6 @@ pub fn app_with_cache_for_principal( litellm_gateway_auth::Permissions::All, Some(cache), Some(principal), - Fallback::Reject, ) } @@ -92,7 +75,6 @@ fn configured_app( permissions: litellm_gateway_auth::Permissions, cache: Option>, principal: Option, - fallback: Fallback, ) -> Router { let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver))); let http = Resolution::from(&HttpSettings::default()).config; @@ -120,7 +102,7 @@ fn configured_app( Some(cache) => gateway.with_cache(cache), None => gateway, }; - router(Arc::new(gateway), fallback).layer(axum::middleware::from_fn_with_state( + router(Arc::new(gateway)).layer(axum::middleware::from_fn_with_state( (permissions, principal), test_identity, )) diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index e8edf2b114a..28eab1466a1 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -6,9 +6,12 @@ license.workspace = true repository.workspace = true [dependencies] -axum.workspace = true +axum = { workspace = true, features = ["original-uri"] } +bytes.workspace = true envy = "0.4.2" http-body-util = "0.1" +hyper-util = { workspace = true, features = ["client-legacy", "http1", "tokio"] } +litellm-core-utils.workspace = true litellm-inference.workspace = true litellm-gateway-inference.workspace = true litellm-gateway-mcp.workspace = true diff --git a/litellm-rust/crates/gateway-inference/src/dispatch/forward.rs b/litellm-rust/crates/gateway/src/dispatch/forward.rs similarity index 96% rename from litellm-rust/crates/gateway-inference/src/dispatch/forward.rs rename to litellm-rust/crates/gateway/src/dispatch/forward.rs index dc24f19b4a0..664c4bfe121 100644 --- a/litellm-rust/crates/gateway-inference/src/dispatch/forward.rs +++ b/litellm-rust/crates/gateway/src/dispatch/forward.rs @@ -20,7 +20,7 @@ use hyper_util::{ use litellm_core_utils::decline::Decline; /// Python reads it to log why the gateway handed the request over. -pub const DECLINED_HEADER: HeaderName = HeaderName::from_static("x-litellm-gateway-declined"); +const DECLINED_HEADER: HeaderName = HeaderName::from_static("x-litellm-gateway-declined"); const X_FORWARDED_FOR: HeaderName = HeaderName::from_static("x-forwarded-for"); @@ -64,7 +64,7 @@ impl Forward { }) } - pub async fn send(&self, request: Request, decline: Decline) -> Response { + pub(super) async fn send(&self, request: Request, decline: Decline) -> Response { let (mut parts, body) = request.into_parts(); let path = parts .extensions diff --git a/litellm-rust/crates/gateway/src/dispatch/mod.rs b/litellm-rust/crates/gateway/src/dispatch/mod.rs new file mode 100644 index 00000000000..dab89790257 --- /dev/null +++ b/litellm-rust/crates/gateway/src/dispatch/mod.rs @@ -0,0 +1,52 @@ +//! What the gateway does with a request it does not serve natively. +//! +//! A mounted router declines by putting a [`Decline`] in its response's extensions. Nothing +//! else crosses this boundary: only this module records requests and talks to Python. + +mod forward; +mod replay; + +pub use forward::{Forward, InvalidUpstream}; + +use axum::{Router, extract::Request}; +use litellm_core_utils::decline::Decline; + +const NOT_NATIVE: Decline = Decline::new("route not native"); + +#[derive(Clone)] +pub enum Fallback { + /// Standalone: an unserved path is a 404, a decline is the mounted router's own error. + Reject, + /// Next to the Python proxy: both go to Python unchanged. + Forward(Forward), +} + +/// Forwards declined requests on the routes `served` holds so far. `record_limit` must cover +/// the largest body those routes accept, or a declined request that size keeps its error. +pub(crate) fn replay_declines(served: Router, fallback: &Fallback, record_limit: usize) -> Router { + match fallback { + Fallback::Reject => served, + Fallback::Forward(forward) => served.route_layer(axum::middleware::from_fn_with_state( + replay::Replay { + forward: forward.clone(), + record_limit, + }, + replay::replay, + )), + } +} + +pub(crate) fn forward_unmatched(app: Router, fallback: Fallback) -> Router { + match fallback { + Fallback::Reject => app, + Fallback::Forward(forward) => { + let unmatched = forward.clone(); + app.fallback(move |request: Request| async move { + unmatched.send(request, NOT_NATIVE).await + }) + .method_not_allowed_fallback(move |request: Request| async move { + forward.send(request, NOT_NATIVE).await + }) + } + } +} diff --git a/litellm-rust/crates/gateway-inference/src/dispatch/replay.rs b/litellm-rust/crates/gateway/src/dispatch/replay.rs similarity index 81% rename from litellm-rust/crates/gateway-inference/src/dispatch/replay.rs rename to litellm-rust/crates/gateway/src/dispatch/replay.rs index 5c4ff545013..0d371cd7902 100644 --- a/litellm-rust/crates/gateway-inference/src/dispatch/replay.rs +++ b/litellm-rust/crates/gateway/src/dispatch/replay.rs @@ -11,18 +11,26 @@ use bytes::BytesMut; use futures_util::StreamExt; use litellm_core_utils::decline::Decline; -use super::Fallback; -use crate::request::MAX_BODY_BYTES; +use super::Forward; + +#[derive(Clone)] +pub(super) struct Replay { + pub forward: Forward, + pub record_limit: usize, +} /// Records the request before any handler reads it, and forwards the recording when the /// response says the request was declined. -pub(crate) async fn replay( - State(fallback): State, +pub(super) async fn replay( + State(Replay { + forward, + record_limit, + }): State, request: Request, next: Next, ) -> Response { let (parts, body) = request.into_parts(); - let body = match buffer(body).await { + let body = match buffer(body, record_limit).await { Ok(body) => body, Err(body) => return next.run(Request::from_parts(parts, body)).await, }; @@ -31,20 +39,17 @@ pub(crate) async fn replay( let Some(&decline) = response.extensions().get::() else { return response; }; - match fallback { - Fallback::Reject => response, - Fallback::Forward(forward) => forward.send(recorded.map(Body::from), decline).await, - } + forward.send(recorded.map(Body::from), decline).await } /// The whole body, or, when it is too large or unreadable, a body that yields the same bytes /// and the same failure, so the handler reports it in its own error envelope. -async fn buffer(body: Body) -> Result { +async fn buffer(body: Body, limit: usize) -> Result { let mut stream = body.into_data_stream(); let mut seen = BytesMut::new(); while let Some(chunk) = stream.next().await { match chunk { - Ok(chunk) if seen.len() + chunk.len() <= MAX_BODY_BYTES => { + Ok(chunk) if seen.len() + chunk.len() <= limit => { seen.extend_from_slice(&chunk); } failed => { diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 44de2ff9570..d74fc670c09 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -1,6 +1,8 @@ mod auth; +mod dispatch; mod error; mod secrets; +pub use dispatch::{Fallback, Forward, InvalidUpstream}; pub use error::Error; use std::{sync::Arc, time::Instant}; @@ -16,10 +18,7 @@ use http_body_util::BodyExt; use litellm_config::Config; use litellm_gateway_auth::Auth; -use litellm_gateway_inference::{ - Gateway, ModelRouter, - dispatch::{Fallback, NOT_NATIVE}, -}; +use litellm_gateway_inference::{Gateway, ModelRouter}; use litellm_http::{ ClientVariant, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify, media::PublicDnsResolver, @@ -124,28 +123,17 @@ pub fn router( fallback: Fallback, ) -> Router { let auth = Auth::from_config(config, inference.secrets.clone()); - let served = litellm_gateway_inference::router(inference, fallback.clone()) - .merge(mcp.unwrap_or_default()) - .route_layer(axum::middleware::from_fn(auth::bind_session_owner)) - .route_layer(axum::middleware::from_fn_with_state( - auth, - litellm_gateway_auth::authenticate, - )) - .layer(axum::middleware::from_fn(log_request)) - .merge(ui.unwrap_or_default()); - match fallback { - Fallback::Reject => served, - Fallback::Forward(forward) => { - let unmatched = forward.clone(); - served - .fallback(move |request: Request| async move { - unmatched.send(request, NOT_NATIVE).await - }) - .method_not_allowed_fallback(move |request: Request| async move { - forward.send(request, NOT_NATIVE).await - }) - } - } + let routes = litellm_gateway_inference::router(inference).merge(mcp.unwrap_or_default()); + let served = + dispatch::replay_declines(routes, &fallback, litellm_gateway_inference::MAX_BODY_BYTES) + .route_layer(axum::middleware::from_fn(auth::bind_session_owner)) + .route_layer(axum::middleware::from_fn_with_state( + auth, + litellm_gateway_auth::authenticate, + )) + .layer(axum::middleware::from_fn(log_request)) + .merge(ui.unwrap_or_default()); + dispatch::forward_unmatched(served, fallback) } async fn log_request(request: Request, next: Next) -> Response { diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index 81fc214bb4b..ff489448b85 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -4,8 +4,8 @@ use std::{ }; use litellm_config::Config; +use litellm_gateway::{Fallback, Forward}; use litellm_gateway_auth::UiBackend; -use litellm_gateway_inference::dispatch::{Fallback, Forward}; use litellm_tracing::{Level, Logger, Metadata, Record, Sink}; use serde_json::json; diff --git a/litellm-rust/crates/gateway-inference/tests/dispatch.rs b/litellm-rust/crates/gateway/tests/dispatch.rs similarity index 74% rename from litellm-rust/crates/gateway-inference/tests/dispatch.rs rename to litellm-rust/crates/gateway/tests/dispatch.rs index b546a1a89a6..a063d81cf48 100644 --- a/litellm-rust/crates/gateway-inference/tests/dispatch.rs +++ b/litellm-rust/crates/gateway/tests/dispatch.rs @@ -1,17 +1,41 @@ -mod support; - use std::net::SocketAddr; use axum::{ + Router, body::{Body, to_bytes}, extract::ConnectInfo, http::Request, + response::Response, }; +use litellm_config::Config; +use litellm_gateway::{Fallback, Forward}; use rstest::rstest; -use serde_json::json; +use serde_json::{Value, json}; use tower::ServiceExt; use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any}; +fn app(python: &str, model: &str, provider: &str) -> Router { + let config = Config::from_yaml(&format!( + "general_settings:\n master_key: gateway-key\nmodel_list:\n - model_name: public/model\n litellm_params:\n model: {model}\n api_base: {provider}\n api_key: test-key\n" + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let fallback = Fallback::Forward(Forward::new(python).unwrap()); + litellm_gateway::router(inference, &config, None, None, fallback) +} + +async fn send(app: Router, path: &str, body: Value) -> Response { + app.oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", "Bearer gateway-key") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap() +} + async fn python() -> MockServer { let python = MockServer::start().await; Mock::given(any()) @@ -62,13 +86,13 @@ async fn a_core_decline_reaches_python_as_the_client_sent_it( ) { let python = python().await; let provider = provider().await; - let app = support::app_forwarding_to(&python.uri(), model, &provider.uri()); + let app = app(&python.uri(), model, &provider.uri()); let response = app .oneshot( Request::post(path) .header("content-type", "application/json") - .header("authorization", "Bearer client-key") + .header("authorization", "Bearer gateway-key") .header("x-client-header", "kept") .header("connection", "te") .header("te", "trailers") @@ -95,44 +119,22 @@ async fn a_core_decline_reaches_python_as_the_client_sent_it( }; assert_eq!(forwarded_path, path); assert_eq!(forwarded.body, body.as_bytes()); - assert_eq!(forwarded.headers["authorization"], "Bearer client-key"); + assert_eq!(forwarded.headers["authorization"], "Bearer gateway-key"); assert_eq!(forwarded.headers["x-client-header"], "kept"); assert_eq!(forwarded.headers["x-litellm-gateway-declined"], reason); assert!(!forwarded.headers.contains_key("te")); } -#[rstest] -#[tokio::test] -async fn a_standalone_gateway_answers_a_decline_with_its_own_501() { - let provider = provider().await; - let response = support::post( - support::app("gemini/test-model", &provider.uri()), - "/v1/responses", - json!({"model": "public/model", "input": "hi"}), - ) - .await; - - assert_eq!(response.status(), 501); - let body = support::json(response).await; - assert!( - body["error"]["message"] - .as_str() - .unwrap() - .contains("native HTTP responses provider") - ); - assert!(provider.received_requests().await.unwrap().is_empty()); -} - #[rstest] #[case::invalid_request(json!({"model": "public/model", "messages": "hi"}), 400)] #[case::unknown_model(json!({"model": "other", "messages": [{"role": "user", "content": "hi"}]}), 400)] #[case::provider_failure(json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), 500)] #[tokio::test] -async fn a_terminal_error_is_never_forwarded(#[case] body: serde_json::Value, #[case] status: u16) { +async fn a_terminal_error_is_never_forwarded(#[case] body: Value, #[case] status: u16) { let python = python().await; let provider = provider().await; - let response = support::post( - support::app_forwarding_to(&python.uri(), "anthropic/test-model", &provider.uri()), + let response = send( + app(&python.uri(), "anthropic/test-model", &provider.uri()), "/v1/chat/completions", body, ) @@ -152,7 +154,9 @@ async fn a_forwarded_decline_names_the_client_in_x_forwarded_for( ) { let python = python().await; let provider = provider().await; - let request = Request::post("/v1/chat/completions").header("content-type", "application/json"); + let request = Request::post("/v1/chat/completions") + .header("content-type", "application/json") + .header("authorization", "Bearer gateway-key"); let request = match incoming { Some(chain) => request.header("x-forwarded-for", chain), None => request, @@ -167,7 +171,7 @@ async fn a_forwarded_decline_names_the_client_in_x_forwarded_for( "203.0.113.7:50000".parse::().unwrap(), )); - support::app_forwarding_to(&python.uri(), "anthropic/test-model", &provider.uri()) + app(&python.uri(), "anthropic/test-model", &provider.uri()) .oneshot(request) .await .unwrap(); diff --git a/litellm-rust/crates/gateway/tests/mcp.rs b/litellm-rust/crates/gateway/tests/mcp.rs index a017a440efc..69212512f83 100644 --- a/litellm-rust/crates/gateway/tests/mcp.rs +++ b/litellm-rust/crates/gateway/tests/mcp.rs @@ -6,7 +6,7 @@ use axum::{ http::Request, }; use litellm_config::Config; -use litellm_gateway_inference::dispatch::Fallback; +use litellm_gateway::Fallback; use litellm_gateway_mcp::{ HttpConfig, rmcp::{ diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index 24093640f5e..b4ccb90fa23 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -6,10 +6,8 @@ use std::{ use axum::{body::Body, http::Request}; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; use litellm_config::Config; -use litellm_gateway_inference::{ - Error, Gateway, - dispatch::{Fallback, Forward}, -}; +use litellm_gateway::{Fallback, Forward}; +use litellm_gateway_inference::{Error, Gateway}; use litellm_http::ClientVariant; use litellm_tracing::{Logger, Metadata, Record, Sink}; use rstest::{fixture, rstest};