feat(gateway): route OpenAI image intents through image bridge

This commit is contained in:
ZheFox
2026-05-17 21:19:17 +08:00
parent d9c8119bda
commit d6c8c14de7
15 changed files with 1849 additions and 25 deletions

View File

@@ -101,7 +101,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
planner_state,
spec_metadata.api_format,
&input.requested_model,
spec_metadata.require_streaming,
false,
input.required_capabilities.as_ref(),
&input.auth_snapshot,
input.client_session_affinity.as_ref(),

View File

@@ -13,7 +13,9 @@ mod gemini;
mod normalize;
mod openai;
pub(crate) use self::codex::apply_codex_openai_responses_special_headers;
pub(crate) use self::codex::{
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
};
pub(crate) use self::family::{
build_local_stream_attempt_source, build_local_stream_plan_and_reports,
build_local_sync_attempt_source, build_local_sync_plan_and_reports,

View File

@@ -9,6 +9,7 @@ pub(super) use self::payload::maybe_build_local_openai_chat_decision_payload_for
pub(super) use self::support::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_image_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts, LocalOpenAiChatCandidateAttempt,
LocalOpenAiChatCandidateAttemptSource, LocalOpenAiChatDecisionInput,
};

View File

@@ -85,7 +85,32 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
&mut extra_fields,
resolved.transport.provider.provider_type.as_str(),
);
if let Some(image_request_summary) = resolved.image_request_summary.as_ref() {
extra_fields.insert("image_request".to_string(), image_request_summary.clone());
}
if resolved
.provider_api_format
.eq_ignore_ascii_case("openai:image")
&& resolved
.transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("chatgpt_web")
{
extra_fields.insert("chatgpt_web_image".to_string(), serde_json::json!(true));
extra_fields.insert(
"local_failover_policy".to_string(),
serde_json::json!({
"stop_status_codes": [400, 401, 403, 429, 500, 502, 503, 504],
"error_stop_patterns": [
{ "pattern": ".*" }
]
}),
);
}
let super::request::LocalOpenAiChatCandidatePayloadParts {
client_api_format,
auth_header,
auth_value,
mapped_model,
@@ -99,6 +124,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
envelope_name,
transport,
request_redacted,
image_request_summary: _,
} = resolved;
let original_request_body_json = if request_redacted {
Some(&provider_request_body)
@@ -122,7 +148,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: &provider_api_format,
client_api_format: "openai:chat",
client_api_format: &client_api_format,
mapped_model: Some(&mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
pool_key_lease: eligible.orchestration.pool_key_lease.as_ref(),
@@ -154,7 +180,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
}),
execution_strategy,
conversion_mode,
"openai:chat",
client_api_format.as_str(),
candidate.endpoint_api_format.as_str(),
),
&transport,
@@ -178,7 +204,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
auth_header: Some(auth_header),
auth_value: Some(auth_value),
provider_api_format,
client_api_format: "openai:chat".to_string(),
client_api_format,
model_name: input.requested_model.clone(),
mapped_model,
prompt_cache_key,

View File

@@ -3,7 +3,7 @@ use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use serde_json::Value;
use serde_json::{json, Value};
use crate::ai_serving::planner::candidate_preparation::{
prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth,
@@ -15,9 +15,10 @@ use crate::ai_serving::planner::common::{
request_requires_body_stream_field, OPENAI_CHAT_STREAM_PLAN_KIND,
};
use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, build_cross_format_openai_chat_request_body,
build_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body,
build_local_openai_chat_upstream_url, request_body_build_failure_extra_data,
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
build_cross_format_openai_chat_request_body, build_cross_format_openai_chat_upstream_url,
build_local_openai_chat_request_body, build_local_openai_chat_upstream_url,
request_body_build_failure_extra_data,
};
use crate::ai_serving::transport::auth::resolve_local_openai_bearer_auth;
use crate::ai_serving::transport::kiro::{
@@ -27,8 +28,10 @@ use crate::ai_serving::transport::kiro::{
};
use crate::ai_serving::transport::local_openai_chat_transport_unsupported_reason;
use crate::ai_serving::transport::{
build_kiro_cross_format_upstream_url, build_standard_provider_request_headers,
StandardProviderRequestHeadersInput,
build_kiro_cross_format_upstream_url, build_openai_image_headers,
build_openai_image_upstream_url, build_standard_provider_request_headers,
openai_image_transport_unsupported_reason, resolve_openai_image_auth,
ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput,
};
use crate::ai_serving::{
ai_local_execution_contract_for_formats, request_conversion_direct_auth,
@@ -51,6 +54,7 @@ use super::support::{
};
pub(crate) struct LocalOpenAiChatCandidatePayloadParts {
pub(super) client_api_format: String,
pub(super) auth_header: String,
pub(super) auth_value: String,
pub(super) mapped_model: String,
@@ -64,6 +68,7 @@ pub(crate) struct LocalOpenAiChatCandidatePayloadParts {
pub(super) envelope_name: Option<&'static str>,
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
pub(super) request_redacted: bool,
pub(super) image_request_summary: Option<Value>,
}
fn request_identity_response_encoding_when_redacted(
@@ -147,6 +152,7 @@ async fn resolve_chat_pii_redaction_feature_settings(
.await
.map_err(|err| {
warn!(
error = ?err,
"gateway failed to read api key chat pii redaction feature settings"
);
@@ -338,6 +344,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
);
return Ok(Some(LocalOpenAiChatCandidatePayloadParts {
client_api_format: "openai:chat".to_string(),
auth_header: resolved_headers.auth_header,
auth_value: resolved_headers.auth_value,
mapped_model: prepared_candidate.mapped_model,
@@ -351,10 +358,26 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
envelope_name: None,
transport: Arc::clone(transport),
request_redacted: redaction.redacted,
image_request_summary: None,
}));
};
let provider_api_format = provider_api_format.trim().to_ascii_lowercase();
if provider_api_format == "openai:image" {
return resolve_openai_chat_to_openai_image_payload_parts(
state,
parts,
trace_id,
body_json,
input,
eligible,
candidate_index,
candidate_id,
upstream_is_stream,
)
.await;
}
let Some(conversion_kind) =
request_conversion_kind("openai:chat", provider_api_format.as_str())
else {
@@ -626,6 +649,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
ai_local_execution_contract_for_formats("openai:chat", provider_api_format.as_str());
Ok(Some(LocalOpenAiChatCandidatePayloadParts {
client_api_format: "openai:chat".to_string(),
auth_header: resolved_headers.auth_header,
auth_value: resolved_headers.auth_value,
mapped_model: prepared_candidate.mapped_model,
@@ -639,9 +663,487 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
envelope_name: None,
transport: Arc::clone(transport),
request_redacted: redaction.redacted,
image_request_summary: None,
}))
}
#[allow(clippy::too_many_arguments)]
async fn resolve_openai_chat_to_openai_image_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput,
eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32,
candidate_id: &str,
upstream_is_stream: bool,
) -> Result<Option<LocalOpenAiChatCandidatePayloadParts>, GatewayError> {
let candidate = &eligible.candidate;
let transport = &eligible.transport;
let provider_api_format = "openai:image";
if let Some(skip_reason) =
openai_image_transport_unsupported_reason(transport, provider_api_format)
{
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return Ok(None);
}
let prepared_candidate = match prepare_header_authenticated_candidate(
crate::ai_serving::PlannerAppState::new(state),
transport,
candidate,
resolve_openai_image_auth(transport),
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "openai_chat_image_bridge",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return Ok(None);
}
};
let is_chatgpt_web = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("chatgpt_web");
let Some((mut provider_request_body, image_request_summary)) = (if is_chatgpt_web {
build_chatgpt_web_image_provider_body_from_openai_chat_body(
body_json,
&input.requested_model,
)
} else {
build_openai_image_provider_body_from_openai_chat_body(
body_json,
&input.requested_model,
upstream_is_stream,
)
}) else {
mark_skipped_local_openai_chat_candidate_with_extra_data(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_build_failed",
request_body_build_failure_extra_data(body_json, "openai:chat", provider_api_format),
)
.await;
return Ok(None);
};
if !is_chatgpt_web {
apply_codex_openai_responses_special_body_edits(
&mut provider_request_body,
transport.provider.provider_type.as_str(),
provider_api_format,
transport.endpoint.body_rules.as_ref(),
Some(candidate.key_id.as_str()),
);
}
let upstream_url = if is_chatgpt_web {
chatgpt_web_image_internal_url(&transport.endpoint.base_url)
} else {
build_openai_image_upstream_url(transport, parts.uri.query())
};
let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: body_json,
})
else {
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
"openai:chat",
provider_api_format,
"openai_chat_image_bridge_headers",
),
)
.await;
return Ok(None);
};
if is_chatgpt_web {
provider_request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string());
} else {
apply_codex_openai_responses_special_headers(
&mut provider_request_headers,
&provider_request_body,
&parts.headers,
transport.provider.provider_type.as_str(),
provider_api_format,
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
}
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:image", provider_api_format);
Ok(Some(LocalOpenAiChatCandidatePayloadParts {
client_api_format: "openai:image".to_string(),
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
mapped_model: prepared_candidate.mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body,
provider_request_headers,
upstream_url,
execution_strategy,
conversion_mode,
report_kind: "openai_image_stream_success".to_string(),
envelope_name: None,
transport: Arc::clone(transport),
request_redacted: false,
image_request_summary: Some(image_request_summary),
}))
}
fn build_openai_image_provider_body_from_openai_chat_body(
body_json: &Value,
requested_model: &str,
upstream_is_stream: bool,
) -> Option<(Value, Value)> {
let (prompt, images) = collect_openai_chat_image_prompt_and_images(body_json)?;
let operation = if images.is_empty() {
"generate"
} else {
"edit"
};
let mut tool = serde_json::Map::new();
tool.insert(
"type".to_string(),
Value::String("image_generation".to_string()),
);
tool.insert("action".to_string(), Value::String(operation.to_string()));
copy_openai_chat_image_tool_option(body_json, &mut tool, "size");
copy_openai_chat_image_tool_option(body_json, &mut tool, "quality");
copy_openai_chat_image_tool_option(body_json, &mut tool, "background");
copy_openai_chat_image_tool_option(body_json, &mut tool, "output_format");
copy_openai_chat_image_tool_option(body_json, &mut tool, "output_compression");
copy_openai_chat_image_tool_option(body_json, &mut tool, "moderation");
copy_openai_chat_image_tool_option(body_json, &mut tool, "input_fidelity");
copy_openai_chat_image_tool_option(body_json, &mut tool, "partial_images");
let input = if images.is_empty() {
serde_json::json!([{
"role": "user",
"content": prompt,
}])
} else {
let mut content = vec![serde_json::json!({
"type": "input_text",
"text": prompt,
})];
content.extend(images);
serde_json::json!([{
"role": "user",
"content": content,
}])
};
let mut body = serde_json::Map::new();
if let Some(model) = body_json
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| {
let requested_model = requested_model.trim();
(!requested_model.is_empty()).then_some(requested_model)
})
{
body.insert("model".to_string(), Value::String(model.to_string()));
}
body.insert("input".to_string(), input);
body.insert(
"tools".to_string(),
Value::Array(vec![Value::Object(tool.clone())]),
);
if upstream_is_stream {
body.insert("stream".to_string(), Value::Bool(true));
}
if let Some(user) = body_json
.get("user")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
body.insert("user".to_string(), Value::String(user.to_string()));
}
let mut summary = serde_json::Map::new();
summary.insert(
"operation".to_string(),
Value::String(operation.to_string()),
);
for key in ["output_format", "partial_images"] {
if let Some(value) = tool.get(key) {
summary.insert(key.to_string(), value.clone());
}
}
Some((Value::Object(body), Value::Object(summary)))
}
fn build_chatgpt_web_image_provider_body_from_openai_chat_body(
body_json: &Value,
requested_model: &str,
) -> Option<(Value, Value)> {
let (prompt, images) = collect_openai_chat_image_prompt_and_images(body_json)?;
let operation = if images.is_empty() {
"generate"
} else {
"edit"
};
let size = body_json
.get("size")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("1024x1024");
let output_format = body_json
.get("output_format")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("png");
let model = body_json
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or_else(|| requested_model.trim());
let web_model = body_json
.get("web_model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("gpt-5-5-thinking");
let image_urls = openai_image_inputs_as_urls(&images);
let body = json!({
"operation": operation,
"model": if model.is_empty() { "gpt-image-2" } else { model },
"web_model": web_model,
"prompt": prompt,
"size": size,
"ratio": chatgpt_web_ratio_for_size(size),
"output_format": output_format,
"images": image_urls,
});
let summary = json!({
"operation": operation,
"output_format": output_format,
});
Some((body, summary))
}
fn copy_openai_chat_image_tool_option(
body_json: &Value,
tool: &mut serde_json::Map<String, Value>,
key: &str,
) {
if let Some(value) = body_json.get(key) {
tool.insert(key.to_string(), value.clone());
}
}
fn collect_openai_chat_image_prompt_and_images(body_json: &Value) -> Option<(String, Vec<Value>)> {
let messages = body_json.get("messages").and_then(Value::as_array)?;
let mut prompt_parts = Vec::new();
let mut images = Vec::new();
for message in messages.iter().filter_map(Value::as_object) {
let role = message
.get("role")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
let content = message.get("content");
if matches!(role, "system" | "developer" | "user") {
if let Some(text) = crate::ai_serving::extract_openai_text_content(content)
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
{
prompt_parts.push(text);
}
}
if role == "user" {
collect_openai_chat_image_inputs(content, &mut images);
}
}
let prompt = prompt_parts.join("\n").trim().to_string();
(!prompt.is_empty()).then_some((prompt, images))
}
fn collect_openai_chat_image_inputs(content: Option<&Value>, images: &mut Vec<Value>) {
let Some(parts) = content.and_then(Value::as_array) else {
return;
};
for part in parts.iter().filter_map(Value::as_object) {
let part_type = part
.get("type")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if matches!(part_type, "image_url" | "input_image") {
if let Some(url) = part
.get("image_url")
.and_then(|value| {
value
.as_str()
.or_else(|| value.get("url").and_then(Value::as_str))
})
.map(str::trim)
.filter(|value| !value.is_empty())
{
images.push(serde_json::json!({
"type": "input_image",
"image_url": url,
}));
} else if let Some(file_id) = part
.get("file_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
images.push(serde_json::json!({
"type": "input_image",
"file_id": file_id,
}));
}
}
}
}
fn openai_image_inputs_as_urls(images: &[Value]) -> Vec<Value> {
images
.iter()
.filter_map(|image| {
image
.get("image_url")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| Value::String(value.to_string()))
})
.collect()
}
fn chatgpt_web_ratio_for_size(size: &str) -> String {
let Some((width, height)) = size.split_once('x') else {
return "1:1".to_string();
};
let Ok(width) = width.trim().parse::<u64>() else {
return "1:1".to_string();
};
let Ok(height) = height.trim().parse::<u64>() else {
return "1:1".to_string();
};
if width == 0 || height == 0 {
return "1:1".to_string();
}
let divisor = gcd(width, height);
format!("{}:{}", width / divisor, height / divisor)
}
fn gcd(mut left: u64, mut right: u64) -> u64 {
while right != 0 {
let next = left % right;
left = right;
right = next;
}
left.max(1)
}
fn chatgpt_web_image_internal_url(base_url: &str) -> String {
let base_url = base_url.trim().trim_end_matches('/');
let base_url = if base_url.is_empty() {
"https://chatgpt.com"
} else {
base_url
};
format!("{base_url}/__aether/chatgpt-web-image")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn chatgpt_web_chat_image_bridge_body_uses_internal_web_shape() {
let body_json = json!({
"model": "gpt-image-2",
"messages": [
{"role": "system", "content": "Use crisp vector-like shapes."},
{
"role": "user",
"content": [
{"type": "text", "text": "Draw a glass city"},
{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}}
]
}
],
"size": "1536x1024",
"output_format": "webp",
"web_model": "gpt-5-image-test"
});
let (provider_body, summary) =
build_chatgpt_web_image_provider_body_from_openai_chat_body(&body_json, "gpt-image-2")
.expect("chat image body should convert");
assert_eq!(provider_body["operation"], "edit");
assert_eq!(provider_body["model"], "gpt-image-2");
assert_eq!(provider_body["web_model"], "gpt-5-image-test");
assert_eq!(
provider_body["prompt"],
"Use crisp vector-like shapes.\nDraw a glass city"
);
assert_eq!(provider_body["size"], "1536x1024");
assert_eq!(provider_body["ratio"], "3:2");
assert_eq!(provider_body["output_format"], "webp");
assert_eq!(provider_body["images"][0], "https://example.com/ref.png");
assert_eq!(summary["operation"], "edit");
assert_eq!(summary["output_format"], "webp");
}
}
#[allow(clippy::too_many_arguments)]
async fn build_kiro_openai_chat_cross_format_payload_parts(
state: &AppState,
@@ -762,6 +1264,7 @@ async fn build_kiro_openai_chat_cross_format_payload_parts(
ai_local_execution_contract_for_formats("openai:chat", provider_api_format);
Some(LocalOpenAiChatCandidatePayloadParts {
client_api_format: "openai:chat".to_string(),
auth_header,
auth_value,
mapped_model,
@@ -775,6 +1278,7 @@ async fn build_kiro_openai_chat_cross_format_payload_parts(
envelope_name: Some(KIRO_ENVELOPE_NAME),
transport: Arc::clone(transport),
request_redacted,
image_request_summary: None,
})
}

View File

@@ -14,7 +14,10 @@ use crate::ai_serving::planner::candidate_metadata::{
LocalExecutionCandidateMetadataParts,
};
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
use crate::ai_serving::planner::candidate_source::LocalCandidatePreselectionKeyMode;
use crate::ai_serving::planner::candidate_source::{
preselect_local_execution_candidates_for_api_formats_with_serving,
LocalCandidatePreselectionKeyMode,
};
use crate::ai_serving::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
@@ -23,7 +26,7 @@ use crate::ai_serving::{
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
ExecutionRuntimeAuthContext, PlannerAppState,
};
use crate::AppState;
use crate::{AppState, GatewayError};
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalOpenAiChatCandidateAttempt;
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttemptSource as LocalOpenAiChatCandidateAttemptSource;
@@ -351,3 +354,93 @@ pub(crate) async fn build_lazy_local_openai_chat_candidate_attempt_source<'a>(
)
.await
}
pub(crate) async fn build_local_openai_chat_image_candidate_attempt_source<'a>(
state: &'a AppState,
trace_id: &str,
input: &LocalOpenAiChatDecisionInput,
body_json: &serde_json::Value,
) -> Result<(LocalOpenAiChatCandidateAttemptSource<'a>, usize), GatewayError> {
let planner_state = PlannerAppState::new(state);
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
let persistence_policy = build_local_candidate_persistence_policy(
auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiChatDecision,
);
let preselection = preselect_local_execution_candidates_for_api_formats_with_serving(
planner_state,
"openai:chat",
&input.requested_model,
false,
input.required_capabilities.as_ref(),
&input.auth_snapshot,
input.client_session_affinity.as_ref(),
false,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
vec!["openai:image".to_string()],
)
.await?;
Ok(build_local_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
"openai:image",
Some(&input.requested_model),
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
preselection.candidates,
preselection.skipped_candidates,
LocalCandidateResolutionMode::WithoutTransportPairGate,
|eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:image", &provider_api_format);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: "openai:image",
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
eligible.candidate.endpoint_api_format.trim(),
))
},
|mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:image", &provider_api_format);
skipped_candidate.extra_data = Some(
build_local_execution_candidate_contract_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
provider_api_format.as_str(),
"openai:image",
serde_json::Map::new(),
execution_strategy,
conversion_mode,
provider_api_format.as_str(),
),
);
skipped_candidate
},
)
.await)
}

View File

@@ -11,6 +11,7 @@ mod plans;
use self::decision::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_image_candidate_attempt_source,
maybe_build_local_openai_chat_decision_payload_for_candidate, LocalOpenAiChatCandidateAttempt,
LocalOpenAiChatCandidateAttemptSource, LocalOpenAiChatDecisionInput,
};

View File

@@ -1,8 +1,10 @@
use async_trait::async_trait;
use tracing::warn;
use super::super::super::openai_request_is_image_generation_intent;
use super::super::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_image_candidate_attempt_source,
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
GatewayError, LocalOpenAiChatCandidateAttempt, LocalOpenAiChatCandidateAttemptSource,
LocalOpenAiChatDecisionInput,
@@ -48,10 +50,39 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
return Ok(None);
};
let (candidates, candidate_count) = build_lazy_local_openai_chat_candidate_attempt_source(
state, trace_id, &input, body_json, true,
)
.await;
let image_generation_intent =
openai_request_is_image_generation_intent(&input.requested_model, body_json);
let (mut candidates, mut candidate_count) = if image_generation_intent {
let (image_candidates, image_candidate_count) =
build_local_openai_chat_image_candidate_attempt_source(
state, trace_id, &input, body_json,
)
.await?;
if image_candidate_count > 0 {
(image_candidates, image_candidate_count)
} else {
build_lazy_local_openai_chat_candidate_attempt_source(
state, trace_id, &input, body_json, true,
)
.await
}
} else {
build_lazy_local_openai_chat_candidate_attempt_source(
state, trace_id, &input, body_json, true,
)
.await
};
if !image_generation_intent && candidate_count == 0 {
let (image_candidates, image_candidate_count) =
build_local_openai_chat_image_candidate_attempt_source(
state, trace_id, &input, body_json,
)
.await?;
if image_candidate_count > 0 {
candidates = image_candidates;
candidate_count = image_candidate_count;
}
}
if candidate_count == 0 {
set_local_openai_chat_candidate_evaluation_diagnostic(
state,

View File

@@ -0,0 +1,84 @@
pub(crate) fn openai_request_is_image_generation_intent(
requested_model: &str,
body_json: &serde_json::Value,
) -> bool {
openai_model_is_image_generation(requested_model)
|| body_json
.get("model")
.and_then(serde_json::Value::as_str)
.is_some_and(openai_model_is_image_generation)
|| openai_tools_contain_image_generation(body_json.get("tools"))
|| openai_tool_choice_selects_image_generation(body_json.get("tool_choice"))
}
fn openai_model_is_image_generation(model: &str) -> bool {
model.trim().to_ascii_lowercase().starts_with("gpt-image-")
}
fn openai_tools_contain_image_generation(tools: Option<&serde_json::Value>) -> bool {
tools
.and_then(serde_json::Value::as_array)
.is_some_and(|items| {
items.iter().any(|item| {
item.get("type")
.and_then(serde_json::Value::as_str)
.is_some_and(|value| value.trim().eq_ignore_ascii_case("image_generation"))
})
})
}
fn openai_tool_choice_selects_image_generation(choice: Option<&serde_json::Value>) -> bool {
let Some(choice) = choice else {
return false;
};
if let Some(value) = choice.as_str() {
return value.trim().eq_ignore_ascii_case("image_generation");
}
let Some(object) = choice.as_object() else {
return false;
};
object
.get("type")
.and_then(serde_json::Value::as_str)
.is_some_and(|value| value.trim().eq_ignore_ascii_case("image_generation"))
|| object
.get("tool")
.and_then(|value| value.get("type"))
.and_then(serde_json::Value::as_str)
.is_some_and(|value| value.trim().eq_ignore_ascii_case("image_generation"))
|| object
.get("function")
.and_then(|value| value.get("name"))
.and_then(serde_json::Value::as_str)
.is_some_and(|value| value.trim().eq_ignore_ascii_case("image_generation"))
}
#[cfg(test)]
mod tests {
use super::openai_request_is_image_generation_intent;
use serde_json::json;
#[test]
fn detects_openai_image_generation_intent_like_compat_proxies() {
assert!(openai_request_is_image_generation_intent(
"GPT-IMAGE-2",
&json!({})
));
assert!(openai_request_is_image_generation_intent(
"gpt-5",
&json!({"model":"gpt-image-2"})
));
assert!(openai_request_is_image_generation_intent(
"gpt-5",
&json!({"tools":[{"type":"image_generation"}]})
));
assert!(openai_request_is_image_generation_intent(
"gpt-5",
&json!({"tool_choice":{"function":{"name":"image_generation"}}})
));
assert!(!openai_request_is_image_generation_intent(
"gpt-5",
&json!({"messages":[{"role":"user","content":"hello"}]})
));
}
}

View File

@@ -1,4 +1,5 @@
mod chat;
mod image_intent;
mod responses;
pub(crate) use crate::ai_serving::{
@@ -14,6 +15,7 @@ pub(crate) use chat::{
maybe_build_stream_local_decision_payload, maybe_build_sync_local_decision_payload,
set_local_openai_chat_execution_exhausted_diagnostic,
};
pub(super) use image_intent::openai_request_is_image_generation_intent;
pub(crate) use responses::{
build_local_openai_responses_stream_attempt_source_for_kind,
build_local_openai_responses_stream_plan_and_reports_for_kind,

View File

@@ -74,6 +74,30 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
if let Some(envelope_name) = resolved.envelope_name {
extra_fields.insert("envelope_name".to_string(), json!(envelope_name));
}
if let Some(image_request_summary) = resolved.image_request_summary.as_ref() {
extra_fields.insert("image_request".to_string(), image_request_summary.clone());
}
if resolved
.provider_api_format
.eq_ignore_ascii_case("openai:image")
&& resolved
.transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("chatgpt_web")
{
extra_fields.insert("chatgpt_web_image".to_string(), json!(true));
extra_fields.insert(
"local_failover_policy".to_string(),
json!({
"stop_status_codes": [400, 401, 403, 429, 500, 502, 503, 504],
"error_stop_patterns": [
{ "pattern": ".*" }
]
}),
);
}
insert_provider_stream_event_api_format(
&mut extra_fields,
resolved.transport.provider.provider_type.as_str(),
@@ -170,6 +194,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
envelope_name: _,
upstream_is_stream,
transport,
image_request_summary: _,
} = resolved;
Some(build_ai_execution_decision_response(

View File

@@ -1,7 +1,7 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use serde_json::Value;
use serde_json::{json, Value};
use tracing::debug;
use crate::ai_serving::planner::candidate_preparation::{
@@ -15,7 +15,8 @@ use crate::ai_serving::planner::common::{
};
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, build_cross_format_openai_responses_request_body,
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
build_cross_format_openai_responses_request_body,
build_cross_format_openai_responses_upstream_url, build_local_openai_responses_request_body,
build_local_openai_responses_upstream_url, request_body_build_failure_extra_data,
};
@@ -35,8 +36,11 @@ use crate::ai_serving::transport::kiro::{
KiroRequestAuth, KIRO_ENVELOPE_NAME,
};
use crate::ai_serving::transport::{
build_kiro_cross_format_upstream_url, build_standard_provider_request_headers,
local_standard_transport_unsupported_reason_with_network, StandardProviderRequestHeadersInput,
build_kiro_cross_format_upstream_url, build_openai_image_headers,
build_openai_image_upstream_url, build_standard_provider_request_headers,
local_standard_transport_unsupported_reason_with_network,
openai_image_transport_unsupported_reason, resolve_openai_image_auth,
ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput,
};
use crate::ai_serving::{
ai_local_execution_contract_for_formats, request_conversion_direct_auth,
@@ -70,6 +74,7 @@ pub(crate) struct LocalOpenAiResponsesCandidatePayloadParts {
pub(super) envelope_name: Option<&'static str>,
pub(super) upstream_is_stream: bool,
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
pub(super) image_request_summary: Option<Value>,
}
#[allow(clippy::too_many_arguments)]
@@ -93,6 +98,21 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
let is_antigravity = is_antigravity_provider_transport(transport);
let is_kiro_claude_cli = is_kiro_claude_messages_transport(transport, provider_api_format);
if provider_api_format.eq_ignore_ascii_case("openai:image") {
return resolve_openai_responses_to_openai_image_payload_parts(
state,
parts,
trace_id,
body_json,
input,
eligible,
candidate_index,
candidate_id,
spec,
)
.await;
}
let same_format = api_format_alias_matches(provider_api_format, &client_api_format);
let conversion_kind = request_conversion_kind(spec_metadata.api_format, provider_api_format);
let transport_unsupported_reason = if same_format && is_kiro_claude_cli {
@@ -516,6 +536,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
},
upstream_is_stream,
transport: Arc::clone(transport),
image_request_summary: None,
})
}
@@ -523,6 +544,459 @@ fn api_format_alias_matches(left: &str, right: &str) -> bool {
crate::ai_serving::api_format_alias_matches(left, right)
}
#[allow(clippy::too_many_arguments)]
async fn resolve_openai_responses_to_openai_image_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiResponsesDecisionInput,
eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32,
candidate_id: &str,
spec: LocalOpenAiResponsesSpec,
) -> Option<LocalOpenAiResponsesCandidatePayloadParts> {
let spec_metadata = local_openai_responses_spec_metadata(spec);
let candidate = &eligible.candidate;
let transport = &eligible.transport;
let provider_api_format = "openai:image";
if let Some(skip_reason) =
openai_image_transport_unsupported_reason(transport, provider_api_format)
{
mark_skipped_local_openai_responses_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return None;
}
let prepared_candidate = match prepare_header_authenticated_candidate(
PlannerAppState::new(state),
transport,
candidate,
resolve_openai_image_auth(transport),
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "openai_responses_image_bridge",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_openai_responses_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return None;
}
};
let is_chatgpt_web = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("chatgpt_web");
let upstream_is_stream = resolve_upstream_is_stream_for_provider(
transport.endpoint.config.as_ref(),
transport.provider.provider_type.as_str(),
provider_api_format,
spec_metadata.require_streaming,
false,
);
let Some((mut provider_request_body, image_request_summary)) = (if is_chatgpt_web {
build_chatgpt_web_image_provider_body_from_openai_responses_body(
body_json,
&input.requested_model,
)
} else {
build_openai_image_provider_body_from_openai_responses_body(
body_json,
&input.requested_model,
upstream_is_stream,
)
}) else {
mark_skipped_local_openai_responses_candidate_with_extra_data(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_build_failed",
request_body_build_failure_extra_data(
body_json,
spec_metadata.api_format,
provider_api_format,
),
)
.await;
return None;
};
if !is_chatgpt_web {
apply_codex_openai_responses_special_body_edits(
&mut provider_request_body,
transport.provider.provider_type.as_str(),
provider_api_format,
transport.endpoint.body_rules.as_ref(),
Some(candidate.key_id.as_str()),
);
}
let upstream_url = if is_chatgpt_web {
chatgpt_web_image_internal_url(&transport.endpoint.base_url)
} else {
build_openai_image_upstream_url(transport, parts.uri.query())
};
let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: body_json,
})
else {
mark_skipped_local_openai_responses_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
spec_metadata.api_format,
provider_api_format,
"openai_responses_image_bridge_headers",
),
)
.await;
return None;
};
if is_chatgpt_web {
provider_request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string());
} else {
apply_codex_openai_responses_special_headers(
&mut provider_request_headers,
&provider_request_body,
&parts.headers,
transport.provider.provider_type.as_str(),
provider_api_format,
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
}
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(spec_metadata.api_format, provider_api_format);
Some(LocalOpenAiResponsesCandidatePayloadParts {
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
mapped_model: prepared_candidate.mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body,
provider_request_headers,
upstream_url,
execution_strategy,
conversion_mode,
is_antigravity: false,
envelope_name: None,
upstream_is_stream,
transport: Arc::clone(transport),
image_request_summary: Some(image_request_summary),
})
}
fn build_openai_image_provider_body_from_openai_responses_body(
body_json: &Value,
requested_model: &str,
upstream_is_stream: bool,
) -> Option<(Value, Value)> {
let object = body_json.as_object()?;
let input = object.get("input")?.clone();
let mut tool = openai_responses_image_generation_tool(object).unwrap_or_else(|| {
serde_json::Map::from_iter([("type".to_string(), json!("image_generation"))])
});
tool.entry("type".to_string())
.or_insert_with(|| json!("image_generation"));
tool.entry("action".to_string())
.or_insert_with(|| json!("generate"));
let mut body = serde_json::Map::new();
body.insert("input".to_string(), input);
body.insert(
"tools".to_string(),
Value::Array(vec![Value::Object(tool.clone())]),
);
if let Some(model) = object
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| {
let requested_model = requested_model.trim();
(!requested_model.is_empty()).then_some(requested_model)
})
{
body.insert("model".to_string(), Value::String(model.to_string()));
}
for key in [
"user",
"metadata",
"include",
"parallel_tool_calls",
"store",
] {
if let Some(value) = object.get(key) {
body.insert(key.to_string(), value.clone());
}
}
if upstream_is_stream {
body.insert("stream".to_string(), Value::Bool(true));
} else if let Some(value) = object.get("stream") {
body.insert("stream".to_string(), value.clone());
}
let mut summary = serde_json::Map::new();
summary.insert(
"operation".to_string(),
tool.get("action")
.cloned()
.unwrap_or_else(|| json!("generate")),
);
for key in ["output_format", "partial_images", "size", "quality"] {
if let Some(value) = tool.get(key).or_else(|| object.get(key)) {
summary.insert(key.to_string(), value.clone());
}
}
Some((Value::Object(body), Value::Object(summary)))
}
fn openai_responses_image_generation_tool(
object: &serde_json::Map<String, Value>,
) -> Option<serde_json::Map<String, Value>> {
object
.get("tools")
.and_then(Value::as_array)?
.iter()
.filter_map(Value::as_object)
.find(|tool| {
tool.get("type")
.and_then(Value::as_str)
.is_some_and(|value| value.trim().eq_ignore_ascii_case("image_generation"))
})
.cloned()
}
fn build_chatgpt_web_image_provider_body_from_openai_responses_body(
body_json: &Value,
requested_model: &str,
) -> Option<(Value, Value)> {
let object = body_json.as_object()?;
let (prompt, images) = collect_openai_responses_image_prompt_and_images(object.get("input"))?;
let operation = if images.is_empty() {
"generate"
} else {
"edit"
};
let tool = openai_responses_image_generation_tool(object);
let size = image_option_string(tool.as_ref(), object, "size").unwrap_or("1024x1024");
let output_format =
image_option_string(tool.as_ref(), object, "output_format").unwrap_or("png");
let model = object
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or_else(|| requested_model.trim());
let web_model = image_option_string(tool.as_ref(), object, "web_model")
.or_else(|| image_option_string(tool.as_ref(), object, "model"))
.unwrap_or("gpt-5-5-thinking");
let image_urls = openai_image_inputs_as_urls(&images);
let body = json!({
"operation": operation,
"model": if model.is_empty() { "gpt-image-2" } else { model },
"web_model": web_model,
"prompt": prompt,
"size": size,
"ratio": chatgpt_web_ratio_for_size(size),
"output_format": output_format,
"images": image_urls,
});
let summary = json!({
"operation": operation,
"output_format": output_format,
});
Some((body, summary))
}
fn image_option_string<'a>(
tool: Option<&'a serde_json::Map<String, Value>>,
object: &'a serde_json::Map<String, Value>,
key: &str,
) -> Option<&'a str> {
tool.and_then(|tool| tool.get(key))
.or_else(|| object.get(key))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn collect_openai_responses_image_prompt_and_images(
input: Option<&Value>,
) -> Option<(String, Vec<Value>)> {
let input = input?;
let mut prompt_parts = Vec::new();
let mut images = Vec::new();
collect_openai_responses_image_input(input, &mut prompt_parts, &mut images);
let prompt = prompt_parts.join("\n").trim().to_string();
(!prompt.is_empty()).then_some((prompt, images))
}
fn collect_openai_responses_image_input(
value: &Value,
prompt_parts: &mut Vec<String>,
images: &mut Vec<Value>,
) {
match value {
Value::String(text) => {
let text = text.trim();
if !text.is_empty() {
prompt_parts.push(text.to_string());
}
}
Value::Array(items) => {
for item in items {
collect_openai_responses_image_input(item, prompt_parts, images);
}
}
Value::Object(object) => {
let item_type = object
.get("type")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if matches!(item_type, "input_text" | "text") {
if let Some(text) = object
.get("text")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
prompt_parts.push(text.to_string());
}
} else if matches!(item_type, "input_image" | "image_url") {
collect_openai_image_input_object(object, images);
}
if let Some(content) = object.get("content") {
collect_openai_responses_image_input(content, prompt_parts, images);
}
}
_ => {}
}
}
fn collect_openai_image_input_object(
object: &serde_json::Map<String, Value>,
images: &mut Vec<Value>,
) {
if let Some(url) = object
.get("image_url")
.and_then(|value| {
value
.as_str()
.or_else(|| value.get("url").and_then(Value::as_str))
})
.or_else(|| object.get("url").and_then(Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())
{
images.push(json!({
"type": "input_image",
"image_url": url,
}));
} else if let Some(file_id) = object
.get("file_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
images.push(json!({
"type": "input_image",
"file_id": file_id,
}));
}
}
fn openai_image_inputs_as_urls(images: &[Value]) -> Vec<Value> {
images
.iter()
.filter_map(|image| {
image
.get("image_url")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| Value::String(value.to_string()))
})
.collect()
}
fn chatgpt_web_ratio_for_size(size: &str) -> String {
let Some((width, height)) = size.split_once('x') else {
return "1:1".to_string();
};
let Ok(width) = width.trim().parse::<u64>() else {
return "1:1".to_string();
};
let Ok(height) = height.trim().parse::<u64>() else {
return "1:1".to_string();
};
if width == 0 || height == 0 {
return "1:1".to_string();
}
let divisor = gcd(width, height);
format!("{}:{}", width / divisor, height / divisor)
}
fn gcd(mut left: u64, mut right: u64) -> u64 {
while right != 0 {
let next = left % right;
left = right;
right = next;
}
left.max(1)
}
fn chatgpt_web_image_internal_url(base_url: &str) -> String {
let base_url = base_url.trim().trim_end_matches('/');
let base_url = if base_url.is_empty() {
"https://chatgpt.com"
} else {
base_url
};
format!("{base_url}/__aether/chatgpt-web-image")
}
#[allow(clippy::too_many_arguments)]
async fn build_kiro_openai_responses_payload_parts(
state: &AppState,
@@ -666,5 +1140,6 @@ async fn build_kiro_openai_responses_payload_parts(
envelope_name: Some(KIRO_ENVELOPE_NAME),
upstream_is_stream,
transport: Arc::clone(transport),
image_request_summary: None,
})
}

View File

@@ -15,6 +15,7 @@ use crate::ai_serving::planner::candidate_metadata::{
LocalExecutionCandidateMetadataParts,
};
use crate::ai_serving::planner::candidate_source::{
preselect_local_execution_candidates_for_api_formats_with_serving,
preselect_local_execution_candidates_with_serving, LocalCandidatePreselectionKeyMode,
};
use crate::ai_serving::planner::common::extract_standard_requested_model;
@@ -35,6 +36,7 @@ use crate::ai_serving::{
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::{AppState, GatewayError};
use super::super::super::openai_request_is_image_generation_intent;
use super::LocalOpenAiResponsesSpec;
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalOpenAiResponsesCandidateAttempt;
@@ -246,6 +248,16 @@ pub(crate) async fn build_local_openai_responses_candidate_attempt_source<'a>(
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiResponsesDecision,
);
if openai_request_is_image_generation_intent(&input.requested_model, body_json) {
let (image_candidates, image_candidate_count) =
build_local_openai_responses_image_candidate_attempt_source(
state, trace_id, input, body_json, spec,
)
.await?;
if image_candidate_count > 0 {
return Ok((image_candidates, image_candidate_count));
}
}
Ok(
build_lazy_requested_model_execution_candidate_attempt_source_with_serving(
planner_state,
@@ -315,6 +327,102 @@ pub(crate) async fn build_local_openai_responses_candidate_attempt_source<'a>(
)
}
pub(crate) async fn build_local_openai_responses_image_candidate_attempt_source<'a>(
state: &'a AppState,
trace_id: &str,
input: &LocalOpenAiResponsesDecisionInput,
body_json: &serde_json::Value,
spec: LocalOpenAiResponsesSpec,
) -> Result<(LocalOpenAiResponsesCandidateAttemptSource<'a>, usize), GatewayError> {
let spec_metadata = local_openai_responses_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
let persistence_policy = build_local_candidate_persistence_policy(
auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiResponsesDecision,
);
let preselection = preselect_local_execution_candidates_for_api_formats_with_serving(
planner_state,
spec_metadata.api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
&input.auth_snapshot,
input.client_session_affinity.as_ref(),
true,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
vec!["openai:image".to_string()],
)
.await?;
Ok(build_local_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
spec_metadata.api_format,
Some(&input.requested_model),
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
preselection.candidates,
preselection.skipped_candidates,
LocalCandidateResolutionMode::WithoutTransportPairGate,
move |eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
&provider_api_format,
);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
eligible.candidate.endpoint_api_format.as_str(),
))
},
move |mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
&provider_api_format,
);
skipped_candidate.extra_data = Some(
build_local_execution_candidate_contract_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
provider_api_format.as_str(),
spec_metadata.api_format,
serde_json::Map::new(),
execution_strategy,
conversion_mode,
provider_api_format.as_str(),
),
);
skipped_candidate
},
)
.await)
}
pub(crate) async fn mark_skipped_local_openai_responses_candidate(
state: &AppState,
input: &LocalOpenAiResponsesDecisionInput,