mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
feat(gateway): normalize compressed request bodies
This commit is contained in:
@@ -7,7 +7,7 @@ use url::form_urlencoded;
|
||||
|
||||
use crate::{
|
||||
ai_serving::extract_gemini_model_from_path,
|
||||
headers::{header_value_str, is_json_request},
|
||||
headers::{decoded_request_body_bytes, header_value_str, is_json_request},
|
||||
};
|
||||
|
||||
use super::super::GatewayControlDecision;
|
||||
@@ -31,7 +31,8 @@ pub(crate) fn extract_requested_model(
|
||||
if !is_json_request(headers) || body.is_empty() {
|
||||
return None;
|
||||
}
|
||||
serde_json::from_slice::<serde_json::Value>(body)
|
||||
let body = decoded_request_body_bytes(headers, body.as_ref()).ok()?;
|
||||
serde_json::from_slice::<serde_json::Value>(body.as_ref())
|
||||
.ok()
|
||||
.and_then(|payload| {
|
||||
payload
|
||||
@@ -356,15 +357,50 @@ pub(super) fn current_unix_secs() -> u64 {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
build_auth_context_cache_key, extract_request_credentials, GatewayCredentialCarrier,
|
||||
GatewayPrimaryCredential, GatewayTrustedAdminHeaders, GatewayTrustedAuthHeaders,
|
||||
build_auth_context_cache_key, extract_request_credentials, extract_requested_model,
|
||||
GatewayCredentialCarrier, GatewayPrimaryCredential, GatewayTrustedAdminHeaders,
|
||||
GatewayTrustedAuthHeaders,
|
||||
};
|
||||
use crate::control::GatewayControlDecision;
|
||||
use axum::body::Bytes;
|
||||
use axum::http::{self, Uri};
|
||||
|
||||
fn uri(path: &str) -> Uri {
|
||||
path.parse().expect("uri should parse")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_requested_model_reads_zstd_encoded_json_body() {
|
||||
let decision = GatewayControlDecision::synthetic(
|
||||
"/v1/responses",
|
||||
Some("ai_public".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("responses".to_string()),
|
||||
Some("openai:responses".to_string()),
|
||||
);
|
||||
let mut headers = http::HeaderMap::new();
|
||||
headers.insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
http::HeaderValue::from_static("application/json"),
|
||||
);
|
||||
headers.insert(
|
||||
http::header::CONTENT_ENCODING,
|
||||
http::HeaderValue::from_static("zstd"),
|
||||
);
|
||||
let encoded =
|
||||
zstd::stream::encode_all(br#"{"model":"gpt-5.4","input":"hello"}"#.as_slice(), 0)
|
||||
.expect("zstd body should encode");
|
||||
|
||||
let requested_model = extract_requested_model(
|
||||
&decision,
|
||||
&uri("/v1/responses"),
|
||||
&headers,
|
||||
&Bytes::from(encoded),
|
||||
);
|
||||
|
||||
assert_eq!(requested_model.as_deref(), Some("gpt-5.4"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selects_openai_bearer_as_provider_api_key() {
|
||||
let mut headers = http::HeaderMap::new();
|
||||
|
||||
@@ -94,6 +94,7 @@ pub(crate) async fn request_model_local_rejection(
|
||||
decision,
|
||||
auth_context,
|
||||
requested_model.as_deref(),
|
||||
headers,
|
||||
body,
|
||||
)
|
||||
.await
|
||||
@@ -104,6 +105,7 @@ async fn balance_capacity_rejection(
|
||||
decision: &GatewayControlDecision,
|
||||
auth_context: &GatewayControlAuthContext,
|
||||
requested_model: Option<&str>,
|
||||
headers: &http::HeaderMap,
|
||||
body: &Bytes,
|
||||
) -> Result<Option<GatewayLocalAuthRejection>, GatewayError> {
|
||||
if auth_context.api_key_is_standalone {
|
||||
@@ -146,7 +148,8 @@ async fn balance_capacity_rejection(
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(estimated_cost_usd) =
|
||||
estimate_request_cost_upper_bound_usd(state, decision, requested_model, body).await?
|
||||
estimate_request_cost_upper_bound_usd(state, decision, requested_model, headers, body)
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -173,6 +176,7 @@ async fn estimate_request_cost_upper_bound_usd(
|
||||
state: &AppState,
|
||||
decision: &GatewayControlDecision,
|
||||
requested_model: &str,
|
||||
headers: &http::HeaderMap,
|
||||
body: &Bytes,
|
||||
) -> Result<Option<f64>, GatewayError> {
|
||||
let Some(api_format) = decision
|
||||
@@ -183,7 +187,11 @@ async fn estimate_request_cost_upper_bound_usd(
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let body_json = serde_json::from_slice::<serde_json::Value>(body).ok();
|
||||
let body = crate::headers::decoded_request_body_bytes(headers, body.as_ref()).ok();
|
||||
let Some(body) = body else {
|
||||
return Ok(None);
|
||||
};
|
||||
let body_json = serde_json::from_slice::<serde_json::Value>(body.as_ref()).ok();
|
||||
let Some(input_tokens) = body_json
|
||||
.as_ref()
|
||||
.map(estimate_json_tokens)
|
||||
|
||||
Reference in New Issue
Block a user