mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
feat(gateway): forward declined and unserved requests to Python
The gateway now acts on the core decline. A replay route_layer records each native request before its handler runs. When the error response carries a Decline it forwards the recording to Python as the client sent it, or keeps the 501 when the gateway runs standalone. Paths and methods with no native route go to Python through the merged router's fallbacks, before authentication LITELLM_PYTHON_UPSTREAM picks the Fallback at startup. Forward uses its own hyper client to the loopback proxy, strips hop-by-hop headers and adds x-litellm-gateway-declined with the reason. It appends the client address to X-Forwarded-For, because Python otherwise sees every forwarded request coming from the gateway. Error responses attach the decline in one place, so the messages handler now builds its Anthropic envelope through Error::anthropic_response instead of a tuple that dropped it. The 501 stubs for embeddings and completions are gone, since a matched stub would never reach Python
This commit is contained in:
Yujong Lee
GitHub
parent
a3ae1be8f2
commit
213740863b
19 changed files with 606 additions and 57 deletions
4
litellm-rust/Cargo.lock
generated
4
litellm-rust/Cargo.lock
generated
|
|
@ -3757,6 +3757,7 @@ dependencies = [
|
|||
"tower-sessions-moka-store",
|
||||
"tracing",
|
||||
"uuid",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3788,9 +3789,11 @@ dependencies = [
|
|||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"hyper-util",
|
||||
"litellm-auth",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-response",
|
||||
"litellm-core-utils",
|
||||
"litellm-gateway-auth",
|
||||
"litellm-host",
|
||||
"litellm-host-http",
|
||||
|
|
@ -3811,6 +3814,7 @@ dependencies = [
|
|||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tower",
|
||||
"tracing",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,10 @@ 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
|
||||
litellm-inference.workspace = true
|
||||
litellm-inference-transcription.workspace = true
|
||||
|
|
@ -28,10 +31,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"] }
|
||||
|
|
|
|||
138
litellm-rust/crates/gateway-inference/src/dispatch/forward.rs
Normal file
138
litellm-rust/crates/gateway-inference/src/dispatch/forward.rs
Normal file
|
|
@ -0,0 +1,138 @@
|
|||
use std::{net::SocketAddr, sync::Arc};
|
||||
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::{ConnectInfo, OriginalUri, Request},
|
||||
http::{
|
||||
HeaderMap, HeaderName, HeaderValue, StatusCode, Uri, Version,
|
||||
header::{
|
||||
CONNECTION, PROXY_AUTHENTICATE, PROXY_AUTHORIZATION, TE, TRAILER, TRANSFER_ENCODING,
|
||||
UPGRADE,
|
||||
},
|
||||
uri::{Authority, PathAndQuery, Scheme},
|
||||
},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use hyper_util::{
|
||||
client::legacy::{Client, connect::HttpConnector},
|
||||
rt::TokioExecutor,
|
||||
};
|
||||
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 X_FORWARDED_FOR: HeaderName = HeaderName::from_static("x-forwarded-for");
|
||||
|
||||
const HOP_BY_HOP: [HeaderName; 8] = [
|
||||
CONNECTION,
|
||||
HeaderName::from_static("keep-alive"),
|
||||
PROXY_AUTHENTICATE,
|
||||
PROXY_AUTHORIZATION,
|
||||
TE,
|
||||
TRAILER,
|
||||
TRANSFER_ENCODING,
|
||||
UPGRADE,
|
||||
];
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("the Python upstream must be an http:// origin, got {0}")]
|
||||
pub struct InvalidUpstream(String);
|
||||
|
||||
/// Reverse proxy to the Python proxy on loopback, with its own client: provider clients carry
|
||||
/// proxy and TLS settings that must not apply here.
|
||||
#[derive(Clone)]
|
||||
pub struct Forward {
|
||||
client: Arc<Client<HttpConnector, Body>>,
|
||||
authority: Authority,
|
||||
}
|
||||
|
||||
impl Forward {
|
||||
pub fn new(upstream: &str) -> Result<Self, InvalidUpstream> {
|
||||
let invalid = || InvalidUpstream(upstream.to_owned());
|
||||
let uri: Uri = upstream.parse().map_err(|_| invalid())?;
|
||||
let origin_only = uri.path_and_query().is_none_or(|path| path.as_str() == "/");
|
||||
let authority = match (uri.scheme(), uri.authority()) {
|
||||
(Some(scheme), Some(authority)) if *scheme == Scheme::HTTP && origin_only => {
|
||||
authority.clone()
|
||||
}
|
||||
_ => return Err(invalid()),
|
||||
};
|
||||
Ok(Self {
|
||||
client: Arc::new(Client::builder(TokioExecutor::new()).build(HttpConnector::new())),
|
||||
authority,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn send(&self, request: Request, decline: Decline) -> Response {
|
||||
let (mut parts, body) = request.into_parts();
|
||||
let path = parts
|
||||
.extensions
|
||||
.get::<OriginalUri>()
|
||||
.map_or(&parts.uri, |original| &original.0)
|
||||
.path_and_query()
|
||||
.cloned()
|
||||
.unwrap_or_else(|| PathAndQuery::from_static("/"));
|
||||
parts.uri = Uri::from_parts({
|
||||
let mut uri = axum::http::uri::Parts::default();
|
||||
uri.scheme = Some(Scheme::HTTP);
|
||||
uri.authority = Some(self.authority.clone());
|
||||
uri.path_and_query = Some(path);
|
||||
uri
|
||||
})
|
||||
.expect("scheme, authority and path make an absolute uri");
|
||||
parts.version = Version::HTTP_11;
|
||||
strip_hop_by_hop(&mut parts.headers);
|
||||
if let Some(ConnectInfo(client)) = parts.extensions.get::<ConnectInfo<SocketAddr>>() {
|
||||
append_forwarded_for(&mut parts.headers, *client);
|
||||
}
|
||||
parts
|
||||
.headers
|
||||
.insert(DECLINED_HEADER, HeaderValue::from_static(decline.reason));
|
||||
tracing::info!(
|
||||
reason = decline.reason,
|
||||
path = parts.uri.path(),
|
||||
"forwarded to python"
|
||||
);
|
||||
match self.client.request(Request::from_parts(parts, body)).await {
|
||||
Ok(response) => {
|
||||
let (mut parts, body) = response.into_parts();
|
||||
strip_hop_by_hop(&mut parts.headers);
|
||||
Response::from_parts(parts, Body::new(body))
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "python upstream unreachable");
|
||||
StatusCode::BAD_GATEWAY.into_response()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Python sees the gateway as its peer, so the client's address goes on the end of the chain
|
||||
/// any proxy in front of the gateway already started.
|
||||
fn append_forwarded_for(headers: &mut HeaderMap, client: SocketAddr) {
|
||||
let client = client.ip().to_string();
|
||||
let chain = headers
|
||||
.get_all(&X_FORWARDED_FOR)
|
||||
.iter()
|
||||
.filter_map(|value| value.to_str().ok())
|
||||
.chain([client.as_str()])
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
if let Ok(chain) = HeaderValue::from_str(&chain) {
|
||||
headers.insert(X_FORWARDED_FOR, chain);
|
||||
}
|
||||
}
|
||||
|
||||
fn strip_hop_by_hop(headers: &mut HeaderMap) {
|
||||
let named: Vec<HeaderName> = headers
|
||||
.get_all(CONNECTION)
|
||||
.iter()
|
||||
.filter_map(|value| value.to_str().ok())
|
||||
.flat_map(|value| value.split(','))
|
||||
.filter_map(|name| HeaderName::try_from(name.trim()).ok())
|
||||
.collect();
|
||||
for name in HOP_BY_HOP.iter().chain(&named) {
|
||||
headers.remove(name);
|
||||
}
|
||||
}
|
||||
24
litellm-rust/crates/gateway-inference/src/dispatch/mod.rs
Normal file
24
litellm-rust/crates/gateway-inference/src/dispatch/mod.rs
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
//! 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),
|
||||
}
|
||||
73
litellm-rust/crates/gateway-inference/src/dispatch/replay.rs
Normal file
73
litellm-rust/crates/gateway-inference/src/dispatch/replay.rs
Normal file
|
|
@ -0,0 +1,73 @@
|
|||
use std::net::SocketAddr;
|
||||
|
||||
use axum::{
|
||||
body::{Body, Bytes},
|
||||
extract::{ConnectInfo, OriginalUri, Request, State},
|
||||
http::request::Parts,
|
||||
middleware::Next,
|
||||
response::Response,
|
||||
};
|
||||
use bytes::BytesMut;
|
||||
use futures_util::StreamExt;
|
||||
use litellm_core_utils::decline::Decline;
|
||||
|
||||
use super::Fallback;
|
||||
use crate::request::MAX_BODY_BYTES;
|
||||
|
||||
/// 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<Fallback>,
|
||||
request: Request,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
let (parts, body) = request.into_parts();
|
||||
let body = match buffer(body).await {
|
||||
Ok(body) => body,
|
||||
Err(body) => return next.run(Request::from_parts(parts, body)).await,
|
||||
};
|
||||
let recorded = record(&parts, body.clone());
|
||||
let response = next.run(Request::from_parts(parts, Body::from(body))).await;
|
||||
let Some(&decline) = response.extensions().get::<Decline>() else {
|
||||
return response;
|
||||
};
|
||||
match fallback {
|
||||
Fallback::Reject => response,
|
||||
Fallback::Forward(forward) => 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<Bytes, Body> {
|
||||
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 => {
|
||||
seen.extend_from_slice(&chunk);
|
||||
}
|
||||
failed => {
|
||||
let replayed = futures_util::stream::iter([Ok(seen.freeze()), failed]);
|
||||
return Err(Body::from_stream(replayed.chain(stream)));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(seen.freeze())
|
||||
}
|
||||
|
||||
fn record(parts: &Parts, body: Bytes) -> Request<Bytes> {
|
||||
let uri = parts
|
||||
.extensions
|
||||
.get::<OriginalUri>()
|
||||
.map_or_else(|| parts.uri.clone(), |original| original.0.clone());
|
||||
let mut recorded = Request::new(body);
|
||||
*recorded.method_mut() = parts.method.clone();
|
||||
*recorded.uri_mut() = uri;
|
||||
*recorded.version_mut() = parts.version;
|
||||
*recorded.headers_mut() = parts.headers.clone();
|
||||
if let Some(&client) = parts.extensions.get::<ConnectInfo<SocketAddr>>() {
|
||||
recorded.extensions_mut().insert(client);
|
||||
}
|
||||
recorded
|
||||
}
|
||||
|
|
@ -3,6 +3,7 @@ use axum::{
|
|||
Json,
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use litellm_core_utils::decline::Decline;
|
||||
use litellm_http::transport::Error as TransportError;
|
||||
use litellm_inference::Error as InferenceError;
|
||||
use litellm_llms::base_llm::ocr::error::Error as OcrError;
|
||||
|
|
@ -22,8 +23,6 @@ pub enum Error {
|
|||
Route(#[from] InferenceError),
|
||||
#[error(transparent)]
|
||||
Ocr(#[from] OcrError),
|
||||
#[error("{0} is not implemented by the Rust gateway")]
|
||||
Unsupported(String),
|
||||
#[error("request body exceeds the size limit")]
|
||||
BodyTooLarge,
|
||||
#[error("{0}")]
|
||||
|
|
@ -49,7 +48,6 @@ impl Error {
|
|||
pub fn status(&self) -> StatusCode {
|
||||
match self {
|
||||
Self::Auth(error) => error.status(),
|
||||
Self::Unsupported(_) => StatusCode::NOT_IMPLEMENTED,
|
||||
Self::Route(error) if error.declined().is_some() => StatusCode::NOT_IMPLEMENTED,
|
||||
Self::BodyTooLarge => StatusCode::PAYLOAD_TOO_LARGE,
|
||||
Self::Ocr(
|
||||
|
|
@ -74,22 +72,44 @@ impl Error {
|
|||
}
|
||||
}
|
||||
|
||||
pub fn declined(&self) -> Option<&Decline> {
|
||||
match self {
|
||||
Self::Route(error) => error.declined(),
|
||||
Self::Auth(_)
|
||||
| Self::InvalidBody(_)
|
||||
| Self::UnknownModel(_)
|
||||
| Self::Ocr(_)
|
||||
| Self::BodyTooLarge
|
||||
| Self::Internal(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn openai_response(self) -> Response {
|
||||
let status = self.status();
|
||||
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()
|
||||
self.respond(json!({"error": {
|
||||
"message": message,
|
||||
"type": error_type(status),
|
||||
"param": null,
|
||||
"code": status.as_u16(),
|
||||
}}))
|
||||
}
|
||||
|
||||
pub fn anthropic_response(self, request_id: Option<&str>) -> Response {
|
||||
let body = self.body(request_id);
|
||||
self.respond(body)
|
||||
}
|
||||
|
||||
/// A decline rides on the response, so the dispatch layer can hand the request on.
|
||||
fn respond(&self, body: Value) -> Response {
|
||||
let mut response = (self.status(), Json(body)).into_response();
|
||||
if let Some(decline) = self.declined() {
|
||||
response.extensions_mut().insert(*decline);
|
||||
}
|
||||
response
|
||||
}
|
||||
|
||||
/// The Anthropic error envelope Python's `AnthropicExceptionMapping` builds: an upstream
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@
|
|||
mod audio_transcription;
|
||||
mod caching;
|
||||
mod chat_completions;
|
||||
pub mod dispatch;
|
||||
mod error;
|
||||
pub mod messages;
|
||||
mod ocr;
|
||||
|
|
@ -88,7 +89,7 @@ impl Gateway {
|
|||
}
|
||||
}
|
||||
|
||||
pub fn router(gateway: Arc<Gateway>) -> Router {
|
||||
pub fn router(gateway: Arc<Gateway>, fallback: dispatch::Fallback) -> Router {
|
||||
Router::new()
|
||||
.route("/v1/messages", post(messages::create))
|
||||
.route("/ocr", post(ocr::create))
|
||||
|
|
@ -104,10 +105,10 @@ pub fn router(gateway: Arc<Gateway>) -> Router {
|
|||
)
|
||||
.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))
|
||||
.route("/v1/completions", post(request::unsupported))
|
||||
.route_layer(axum::middleware::from_fn_with_state(
|
||||
fallback,
|
||||
dispatch::replay,
|
||||
))
|
||||
.layer(axum::extract::DefaultBodyLimit::max(
|
||||
request::MAX_BODY_BYTES,
|
||||
))
|
||||
|
|
@ -115,11 +116,8 @@ pub fn router(gateway: Arc<Gateway>) -> Router {
|
|||
}
|
||||
|
||||
fn model_routes() -> Router<Arc<Gateway>> {
|
||||
Router::new()
|
||||
.route(
|
||||
"/chat/completions",
|
||||
post(chat_completions::create_from_model_path),
|
||||
)
|
||||
.route("/embeddings", post(request::unsupported))
|
||||
.route("/completions", post(request::unsupported))
|
||||
Router::new().route(
|
||||
"/chat/completions",
|
||||
post(chat_completions::create_from_model_path),
|
||||
)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -37,7 +37,7 @@ pub async fn create(
|
|||
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()))))
|
||||
result.map_err(|error| error.anthropic_response(request_id.as_deref()))
|
||||
}
|
||||
|
||||
async fn handle(
|
||||
|
|
|
|||
|
|
@ -1,8 +1,7 @@
|
|||
use axum::{
|
||||
body::Bytes,
|
||||
extract::{FromRequest, FromRequestParts, Multipart, OriginalUri, Request},
|
||||
extract::{FromRequest, FromRequestParts, Multipart, Request},
|
||||
http::request::Parts,
|
||||
response::IntoResponse,
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
|
|
@ -162,10 +161,6 @@ fn multipart_error(error: axum::extract::multipart::MultipartError) -> Error {
|
|||
Error::InvalidBody(error.to_string())
|
||||
}
|
||||
|
||||
pub(crate) async fn unsupported(OriginalUri(uri): OriginalUri) -> impl IntoResponse {
|
||||
Error::Unsupported(uri.path().to_owned())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use axum::{
|
||||
|
|
|
|||
177
litellm-rust/crates/gateway-inference/tests/dispatch.rs
Normal file
177
litellm-rust/crates/gateway-inference/tests/dispatch.rs
Normal file
|
|
@ -0,0 +1,177 @@
|
|||
mod support;
|
||||
|
||||
use std::net::SocketAddr;
|
||||
|
||||
use axum::{
|
||||
body::{Body, to_bytes},
|
||||
extract::ConnectInfo,
|
||||
http::Request,
|
||||
};
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
use tower::ServiceExt;
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any};
|
||||
|
||||
async fn python() -> MockServer {
|
||||
let python = MockServer::start().await;
|
||||
Mock::given(any())
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("x-served-by", "python")
|
||||
.set_body_string("from python"),
|
||||
)
|
||||
.mount(&python)
|
||||
.await;
|
||||
python
|
||||
}
|
||||
|
||||
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?trace=1",
|
||||
r#"{"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",
|
||||
r#"{"model":"public/model","max_tokens":16,"messages":[{"role":"user","content":"hi"}]}"#,
|
||||
"Bedrock invoke messages request shaping"
|
||||
)]
|
||||
#[case::responses_provider(
|
||||
"gemini/test-model",
|
||||
"/v1/responses",
|
||||
r#"{"model":"public/model","input":"hi"}"#,
|
||||
"native HTTP responses provider"
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn a_core_decline_reaches_python_as_the_client_sent_it(
|
||||
#[case] model: &str,
|
||||
#[case] path: &str,
|
||||
#[case] body: &'static str,
|
||||
#[case] reason: &str,
|
||||
) {
|
||||
let python = python().await;
|
||||
let provider = provider().await;
|
||||
let app = support::app_forwarding_to(&python.uri(), model, &provider.uri());
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::post(path)
|
||||
.header("content-type", "application/json")
|
||||
.header("authorization", "Bearer client-key")
|
||||
.header("x-client-header", "kept")
|
||||
.header("connection", "te")
|
||||
.header("te", "trailers")
|
||||
.body(Body::from(body))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), 200);
|
||||
assert_eq!(response.headers()["x-served-by"], "python");
|
||||
assert_eq!(
|
||||
to_bytes(response.into_body(), 1024).await.unwrap(),
|
||||
"from python"
|
||||
);
|
||||
assert!(provider.received_requests().await.unwrap().is_empty());
|
||||
let forwarded = python.received_requests().await.unwrap();
|
||||
assert_eq!(forwarded.len(), 1);
|
||||
let forwarded = &forwarded[0];
|
||||
assert_eq!(forwarded.method.as_str(), "POST");
|
||||
let forwarded_path = match forwarded.url.query() {
|
||||
Some(query) => format!("{}?{query}", forwarded.url.path()),
|
||||
None => forwarded.url.path().to_owned(),
|
||||
};
|
||||
assert_eq!(forwarded_path, path);
|
||||
assert_eq!(forwarded.body, body.as_bytes());
|
||||
assert_eq!(forwarded.headers["authorization"], "Bearer client-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) {
|
||||
let python = python().await;
|
||||
let provider = provider().await;
|
||||
let response = support::post(
|
||||
support::app_forwarding_to(&python.uri(), "anthropic/test-model", &provider.uri()),
|
||||
"/v1/chat/completions",
|
||||
body,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), status);
|
||||
assert!(python.received_requests().await.unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::direct_client(None, "203.0.113.7")]
|
||||
#[case::behind_a_load_balancer(Some("198.51.100.1"), "198.51.100.1, 203.0.113.7")]
|
||||
#[tokio::test]
|
||||
async fn a_forwarded_decline_names_the_client_in_x_forwarded_for(
|
||||
#[case] incoming: Option<&str>,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let python = python().await;
|
||||
let provider = provider().await;
|
||||
let request = Request::post("/v1/chat/completions").header("content-type", "application/json");
|
||||
let request = match incoming {
|
||||
Some(chain) => request.header("x-forwarded-for", chain),
|
||||
None => request,
|
||||
};
|
||||
let mut request = request
|
||||
.body(Body::from(
|
||||
json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], "stream": true})
|
||||
.to_string(),
|
||||
))
|
||||
.unwrap();
|
||||
request.extensions_mut().insert(ConnectInfo(
|
||||
"203.0.113.7:50000".parse::<SocketAddr>().unwrap(),
|
||||
));
|
||||
|
||||
support::app_forwarding_to(&python.uri(), "anthropic/test-model", &provider.uri())
|
||||
.oneshot(request)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let forwarded = python.received_requests().await.unwrap();
|
||||
assert_eq!(forwarded[0].headers["x-forwarded-for"], expected);
|
||||
}
|
||||
|
|
@ -157,18 +157,14 @@ async fn missing_inference_fields_use_the_same_validation_as_null(
|
|||
#[case("/engines/public%2Fmodel/embeddings")]
|
||||
#[case("/openai/deployments/public%2Fmodel/completions")]
|
||||
#[tokio::test]
|
||||
async fn unimplemented_routes_return_an_explicit_error(#[case] path: &str) {
|
||||
async fn routes_without_a_native_handler_are_left_to_the_mounting_server(#[case] path: &str) {
|
||||
let response = support::post(
|
||||
support::app("anthropic/test-model", "http://127.0.0.1:1"),
|
||||
path,
|
||||
json!({}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), 501);
|
||||
let body = support::json(response).await;
|
||||
let message = body["error"]["message"].as_str().unwrap();
|
||||
assert!(message.contains("not implemented"));
|
||||
assert!(message.contains(path));
|
||||
assert_eq!(response.status(), 404);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
|
|
@ -285,7 +281,7 @@ async fn json_extraction_rejections_use_the_endpoint_error_envelope(
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unsupported("/engines/public%2Fmodel/embeddings", 501)]
|
||||
#[case::unserved("/engines/public%2Fmodel/embeddings", 404)]
|
||||
#[case::unknown("/engines/public%2Fmodel/unknown", 404)]
|
||||
#[case::unescaped_model("/engines/public/model/chat/completions", 404)]
|
||||
#[case::extra_segment("/engines/public%2Fmodel/extra/chat/completions", 404)]
|
||||
|
|
|
|||
|
|
@ -10,7 +10,11 @@ use axum::{
|
|||
response::Response,
|
||||
};
|
||||
use futures_util::future::BoxFuture;
|
||||
use litellm_gateway_inference::{Deployment, Gateway, router};
|
||||
use litellm_gateway_inference::{
|
||||
Deployment, Gateway,
|
||||
dispatch::{Fallback, Forward},
|
||||
router,
|
||||
};
|
||||
use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver};
|
||||
use litellm_inference::resources::CoreResources;
|
||||
use litellm_secrets::{SecretValue, source::SecretSource};
|
||||
|
|
@ -37,7 +41,18 @@ pub fn app_with_permissions(
|
|||
api_base: &str,
|
||||
permissions: litellm_gateway_auth::Permissions,
|
||||
) -> Router {
|
||||
configured_app(model, api_base, permissions, None, None)
|
||||
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()),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn app_with_cache(
|
||||
|
|
@ -51,6 +66,7 @@ pub fn app_with_cache(
|
|||
litellm_gateway_auth::Permissions::All,
|
||||
Some(cache),
|
||||
None,
|
||||
Fallback::Reject,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -66,6 +82,7 @@ pub fn app_with_cache_for_principal(
|
|||
litellm_gateway_auth::Permissions::All,
|
||||
Some(cache),
|
||||
Some(principal),
|
||||
Fallback::Reject,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -75,6 +92,7 @@ fn configured_app(
|
|||
permissions: litellm_gateway_auth::Permissions,
|
||||
cache: Option<Arc<dyn litellm_cache_response::ResponseCacheService>>,
|
||||
principal: Option<litellm_gateway_auth::Principal>,
|
||||
fallback: Fallback,
|
||||
) -> Router {
|
||||
let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver)));
|
||||
let http = Resolution::from(&HttpSettings::default()).config;
|
||||
|
|
@ -102,7 +120,7 @@ fn configured_app(
|
|||
Some(cache) => gateway.with_cache(cache),
|
||||
None => gateway,
|
||||
};
|
||||
router(Arc::new(gateway)).layer(axum::middleware::from_fn_with_state(
|
||||
router(Arc::new(gateway), fallback).layer(axum::middleware::from_fn_with_state(
|
||||
(permissions, principal),
|
||||
test_identity,
|
||||
))
|
||||
|
|
|
|||
|
|
@ -31,6 +31,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
|
||||
|
|
|
|||
|
|
@ -41,3 +41,11 @@ Each MCP entry supports HTTP or stdio transport, an optional pinned `server_id`,
|
|||
The gateway initializes upstreams before opening its listener. Startup fails when credentials cannot be resolved or an upstream cannot initialize. Shutdown cancels upstream connections and terminates owned stdio children
|
||||
|
||||
The config schema preserves Python sections, but parsing a section does not implement its service. MCP OAuth, database-backed server management, guardrails, spend accounting, local Python tools, and legacy SSE upstream connections remain unimplemented. Unsupported MCP settings and configured guardrail policies fail at startup. Incoming legacy SSE clients remain supported
|
||||
|
||||
## Python fallback
|
||||
|
||||
Set `LITELLM_PYTHON_UPSTREAM` to the Python proxy's `http://` origin, usually a loopback address, and the gateway forwards every request it does not serve natively. That covers paths and methods with no native route, which reach Python without gateway authentication, and native requests that Rust core declines before calling the provider. A forwarded request keeps its method, path, query, headers and body bytes, and carries `x-litellm-gateway-declined` with the reason. Errors after the provider call are never forwarded
|
||||
|
||||
The gateway appends the client address to `X-Forwarded-For`, since Python otherwise sees every forwarded request coming from the gateway. Set `general_settings.use_x_forwarded_for: true` on the Python proxy so `allowed_ips`, spend log IPs and the other IP checks read it, and include the gateway's address in its trusted proxy ranges
|
||||
|
||||
Without it the gateway runs standalone. An unserved path is a 404 and a decline is a 501 that names the reason
|
||||
|
|
|
|||
|
|
@ -16,7 +16,10 @@ use http_body_util::BodyExt;
|
|||
|
||||
use litellm_config::Config;
|
||||
use litellm_gateway_auth::Auth;
|
||||
use litellm_gateway_inference::{Gateway, ModelRouter};
|
||||
use litellm_gateway_inference::{
|
||||
Gateway, ModelRouter,
|
||||
dispatch::{Fallback, NOT_NATIVE},
|
||||
};
|
||||
use litellm_http::{
|
||||
ClientVariant, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify,
|
||||
media::PublicDnsResolver,
|
||||
|
|
@ -118,19 +121,30 @@ pub fn router(
|
|||
config: &Config,
|
||||
ui: Option<Router>,
|
||||
mcp: Option<Router>,
|
||||
fallback: Fallback,
|
||||
) -> Router {
|
||||
let auth = Auth::from_config(config, inference.secrets.clone());
|
||||
let inference = litellm_gateway_inference::router(inference)
|
||||
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));
|
||||
match ui {
|
||||
Some(ui) => inference.merge(ui),
|
||||
None => inference,
|
||||
.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
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ use std::{
|
|||
|
||||
use litellm_config::Config;
|
||||
use litellm_gateway_auth::UiBackend;
|
||||
use litellm_gateway_inference::dispatch::{Fallback, Forward};
|
||||
use litellm_tracing::{Level, Logger, Metadata, Record, Sink};
|
||||
use serde_json::json;
|
||||
|
||||
|
|
@ -46,6 +47,10 @@ async fn main() -> Result<(), Box<dyn Error>> {
|
|||
Logger::new(StderrSink { level }).install_global()?;
|
||||
let config = Config::load(settings.litellm_config)?;
|
||||
let inference = litellm_gateway::build_inference(&config)?;
|
||||
let fallback = match settings.litellm_python_upstream.as_deref() {
|
||||
Some(upstream) => Fallback::Forward(Forward::new(upstream)?),
|
||||
None => Fallback::Reject,
|
||||
};
|
||||
let ui = match std::env::var_os("LITELLM_UI_PATH") {
|
||||
Some(directory) => {
|
||||
let ui_settings: UiSettings = settings::from_iter(std::env::vars_os())?;
|
||||
|
|
@ -97,11 +102,12 @@ async fn main() -> Result<(), Box<dyn Error>> {
|
|||
});
|
||||
let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?;
|
||||
|
||||
tracing::info!(address = %listener.local_addr()?, models = config.model_list.len(), log_level = %level, "gateway listening");
|
||||
tracing::info!(address = %listener.local_addr()?, models = config.model_list.len(), log_level = %level, python_upstream = settings.litellm_python_upstream.as_deref(), "gateway listening");
|
||||
|
||||
let result = axum::serve(
|
||||
listener,
|
||||
litellm_gateway::router(inference, &config, ui, mcp_router),
|
||||
litellm_gateway::router(inference, &config, ui, mcp_router, fallback)
|
||||
.into_make_service_with_connect_info::<std::net::SocketAddr>(),
|
||||
)
|
||||
.with_graceful_shutdown(shutdown.clone().cancelled_owned())
|
||||
.await;
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ pub(super) struct Settings {
|
|||
pub host: String,
|
||||
pub port: u16,
|
||||
pub litellm_config: String,
|
||||
pub litellm_python_upstream: Option<String>,
|
||||
rust_log: Option<String>,
|
||||
}
|
||||
|
||||
|
|
@ -19,6 +20,7 @@ impl Default for Settings {
|
|||
host: "0.0.0.0".into(),
|
||||
port: 4000,
|
||||
litellm_config: "config.yaml".into(),
|
||||
litellm_python_upstream: None,
|
||||
rust_log: None,
|
||||
}
|
||||
}
|
||||
|
|
@ -81,6 +83,7 @@ mod tests {
|
|||
assert_eq!(settings.host, "0.0.0.0");
|
||||
assert_eq!(settings.port, 4000);
|
||||
assert_eq!(settings.litellm_config, "config.yaml");
|
||||
assert_eq!(settings.litellm_python_upstream, None);
|
||||
assert_eq!(settings.log_level(), Level::INFO);
|
||||
}
|
||||
|
||||
|
|
@ -90,6 +93,7 @@ mod tests {
|
|||
("HOST", "127.0.0.1"),
|
||||
("PORT", "8080"),
|
||||
("LITELLM_CONFIG", "/tmp/custom.yaml"),
|
||||
("LITELLM_PYTHON_UPSTREAM", "http://127.0.0.1:4001"),
|
||||
("RUST_LOG", "debug"),
|
||||
("LITELLM_UI_SECURE_COOKIES", "invalid-but-ui-disabled"),
|
||||
])
|
||||
|
|
@ -98,6 +102,10 @@ mod tests {
|
|||
assert_eq!(settings.host, "127.0.0.1");
|
||||
assert_eq!(settings.port, 8080);
|
||||
assert_eq!(settings.litellm_config, "/tmp/custom.yaml");
|
||||
assert_eq!(
|
||||
settings.litellm_python_upstream.as_deref(),
|
||||
Some("http://127.0.0.1:4001")
|
||||
);
|
||||
assert_eq!(settings.log_level(), Level::DEBUG);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ use axum::{
|
|||
http::Request,
|
||||
};
|
||||
use litellm_config::Config;
|
||||
use litellm_gateway_inference::dispatch::Fallback;
|
||||
use litellm_gateway_mcp::{
|
||||
HttpConfig,
|
||||
rmcp::{
|
||||
|
|
@ -145,7 +146,7 @@ mcp_servers:
|
|||
..Default::default()
|
||||
},
|
||||
);
|
||||
let app = litellm_gateway::router(inference, &config, None, Some(mcp_router));
|
||||
let app = litellm_gateway::router(inference, &config, None, Some(mcp_router), Fallback::Reject);
|
||||
|
||||
for (path, method) in [
|
||||
("/mcp", "POST"),
|
||||
|
|
|
|||
|
|
@ -6,7 +6,10 @@ 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};
|
||||
use litellm_gateway_inference::{
|
||||
Error, Gateway,
|
||||
dispatch::{Fallback, Forward},
|
||||
};
|
||||
use litellm_http::ClientVariant;
|
||||
use litellm_tracing::{Logger, Metadata, Record, Sink};
|
||||
use rstest::{fixture, rstest};
|
||||
|
|
@ -74,7 +77,7 @@ async fn authenticates_before_serving_mounted_inference_routes(
|
|||
let server = tokio::spawn(async move {
|
||||
axum::serve(
|
||||
listener,
|
||||
litellm_gateway::router(inference, &config, None, None),
|
||||
litellm_gateway::router(inference, &config, None, None, Fallback::Reject),
|
||||
)
|
||||
.with_graceful_shutdown(async move {
|
||||
let _ = stopped.await;
|
||||
|
|
@ -129,7 +132,10 @@ async fn logs_request_outcome_without_credentials_or_query(inference: Arc<Gatewa
|
|||
let logger = Logger::new(LogSink(sender));
|
||||
|
||||
let response = logger
|
||||
.instrument(litellm_gateway::router(inference, &config, None, None).oneshot(request))
|
||||
.instrument(
|
||||
litellm_gateway::router(inference, &config, None, None, Fallback::Reject)
|
||||
.oneshot(request),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
|
|
@ -165,7 +171,7 @@ async fn mounts_ui_without_exposing_credentials_or_authorizing_inference(inferen
|
|||
false,
|
||||
)
|
||||
.merge(litellm_gateway_ui::dashboard_assets(assets.path()));
|
||||
let app = litellm_gateway::router(inference, &config, Some(ui), None);
|
||||
let app = litellm_gateway::router(inference, &config, Some(ui), None, Fallback::Reject);
|
||||
let page = app
|
||||
.clone()
|
||||
.oneshot(Request::get("/ui/").body(Body::empty()).unwrap())
|
||||
|
|
@ -275,7 +281,7 @@ async fn ui_routes_are_absent_when_not_mounted(
|
|||
let config =
|
||||
Config::from_yaml("model_list: []\ngeneral_settings:\n master_key: gateway-key\n")
|
||||
.unwrap();
|
||||
let response = litellm_gateway::router(inference, &config, None, None)
|
||||
let response = litellm_gateway::router(inference, &config, None, None, Fallback::Reject)
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method(method)
|
||||
|
|
@ -287,3 +293,62 @@ async fn ui_routes_are_absent_when_not_mounted(
|
|||
.unwrap();
|
||||
assert_eq!(response.status().as_u16(), 404);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unserved_path("POST", "/v1/embeddings", None, true)]
|
||||
#[case::management_api("GET", "/key/info?key=sk-1", None, true)]
|
||||
#[case::wrong_method_on_a_native_path("GET", "/v1/messages", None, true)]
|
||||
#[case::bad_key_on_a_native_route("POST", "/v1/messages", Some("Bearer wrong"), false)]
|
||||
#[tokio::test]
|
||||
async fn python_serves_every_request_the_gateway_does_not_route(
|
||||
inference: Arc<Gateway>,
|
||||
#[case] method: &str,
|
||||
#[case] path: &str,
|
||||
#[case] authorization: Option<&str>,
|
||||
#[case] forwarded: bool,
|
||||
) {
|
||||
let python = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::any())
|
||||
.respond_with(wiremock::ResponseTemplate::new(200).set_body_string("from python"))
|
||||
.mount(&python)
|
||||
.await;
|
||||
let config =
|
||||
Config::from_yaml("model_list: []\ngeneral_settings:\n master_key: gateway-key\n")
|
||||
.unwrap();
|
||||
let fallback = Fallback::Forward(Forward::new(&python.uri()).unwrap());
|
||||
let request = Request::builder().method(method).uri(path);
|
||||
let request = match authorization {
|
||||
Some(value) => request.header("authorization", value),
|
||||
None => request,
|
||||
};
|
||||
|
||||
let mut request = request.body(Body::from("{}")).unwrap();
|
||||
request.extensions_mut().insert(axum::extract::ConnectInfo(
|
||||
"203.0.113.7:50000".parse::<std::net::SocketAddr>().unwrap(),
|
||||
));
|
||||
|
||||
let response = litellm_gateway::router(inference, &config, None, None, fallback)
|
||||
.oneshot(request)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let received = python.received_requests().await.unwrap();
|
||||
assert_eq!(received.len(), usize::from(forwarded));
|
||||
if forwarded {
|
||||
assert_eq!(response.status().as_u16(), 200);
|
||||
assert_eq!(received[0].method.as_str(), method);
|
||||
let url = &received[0].url;
|
||||
let forwarded_path = match url.query() {
|
||||
Some(query) => format!("{}?{query}", url.path()),
|
||||
None => url.path().to_owned(),
|
||||
};
|
||||
assert_eq!(forwarded_path, path);
|
||||
assert_eq!(
|
||||
received[0].headers["x-litellm-gateway-declined"],
|
||||
"route not native"
|
||||
);
|
||||
assert_eq!(received[0].headers["x-forwarded-for"], "203.0.113.7");
|
||||
} else {
|
||||
assert_eq!(response.status().as_u16(), 401);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue