fix(gateway): route openai image api requests with mapped models

This commit is contained in:
ZheFox
2026-05-20 17:16:38 +08:00
parent cde2062618
commit de7be4f15b
5 changed files with 118 additions and 11 deletions

View File

@@ -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(

View File

@@ -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,

View File

@@ -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(

View File

@@ -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,

View File

@@ -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"));