From e4511b7aa8aaad6d72c176610fe1955f4371bebd Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 6 Oct 2026 15:45:10 -0700 Subject: [PATCH] refactor(rust): add gateway and router hook seams --- litellm-rust/Cargo.lock | 3 + .../src/audio_transcription.rs | 16 +- .../crates/gateway-inference/src/caching.rs | 33 +++- .../gateway-inference/src/chat_completions.rs | 25 ++- .../crates/gateway-inference/src/lib.rs | 84 +++++--- .../crates/gateway-inference/src/messages.rs | 15 +- .../crates/gateway-inference/src/ocr.rs | 16 +- .../crates/gateway-inference/src/request.rs | 7 +- .../crates/gateway-inference/src/responses.rs | 14 +- litellm-rust/crates/gateway/Cargo.toml | 2 + litellm-rust/crates/gateway/src/hooks.rs | 23 +++ litellm-rust/crates/gateway/src/lib.rs | 33 +++- litellm-rust/crates/gateway/tests/hooks.rs | 183 ++++++++++++++++++ litellm-rust/crates/router/Cargo.toml | 1 + litellm-rust/crates/router/src/hooks.rs | 15 ++ litellm-rust/crates/router/src/lib.rs | 34 +++- litellm-rust/crates/router/tests/hooks.rs | 54 ++++++ 17 files changed, 487 insertions(+), 71 deletions(-) create mode 100644 litellm-rust/crates/gateway/src/hooks.rs create mode 100644 litellm-rust/crates/gateway/tests/hooks.rs create mode 100644 litellm-rust/crates/router/src/hooks.rs create mode 100644 litellm-rust/crates/router/tests/hooks.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index a3ec1f5150c..1690f825434 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3781,6 +3781,7 @@ dependencies = [ "litellm-gateway-inference", "litellm-gateway-mcp", "litellm-gateway-ui", + "litellm-host", "litellm-http", "litellm-inference", "litellm-llms", @@ -3797,6 +3798,7 @@ dependencies = [ "tower-sessions-moka-store", "tracing", "uuid", + "wiremock", ] [[package]] @@ -4312,6 +4314,7 @@ dependencies = [ "litellm-inference", "litellm-inference-messages", "rstest", + "tokio", ] [[package]] diff --git a/litellm-rust/crates/gateway-inference/src/audio_transcription.rs b/litellm-rust/crates/gateway-inference/src/audio_transcription.rs index 8d571eccc5b..909429845f5 100644 --- a/litellm-rust/crates/gateway-inference/src/audio_transcription.rs +++ b/litellm-rust/crates/gateway-inference/src/audio_transcription.rs @@ -1,4 +1,7 @@ use litellm_gateway_auth::AuthenticatedRequest; +use litellm_host::interceptors::Interceptors; +use litellm_inference::RouteError; +use litellm_router::RouterHooks; use std::{path::Path, sync::Arc}; use axum::{Json, extract::State, response::IntoResponse}; @@ -11,23 +14,26 @@ use crate::{ request::{self, InferenceBody}, }; -pub(crate) async fn create( - State(gateway): State>, +pub(crate) async fn create< + R: RouterHooks + 'static, + I: Interceptors + Clone + 'static, +>( + State(gateway): State>>, identity: AuthenticatedRequest, body: InferenceBody, ) -> Result { handle(&gateway, &identity, body).await.map(Json) } -async fn handle( - gateway: &Gateway, +async fn handle + Clone + 'static>( + gateway: &Gateway, identity: &AuthenticatedRequest, InferenceBody { fields: body, upload, }: InferenceBody, ) -> Result { - let deployment = request::resolve_deployment(gateway, &body)?; + let deployment = request::resolve_deployment(gateway, &body).await?; request::authorize_model(identity, deployment, &body).await?; let audio = match upload { Some(upload) => { diff --git a/litellm-rust/crates/gateway-inference/src/caching.rs b/litellm-rust/crates/gateway-inference/src/caching.rs index 1abea61b164..b43fa07515c 100644 --- a/litellm-rust/crates/gateway-inference/src/caching.rs +++ b/litellm-rust/crates/gateway-inference/src/caching.rs @@ -69,39 +69,52 @@ fn duration(seconds: f64) -> Result { .ok_or_else(|| Error::InvalidBody("cache durations must be finite and positive".into())) } -#[derive(Clone, Default)] -pub(crate) struct CacheHeaders(std::sync::Arc>); +#[derive(Clone)] +pub(crate) struct CacheHeaders { + key: std::sync::Arc>, + hooks: H, +} -impl litellm_host::interceptors::Interceptors for CacheHeaders { +impl> + litellm_host::interceptors::Interceptors for CacheHeaders +{ async fn before_provider_request( &self, wire: litellm_host::interceptors::WireRequest, - _: litellm_host::interceptors::RequestContext, + context: litellm_host::interceptors::RequestContext, ) -> Result { - Ok(wire) + self.hooks.before_provider_request(wire, context).await } async fn after_provider_response( &self, - _: litellm_host::interceptors::RawResponse, + raw: litellm_host::interceptors::RawResponse, ) -> Result<(), litellm_inference::RouteError> { - Ok(()) + self.hooks.after_provider_response(raw).await } async fn result_ready( &self, facts: litellm_host::interceptors::ExecutionFacts, ) -> Result<(), litellm_inference::RouteError> { + self.hooks.result_ready(facts.clone()).await?; if let litellm_host::interceptors::ResultSource::Cache { key } = facts.source { - let _ = self.0.set(key); + let _ = self.key.set(key); } Ok(()) } } -impl CacheHeaders { +impl CacheHeaders { + pub(crate) fn new(hooks: H) -> Self { + Self { + key: Default::default(), + hooks, + } + } + pub(crate) fn apply(&self, mut response: axum::response::Response) -> axum::response::Response { - if let Some(key) = self.0.get() + if let Some(key) = self.key.get() && let Ok(value) = axum::http::HeaderValue::from_str(key) { response.headers_mut().insert("x-litellm-cache-key", value); diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index 742a373b286..a560d85b20f 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -1,4 +1,7 @@ use litellm_gateway_auth::AuthenticatedRequest; +use litellm_host::interceptors::Interceptors; +use litellm_inference::RouteError; +use litellm_router::RouterHooks; use std::sync::Arc; use axum::{ @@ -11,16 +14,22 @@ use serde_json::{Map, Value}; use crate::{Error, Gateway, JsonObject, request}; -pub(crate) async fn create( - State(gateway): State>, +pub(crate) async fn create< + R: RouterHooks + 'static, + I: Interceptors + Clone + 'static, +>( + State(gateway): State>>, identity: AuthenticatedRequest, JsonObject(body): JsonObject, ) -> Result { handle(&gateway, &identity, body).await } -pub(crate) async fn create_from_model_path( - State(gateway): State>, +pub(crate) async fn create_from_model_path< + R: RouterHooks + 'static, + I: Interceptors + Clone + 'static, +>( + State(gateway): State>>, identity: AuthenticatedRequest, Path(model): Path, JsonObject(body): JsonObject, @@ -35,12 +44,12 @@ pub(crate) async fn create_from_model_path( handle(&gateway, &identity, body).await } -async fn handle( - gateway: &Gateway, +async fn handle + Clone + 'static>( + gateway: &Gateway, identity: &AuthenticatedRequest, body: Map, ) -> Result { - let deployment = request::resolve_deployment(gateway, &body)?; + let deployment = request::resolve_deployment(gateway, &body).await?; request::authorize_model(identity, deployment, &body).await?; let (body, cache_options) = crate::caching::prepare(identity, body)?; let route = gateway.chat_completions.clone(); @@ -53,7 +62,7 @@ async fn handle( }; let messages = body.get("messages").cloned().unwrap_or_default(); - let headers = crate::caching::CacheHeaders::default(); + let headers = crate::caching::CacheHeaders::new(gateway.interceptors.clone()); let response = litellm_host_http::serve_unary( route.machine( ChatCompletionsCall { diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index 4663151f1de..fcd3569d553 100644 --- a/litellm-rust/crates/gateway-inference/src/lib.rs +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -26,30 +26,26 @@ use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings}; use litellm_secrets::source::SecretSource; pub use error::Error; -pub use litellm_router::{Deployment, Router as ModelRouter}; +use litellm_host::interceptors::Interceptors; +use litellm_inference::RouteError; +pub use litellm_router::{Deployment, Router as ModelRouter, RouterHooks}; pub use request::{JsonObject, RequestId}; -pub struct Gateway { +pub struct Gateway { cache: Option>, pub audio_transcription: AudioTranscriptionRoute, pub chat_completions: ChatCompletionsRoute, pub messages: MessagesRoute, pub ocr: OcrRoute, pub responses: ResponsesRoute, - pub models: ModelRouter, + pub models: ModelRouter, + pub interceptors: I, pub secrets: Arc, pub resources: CoreResources, pub http: HttpClientConfig, } impl Gateway { - pub fn with_cache(self, cache: Arc) -> Self { - Self { - cache: Some(cache), - ..self - } - } - pub fn new( resources: CoreResources, http: HttpClientConfig, @@ -60,6 +56,7 @@ impl Gateway { let auth = resources.auth.clone(); Ok(Self { cache: None, + interceptors: (), audio_transcription: AudioTranscriptionRoute::new( provider.clone(), auth.clone(), @@ -88,22 +85,57 @@ impl Gateway { } } -pub fn router(gateway: Arc) -> Router { +impl Gateway { + pub fn with_cache(self, cache: Arc) -> Self { + Self { + cache: Some(cache), + ..self + } + } + + pub fn with_hooks(self, router_hooks: T, interceptors: H) -> Gateway { + Gateway { + cache: self.cache, + audio_transcription: self.audio_transcription, + chat_completions: self.chat_completions, + messages: self.messages, + ocr: self.ocr, + responses: self.responses, + models: self.models.with_hooks(router_hooks), + interceptors, + secrets: self.secrets, + resources: self.resources, + http: self.http, + } + } +} + +pub fn router(gateway: Arc>) -> Router +where + R: RouterHooks + 'static, + I: Interceptors + Clone + 'static, +{ Router::new() - .route("/v1/messages", post(messages::create)) - .route("/ocr", post(ocr::create)) - .route("/v1/ocr", post(ocr::create)) - .route("/chat/completions", post(chat_completions::create)) - .route("/v1/chat/completions", post(chat_completions::create)) - .nest("/engines/{model}", model_routes()) - .nest("/openai/deployments/{model}", model_routes()) - .route("/audio/transcriptions", post(audio_transcription::create)) + .route("/v1/messages", post(messages::create::)) + .route("/ocr", post(ocr::create::)) + .route("/v1/ocr", post(ocr::create::)) + .route("/chat/completions", post(chat_completions::create::)) + .route( + "/v1/chat/completions", + post(chat_completions::create::), + ) + .nest("/engines/{model}", model_routes::()) + .nest("/openai/deployments/{model}", model_routes::()) + .route( + "/audio/transcriptions", + post(audio_transcription::create::), + ) .route( "/v1/audio/transcriptions", - post(audio_transcription::create), + post(audio_transcription::create::), ) - .route("/responses", post(responses::create)) - .route("/v1/responses", post(responses::create)) + .route("/responses", post(responses::create::)) + .route("/v1/responses", post(responses::create::)) .route("/embeddings", post(request::unsupported)) .route("/v1/embeddings", post(request::unsupported)) .route("/completions", post(request::unsupported)) @@ -114,11 +146,15 @@ pub fn router(gateway: Arc) -> Router { .with_state(gateway) } -fn model_routes() -> Router> { +fn model_routes() -> Router>> +where + R: RouterHooks + 'static, + I: Interceptors + Clone + 'static, +{ Router::new() .route( "/chat/completions", - post(chat_completions::create_from_model_path), + post(chat_completions::create_from_model_path::), ) .route("/embeddings", post(request::unsupported)) .route("/completions", post(request::unsupported)) diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index 8969ad80934..f19b1237690 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -1,6 +1,9 @@ //! `POST /v1/messages`, as the Python proxy's `anthropic_response` serves it. use litellm_gateway_auth::AuthenticatedRequest; +use litellm_host::interceptors::Interceptors; +use litellm_inference::RouteError; +use litellm_router::RouterHooks; use std::sync::Arc; use axum::{ @@ -21,8 +24,8 @@ use crate::{Deployment, Error, Gateway, JsonObject, RequestId, request}; const ANTHROPIC_API_HEADERS: [&str; 2] = ["anthropic-version", "anthropic-beta"]; const ANTHROPIC_API_HEADER_PROVIDERS: &str = "anthropic,bedrock,bedrock_mantle,vertex_ai"; -pub async fn create( - State(gateway): State>, +pub async fn create + Clone + 'static>( + State(gateway): State>>, identity: AuthenticatedRequest, RequestId(request_id): RequestId, headers: HeaderMap, @@ -35,13 +38,13 @@ pub async fn create( result.map_err(|error| (error.status(), Json(error.body(request_id.as_deref())))) } -async fn handle( - gateway: &Gateway, +async fn handle + Clone + 'static>( + gateway: &Gateway, identity: &AuthenticatedRequest, headers: &HeaderMap, body: Map, ) -> Result { - let deployment = request::resolve_deployment(gateway, &body)?; + let deployment = request::resolve_deployment(gateway, &body).await?; request::authorize_model(identity, deployment, &body).await?; let (body, cache_options) = crate::caching::prepare(identity, body)?; let route = gateway.messages.clone(); @@ -57,7 +60,7 @@ async fn handle( let machine = route.machine(call, cache_options.policy); let stream = Sse::::new(Json, |error| Bytes::from(Error::from(error).sse_frame())); - let headers = crate::caching::CacheHeaders::default(); + let headers = crate::caching::CacheHeaders::new(gateway.interceptors.clone()); let response = litellm_host_http::serve(machine, (), headers.clone(), stream, None).await?; Ok(headers.apply(response)) } diff --git a/litellm-rust/crates/gateway-inference/src/ocr.rs b/litellm-rust/crates/gateway-inference/src/ocr.rs index 03d064b6bae..c78296a8127 100644 --- a/litellm-rust/crates/gateway-inference/src/ocr.rs +++ b/litellm-rust/crates/gateway-inference/src/ocr.rs @@ -1,4 +1,7 @@ use litellm_gateway_auth::AuthenticatedRequest; +use litellm_host::interceptors::Interceptors; +use litellm_inference::RouteError; +use litellm_router::RouterHooks; use std::sync::Arc; use axum::{Json, extract::State, http::HeaderMap, response::IntoResponse}; @@ -13,8 +16,11 @@ use crate::{ request::{self, InferenceBody}, }; -pub(crate) async fn create( - State(gateway): State>, +pub(crate) async fn create< + R: RouterHooks + 'static, + I: Interceptors + Clone + 'static, +>( + State(gateway): State>>, identity: AuthenticatedRequest, headers: HeaderMap, body: InferenceBody, @@ -22,8 +28,8 @@ pub(crate) async fn create( handle(&gateway, &identity, &headers, body).await.map(Json) } -async fn handle( - gateway: &Gateway, +async fn handle + Clone + 'static>( + gateway: &Gateway, identity: &AuthenticatedRequest, headers: &HeaderMap, InferenceBody { @@ -35,7 +41,7 @@ async fn handle( .get("x-req-format") .and_then(|value| value.to_str().ok()) .map(str::to_owned); - let deployment = request::resolve_deployment(gateway, &body)?; + let deployment = request::resolve_deployment(gateway, &body).await?; request::authorize_model(identity, deployment, &body).await?; let document = match upload { Some(upload) => OcrDocumentInput::Bytes { diff --git a/litellm-rust/crates/gateway-inference/src/request.rs b/litellm-rust/crates/gateway-inference/src/request.rs index b398da2013b..e899592ca4a 100644 --- a/litellm-rust/crates/gateway-inference/src/request.rs +++ b/litellm-rust/crates/gateway-inference/src/request.rs @@ -110,8 +110,8 @@ pub(crate) async fn authorize_model( Ok(()) } -pub(crate) fn resolve_deployment<'a>( - gateway: &'a Gateway, +pub(crate) async fn resolve_deployment<'a, R: crate::RouterHooks, I>( + gateway: &'a Gateway, body: &Map, ) -> Result<&'a Deployment, Error> { let model = body @@ -120,7 +120,8 @@ pub(crate) fn resolve_deployment<'a>( .ok_or_else(|| Error::InvalidBody("model is required".into()))?; gateway .models - .get(model) + .select(model) + .await .ok_or_else(|| Error::UnknownModel(model.to_owned())) } diff --git a/litellm-rust/crates/gateway-inference/src/responses.rs b/litellm-rust/crates/gateway-inference/src/responses.rs index 77bebf610e1..e0164155f42 100644 --- a/litellm-rust/crates/gateway-inference/src/responses.rs +++ b/litellm-rust/crates/gateway-inference/src/responses.rs @@ -1,3 +1,6 @@ +use litellm_host::interceptors::Interceptors; +use litellm_inference::RouteError; +use litellm_router::RouterHooks; use std::sync::Arc; use axum::{Json, body::Bytes, extract::State, response::Response}; @@ -8,12 +11,15 @@ use serde_json::json; use crate::{Error, Gateway, JsonObject, request}; -pub(crate) async fn create( - State(gateway): State>, +pub(crate) async fn create< + R: RouterHooks + 'static, + I: Interceptors + Clone + 'static, +>( + State(gateway): State>>, identity: AuthenticatedRequest, JsonObject(body): JsonObject, ) -> Result { - let deployment = request::resolve_deployment(&gateway, &body)?; + let deployment = request::resolve_deployment(&gateway, &body).await?; request::authorize_model(&identity, deployment, &body).await?; let (body, cache_options) = crate::caching::prepare(&identity, body)?; let route = gateway.responses.clone(); @@ -46,7 +52,7 @@ pub(crate) async fn create( json!({"type": "error", "code": error.status().as_u16().to_string(), "message": error.to_string(), "param": null}) )) }); - let headers = crate::caching::CacheHeaders::default(); + let headers = crate::caching::CacheHeaders::new(gateway.interceptors.clone()); let response = litellm_host_http::serve(machine, (), headers.clone(), stream, None).await?; Ok(headers.apply(response)) } diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 387bcd43231..e4471fe28d7 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -10,6 +10,7 @@ axum.workspace = true envy = "0.4.2" http-body-util = "0.1" litellm-inference.workspace = true +litellm-host.workspace = true litellm-gateway-inference.workspace = true litellm-gateway-mcp.workspace = true tokio-util = "0.7" @@ -31,6 +32,7 @@ uuid.workspace = true tower-sessions-moka-store = "0.15.0" [dev-dependencies] +wiremock.workspace = true base64.workspace = true rstest.workspace = true tempfile.workspace = true diff --git a/litellm-rust/crates/gateway/src/hooks.rs b/litellm-rust/crates/gateway/src/hooks.rs new file mode 100644 index 00000000000..a88c1921aea --- /dev/null +++ b/litellm-rust/crates/gateway/src/hooks.rs @@ -0,0 +1,23 @@ +use std::future::Future; + +use axum::{extract::Request, middleware::Next, response::Response}; + +pub trait GatewayHooks: Send + Sync { + fn pre_call(&self, request: Request) -> impl Future> + Send { + async { Ok(request) } + } + + fn post_call(&self, response: Response) -> impl Future + Send { + async { response } + } +} + +impl GatewayHooks for () {} + +pub(crate) async fn run(hooks: &H, request: Request, next: Next) -> Response { + let request = match hooks.pre_call(request).await { + Ok(request) => request, + Err(response) => return response, + }; + hooks.post_call(next.run(request).await).await +} diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index d711eb78caa..046d692a46d 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -1,5 +1,6 @@ mod auth; mod error; +pub mod hooks; mod secrets; pub use error::Error; @@ -113,14 +114,40 @@ pub async fn build_mcp( )) } -pub fn router( - inference: Arc, +pub fn router( + inference: Arc>, config: &Config, ui: Option, mcp: Option, -) -> Router { +) -> Router +where + R: litellm_gateway_inference::RouterHooks + 'static, + I: litellm_host::interceptors::Interceptors + Clone + 'static, +{ + router_with_hooks(inference, config, ui, mcp, ()) +} + +pub fn router_with_hooks( + inference: Arc>, + config: &Config, + ui: Option, + mcp: Option, + hooks: H, +) -> Router +where + R: litellm_gateway_inference::RouterHooks + 'static, + I: litellm_host::interceptors::Interceptors + Clone + 'static, + H: hooks::GatewayHooks + 'static, +{ + let hooks = Arc::new(hooks); let auth = Auth::from_config(config, inference.secrets.clone()); let inference = litellm_gateway_inference::router(inference) + .route_layer(axum::middleware::from_fn( + move |request: Request, next: Next| { + let hooks = hooks.clone(); + async move { hooks::run(hooks.as_ref(), request, next).await } + }, + )) .merge(mcp.unwrap_or_default()) .route_layer(axum::middleware::from_fn(auth::bind_session_owner)) .route_layer(axum::middleware::from_fn_with_state( diff --git a/litellm-rust/crates/gateway/tests/hooks.rs b/litellm-rust/crates/gateway/tests/hooks.rs new file mode 100644 index 00000000000..f2cb63a92eb --- /dev/null +++ b/litellm-rust/crates/gateway/tests/hooks.rs @@ -0,0 +1,183 @@ +use std::sync::{Arc, mpsc}; + +use axum::{ + body::{Body, to_bytes}, + extract::Request, + http::StatusCode, + response::{IntoResponse, Response}, +}; +use litellm_config::Config; +use litellm_gateway::hooks::GatewayHooks; +use litellm_gateway_inference::{Deployment, RouterHooks}; +use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest}; +use litellm_inference::RouteError; +use rstest::{fixture, rstest}; +use serde_json::{Value, json}; +use tower::ServiceExt; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{header, method, path}, +}; + +#[derive(Clone)] +struct Recorder { + phases: mpsc::Sender<&'static str>, + reject_gateway: bool, + reject_router: bool, +} + +impl GatewayHooks for Recorder { + async fn pre_call(&self, request: Request) -> Result { + self.phases.send("gateway.pre_call").unwrap(); + if self.reject_gateway { + return Err((StatusCode::FORBIDDEN, "rejected by gateway hook").into_response()); + } + Ok(request) + } + + async fn post_call(&self, response: Response) -> Response { + self.phases.send("gateway.post_call").unwrap(); + ([("x-hook-completed", "true")], response).into_response() + } +} + +impl RouterHooks for Recorder { + async fn filter_deployments<'a>( + &self, + model: &str, + deployments: Vec<&'a Deployment>, + ) -> Vec<&'a Deployment> { + self.phases.send("router.filter_deployments").unwrap(); + assert_eq!(model, "public-model"); + assert_eq!(deployments.len(), 1); + if self.reject_router { + return Vec::new(); + } + deployments + } +} + +impl Interceptors for Recorder { + async fn before_provider_request( + &self, + wire: WireRequest, + context: RequestContext, + ) -> Result { + self.phases + .send("inference.before_provider_request") + .unwrap(); + assert_eq!(context.model, wire.body["model"]); + Ok(WireRequest { + headers: wire + .headers + .into_iter() + .chain([("x-inference-hook".into(), "called".into())]) + .collect(), + ..wire + }) + } + + async fn after_provider_response(&self, raw: RawResponse) -> Result<(), RouteError> { + self.phases + .send("inference.after_provider_response") + .unwrap(); + let response: Value = serde_json::from_str(&raw.body).unwrap(); + assert_eq!(response["content"][0]["text"], "provider reply"); + Ok(()) + } +} + +#[fixture] +fn phases() -> (mpsc::Sender<&'static str>, mpsc::Receiver<&'static str>) { + mpsc::channel() +} + +#[rstest] +#[case::all_layers(false, false, 200, &[ + "gateway.pre_call", + "router.filter_deployments", + "inference.before_provider_request", + "inference.after_provider_response", + "gateway.post_call", +])] +#[case::gateway_rejection(true, false, 403, &["gateway.pre_call"])] +#[case::router_filters_every_candidate(false, true, 400, &[ + "gateway.pre_call", + "router.filter_deployments", + "gateway.post_call", +])] +#[tokio::test] +async fn request_hooks_compose_and_rejections_stop_inner_layers( + phases: (mpsc::Sender<&'static str>, mpsc::Receiver<&'static str>), + #[case] reject_gateway: bool, + #[case] reject_router: bool, + #[case] status: u16, + #[case] expected_phases: &[&str], +) { + let upstream = MockServer::start().await; + let message = json!({ + "id": "msg_test", "type": "message", "role": "assistant", "model": "test-model", + "content": [{"type": "text", "text": "provider reply"}], + "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1}, + }); + Mock::given(method("POST")) + .and(path("/v1/messages")) + .and(header("x-inference-hook", "called")) + .respond_with(ResponseTemplate::new(200).set_body_json(&message)) + .expect(if reject_gateway || reject_router { + 0 + } else { + 1 + }) + .mount(&upstream) + .await; + let config = Config::from_yaml(&format!( + "model_list:\n - model_name: public-model\n litellm_params:\n model: anthropic/test-model\n api_key: test-key\n api_base: {}\ngeneral_settings:\n master_key: gateway-key\n", + upstream.uri(), + )).unwrap(); + let (sender, receiver) = phases; + let recorder = Recorder { + phases: sender, + reject_gateway, + reject_router, + }; + let inference = Arc::into_inner(litellm_gateway::build_inference(&config).unwrap()) + .unwrap() + .with_hooks(recorder.clone(), recorder.clone()); + let app = + litellm_gateway::router_with_hooks(Arc::new(inference), &config, None, None, recorder); + let response = app + .oneshot( + Request::post("/v1/messages") + .header("content-type", "application/json") + .header("authorization", "Bearer gateway-key") + .body(Body::from( + json!({ + "model": "public-model", "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 16, + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status().as_u16(), status); + assert_eq!( + response.headers().contains_key("x-hook-completed"), + !reject_gateway + ); + let body = to_bytes(response.into_body(), 4096).await.unwrap(); + if reject_gateway { + assert_eq!(body, "rejected by gateway hook"); + } else if !reject_router { + let response: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(response["content"], message["content"]); + assert_eq!(response["usage"], message["usage"]); + } + assert_eq!(receiver.try_iter().collect::>(), expected_phases); + assert_eq!( + upstream.received_requests().await.unwrap().len(), + usize::from(!reject_gateway && !reject_router) + ); +} diff --git a/litellm-rust/crates/router/Cargo.toml b/litellm-rust/crates/router/Cargo.toml index 8c01c1e5a99..0a7387952ad 100644 --- a/litellm-rust/crates/router/Cargo.toml +++ b/litellm-rust/crates/router/Cargo.toml @@ -11,4 +11,5 @@ litellm-inference.workspace = true litellm-inference-messages.workspace = true [dev-dependencies] +tokio.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/router/src/hooks.rs b/litellm-rust/crates/router/src/hooks.rs new file mode 100644 index 00000000000..8123e3dd543 --- /dev/null +++ b/litellm-rust/crates/router/src/hooks.rs @@ -0,0 +1,15 @@ +use std::future::Future; + +use crate::Deployment; + +pub trait RouterHooks: Send + Sync { + fn filter_deployments<'a>( + &self, + _model: &str, + deployments: Vec<&'a Deployment>, + ) -> impl Future> + Send { + async { deployments } + } +} + +impl RouterHooks for () {} diff --git a/litellm-rust/crates/router/src/lib.rs b/litellm-rust/crates/router/src/lib.rs index da33bfb04bd..640af434cab 100644 --- a/litellm-rust/crates/router/src/lib.rs +++ b/litellm-rust/crates/router/src/lib.rs @@ -1,4 +1,6 @@ mod deployment; +mod hooks; +pub use hooks::RouterHooks; use std::collections::HashMap; @@ -7,7 +9,10 @@ use litellm_config::Model; pub use deployment::Deployment; #[derive(Clone, Debug, Default)] -pub struct Router(HashMap); +pub struct Router { + deployments: HashMap, + hooks: H, +} impl Router { pub fn from_model_list(model_list: &[Model]) -> Self { @@ -33,12 +38,35 @@ impl Router { } pub fn get(&self, model_name: &str) -> Option<&Deployment> { - self.0.get(model_name) + self.deployments.get(model_name) + } +} + +impl Router { + pub fn with_hooks(self, hooks: T) -> Router { + Router { + deployments: self.deployments, + hooks, + } + } +} + +impl Router { + pub async fn select(&self, model_name: &str) -> Option<&Deployment> { + let candidates = self.deployments.get(model_name).into_iter().collect(); + self.hooks + .filter_deployments(model_name, candidates) + .await + .into_iter() + .next() } } impl FromIterator<(String, Deployment)> for Router { fn from_iter>(entries: I) -> Self { - Self(entries.into_iter().collect()) + Self { + deployments: entries.into_iter().collect(), + hooks: (), + } } } diff --git a/litellm-rust/crates/router/tests/hooks.rs b/litellm-rust/crates/router/tests/hooks.rs new file mode 100644 index 00000000000..660ba956397 --- /dev/null +++ b/litellm-rust/crates/router/tests/hooks.rs @@ -0,0 +1,54 @@ +use litellm_router::{Deployment, Router, RouterHooks}; +use rstest::{fixture, rstest}; + +#[fixture] +fn router() -> Router { + Router::from_iter([( + "public-model".into(), + Deployment { + model: "provider/model".into(), + ..Default::default() + }, + )]) +} + +struct Reject; + +impl RouterHooks for Reject { + async fn filter_deployments<'a>( + &self, + model: &str, + deployments: Vec<&'a Deployment>, + ) -> Vec<&'a Deployment> { + assert_eq!(model, "public-model"); + assert_eq!(deployments.len(), 1); + assert_eq!(deployments[0].model, "provider/model"); + Vec::new() + } +} + +#[rstest] +#[case::known("public-model")] +#[case::unknown("missing")] +#[tokio::test] +async fn default_hooks_preserve_model_lookup(router: Router, #[case] model: &str) { + assert_eq!( + router + .select(model) + .await + .map(|deployment| &deployment.model), + router.get(model).map(|deployment| &deployment.model), + ); +} + +#[rstest] +#[tokio::test] +async fn selection_uses_filtered_candidates(router: Router) { + assert!( + router + .with_hooks(Reject) + .select("public-model") + .await + .is_none() + ); +}