mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
fix(gateway): route openai image api requests with mapped models
This commit is contained in:
@@ -19,9 +19,10 @@ use crate::ai_serving::{
|
||||
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
|
||||
build_chatgpt_web_image_request_body,
|
||||
build_gemini_image_request_body_from_openai_image_request,
|
||||
build_openai_image_provider_request_body, default_model_for_openai_image_operation,
|
||||
normalize_openai_image_request, request_conversion_direct_auth, CandidateFailureDiagnostic,
|
||||
GatewayProviderTransportSnapshot, PlannerAppState, RequestConversionKind,
|
||||
build_openai_image_api_provider_request_body, build_openai_image_provider_request_body,
|
||||
default_model_for_openai_image_operation, normalize_openai_image_request,
|
||||
request_conversion_direct_auth, CandidateFailureDiagnostic, GatewayProviderTransportSnapshot,
|
||||
PlannerAppState, RequestConversionKind,
|
||||
};
|
||||
use crate::image_capabilities::openai_image_normalize_options_for_provider;
|
||||
use crate::AppState;
|
||||
@@ -160,6 +161,11 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("grok");
|
||||
let is_codex = transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("codex");
|
||||
let transport_profile = crate::ai_serving::transport::resolve_transport_profile(transport);
|
||||
let upstream_url = if is_chatgpt_web {
|
||||
chatgpt_web_image_internal_url(&transport.endpoint.base_url)
|
||||
@@ -173,8 +179,13 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
|
||||
Ok(body) => body,
|
||||
Err(err) => err.to_error_json(),
|
||||
}
|
||||
} else {
|
||||
} else if is_codex || is_grok {
|
||||
build_openai_image_provider_request_body(&normalized_request)
|
||||
} else {
|
||||
build_openai_image_api_provider_request_body(
|
||||
&normalized_request,
|
||||
Some(prepared_candidate.mapped_model.as_str()),
|
||||
)
|
||||
};
|
||||
if !is_chatgpt_web {
|
||||
apply_codex_openai_responses_special_body_edits(
|
||||
|
||||
@@ -20,6 +20,7 @@ pub(crate) use aether_ai_formats::api::{
|
||||
build_local_openai_responses_request_body,
|
||||
build_local_openai_responses_request_body_with_model_directives,
|
||||
build_local_success_background_report, build_local_success_conversion_background_report,
|
||||
build_openai_image_api_provider_request_body,
|
||||
build_openai_image_provider_body_from_response_stream_sync_body,
|
||||
build_openai_image_provider_request_body,
|
||||
build_openai_image_request_body_from_gemini_image_request,
|
||||
|
||||
@@ -1944,13 +1944,23 @@ async fn provider_query_execute_openai_image_test_candidate(
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("grok");
|
||||
let is_codex = transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("codex");
|
||||
let mut provider_request_body = if is_chatgpt_web {
|
||||
match crate::ai_serving::build_chatgpt_web_image_request_body(&parts, &request_body, None) {
|
||||
Ok(body) => body,
|
||||
Err(err) => err.to_error_json(),
|
||||
}
|
||||
} else {
|
||||
} else if is_codex || is_grok {
|
||||
crate::ai_serving::build_openai_image_provider_request_body(&normalized_request)
|
||||
} else {
|
||||
crate::ai_serving::build_openai_image_api_provider_request_body(
|
||||
&normalized_request,
|
||||
Some(candidate.effective_model.as_str()),
|
||||
)
|
||||
};
|
||||
if !is_chatgpt_web {
|
||||
crate::ai_serving::apply_codex_openai_responses_special_body_edits(
|
||||
|
||||
@@ -185,12 +185,13 @@ pub use crate::formats::{
|
||||
},
|
||||
openai::image::{
|
||||
request::{
|
||||
build_chatgpt_web_image_request_body, build_openai_image_provider_request_body,
|
||||
default_model_for_openai_image_operation, is_openai_image_stream_request,
|
||||
normalize_openai_image_request, normalize_openai_image_request_with_options,
|
||||
openai_image_operation_from_path, resolve_requested_openai_image_model_for_request,
|
||||
ChatGptWebImageRequestError, NormalizedOpenAiImageRequest, OpenAiImageNormalizeOptions,
|
||||
OpenAiImageOperation, OpenAiImageResponseFormat,
|
||||
build_chatgpt_web_image_request_body, build_openai_image_api_provider_request_body,
|
||||
build_openai_image_provider_request_body, default_model_for_openai_image_operation,
|
||||
is_openai_image_stream_request, normalize_openai_image_request,
|
||||
normalize_openai_image_request_with_options, openai_image_operation_from_path,
|
||||
resolve_requested_openai_image_model_for_request, ChatGptWebImageRequestError,
|
||||
NormalizedOpenAiImageRequest, OpenAiImageNormalizeOptions, OpenAiImageOperation,
|
||||
OpenAiImageResponseFormat,
|
||||
},
|
||||
spec::{
|
||||
resolve_stream_spec as resolve_local_image_stream_spec,
|
||||
|
||||
@@ -44,6 +44,7 @@ pub struct NormalizedOpenAiImageRequest {
|
||||
images: Vec<Value>,
|
||||
tool: Map<String, Value>,
|
||||
image_count: Option<u64>,
|
||||
stream: Option<bool>,
|
||||
user: Option<String>,
|
||||
}
|
||||
|
||||
@@ -521,6 +522,54 @@ pub fn build_openai_image_provider_request_body(request: &NormalizedOpenAiImageR
|
||||
Value::Object(body)
|
||||
}
|
||||
|
||||
pub fn build_openai_image_api_provider_request_body(
|
||||
request: &NormalizedOpenAiImageRequest,
|
||||
mapped_model: Option<&str>,
|
||||
) -> Value {
|
||||
let model = mapped_model
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.or(request.requested_model.as_deref())
|
||||
.unwrap_or_else(|| default_model_for_openai_image_operation(request.operation));
|
||||
let mut body = Map::new();
|
||||
body.insert("model".to_string(), Value::String(model.to_string()));
|
||||
if let Some(prompt) = request.prompt.as_ref() {
|
||||
body.insert("prompt".to_string(), Value::String(prompt.clone()));
|
||||
}
|
||||
if let Some(image_count) = request.image_count {
|
||||
body.insert("n".to_string(), Value::Number(Number::from(image_count)));
|
||||
}
|
||||
if let Some(user) = request.user.as_ref() {
|
||||
body.insert("user".to_string(), Value::String(user.clone()));
|
||||
}
|
||||
if let Some(stream) = request.stream {
|
||||
body.insert("stream".to_string(), Value::Bool(stream));
|
||||
}
|
||||
for (key, value) in &request.tool {
|
||||
match key.as_str() {
|
||||
"type" | "action" => {}
|
||||
"input_image_mask" => {
|
||||
body.insert("mask".to_string(), value.clone());
|
||||
}
|
||||
_ => {
|
||||
body.insert(key.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(response_format) = request.summary_json.get("response_format") {
|
||||
body.entry("response_format".to_string())
|
||||
.or_insert_with(|| response_format.clone());
|
||||
}
|
||||
if !request.images.is_empty() {
|
||||
if request.images.len() == 1 {
|
||||
body.insert("image".to_string(), request.images[0].clone());
|
||||
} else {
|
||||
body.insert("images".to_string(), Value::Array(request.images.clone()));
|
||||
}
|
||||
}
|
||||
Value::Object(body)
|
||||
}
|
||||
|
||||
fn normalize_openai_image_json_request(
|
||||
body_json: &Value,
|
||||
operation: OpenAiImageOperation,
|
||||
@@ -549,6 +598,7 @@ fn normalize_openai_image_json_request(
|
||||
let output_format =
|
||||
normalize_output_format(object.get("output_format").and_then(Value::as_str))?;
|
||||
let partial_images = normalize_partial_images(object.get("partial_images"))?;
|
||||
let stream = object.get("stream").and_then(value_as_bool);
|
||||
let user = object
|
||||
.get("user")
|
||||
.and_then(Value::as_str)
|
||||
@@ -579,6 +629,7 @@ fn normalize_openai_image_json_request(
|
||||
images,
|
||||
tool,
|
||||
image_count,
|
||||
stream,
|
||||
user,
|
||||
summary_json: build_image_request_summary_json(
|
||||
operation,
|
||||
@@ -627,6 +678,9 @@ fn normalize_openai_image_multipart_request(
|
||||
.map(Value::String)
|
||||
.as_ref(),
|
||||
)?;
|
||||
let stream = find_multipart_text_field(&multipart_fields, "stream")
|
||||
.as_deref()
|
||||
.and_then(parse_bool_string);
|
||||
let user = find_multipart_text_field(&multipart_fields, "user")
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty());
|
||||
@@ -679,6 +733,7 @@ fn normalize_openai_image_multipart_request(
|
||||
images,
|
||||
tool,
|
||||
image_count,
|
||||
stream,
|
||||
user,
|
||||
summary_json: build_image_request_summary_json(
|
||||
operation,
|
||||
@@ -1478,6 +1533,35 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_image_api_provider_request_body_keeps_images_api_shape() {
|
||||
let parts = request_parts("/v1/images/generations", Some("application/json"));
|
||||
let request = normalize_openai_image_request_with_options(
|
||||
&parts,
|
||||
&json!({
|
||||
"model": "grok-imagine-image-lite",
|
||||
"prompt": "draw a cat",
|
||||
"n": 1,
|
||||
"size": "1024x1024",
|
||||
"stream": true
|
||||
}),
|
||||
None,
|
||||
OpenAiImageNormalizeOptions::with_max_generation_count(4),
|
||||
)
|
||||
.expect("generation request should normalize");
|
||||
|
||||
let provider_request_body =
|
||||
build_openai_image_api_provider_request_body(&request, Some("mapped-image-model"));
|
||||
|
||||
assert_eq!(provider_request_body["model"], "mapped-image-model");
|
||||
assert_eq!(provider_request_body["prompt"], "draw a cat");
|
||||
assert_eq!(provider_request_body["n"], 1);
|
||||
assert_eq!(provider_request_body["size"], "1024x1024");
|
||||
assert_eq!(provider_request_body["stream"], true);
|
||||
assert!(provider_request_body.get("input").is_none());
|
||||
assert!(provider_request_body.get("tools").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chatgpt_web_accepts_1k_tier_and_1024_size() {
|
||||
let parts = request_parts("/v1/images/generations", Some("application/json"));
|
||||
|
||||
Reference in New Issue
Block a user