mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
Merge remote-tracking branch 'origin/pr-483' into merge-pr-483
# Conflicts: # apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs # apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs # apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs # apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs # apps/aether-gateway/src/execution_runtime/chatgpt_web_image.rs
This commit is contained in:
@@ -121,7 +121,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.routing_policy.as_ref(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -89,7 +89,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,
|
||||
@@ -104,6 +129,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
||||
transport,
|
||||
request_redacted,
|
||||
transport_profile: _,
|
||||
image_request_summary: _,
|
||||
} = resolved;
|
||||
let original_request_body_json = if request_redacted {
|
||||
Some(&provider_request_body)
|
||||
@@ -128,7 +154,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(),
|
||||
@@ -160,7 +186,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,
|
||||
|
||||
@@ -4,7 +4,7 @@ use std::sync::Arc;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_contracts::ResolvedTransportProfile;
|
||||
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,
|
||||
@@ -16,9 +16,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::{
|
||||
@@ -29,8 +30,10 @@ use crate::ai_serving::transport::kiro::{
|
||||
use crate::ai_serving::transport::local_openai_chat_transport_unsupported_reason;
|
||||
use crate::ai_serving::transport::{
|
||||
build_grok_browser_headers, build_grok_upstream_url, build_kiro_cross_format_upstream_url,
|
||||
build_standard_provider_request_headers, GrokHeaderInput, StandardProviderRequestHeadersInput,
|
||||
GROK_CHAT_PATH,
|
||||
build_openai_image_headers, build_openai_image_upstream_url,
|
||||
build_standard_provider_request_headers, openai_image_transport_unsupported_reason,
|
||||
resolve_openai_image_auth, GrokHeaderInput, ProviderOpenAiImageHeadersInput,
|
||||
StandardProviderRequestHeadersInput, GROK_CHAT_PATH,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
ai_local_execution_contract_for_formats, request_conversion_direct_auth,
|
||||
@@ -53,6 +56,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,
|
||||
@@ -67,6 +71,7 @@ pub(crate) struct LocalOpenAiChatCandidatePayloadParts {
|
||||
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
pub(super) request_redacted: bool,
|
||||
pub(super) transport_profile: Option<ResolvedTransportProfile>,
|
||||
pub(super) image_request_summary: Option<Value>,
|
||||
}
|
||||
|
||||
fn is_grok_text_provider_api_format(provider_api_format: &str) -> bool {
|
||||
@@ -157,6 +162,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"
|
||||
);
|
||||
@@ -309,6 +315,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: prepared_candidate.auth_header,
|
||||
auth_value: prepared_candidate.auth_value,
|
||||
mapped_model: prepared_candidate.mapped_model,
|
||||
@@ -323,6 +330,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
transport: Arc::clone(transport),
|
||||
request_redacted: redaction.redacted,
|
||||
transport_profile,
|
||||
image_request_summary: None,
|
||||
}));
|
||||
}
|
||||
|
||||
@@ -473,6 +481,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,
|
||||
@@ -487,10 +496,26 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
transport: Arc::clone(transport),
|
||||
request_redacted: redaction.redacted,
|
||||
transport_profile,
|
||||
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 {
|
||||
@@ -762,6 +787,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,
|
||||
@@ -776,9 +802,496 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
transport: Arc::clone(transport),
|
||||
request_redacted: redaction.redacted,
|
||||
transport_profile: None,
|
||||
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:chat", provider_api_format);
|
||||
|
||||
Ok(Some(LocalOpenAiChatCandidatePayloadParts {
|
||||
client_api_format: "openai:chat".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_chat_stream_success".to_string(),
|
||||
envelope_name: None,
|
||||
transport: Arc::clone(transport),
|
||||
request_redacted: false,
|
||||
transport_profile: None,
|
||||
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", "size", "quality"] {
|
||||
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 quality = body_json
|
||||
.get("quality")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("medium");
|
||||
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,
|
||||
"size": size,
|
||||
"quality": quality,
|
||||
});
|
||||
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,
|
||||
@@ -900,6 +1413,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,
|
||||
@@ -914,6 +1428,7 @@ async fn build_kiro_openai_chat_cross_format_payload_parts(
|
||||
transport: Arc::clone(transport),
|
||||
request_redacted,
|
||||
transport_profile: None,
|
||||
image_request_summary: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -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;
|
||||
@@ -354,3 +357,95 @@ 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.routing_policy.as_ref(),
|
||||
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:chat",
|
||||
Some(&input.requested_model),
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.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:chat", &provider_api_format);
|
||||
Some(build_local_execution_candidate_contract_metadata(
|
||||
LocalExecutionCandidateMetadataParts {
|
||||
eligible,
|
||||
provider_api_format: provider_api_format.as_str(),
|
||||
client_api_format: "openai:chat",
|
||||
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:chat", &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:chat",
|
||||
serde_json::Map::new(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
provider_api_format.as_str(),
|
||||
),
|
||||
);
|
||||
skipped_candidate
|
||||
},
|
||||
)
|
||||
.await)
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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,
|
||||
@@ -49,14 +51,53 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
|
||||
};
|
||||
let effective_body_json = input.effective_body_json(body_json).clone();
|
||||
|
||||
let (candidates, candidate_count) = build_lazy_local_openai_chat_candidate_attempt_source(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
&effective_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,
|
||||
&effective_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,
|
||||
&effective_body_json,
|
||||
true,
|
||||
)
|
||||
.await
|
||||
}
|
||||
} else {
|
||||
build_lazy_local_openai_chat_candidate_attempt_source(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
&effective_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,
|
||||
&effective_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,
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
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_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_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!({"tool_choice":{"function":{"name":"image_generation"}}})
|
||||
));
|
||||
assert!(openai_request_is_image_generation_intent(
|
||||
"gpt-5",
|
||||
&json!({"tool_choice":{"type":"image_generation"}})
|
||||
));
|
||||
assert!(!openai_request_is_image_generation_intent(
|
||||
"gpt-5",
|
||||
&json!({"tools":[{"type":"image_generation"}]})
|
||||
));
|
||||
assert!(!openai_request_is_image_generation_intent(
|
||||
"gpt-5",
|
||||
&json!({"messages":[{"role":"user","content":"hello"}]})
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -81,6 +81,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(),
|
||||
@@ -179,6 +203,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
|
||||
upstream_is_stream,
|
||||
transport,
|
||||
transport_profile: _,
|
||||
image_request_summary: _,
|
||||
} = resolved;
|
||||
|
||||
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
|
||||
|
||||
@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_contracts::ResolvedTransportProfile;
|
||||
use serde_json::Value;
|
||||
use serde_json::{json, Value};
|
||||
use tracing::debug;
|
||||
|
||||
use crate::ai_serving::planner::candidate_preparation::{
|
||||
@@ -16,7 +16,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,
|
||||
};
|
||||
@@ -37,9 +38,11 @@ use crate::ai_serving::transport::kiro::{
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
build_grok_browser_headers, build_grok_upstream_url, 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, GrokHeaderInput,
|
||||
StandardProviderRequestHeadersInput, GROK_CHAT_PATH,
|
||||
local_standard_transport_unsupported_reason_with_network,
|
||||
openai_image_transport_unsupported_reason, resolve_openai_image_auth, GrokHeaderInput,
|
||||
ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput, GROK_CHAT_PATH,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
ai_local_execution_contract_for_formats, request_conversion_direct_auth,
|
||||
@@ -81,6 +84,7 @@ pub(crate) struct LocalOpenAiResponsesCandidatePayloadParts {
|
||||
pub(super) upstream_is_stream: bool,
|
||||
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
pub(super) transport_profile: Option<ResolvedTransportProfile>,
|
||||
pub(super) image_request_summary: Option<Value>,
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
@@ -110,6 +114,21 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("grok");
|
||||
|
||||
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 is_grok
|
||||
@@ -596,6 +615,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
|
||||
upstream_is_stream,
|
||||
transport: Arc::clone(transport),
|
||||
transport_profile,
|
||||
image_request_summary: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -603,6 +623,463 @@ 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),
|
||||
transport_profile: None,
|
||||
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 quality = image_option_string(tool.as_ref(), object, "quality").unwrap_or("medium");
|
||||
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,
|
||||
"size": size,
|
||||
"quality": quality,
|
||||
});
|
||||
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,
|
||||
@@ -748,5 +1225,6 @@ async fn build_kiro_openai_responses_payload_parts(
|
||||
upstream_is_stream,
|
||||
transport: Arc::clone(transport),
|
||||
transport_profile: None,
|
||||
image_request_summary: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
@@ -36,6 +37,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;
|
||||
@@ -265,6 +267,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,
|
||||
@@ -335,6 +347,104 @@ 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.routing_policy.as_ref(),
|
||||
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(),
|
||||
input.routing_policy.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,
|
||||
|
||||
@@ -3,20 +3,22 @@ use std::io::Error as IoError;
|
||||
use std::time::Instant;
|
||||
|
||||
use aether_contracts::{
|
||||
ExecutionPlan, ExecutionResult, ExecutionTelemetry, RequestBody, ResolvedTransportProfile,
|
||||
ResponseBody, StreamFrame, StreamFramePayload, StreamFrameType,
|
||||
EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
|
||||
TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_HTTP_MODE_AUTO, TRANSPORT_POOL_SCOPE_KEY,
|
||||
ExecutionPlan, ExecutionResult, ExecutionStreamTerminalSummary, ExecutionTelemetry,
|
||||
RequestBody, ResolvedTransportProfile, ResponseBody, StreamFrame, StreamFramePayload,
|
||||
StreamFrameType, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
|
||||
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, TRANSPORT_BACKEND_BROWSER_WREQ,
|
||||
TRANSPORT_HTTP_MODE_AUTO, TRANSPORT_POOL_SCOPE_KEY,
|
||||
};
|
||||
use axum::body::Bytes;
|
||||
use base64::Engine as _;
|
||||
use chrono::{FixedOffset, Utc};
|
||||
use futures_util::stream::{self, BoxStream};
|
||||
use futures_util::StreamExt;
|
||||
use serde_json::{json, Value};
|
||||
use serde_json::{json, Map, Value};
|
||||
use tracing::debug;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_serving::api::StreamingStandardTerminalObserver;
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
|
||||
use crate::execution_runtime::transport::{
|
||||
@@ -111,16 +113,15 @@ pub(crate) async fn maybe_execute_chatgpt_web_image_stream(
|
||||
Err(err) => chatgpt_web_transport_error_execution_result(plan, started_at, &err),
|
||||
};
|
||||
Ok(Some(ChatGptWebImageStream {
|
||||
frame_stream: execution_result_frame_stream(&result),
|
||||
frame_stream: execution_result_frame_stream(plan, &result, report_context),
|
||||
report_context: report_context.cloned(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn is_chatgpt_web_image_plan(plan: &ExecutionPlan, report_context: Option<&Value>) -> bool {
|
||||
if !plan.client_api_format.eq_ignore_ascii_case("openai:image")
|
||||
|| !plan
|
||||
.provider_api_format
|
||||
.eq_ignore_ascii_case("openai:image")
|
||||
if !plan
|
||||
.provider_api_format
|
||||
.eq_ignore_ascii_case("openai:image")
|
||||
{
|
||||
return false;
|
||||
}
|
||||
@@ -1510,9 +1511,12 @@ fn bytes_execution_result(
|
||||
}
|
||||
|
||||
fn execution_result_frame_stream(
|
||||
plan: &ExecutionPlan,
|
||||
result: &ExecutionResult,
|
||||
report_context: Option<&Value>,
|
||||
) -> BoxStream<'static, Result<Bytes, IoError>> {
|
||||
let body = execution_result_body_bytes_lossy(result);
|
||||
let terminal_summary = chatgpt_web_stream_terminal_summary(plan, result, report_context, &body);
|
||||
let mut frames = vec![
|
||||
StreamFrame {
|
||||
frame_type: StreamFrameType::Headers,
|
||||
@@ -1551,7 +1555,7 @@ fn execution_result_frame_stream(
|
||||
}),
|
||||
},
|
||||
});
|
||||
frames.push(StreamFrame::eof());
|
||||
frames.push(StreamFrame::eof_with_summary(terminal_summary));
|
||||
stream::iter(
|
||||
frames
|
||||
.into_iter()
|
||||
@@ -1560,6 +1564,81 @@ fn execution_result_frame_stream(
|
||||
.boxed()
|
||||
}
|
||||
|
||||
fn chatgpt_web_stream_terminal_summary(
|
||||
plan: &ExecutionPlan,
|
||||
result: &ExecutionResult,
|
||||
report_context: Option<&Value>,
|
||||
body: &[u8],
|
||||
) -> Option<ExecutionStreamTerminalSummary> {
|
||||
if !(200..300).contains(&result.status_code) || body.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let observer_context = chatgpt_web_stream_observer_context(plan, report_context);
|
||||
let mut observer = StreamingStandardTerminalObserver::default();
|
||||
let mut line_start = 0usize;
|
||||
for (index, byte) in body.iter().enumerate() {
|
||||
if *byte != b'\n' {
|
||||
continue;
|
||||
}
|
||||
observer
|
||||
.push_line(&observer_context, body[line_start..=index].to_vec())
|
||||
.ok()?;
|
||||
line_start = index.saturating_add(1);
|
||||
}
|
||||
if line_start < body.len() {
|
||||
observer
|
||||
.push_line(&observer_context, body[line_start..].to_vec())
|
||||
.ok()?;
|
||||
}
|
||||
observer.finish(&observer_context).ok().flatten()
|
||||
}
|
||||
|
||||
fn chatgpt_web_stream_observer_context(
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&Value>,
|
||||
) -> Value {
|
||||
let mut context = report_context
|
||||
.cloned()
|
||||
.filter(Value::is_object)
|
||||
.unwrap_or_else(|| json!({}));
|
||||
let object = context
|
||||
.as_object_mut()
|
||||
.expect("observer context should be an object");
|
||||
object
|
||||
.entry("provider_api_format".to_string())
|
||||
.or_insert_with(|| Value::String(plan.provider_api_format.clone()));
|
||||
object
|
||||
.entry("client_api_format".to_string())
|
||||
.or_insert_with(|| Value::String(plan.client_api_format.clone()));
|
||||
object
|
||||
.entry("model".to_string())
|
||||
.or_insert_with(|| Value::String(plan.model_name.clone().unwrap_or_default()));
|
||||
if !object.contains_key("image_request") {
|
||||
if let Some(image_request) = chatgpt_web_image_request_context(plan) {
|
||||
object.insert("image_request".to_string(), image_request);
|
||||
}
|
||||
}
|
||||
context
|
||||
}
|
||||
|
||||
fn chatgpt_web_image_request_context(plan: &ExecutionPlan) -> Option<Value> {
|
||||
let body = plan.body.json_body.as_ref()?.as_object()?;
|
||||
let mut image_request = Map::new();
|
||||
image_request.insert(
|
||||
"operation".to_string(),
|
||||
Value::String("generate".to_string()),
|
||||
);
|
||||
for key in ["model", "size", "quality", "output_format"] {
|
||||
if let Some(value) = body.get(key).and_then(Value::as_str).map(str::trim) {
|
||||
if !value.is_empty() {
|
||||
image_request.insert(key.to_string(), Value::String(value.to_string()));
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(Value::Object(image_request))
|
||||
}
|
||||
|
||||
fn telemetry(started_at: Instant, upstream_bytes: u64) -> ExecutionTelemetry {
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
ExecutionTelemetry {
|
||||
@@ -2280,6 +2359,31 @@ data: [DONE]
|
||||
assert!(decoded_data.contains("\"width\":2"));
|
||||
assert!(decoded_data.contains("\"height\":3"));
|
||||
assert!(text.contains("\"type\":\"eof\""));
|
||||
let eof_frame = text
|
||||
.lines()
|
||||
.filter_map(|line| serde_json::from_str::<Value>(line).ok())
|
||||
.find(|frame| frame.get("type").and_then(Value::as_str) == Some("eof"))
|
||||
.expect("eof frame should exist");
|
||||
assert_eq!(
|
||||
eof_frame
|
||||
.get("payload")
|
||||
.and_then(|payload| payload.get("summary"))
|
||||
.and_then(|summary| summary.get("standardized_usage"))
|
||||
.and_then(|usage| usage.get("dimensions"))
|
||||
.and_then(|dimensions| dimensions.get("image_count"))
|
||||
.and_then(Value::as_u64),
|
||||
Some(1)
|
||||
);
|
||||
assert_eq!(
|
||||
eof_frame
|
||||
.get("payload")
|
||||
.and_then(|payload| payload.get("summary"))
|
||||
.and_then(|summary| summary.get("standardized_usage"))
|
||||
.and_then(|usage| usage.get("dimensions"))
|
||||
.and_then(|dimensions| dimensions.get("image_size"))
|
||||
.and_then(Value::as_str),
|
||||
Some("1024x1024")
|
||||
);
|
||||
|
||||
handle.abort();
|
||||
}
|
||||
@@ -2314,6 +2418,32 @@ data: [DONE]
|
||||
assert_eq!(body["error"]["code"], "chatgpt_web_image_unsupported");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn chatgpt_web_image_executor_accepts_marked_responses_client_plan() {
|
||||
let state = crate::AppState::new().expect("state should build");
|
||||
let mut plan = sample_plan(
|
||||
CHATGPT_WEB_DEFAULT_BASE_URL,
|
||||
json!({
|
||||
"error": {
|
||||
"message": "ChatGPT-Web 不支持该分辨率",
|
||||
"type": "invalid_request_error",
|
||||
"code": "chatgpt_web_image_unsupported"
|
||||
}
|
||||
}),
|
||||
false,
|
||||
);
|
||||
plan.client_api_format = "openai:responses".to_string();
|
||||
|
||||
let result = maybe_execute_chatgpt_web_image_sync(&state, &plan, None)
|
||||
.await
|
||||
.expect("executor should run")
|
||||
.expect("marked image provider plan should be intercepted");
|
||||
|
||||
assert_eq!(result.status_code, 400);
|
||||
let body = execution_result_json(&result).expect("error should be json");
|
||||
assert_eq!(body["error"]["code"], "chatgpt_web_image_unsupported");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn chatgpt_web_image_stream_path_wraps_executor_result_as_ndjson_frames() {
|
||||
let state = crate::AppState::new().expect("state should build");
|
||||
|
||||
@@ -13,6 +13,16 @@ fn sync_plan_kind_disables_local_candidate_failover(plan_kind: &str) -> bool {
|
||||
)
|
||||
}
|
||||
|
||||
fn openai_image_success_disables_local_success_failover(
|
||||
plan: &ExecutionPlan,
|
||||
status_code: u16,
|
||||
) -> bool {
|
||||
status_code == 200
|
||||
&& plan
|
||||
.provider_api_format
|
||||
.eq_ignore_ascii_case("openai:image")
|
||||
}
|
||||
|
||||
pub(crate) async fn should_retry_next_local_candidate_sync(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
@@ -48,6 +58,10 @@ pub(crate) async fn analyze_local_candidate_failover_sync(
|
||||
return LocalFailoverAnalysis::use_default();
|
||||
}
|
||||
|
||||
if openai_image_success_disables_local_success_failover(plan, result.status_code) {
|
||||
return LocalFailoverAnalysis::use_default();
|
||||
}
|
||||
|
||||
resolve_local_failover_analysis_for_attempt(
|
||||
state,
|
||||
plan,
|
||||
@@ -218,6 +232,10 @@ pub(crate) async fn resolve_local_candidate_failover_analysis_stream(
|
||||
status_code: u16,
|
||||
response_text: Option<&str>,
|
||||
) -> LocalFailoverAnalysis {
|
||||
if openai_image_success_disables_local_success_failover(plan, status_code) {
|
||||
return LocalFailoverAnalysis::use_default();
|
||||
}
|
||||
|
||||
resolve_local_failover_analysis_for_attempt(
|
||||
state,
|
||||
plan,
|
||||
@@ -756,6 +774,75 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stream_success_failover_does_not_retry_openai_image_success() {
|
||||
let local_report_context = serde_json::json!({
|
||||
"candidate_index": 0,
|
||||
"retry_index": 0,
|
||||
});
|
||||
let state = build_state_with_provider_config(Some(serde_json::json!({
|
||||
"failover_rules": {
|
||||
"success_failover_patterns": [
|
||||
{"pattern": ".*"}
|
||||
]
|
||||
}
|
||||
})));
|
||||
let mut plan = sample_plan();
|
||||
plan.provider_api_format = "openai:image".to_string();
|
||||
|
||||
assert!(
|
||||
!should_retry_next_local_candidate_stream(
|
||||
&state,
|
||||
&plan,
|
||||
"openai_image_stream",
|
||||
Some(&local_report_context),
|
||||
200,
|
||||
Some("{\"data\":[{\"b64_json\":\"aGVsbG8=\"}]}"),
|
||||
)
|
||||
.await,
|
||||
"successful OpenAI image responses should not be retried by success failover rules"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_success_failover_does_not_retry_openai_image_success() {
|
||||
let local_report_context = serde_json::json!({
|
||||
"candidate_index": 0,
|
||||
"retry_index": 0,
|
||||
});
|
||||
let state = build_state_with_provider_config(Some(serde_json::json!({
|
||||
"failover_rules": {
|
||||
"success_failover_patterns": [
|
||||
{"pattern": ".*"}
|
||||
]
|
||||
}
|
||||
})));
|
||||
let mut plan = sample_plan();
|
||||
plan.provider_api_format = "openai:image".to_string();
|
||||
let result = ExecutionResult {
|
||||
request_id: "req-1".to_string(),
|
||||
candidate_id: None,
|
||||
status_code: 200,
|
||||
headers: Default::default(),
|
||||
body: None,
|
||||
telemetry: None,
|
||||
error: None,
|
||||
};
|
||||
|
||||
assert!(
|
||||
!should_retry_next_local_candidate_sync(
|
||||
&state,
|
||||
&plan,
|
||||
"openai_image_sync",
|
||||
Some(&local_report_context),
|
||||
&result,
|
||||
Some("{\"data\":[{\"b64_json\":\"aGVsbG8=\"}]}")
|
||||
)
|
||||
.await,
|
||||
"successful OpenAI image responses should not be retried by success failover rules"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_local_failover_policy_reads_provider_rules() {
|
||||
let state = build_state_with_provider_config(Some(serde_json::json!({
|
||||
|
||||
@@ -1044,6 +1044,12 @@ fn build_sse_body_stream(
|
||||
}
|
||||
}
|
||||
|
||||
fn stream_chunk_contains_sse_done(chunk: &[u8]) -> bool {
|
||||
std::str::from_utf8(chunk)
|
||||
.ok()
|
||||
.is_some_and(|text| text.lines().any(|line| line.trim() == "data: [DONE]"))
|
||||
}
|
||||
|
||||
async fn next_stream_frame<R>(
|
||||
buffered_frames: &mut VecDeque<StreamFrame>,
|
||||
lines: &mut FramedRead<R, LinesCodec>,
|
||||
@@ -1088,6 +1094,33 @@ fn should_refresh_stream_usage_telemetry(
|
||||
|| (next_elapsed.is_some() && next_elapsed != previous_elapsed)
|
||||
}
|
||||
|
||||
fn build_terminal_stream_telemetry(
|
||||
stream_started_at: Instant,
|
||||
telemetry: Option<&ExecutionTelemetry>,
|
||||
usage_stream_telemetry: Option<&ExecutionTelemetry>,
|
||||
upstream_bytes: u64,
|
||||
) -> ExecutionTelemetry {
|
||||
let current_elapsed_ms = stream_started_at
|
||||
.elapsed()
|
||||
.as_millis()
|
||||
.min(u128::from(u64::MAX)) as u64;
|
||||
let ttfb_ms = telemetry
|
||||
.and_then(|telemetry| telemetry.ttfb_ms)
|
||||
.or_else(|| usage_stream_telemetry.and_then(|telemetry| telemetry.ttfb_ms));
|
||||
let prior_elapsed_ms = telemetry
|
||||
.and_then(|telemetry| telemetry.elapsed_ms)
|
||||
.or_else(|| usage_stream_telemetry.and_then(|telemetry| telemetry.elapsed_ms))
|
||||
.unwrap_or(0);
|
||||
let elapsed_ms = current_elapsed_ms
|
||||
.max(prior_elapsed_ms)
|
||||
.max(ttfb_ms.unwrap_or(0));
|
||||
ExecutionTelemetry {
|
||||
ttfb_ms,
|
||||
elapsed_ms: Some(elapsed_ms),
|
||||
upstream_bytes: Some(upstream_bytes),
|
||||
}
|
||||
}
|
||||
|
||||
fn should_skip_direct_finalize_prefetch(
|
||||
direct_stream_finalize_kind: Option<&str>,
|
||||
content_type: Option<&str>,
|
||||
@@ -2052,6 +2085,8 @@ async fn execute_stream_from_frame_stream(
|
||||
max_stream_body_buffer_bytes,
|
||||
&mut client_body_truncated,
|
||||
);
|
||||
let mut client_visible_stream_completed =
|
||||
stream_chunk_contains_sse_done(&prefetched_body_for_report);
|
||||
let mut usage_stream_telemetry: Option<ExecutionTelemetry> = initial_telemetry.clone();
|
||||
let mut telemetry: Option<ExecutionTelemetry> = initial_telemetry;
|
||||
let reached_eof = initial_reached_eof;
|
||||
@@ -2478,6 +2513,8 @@ async fn execute_stream_from_frame_stream(
|
||||
);
|
||||
let rewritten_chunk_len =
|
||||
u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX);
|
||||
let chunk_completed_stream =
|
||||
stream_chunk_contains_sse_done(&rewritten_chunk);
|
||||
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
|
||||
warn!(
|
||||
event_name = "stream_execution_downstream_disconnected",
|
||||
@@ -2490,6 +2527,7 @@ async fn execute_stream_from_frame_stream(
|
||||
downstream_dropped = true;
|
||||
break;
|
||||
} else {
|
||||
client_visible_stream_completed |= chunk_completed_stream;
|
||||
client_stream_bytes.fetch_add(rewritten_chunk_len, Ordering::Relaxed);
|
||||
last_client_chunk_elapsed_ms.store(
|
||||
stream_started_at_for_report
|
||||
@@ -2604,6 +2642,8 @@ async fn execute_stream_from_frame_stream(
|
||||
);
|
||||
let rewritten_chunk_len =
|
||||
u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX);
|
||||
let chunk_completed_stream =
|
||||
stream_chunk_contains_sse_done(&rewritten_chunk);
|
||||
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
|
||||
warn!(
|
||||
event_name = "stream_execution_downstream_flush_disconnected",
|
||||
@@ -2615,6 +2655,7 @@ async fn execute_stream_from_frame_stream(
|
||||
);
|
||||
downstream_dropped = true;
|
||||
} else {
|
||||
client_visible_stream_completed |= chunk_completed_stream;
|
||||
client_stream_bytes
|
||||
.fetch_add(rewritten_chunk_len, Ordering::Relaxed);
|
||||
last_client_chunk_elapsed_ms.store(
|
||||
@@ -2661,6 +2702,8 @@ async fn execute_stream_from_frame_stream(
|
||||
);
|
||||
let flushed_chunk_len =
|
||||
u64::try_from(flushed_chunk.len()).unwrap_or(u64::MAX);
|
||||
let chunk_completed_stream =
|
||||
stream_chunk_contains_sse_done(&flushed_chunk);
|
||||
if tx.send(Ok(Bytes::from(flushed_chunk))).await.is_err() {
|
||||
warn!(
|
||||
event_name = "stream_execution_downstream_rewrite_flush_disconnected",
|
||||
@@ -2672,6 +2715,7 @@ async fn execute_stream_from_frame_stream(
|
||||
);
|
||||
downstream_dropped = true;
|
||||
} else {
|
||||
client_visible_stream_completed |= chunk_completed_stream;
|
||||
client_stream_bytes.fetch_add(flushed_chunk_len, Ordering::Relaxed);
|
||||
last_client_chunk_elapsed_ms.store(
|
||||
stream_started_at_for_report
|
||||
@@ -2781,6 +2825,18 @@ async fn execute_stream_from_frame_stream(
|
||||
),
|
||||
);
|
||||
|
||||
if downstream_dropped && client_visible_stream_completed && terminal_failure.is_none() {
|
||||
debug!(
|
||||
event_name = "execution_runtime_stream_downstream_closed_after_done",
|
||||
log_type = "debug",
|
||||
trace_id = %trace_id_owned,
|
||||
request_id = %request_id_for_report_log,
|
||||
candidate_id = ?candidate_id_for_report.as_deref(),
|
||||
"gateway treats downstream close after client-visible SSE DONE as completed"
|
||||
);
|
||||
downstream_dropped = false;
|
||||
}
|
||||
|
||||
if downstream_dropped {
|
||||
debug!(
|
||||
event_name = "execution_runtime_stream_report_skipped",
|
||||
@@ -2791,6 +2847,12 @@ async fn execute_stream_from_frame_stream(
|
||||
trace_id = %trace_id_owned,
|
||||
"gateway skipped stream report because downstream disconnected before completion"
|
||||
);
|
||||
let terminal_telemetry = Some(build_terminal_stream_telemetry(
|
||||
stream_started_at_for_report,
|
||||
telemetry.as_ref(),
|
||||
usage_stream_telemetry.as_ref(),
|
||||
provider_stream_bytes.load(Ordering::Relaxed),
|
||||
));
|
||||
let usage_payload = build_stream_usage_payload(
|
||||
trace_id_owned,
|
||||
report_kind_owned.unwrap_or_default(),
|
||||
@@ -2802,7 +2864,7 @@ async fn execute_stream_from_frame_stream(
|
||||
&buffered_body,
|
||||
client_body_truncated,
|
||||
stream_terminal_summary,
|
||||
telemetry,
|
||||
terminal_telemetry,
|
||||
);
|
||||
record_stream_terminal_usage(
|
||||
&state_for_report,
|
||||
@@ -2834,6 +2896,12 @@ async fn execute_stream_from_frame_stream(
|
||||
|
||||
if let Some(failure) = terminal_failure {
|
||||
record_manual_proxy_stream_error(&state_for_report, &plan_for_report).await;
|
||||
let terminal_telemetry = Some(build_terminal_stream_telemetry(
|
||||
stream_started_at_for_report,
|
||||
telemetry.as_ref(),
|
||||
usage_stream_telemetry.as_ref(),
|
||||
provider_stream_bytes.load(Ordering::Relaxed),
|
||||
));
|
||||
submit_midstream_stream_failure(
|
||||
&state_for_report,
|
||||
&trace_id_owned,
|
||||
@@ -2841,7 +2909,7 @@ async fn execute_stream_from_frame_stream(
|
||||
direct_stream_finalize_kind_owned.as_deref(),
|
||||
report_context_owned,
|
||||
headers_for_report,
|
||||
telemetry,
|
||||
terminal_telemetry,
|
||||
&provider_buffered_body,
|
||||
candidate_started_unix_secs_for_report,
|
||||
failure,
|
||||
@@ -2851,6 +2919,12 @@ async fn execute_stream_from_frame_stream(
|
||||
}
|
||||
|
||||
let should_submit_report = report_kind_owned.is_some();
|
||||
let terminal_telemetry = Some(build_terminal_stream_telemetry(
|
||||
stream_started_at_for_report,
|
||||
telemetry.as_ref(),
|
||||
usage_stream_telemetry.as_ref(),
|
||||
provider_stream_bytes.load(Ordering::Relaxed),
|
||||
));
|
||||
let usage_payload = build_stream_usage_payload(
|
||||
trace_id_owned.clone(),
|
||||
report_kind_owned.unwrap_or_default(),
|
||||
@@ -2862,7 +2936,7 @@ async fn execute_stream_from_frame_stream(
|
||||
&buffered_body,
|
||||
client_body_truncated,
|
||||
stream_terminal_summary,
|
||||
telemetry,
|
||||
terminal_telemetry,
|
||||
);
|
||||
apply_local_execution_effect(
|
||||
&state_for_report,
|
||||
@@ -3365,6 +3439,7 @@ mod tests {
|
||||
.await
|
||||
.expect("first business chunk should arrive");
|
||||
assert_eq!(first.as_ref(), b"data: {\"id\":\"first\"}\n\n");
|
||||
tokio::time::sleep(Duration::from_millis(30)).await;
|
||||
drop(body_stream);
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(1), frame_stream_dropped.notified())
|
||||
@@ -3392,6 +3467,172 @@ mod tests {
|
||||
candidates[0].error_type.as_deref(),
|
||||
Some("downstream_disconnect")
|
||||
);
|
||||
|
||||
let stored_usage = tokio::time::timeout(Duration::from_secs(1), async {
|
||||
loop {
|
||||
let usage = usage_repository
|
||||
.find_by_request_id("req-client-drop-cancels-upstream")
|
||||
.await
|
||||
.expect("usage should read");
|
||||
if usage
|
||||
.as_ref()
|
||||
.is_some_and(|usage| usage.status == "cancelled")
|
||||
{
|
||||
break usage.expect("cancelled usage should exist");
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("usage should be marked cancelled");
|
||||
assert_eq!(stored_usage.billing_status, "pending");
|
||||
assert_eq!(stored_usage.status_code, Some(499));
|
||||
let first_byte_time_ms = stored_usage
|
||||
.first_byte_time_ms
|
||||
.expect("cancelled stream should retain first byte time");
|
||||
let response_time_ms = stored_usage
|
||||
.response_time_ms
|
||||
.expect("cancelled stream should record terminal duration");
|
||||
assert!(
|
||||
response_time_ms > first_byte_time_ms,
|
||||
"terminal duration should include time after the first byte"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn image_stream_downstream_close_after_done_is_recorded_success() {
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let state = AppState::new()
|
||||
.expect("app state should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
|
||||
Arc::clone(&request_candidate_repository),
|
||||
Arc::clone(&usage_repository),
|
||||
),
|
||||
)
|
||||
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
||||
enabled: true,
|
||||
..UsageRuntimeConfig::default()
|
||||
});
|
||||
let plan = ExecutionPlan {
|
||||
request_id: "req-image-done-close-success".into(),
|
||||
candidate_id: Some("cand-image-done-close-success".into()),
|
||||
provider_name: Some("openai".into()),
|
||||
provider_id: "prov-1".into(),
|
||||
endpoint_id: "ep-1".into(),
|
||||
key_id: "key-1".into(),
|
||||
method: "POST".into(),
|
||||
url: "https://example.com/v1/images/generations".into(),
|
||||
headers: BTreeMap::from([("accept".into(), "text/event-stream".into())]),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({
|
||||
"model": "gpt-image-2",
|
||||
"prompt": "draw a small image",
|
||||
"stream": true
|
||||
})),
|
||||
stream: true,
|
||||
client_api_format: "openai:chat".into(),
|
||||
provider_api_format: "openai:image".into(),
|
||||
model_name: Some("gpt-image-2".into()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
};
|
||||
let frame_stream = stream! {
|
||||
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
|
||||
b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
|
||||
));
|
||||
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
|
||||
b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"output_index\\\":0,\\\"item\\\":{\\\"id\\\":\\\"ig_1\\\",\\\"type\\\":\\\"image_generation_call\\\",\\\"result\\\":\\\"aGVsbG8=\\\"}}\\n\\nevent: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_1\\\",\\\"model\\\":\\\"gpt-image-2\\\",\\\"status\\\":\\\"completed\\\",\\\"usage\\\":null}}\\n\\n\"}}\n",
|
||||
));
|
||||
std::future::pending::<()>().await;
|
||||
}
|
||||
.boxed();
|
||||
|
||||
let response = execute_stream_from_frame_stream(
|
||||
&state,
|
||||
plan,
|
||||
"trace-image-done-close-success",
|
||||
&test_decision(),
|
||||
"openai_chat_stream",
|
||||
Some("openai_chat_stream_success".to_string()),
|
||||
Some(json!({
|
||||
"request_id": "req-image-done-close-success",
|
||||
"candidate_id": "cand-image-done-close-success",
|
||||
"candidate_index": 0,
|
||||
"retry_index": 0,
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:chat",
|
||||
"image_request": {
|
||||
"size": "1024x1024",
|
||||
"quality": "medium"
|
||||
}
|
||||
})),
|
||||
crate::clock::current_unix_ms(),
|
||||
Instant::now(),
|
||||
frame_stream,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("execution should succeed")
|
||||
.expect("execution should return a client response");
|
||||
|
||||
let mut body_stream = response.into_body().into_data_stream();
|
||||
let mut body = Vec::new();
|
||||
tokio::time::timeout(Duration::from_secs(1), async {
|
||||
while !String::from_utf8_lossy(&body).contains("data: [DONE]") {
|
||||
let chunk = body_stream
|
||||
.next()
|
||||
.await
|
||||
.expect("body should yield until done")
|
||||
.expect("chunk should be ok");
|
||||
body.extend_from_slice(&chunk);
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("final DONE should arrive");
|
||||
drop(body_stream);
|
||||
|
||||
let candidates = tokio::time::timeout(Duration::from_secs(1), async {
|
||||
loop {
|
||||
let candidates = request_candidate_repository
|
||||
.list_by_request_id("req-image-done-close-success")
|
||||
.await
|
||||
.expect("request candidates should read");
|
||||
if candidates
|
||||
.first()
|
||||
.is_some_and(|candidate| candidate.status == RequestCandidateStatus::Success)
|
||||
{
|
||||
break candidates;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("candidate should be marked success");
|
||||
assert_eq!(candidates[0].status_code, Some(200));
|
||||
|
||||
let stored_usage = tokio::time::timeout(Duration::from_secs(1), async {
|
||||
loop {
|
||||
let usage = usage_repository
|
||||
.find_by_request_id("req-image-done-close-success")
|
||||
.await
|
||||
.expect("usage should read");
|
||||
if usage
|
||||
.as_ref()
|
||||
.is_some_and(|usage| usage.status == "completed")
|
||||
{
|
||||
break usage.expect("completed usage should exist");
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("usage should be marked completed");
|
||||
assert_eq!(stored_usage.status_code, Some(200));
|
||||
assert!(stored_usage.total_tokens > 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -107,6 +107,13 @@ pub(crate) async fn build_admin_global_model_routing_payload(
|
||||
let Some(provider) = providers.get(&model.provider_id) else {
|
||||
continue;
|
||||
};
|
||||
let provider_model_mapping_names =
|
||||
provider_model_mapping_names_for_routing(model.provider_model_mappings.as_ref());
|
||||
let key_match_model_names = key_match_model_names_for_routing(
|
||||
&global_model.name,
|
||||
&model.provider_model_name,
|
||||
&provider_model_mapping_names,
|
||||
);
|
||||
let mut endpoint_payloads = Vec::new();
|
||||
let mut active_endpoints = 0usize;
|
||||
for endpoint in endpoints_by_provider
|
||||
@@ -132,7 +139,7 @@ pub(crate) async fn build_admin_global_model_routing_payload(
|
||||
.filter(|key| {
|
||||
key_allowed_models_match_global_model_for_routing(
|
||||
key.allowed_models.as_ref(),
|
||||
&global_model.name,
|
||||
&key_match_model_names,
|
||||
&global_model_mappings,
|
||||
)
|
||||
})
|
||||
@@ -313,7 +320,7 @@ pub(crate) async fn build_admin_global_model_routing_payload(
|
||||
|
||||
fn key_allowed_models_match_global_model_for_routing(
|
||||
raw_allowed_models: Option<&serde_json::Value>,
|
||||
global_model_name: &str,
|
||||
model_names: &[String],
|
||||
global_model_mappings: &[String],
|
||||
) -> bool {
|
||||
// 兼容 Python 预览逻辑:None/[] 视为“不限制”,在链路预览中保留该 Key。
|
||||
@@ -322,14 +329,16 @@ fn key_allowed_models_match_global_model_for_routing(
|
||||
return true;
|
||||
}
|
||||
|
||||
if allowed_models
|
||||
.iter()
|
||||
.any(|value| value == global_model_name)
|
||||
{
|
||||
return true;
|
||||
}
|
||||
|
||||
for allowed_model in &allowed_models {
|
||||
for allowed_model in allowed_models.iter().map(String::as_str).map(str::trim) {
|
||||
if allowed_model.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if model_names
|
||||
.iter()
|
||||
.any(|model_name| model_name.eq_ignore_ascii_case(allowed_model))
|
||||
{
|
||||
return true;
|
||||
}
|
||||
for pattern in global_model_mappings {
|
||||
if matches_model_mapping(pattern, allowed_model) {
|
||||
return true;
|
||||
@@ -340,6 +349,54 @@ fn key_allowed_models_match_global_model_for_routing(
|
||||
false
|
||||
}
|
||||
|
||||
fn provider_model_mapping_names_for_routing(
|
||||
raw_mappings: Option<&serde_json::Value>,
|
||||
) -> Vec<String> {
|
||||
raw_mappings
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.map(|items| {
|
||||
items
|
||||
.iter()
|
||||
.filter_map(|item| {
|
||||
item.as_str()
|
||||
.or_else(|| item.get("name").and_then(serde_json::Value::as_str))
|
||||
})
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn key_match_model_names_for_routing(
|
||||
global_model_name: &str,
|
||||
provider_model_name: &str,
|
||||
provider_model_mapping_names: &[String],
|
||||
) -> Vec<String> {
|
||||
let mut names = Vec::new();
|
||||
push_unique_model_name(&mut names, global_model_name);
|
||||
push_unique_model_name(&mut names, provider_model_name);
|
||||
for mapping_name in provider_model_mapping_names {
|
||||
push_unique_model_name(&mut names, mapping_name);
|
||||
}
|
||||
names
|
||||
}
|
||||
|
||||
fn push_unique_model_name(names: &mut Vec<String>, value: &str) {
|
||||
let value = value.trim();
|
||||
if value.is_empty() {
|
||||
return;
|
||||
}
|
||||
if names
|
||||
.iter()
|
||||
.any(|existing| existing.eq_ignore_ascii_case(value))
|
||||
{
|
||||
return;
|
||||
}
|
||||
names.push(value.to_string());
|
||||
}
|
||||
|
||||
pub(crate) async fn build_admin_assign_global_model_to_providers_payload(
|
||||
state: &AdminAppState<'_>,
|
||||
global_model_id: &str,
|
||||
|
||||
@@ -13,6 +13,7 @@ use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadReposi
|
||||
use aether_data_contracts::repository::candidate_selection::{
|
||||
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
|
||||
};
|
||||
use aether_data_contracts::repository::candidates::RequestCandidateReadRepository;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
@@ -741,3 +742,451 @@ async fn gateway_bridges_codex_image_sync_json_to_streaming_image_sse() {
|
||||
execution_runtime_handle.abort();
|
||||
refresh_handle.abort();
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenImageBridgeExecutionPlan {
|
||||
trace_id: String,
|
||||
client_api_format: String,
|
||||
provider_api_format: String,
|
||||
url: String,
|
||||
plan_stream: bool,
|
||||
auth_header: String,
|
||||
chatgpt_web_marker: String,
|
||||
body_json: serde_json::Value,
|
||||
}
|
||||
|
||||
fn image_bridge_hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn image_bridge_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
Some(serde_json::json!([
|
||||
"openai:chat",
|
||||
"openai:responses",
|
||||
"openai:image"
|
||||
])),
|
||||
Some(serde_json::json!(["gpt-image-2"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800_i64),
|
||||
None,
|
||||
Some(serde_json::json!([
|
||||
"openai:chat",
|
||||
"openai:responses",
|
||||
"openai:image"
|
||||
])),
|
||||
Some(serde_json::json!(["gpt-image-2"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn image_bridge_candidate_row(
|
||||
prefix: &str,
|
||||
provider_name: &str,
|
||||
provider_type: &str,
|
||||
) -> StoredMinimalCandidateSelectionRow {
|
||||
let key_auth_type = if provider_type == "chatgpt_web" {
|
||||
"bearer"
|
||||
} else {
|
||||
"api_key"
|
||||
};
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: format!("provider-{prefix}"),
|
||||
provider_name: provider_name.to_string(),
|
||||
provider_type: provider_type.to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: format!("endpoint-{prefix}"),
|
||||
endpoint_api_format: "openai:image".to_string(),
|
||||
endpoint_api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("image".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: format!("key-{prefix}"),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: key_auth_type.to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["openai:image".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"openai:image": 1})),
|
||||
model_id: format!("model-{prefix}"),
|
||||
global_model_id: format!("global-model-{prefix}"),
|
||||
global_model_name: "gpt-image-2".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(false),
|
||||
model_provider_model_name: "gpt-image-2".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-image-2".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:image".to_string()]),
|
||||
endpoint_ids: None,
|
||||
}]),
|
||||
model_supports_streaming: Some(false),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn image_bridge_provider_catalog_provider(
|
||||
prefix: &str,
|
||||
provider_name: &str,
|
||||
provider_type: &str,
|
||||
base_url: &str,
|
||||
) -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
format!("provider-{prefix}"),
|
||||
provider_name.to_string(),
|
||||
Some(base_url.to_string()),
|
||||
provider_type.to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn image_bridge_provider_catalog_endpoint(
|
||||
prefix: &str,
|
||||
base_url: &str,
|
||||
) -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
format!("endpoint-{prefix}"),
|
||||
format!("provider-{prefix}"),
|
||||
"openai:image".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("image".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
base_url.to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn image_bridge_provider_catalog_key(
|
||||
prefix: &str,
|
||||
provider_type: &str,
|
||||
) -> StoredProviderCatalogKey {
|
||||
let auth_type = if provider_type == "chatgpt_web" {
|
||||
"bearer"
|
||||
} else {
|
||||
"api_key"
|
||||
};
|
||||
StoredProviderCatalogKey::new(
|
||||
format!("key-{prefix}"),
|
||||
format!("provider-{prefix}"),
|
||||
"prod".to_string(),
|
||||
auth_type.to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["openai:image"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-image-bridge")
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"openai:image": 1})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
async fn start_image_bridge_gateway(
|
||||
prefix: &str,
|
||||
provider_name: &str,
|
||||
provider_type: &str,
|
||||
base_url: &str,
|
||||
execution_runtime_url: String,
|
||||
) -> (
|
||||
String,
|
||||
tokio::task::JoinHandle<()>,
|
||||
String,
|
||||
Arc<InMemoryRequestCandidateRepository>,
|
||||
) {
|
||||
let client_api_key = format!("sk-client-{prefix}");
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(image_bridge_hash_api_key(&client_api_key)),
|
||||
image_bridge_auth_snapshot(&format!("api-key-{prefix}"), &format!("user-{prefix}")),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
image_bridge_candidate_row(prefix, provider_name, provider_type),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![image_bridge_provider_catalog_provider(
|
||||
prefix,
|
||||
provider_name,
|
||||
provider_type,
|
||||
base_url,
|
||||
)],
|
||||
vec![image_bridge_provider_catalog_endpoint(prefix, base_url)],
|
||||
vec![image_bridge_provider_catalog_key(prefix, provider_type)],
|
||||
));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
(
|
||||
gateway_url,
|
||||
gateway_handle,
|
||||
client_api_key,
|
||||
request_candidate_repository,
|
||||
)
|
||||
}
|
||||
|
||||
fn capture_image_bridge_execution_plan(
|
||||
parts: http::request::Parts,
|
||||
payload: serde_json::Value,
|
||||
) -> SeenImageBridgeExecutionPlan {
|
||||
SeenImageBridgeExecutionPlan {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
client_api_format: payload
|
||||
.get("client_api_format")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
provider_api_format: payload
|
||||
.get("provider_api_format")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
plan_stream: payload
|
||||
.get("stream")
|
||||
.and_then(|value| value.as_bool())
|
||||
.unwrap_or(false),
|
||||
auth_header: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("authorization"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
chatgpt_web_marker: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-aether-chatgpt-web-image"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
body_json: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.cloned()
|
||||
.unwrap_or(serde_json::Value::Null),
|
||||
}
|
||||
}
|
||||
|
||||
fn image_bridge_execution_runtime(
|
||||
seen_execution_plan: Arc<Mutex<Option<SeenImageBridgeExecutionPlan>>>,
|
||||
) -> Router {
|
||||
Router::new().route(
|
||||
"/v1/execute/stream",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_plan_inner = Arc::clone(&seen_execution_plan);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
|
||||
.expect("execution runtime payload should parse");
|
||||
*seen_execution_plan_inner.lock().expect("mutex should lock") =
|
||||
Some(capture_image_bridge_execution_plan(parts, payload));
|
||||
let frames = concat!(
|
||||
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"output_index\\\":0,\\\"item\\\":{\\\"id\\\":\\\"ig_bridge_123\\\",\\\"type\\\":\\\"image_generation_call\\\",\\\"result\\\":\\\"aGVsbG8=\\\",\\\"output_format\\\":\\\"png\\\"}}\\n\\n\"}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_bridge_123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-image-2\\\",\\\"status\\\":\\\"completed\\\",\\\"output\\\":[]}}\\n\\n\"}}\n",
|
||||
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
|
||||
);
|
||||
let mut response = http::Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from(frames))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
http::HeaderValue::from_static("application/x-ndjson"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_routes_openai_chat_stream_image_intent_to_openai_image_plan_without_streaming_support(
|
||||
) {
|
||||
let seen_execution_plan = Arc::new(Mutex::new(None::<SeenImageBridgeExecutionPlan>));
|
||||
let execution_runtime = image_bridge_execution_runtime(Arc::clone(&seen_execution_plan));
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let (gateway_url, gateway_handle, client_api_key, request_candidate_repository) =
|
||||
start_image_bridge_gateway(
|
||||
"chat-stream-image-bridge",
|
||||
"image-provider",
|
||||
"custom",
|
||||
"https://images.example.com",
|
||||
execution_runtime_url,
|
||||
)
|
||||
.await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(http::header::AUTHORIZATION, format!("Bearer {client_api_key}"))
|
||||
.header(TRACE_ID_HEADER, "trace-chat-stream-image-bridge-123")
|
||||
.body(
|
||||
r#"{"model":"gpt-image-2","messages":[{"role":"user","content":"Draw a city made of glass"}],"stream":true,"size":"1024x1024"}"#,
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let response_text = response.text().await.expect("body should read");
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-chat-stream-image-bridge-123")
|
||||
.await
|
||||
.expect("request candidates should read");
|
||||
assert_eq!(
|
||||
status,
|
||||
StatusCode::OK,
|
||||
"{response_text}\n{stored_candidates:#?}"
|
||||
);
|
||||
assert!(response_text.contains("\"object\":\"chat.completion.chunk\""));
|
||||
assert!(response_text.contains(""));
|
||||
assert!(response_text.contains("data: [DONE]"));
|
||||
assert!(!response_text.contains("image_generation.completed"));
|
||||
|
||||
let seen_plan = seen_execution_plan
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution plan should be captured");
|
||||
assert_eq!(seen_plan.trace_id, "trace-chat-stream-image-bridge-123");
|
||||
assert_eq!(seen_plan.client_api_format, "openai:chat");
|
||||
assert_eq!(seen_plan.provider_api_format, "openai:image");
|
||||
assert_eq!(seen_plan.url, "https://images.example.com/v1/responses");
|
||||
assert!(seen_plan.plan_stream);
|
||||
assert_eq!(seen_plan.auth_header, "Bearer sk-upstream-image-bridge");
|
||||
assert_eq!(seen_plan.chatgpt_web_marker, "");
|
||||
assert_eq!(seen_plan.body_json["model"], "gpt-image-2");
|
||||
assert_eq!(seen_plan.body_json["stream"], true);
|
||||
assert_eq!(
|
||||
seen_plan.body_json["input"][0]["content"],
|
||||
"Draw a city made of glass"
|
||||
);
|
||||
assert_eq!(seen_plan.body_json["tools"][0]["type"], "image_generation");
|
||||
assert_eq!(seen_plan.body_json["tools"][0]["size"], "1024x1024");
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_routes_openai_responses_stream_image_intent_to_openai_image_plan_without_streaming_support(
|
||||
) {
|
||||
let seen_execution_plan = Arc::new(Mutex::new(None::<SeenImageBridgeExecutionPlan>));
|
||||
let execution_runtime = image_bridge_execution_runtime(Arc::clone(&seen_execution_plan));
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let (gateway_url, gateway_handle, client_api_key, _request_candidate_repository) =
|
||||
start_image_bridge_gateway(
|
||||
"responses-stream-image-bridge",
|
||||
"image-provider",
|
||||
"custom",
|
||||
"https://images.example.com",
|
||||
execution_runtime_url,
|
||||
)
|
||||
.await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/responses"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(http::header::AUTHORIZATION, format!("Bearer {client_api_key}"))
|
||||
.header(TRACE_ID_HEADER, "trace-responses-stream-image-bridge-123")
|
||||
.body(
|
||||
r#"{"model":"gpt-image-2","input":"Draw a mountain observatory","tools":[{"type":"image_generation","size":"1024x1024"}],"stream":true}"#,
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let response_text = response.text().await.expect("body should read");
|
||||
assert_eq!(status, StatusCode::OK, "{response_text}");
|
||||
assert!(response_text.contains("response.output_item.done"));
|
||||
assert!(response_text.contains("image_generation_call"));
|
||||
|
||||
let seen_plan = seen_execution_plan
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution plan should be captured");
|
||||
assert_eq!(
|
||||
seen_plan.trace_id,
|
||||
"trace-responses-stream-image-bridge-123"
|
||||
);
|
||||
assert_eq!(seen_plan.client_api_format, "openai:responses");
|
||||
assert_eq!(seen_plan.provider_api_format, "openai:image");
|
||||
assert_eq!(seen_plan.url, "https://images.example.com/v1/responses");
|
||||
assert!(seen_plan.plan_stream);
|
||||
assert_eq!(seen_plan.auth_header, "Bearer sk-upstream-image-bridge");
|
||||
assert_eq!(seen_plan.body_json["stream"], true);
|
||||
assert_eq!(seen_plan.body_json["input"], "Draw a mountain observatory");
|
||||
assert_eq!(seen_plan.body_json["tools"][0]["type"], "image_generation");
|
||||
assert_eq!(seen_plan.body_json["tools"][0]["size"], "1024x1024");
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
@@ -148,6 +148,120 @@ async fn gateway_handles_admin_provider_endpoints_locally_with_trusted_admin_pri
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_counts_fixed_provider_oauth_keys_for_inherited_endpoint_formats() {
|
||||
let mut codex_provider = sample_provider("provider-codex", "codex", 10);
|
||||
codex_provider.provider_type = "codex".to_string();
|
||||
let mut chatgpt_web_provider = sample_provider("provider-chatgpt-web", "chatgpt_web", 20);
|
||||
chatgpt_web_provider.provider_type = "chatgpt_web".to_string();
|
||||
|
||||
let mut codex_key = sample_key(
|
||||
"key-codex-oauth",
|
||||
"provider-codex",
|
||||
"openai:responses:compact",
|
||||
"oauth-token",
|
||||
);
|
||||
codex_key.auth_type = "oauth".to_string();
|
||||
codex_key.api_formats = Some(json!(["legacy:mismatch"]));
|
||||
|
||||
let mut chatgpt_web_key = sample_key(
|
||||
"key-chatgpt-web-oauth",
|
||||
"provider-chatgpt-web",
|
||||
"openai:image",
|
||||
"oauth-token",
|
||||
);
|
||||
chatgpt_web_key.auth_type = "oauth".to_string();
|
||||
chatgpt_web_key.api_formats = Some(json!(["legacy:mismatch"]));
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![codex_provider, chatgpt_web_provider],
|
||||
vec![
|
||||
sample_endpoint(
|
||||
"endpoint-codex-compact",
|
||||
"provider-codex",
|
||||
"openai:responses:compact",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
),
|
||||
sample_endpoint(
|
||||
"endpoint-codex-image",
|
||||
"provider-codex",
|
||||
"openai:image",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
),
|
||||
sample_endpoint(
|
||||
"endpoint-chatgpt-web-image",
|
||||
"provider-chatgpt-web",
|
||||
"openai:image",
|
||||
"https://chatgpt.com",
|
||||
),
|
||||
],
|
||||
vec![codex_key, chatgpt_web_key],
|
||||
));
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let codex_response = client
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-codex/endpoints?skip=0&limit=50"
|
||||
))
|
||||
.header(GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(codex_response.status(), StatusCode::OK);
|
||||
let codex_payload: serde_json::Value = codex_response.json().await.expect("json should parse");
|
||||
let codex_items = codex_payload
|
||||
.as_array()
|
||||
.expect("payload should be an array");
|
||||
for api_format in ["openai:responses:compact", "openai:image"] {
|
||||
let endpoint = codex_items
|
||||
.iter()
|
||||
.find(|item| item["api_format"] == api_format)
|
||||
.expect("endpoint should exist");
|
||||
assert_eq!(endpoint["total_keys"], 1);
|
||||
assert_eq!(endpoint["active_keys"], 1);
|
||||
}
|
||||
|
||||
let chatgpt_web_response = client
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-chatgpt-web/endpoints?skip=0&limit=50"
|
||||
))
|
||||
.header(GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(chatgpt_web_response.status(), StatusCode::OK);
|
||||
let chatgpt_web_payload: serde_json::Value = chatgpt_web_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json should parse");
|
||||
let chatgpt_web_items = chatgpt_web_payload
|
||||
.as_array()
|
||||
.expect("payload should be an array");
|
||||
let chatgpt_web_image = chatgpt_web_items
|
||||
.iter()
|
||||
.find(|item| item["api_format"] == "openai:image")
|
||||
.expect("image endpoint should exist");
|
||||
assert_eq!(chatgpt_web_image["total_keys"], 1);
|
||||
assert_eq!(chatgpt_web_image["active_keys"], 1);
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_counts_keys_with_null_api_formats_for_each_fixed_provider_endpoint() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
@@ -195,7 +309,7 @@ async fn gateway_counts_keys_with_null_api_formats_for_each_fixed_provider_endpo
|
||||
vec![inherited_key],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (_, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
|
||||
@@ -936,6 +936,142 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_global_model_routing_counts_image_provider_keys_by_provider_model_name() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/models/global/global-gpt-image/routing",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut image_provider = sample_provider("provider-image", "image", 10);
|
||||
image_provider.provider_type = "chatgpt_web".to_string();
|
||||
let grok_provider = sample_provider("provider-grok", "grok2api", 20);
|
||||
|
||||
let mut image_key = sample_key(
|
||||
"key-image-routing",
|
||||
"provider-image",
|
||||
"legacy:mismatch",
|
||||
"sk-image-routing-1234",
|
||||
);
|
||||
image_key.name = "image-account".to_string();
|
||||
image_key.auth_type = "oauth".to_string();
|
||||
image_key.allowed_models = Some(json!(["gpt-image-2"]));
|
||||
|
||||
let mut grok_key = sample_key(
|
||||
"key-grok-routing",
|
||||
"provider-grok",
|
||||
"openai:chat",
|
||||
"sk-grok-routing-5678",
|
||||
);
|
||||
grok_key.name = "all".to_string();
|
||||
grok_key.allowed_models = Some(json!(["gpt-image-2"]));
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![image_provider, grok_provider],
|
||||
vec![
|
||||
sample_endpoint(
|
||||
"endpoint-image",
|
||||
"provider-image",
|
||||
"openai:image",
|
||||
"https://chatgpt.example",
|
||||
),
|
||||
sample_endpoint(
|
||||
"endpoint-grok-chat",
|
||||
"provider-grok",
|
||||
"openai:chat",
|
||||
"https://grok.example",
|
||||
),
|
||||
],
|
||||
vec![image_key, grok_key],
|
||||
));
|
||||
let global_model_repository = Arc::new(
|
||||
InMemoryGlobalModelReadRepository::seed(Vec::new())
|
||||
.with_admin_global_models(vec![sample_admin_global_model(
|
||||
"global-gpt-image",
|
||||
"GPT-Image-2",
|
||||
"GPT-Image-2",
|
||||
)])
|
||||
.with_admin_provider_models(vec![
|
||||
sample_admin_provider_model(
|
||||
"model-image-gpt-image",
|
||||
"provider-image",
|
||||
"global-gpt-image",
|
||||
"gpt-image-2",
|
||||
),
|
||||
sample_admin_provider_model(
|
||||
"model-grok-gpt-image",
|
||||
"provider-grok",
|
||||
"global-gpt-image",
|
||||
"gpt-image-2",
|
||||
),
|
||||
]),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_global_model_repository_for_tests(global_model_repository),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/models/global/global-gpt-image/routing"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["global_model_name"], "GPT-Image-2");
|
||||
assert_eq!(payload["total_providers"], 2);
|
||||
assert_eq!(payload["active_providers"], 2);
|
||||
|
||||
let providers = payload["providers"].as_array().expect("providers array");
|
||||
assert_eq!(providers.len(), 2);
|
||||
|
||||
let image_endpoints = providers[0]["endpoints"]
|
||||
.as_array()
|
||||
.expect("image endpoints array");
|
||||
assert_eq!(providers[0]["id"], "provider-image");
|
||||
assert_eq!(image_endpoints[0]["api_format"], "openai:image");
|
||||
assert_eq!(image_endpoints[0]["total_keys"], 1);
|
||||
assert_eq!(image_endpoints[0]["active_keys"], 1);
|
||||
assert_eq!(image_endpoints[0]["keys"][0]["name"], "image-account");
|
||||
|
||||
let grok_endpoints = providers[1]["endpoints"]
|
||||
.as_array()
|
||||
.expect("grok endpoints array");
|
||||
assert_eq!(providers[1]["id"], "provider-grok");
|
||||
assert_eq!(grok_endpoints[0]["api_format"], "openai:chat");
|
||||
assert_eq!(grok_endpoints[0]["total_keys"], 1);
|
||||
assert_eq!(grok_endpoints[0]["active_keys"], 1);
|
||||
assert_eq!(grok_endpoints[0]["keys"][0]["name"], "all");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_creates_admin_global_model_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
|
||||
@@ -113,10 +113,7 @@ struct ExpectedUsagePricing {
|
||||
|
||||
impl ExpectedUsagePricing {
|
||||
fn total_tokens(self) -> u64 {
|
||||
self.input_tokens
|
||||
.saturating_add(self.output_tokens)
|
||||
.saturating_add(self.cache_creation_tokens)
|
||||
.saturating_add(self.cache_read_tokens)
|
||||
self.input_tokens.saturating_add(self.output_tokens)
|
||||
}
|
||||
|
||||
fn cache_creation_uncategorized_tokens(self) -> u64 {
|
||||
|
||||
@@ -1,9 +1,15 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use aether_contracts::{ExecutionStreamTerminalSummary, StandardizedUsage};
|
||||
use base64::Engine as _;
|
||||
use serde_json::Value;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::contracts::OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND;
|
||||
use crate::formats::openai::responses::codex::CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT;
|
||||
use crate::formats::shared::sse::encode_json_sse;
|
||||
use crate::formats::shared::sse::{encode_done_sse, encode_json_sse};
|
||||
use crate::formats::shared::stream_core::common::{
|
||||
build_openai_chat_chunk, build_openai_chat_finish_chunk, build_openai_chat_usage_chunk,
|
||||
};
|
||||
use crate::formats::shared::AiSurfaceFinalizeError;
|
||||
|
||||
#[derive(Default)]
|
||||
@@ -20,6 +26,38 @@ struct OpenAiImageFrame {
|
||||
b64_json: String,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct OpenAiImageChatStreamState {
|
||||
buffered: Vec<u8>,
|
||||
response_id: Option<String>,
|
||||
model: Option<String>,
|
||||
latest_image: Option<OpenAiImageChatFrame>,
|
||||
emitted_image_count: u64,
|
||||
emitted_image_keys: BTreeSet<String>,
|
||||
started: bool,
|
||||
finished: bool,
|
||||
emitted_failure: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct OpenAiImageChatFrame {
|
||||
b64_json: String,
|
||||
output_format: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct OpenAiImageStreamTerminalState {
|
||||
event_name: Option<String>,
|
||||
data_lines: Vec<String>,
|
||||
response_id: Option<String>,
|
||||
model: Option<String>,
|
||||
image_count: u64,
|
||||
image_keys: BTreeSet<String>,
|
||||
usage: Option<Value>,
|
||||
observed_finish: bool,
|
||||
parser_error: Option<String>,
|
||||
}
|
||||
|
||||
impl OpenAiImageStreamState {
|
||||
pub fn push_chunk(
|
||||
&mut self,
|
||||
@@ -229,6 +267,662 @@ impl OpenAiImageStreamState {
|
||||
}
|
||||
}
|
||||
|
||||
impl OpenAiImageChatStreamState {
|
||||
pub fn push_chunk(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
chunk: &[u8],
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
self.buffered.extend_from_slice(chunk);
|
||||
let mut output = Vec::new();
|
||||
while let Some(block_end) = find_sse_block_end(&self.buffered) {
|
||||
let block = self.buffered.drain(..block_end).collect::<Vec<_>>();
|
||||
output.extend(self.transform_block(report_context, &block)?);
|
||||
drain_sse_separator(&mut self.buffered);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
pub fn finish(&mut self, report_context: &Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let mut output = if self.buffered.is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
let block = std::mem::take(&mut self.buffered);
|
||||
self.transform_block(report_context, &block)?
|
||||
};
|
||||
if !self.finished && !self.emitted_failure && self.latest_image.is_some() {
|
||||
output.extend(self.emit_final(report_context, None)?);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn transform_block(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
block: &[u8],
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let text = std::str::from_utf8(block)
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?;
|
||||
let mut event_name = None::<String>;
|
||||
let mut data_lines = Vec::new();
|
||||
for raw_line in text.lines() {
|
||||
let line = raw_line.trim_end_matches('\r');
|
||||
if let Some(value) = line.strip_prefix("event:") {
|
||||
event_name = Some(value.trim().to_string());
|
||||
} else if let Some(value) = line.strip_prefix("data:") {
|
||||
data_lines.push(value.trim().to_string());
|
||||
}
|
||||
}
|
||||
let data = data_lines.join("\n");
|
||||
if data.is_empty() || data == "[DONE]" {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let event: Value = serde_json::from_str(&data)?;
|
||||
let event_type = event
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.or(event_name.as_deref())
|
||||
.unwrap_or_default();
|
||||
match event_type {
|
||||
"error" | "response.failed" | "image_generation.failed" | "image_edit.failed" => {
|
||||
self.handle_failed(report_context, &event)
|
||||
}
|
||||
"response.image_generation_call.partial_image" => {
|
||||
self.emit_empty_progress_chunk(report_context)
|
||||
}
|
||||
"response.output_item.done" => self.handle_output_item_done(report_context, &event),
|
||||
"response.completed" | "response.done" => self.handle_completed(report_context, &event),
|
||||
"image_generation.completed" | "image_edit.completed" => {
|
||||
self.handle_image_completed(report_context, &event)
|
||||
}
|
||||
_ => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_output_item_done(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
event: &Value,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if self.finished || self.emitted_failure {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let Some(item) = event.get("item").and_then(Value::as_object) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
if item.get("type").and_then(Value::as_str) != Some("image_generation_call") {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
if let Some(result) = item.get("result").and_then(Value::as_str).map(str::trim) {
|
||||
if !result.is_empty() {
|
||||
let key = image_chat_output_key(item, result);
|
||||
if self.emitted_image_keys.insert(key) {
|
||||
self.latest_image = Some(OpenAiImageChatFrame {
|
||||
b64_json: result.to_string(),
|
||||
output_format: item
|
||||
.get("output_format")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
});
|
||||
self.emitted_image_count = self.emitted_image_count.saturating_add(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
self.ensure_started(report_context)
|
||||
}
|
||||
|
||||
fn handle_completed(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
event: &Value,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if self.finished || self.emitted_failure {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
if let Some(response) = event.get("response") {
|
||||
self.update_identity_from_response(response);
|
||||
if self.latest_image.is_none() {
|
||||
if let Some(frame) = completed_response_image_chat_frame(response) {
|
||||
self.latest_image = Some(frame);
|
||||
self.emitted_image_count = self.emitted_image_count.saturating_add(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
let usage = event
|
||||
.get("response")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|response| {
|
||||
response
|
||||
.get("tool_usage")
|
||||
.and_then(|value| value.get("image_gen"))
|
||||
.cloned()
|
||||
.or_else(|| response.get("usage").cloned())
|
||||
});
|
||||
self.emit_final(report_context, usage.as_ref())
|
||||
}
|
||||
|
||||
fn handle_image_completed(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
event: &Value,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if self.finished || self.emitted_failure {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
if let Some(result) = event
|
||||
.get("b64_json")
|
||||
.or_else(|| event.get("result"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
self.latest_image = Some(OpenAiImageChatFrame {
|
||||
b64_json: result.to_string(),
|
||||
output_format: event
|
||||
.get("output_format")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
});
|
||||
self.emitted_image_count = self.emitted_image_count.max(1);
|
||||
}
|
||||
self.emit_final(report_context, event.get("usage"))
|
||||
}
|
||||
|
||||
fn handle_failed(
|
||||
&mut self,
|
||||
_report_context: &Value,
|
||||
event: &Value,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if self.emitted_failure {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.emitted_failure = true;
|
||||
self.finished = true;
|
||||
let mut output = encode_json_sse(
|
||||
None,
|
||||
&serde_json::json!({
|
||||
"error": image_failure_error(event),
|
||||
}),
|
||||
)?;
|
||||
output.extend(encode_done_sse());
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn ensure_started(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if self.started {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.emit_empty_progress_chunk(report_context)
|
||||
}
|
||||
|
||||
fn emit_empty_progress_chunk(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
self.started = true;
|
||||
let (response_id, model) = self.identity(report_context);
|
||||
encode_json_sse(
|
||||
None,
|
||||
&build_openai_chat_chunk(&response_id, &model, String::new(), None, None),
|
||||
)
|
||||
}
|
||||
|
||||
fn emit_final(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
usage: Option<&Value>,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if self.finished || self.emitted_failure {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let Some(latest_image) = self.latest_image.clone() else {
|
||||
return self.ensure_started(report_context);
|
||||
};
|
||||
let mut output = self.ensure_started(report_context)?;
|
||||
let (response_id, model) = self.identity(report_context);
|
||||
output.extend(encode_json_sse(
|
||||
None,
|
||||
&build_openai_chat_chunk(
|
||||
&response_id,
|
||||
&model,
|
||||
image_chat_markdown(&latest_image),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
)?);
|
||||
output.extend(encode_json_sse(
|
||||
None,
|
||||
&build_openai_chat_finish_chunk(&response_id, &model, Some("stop")),
|
||||
)?);
|
||||
if let Some((input_tokens, output_tokens, total_tokens, reasoning_tokens)) =
|
||||
openai_image_chat_usage_counts(usage)
|
||||
{
|
||||
output.extend(encode_json_sse(
|
||||
None,
|
||||
&build_openai_chat_usage_chunk(
|
||||
&response_id,
|
||||
&model,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
reasoning_tokens,
|
||||
),
|
||||
)?);
|
||||
}
|
||||
output.extend(encode_done_sse());
|
||||
self.finished = true;
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn update_identity_from_response(&mut self, response: &Value) {
|
||||
if let Some(id) = response
|
||||
.get("id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
self.response_id = Some(id.replace("resp", "chatcmpl"));
|
||||
}
|
||||
if let Some(model) = response
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
self.model = Some(model.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
fn identity(&self, report_context: &Value) -> (String, String) {
|
||||
let response_id = self.response_id.clone().unwrap_or_else(|| {
|
||||
report_context
|
||||
.get("request_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(|value| format!("chatcmpl-image-{value}"))
|
||||
.unwrap_or_else(|| "chatcmpl-image".to_string())
|
||||
});
|
||||
let model = self
|
||||
.model
|
||||
.clone()
|
||||
.or_else(|| {
|
||||
report_context
|
||||
.get("mapped_model")
|
||||
.or_else(|| report_context.get("model"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
.unwrap_or_else(|| "gpt-image".to_string());
|
||||
(response_id, model)
|
||||
}
|
||||
}
|
||||
|
||||
impl OpenAiImageStreamTerminalState {
|
||||
pub fn push_line(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
line: Vec<u8>,
|
||||
) -> Result<Option<ExecutionStreamTerminalSummary>, AiSurfaceFinalizeError> {
|
||||
let text = std::str::from_utf8(&line)
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?;
|
||||
let trimmed = text.trim_matches('\r').trim_matches('\n');
|
||||
if trimmed.is_empty() {
|
||||
self.flush_event(report_context)?;
|
||||
return Ok(self.latest_summary(report_context));
|
||||
}
|
||||
if let Some(value) = trimmed.strip_prefix("event:") {
|
||||
self.event_name = Some(value.trim().to_string());
|
||||
} else if let Some(value) = trimmed.strip_prefix("data:") {
|
||||
self.data_lines.push(value.trim().to_string());
|
||||
}
|
||||
Ok(self.latest_summary(report_context))
|
||||
}
|
||||
|
||||
pub fn finish(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
) -> Result<Option<ExecutionStreamTerminalSummary>, AiSurfaceFinalizeError> {
|
||||
self.flush_event(report_context)?;
|
||||
if self.image_count > 0 && !self.observed_finish {
|
||||
self.observed_finish = true;
|
||||
}
|
||||
Ok(self.latest_summary(report_context))
|
||||
}
|
||||
|
||||
fn flush_event(&mut self, report_context: &Value) -> Result<(), AiSurfaceFinalizeError> {
|
||||
if self.data_lines.is_empty() {
|
||||
self.event_name = None;
|
||||
return Ok(());
|
||||
}
|
||||
let data = std::mem::take(&mut self.data_lines).join("\n");
|
||||
let event_name = self.event_name.take();
|
||||
if data.is_empty() || data == "[DONE]" {
|
||||
return Ok(());
|
||||
}
|
||||
let event = match serde_json::from_str::<Value>(&data) {
|
||||
Ok(event) => event,
|
||||
Err(err) => {
|
||||
self.parser_error.get_or_insert_with(|| err.to_string());
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
let event_type = event
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.or(event_name.as_deref())
|
||||
.unwrap_or_default();
|
||||
match event_type {
|
||||
"response.output_item.done" => self.observe_output_item_done(&event),
|
||||
"response.completed" | "response.done" => self.observe_completed(&event),
|
||||
"image_generation.completed" | "image_edit.completed" => {
|
||||
self.observe_image_completed(&event)
|
||||
}
|
||||
"error" | "response.failed" | "image_generation.failed" | "image_edit.failed" => {
|
||||
self.parser_error
|
||||
.get_or_insert_with(|| image_failure_error(&event).to_string());
|
||||
self.observed_finish = true;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
if self.model.is_none() {
|
||||
self.model = image_bridge_model(Some(report_context));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn observe_output_item_done(&mut self, event: &Value) {
|
||||
let Some(item) = event.get("item").and_then(Value::as_object) else {
|
||||
return;
|
||||
};
|
||||
if item.get("type").and_then(Value::as_str) != Some("image_generation_call") {
|
||||
return;
|
||||
}
|
||||
let Some(result) = item
|
||||
.get("result")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let key = image_chat_output_key(item, result);
|
||||
if self.image_keys.insert(key) {
|
||||
self.image_count = self.image_count.saturating_add(1);
|
||||
}
|
||||
}
|
||||
|
||||
fn observe_completed(&mut self, event: &Value) {
|
||||
self.observed_finish = true;
|
||||
let Some(response) = event.get("response") else {
|
||||
return;
|
||||
};
|
||||
self.update_identity_from_response(response);
|
||||
if self.image_count == 0 {
|
||||
self.image_count = completed_response_image_count(response);
|
||||
}
|
||||
self.usage = response
|
||||
.get("tool_usage")
|
||||
.and_then(|value| value.get("image_gen"))
|
||||
.cloned()
|
||||
.or_else(|| response.get("usage").cloned())
|
||||
.or_else(|| self.usage.clone());
|
||||
}
|
||||
|
||||
fn observe_image_completed(&mut self, event: &Value) {
|
||||
self.observed_finish = true;
|
||||
if self.image_count == 0 {
|
||||
if event
|
||||
.get("b64_json")
|
||||
.or_else(|| event.get("result"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
{
|
||||
self.image_count = 1;
|
||||
}
|
||||
}
|
||||
self.usage = event.get("usage").cloned().or_else(|| self.usage.clone());
|
||||
}
|
||||
|
||||
fn update_identity_from_response(&mut self, response: &Value) {
|
||||
if let Some(id) = response
|
||||
.get("id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
self.response_id = Some(id.to_string());
|
||||
}
|
||||
if let Some(model) = response
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
self.model = Some(model.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
fn latest_summary(&self, report_context: &Value) -> Option<ExecutionStreamTerminalSummary> {
|
||||
if self.image_count == 0
|
||||
&& self.usage.is_none()
|
||||
&& self.response_id.is_none()
|
||||
&& self.model.is_none()
|
||||
&& self.parser_error.is_none()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
Some(ExecutionStreamTerminalSummary {
|
||||
standardized_usage: openai_image_stream_standardized_usage(
|
||||
self.usage.as_ref(),
|
||||
Some(report_context),
|
||||
self.image_count,
|
||||
),
|
||||
finish_reason: self.observed_finish.then(|| "stop".to_string()),
|
||||
response_id: self.response_id.clone(),
|
||||
model: self
|
||||
.model
|
||||
.clone()
|
||||
.or_else(|| image_bridge_model(Some(report_context))),
|
||||
observed_finish: self.observed_finish,
|
||||
unknown_event_count: 0,
|
||||
parser_error: self.parser_error.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn completed_response_image_chat_frame(response: &Value) -> Option<OpenAiImageChatFrame> {
|
||||
response
|
||||
.get("output")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter(|item| item.get("type").and_then(Value::as_str) == Some("image_generation_call"))
|
||||
.find_map(|item| {
|
||||
let result = item.get("result").and_then(Value::as_str)?.trim();
|
||||
if result.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(OpenAiImageChatFrame {
|
||||
b64_json: result.to_string(),
|
||||
output_format: item
|
||||
.get("output_format")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn completed_response_image_count(response: &Value) -> u64 {
|
||||
response
|
||||
.get("output")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter(|item| item.get("type").and_then(Value::as_str) == Some("image_generation_call"))
|
||||
.filter(|item| {
|
||||
item.get("result")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
})
|
||||
.count() as u64
|
||||
}
|
||||
|
||||
fn image_chat_output_key(item: &Map<String, Value>, result: &str) -> String {
|
||||
item.get("id")
|
||||
.or_else(|| item.get("call_id"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| result.to_string())
|
||||
}
|
||||
|
||||
fn openai_image_stream_standardized_usage(
|
||||
usage: Option<&Value>,
|
||||
report_context: Option<&Value>,
|
||||
image_count: u64,
|
||||
) -> Option<StandardizedUsage> {
|
||||
let mut standardized_usage = usage
|
||||
.and_then(openai_image_usage_to_standardized_usage)
|
||||
.unwrap_or_else(StandardizedUsage::new);
|
||||
if image_count > 0 {
|
||||
standardized_usage.request_count = i64::try_from(image_count).unwrap_or(i64::MAX);
|
||||
standardized_usage
|
||||
.dimensions
|
||||
.insert("image_count".to_string(), serde_json::json!(image_count));
|
||||
}
|
||||
if let Some(output_format) = image_request_output_format(report_context) {
|
||||
standardized_usage.dimensions.insert(
|
||||
"image_output_format".to_string(),
|
||||
serde_json::json!(output_format),
|
||||
);
|
||||
}
|
||||
if let Some(size) = image_request_size(report_context) {
|
||||
standardized_usage
|
||||
.dimensions
|
||||
.insert("image_size".to_string(), serde_json::json!(size));
|
||||
}
|
||||
if let Some(quality) = image_request_quality(report_context) {
|
||||
standardized_usage
|
||||
.dimensions
|
||||
.insert("image_quality".to_string(), serde_json::json!(quality));
|
||||
}
|
||||
(standardized_usage.signal_score() > 0).then_some(standardized_usage)
|
||||
}
|
||||
|
||||
fn openai_image_usage_to_standardized_usage(value: &Value) -> Option<StandardizedUsage> {
|
||||
let usage = value.as_object()?;
|
||||
let mut input_tokens = usage
|
||||
.get("input_tokens")
|
||||
.or_else(|| usage.get("prompt_tokens"))
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or(0);
|
||||
let output_tokens = usage
|
||||
.get("output_tokens")
|
||||
.or_else(|| usage.get("completion_tokens"))
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or(0);
|
||||
let cache_creation_tokens = usage
|
||||
.get("cache_creation_input_tokens")
|
||||
.and_then(Value::as_i64)
|
||||
.or_else(|| {
|
||||
usage
|
||||
.get("input_tokens_details")
|
||||
.or_else(|| usage.get("prompt_tokens_details"))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|details| details.get("cached_creation_tokens"))
|
||||
.and_then(Value::as_i64)
|
||||
})
|
||||
.unwrap_or(0);
|
||||
let cache_read_tokens = usage
|
||||
.get("cache_read_input_tokens")
|
||||
.and_then(Value::as_i64)
|
||||
.or_else(|| {
|
||||
usage
|
||||
.get("input_tokens_details")
|
||||
.or_else(|| usage.get("prompt_tokens_details"))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|details| details.get("cached_tokens"))
|
||||
.and_then(Value::as_i64)
|
||||
})
|
||||
.unwrap_or(0);
|
||||
let total_tokens = usage.get("total_tokens").and_then(Value::as_i64).unwrap_or(
|
||||
input_tokens
|
||||
.saturating_add(output_tokens)
|
||||
.saturating_add(cache_creation_tokens)
|
||||
.saturating_add(cache_read_tokens),
|
||||
);
|
||||
if input_tokens == 0 && total_tokens > output_tokens {
|
||||
input_tokens = total_tokens.saturating_sub(output_tokens);
|
||||
}
|
||||
let mut standardized_usage = StandardizedUsage::new();
|
||||
standardized_usage.input_tokens = input_tokens;
|
||||
standardized_usage.output_tokens = output_tokens;
|
||||
standardized_usage.cache_creation_tokens = cache_creation_tokens;
|
||||
standardized_usage.cache_read_tokens = cache_read_tokens;
|
||||
standardized_usage
|
||||
.dimensions
|
||||
.insert("total_tokens".to_string(), serde_json::json!(total_tokens));
|
||||
Some(standardized_usage.normalize_cache_creation_breakdown())
|
||||
}
|
||||
|
||||
fn image_chat_markdown(frame: &OpenAiImageChatFrame) -> String {
|
||||
let mime_type = match frame
|
||||
.output_format
|
||||
.as_deref()
|
||||
.unwrap_or("png")
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
.as_str()
|
||||
{
|
||||
"jpg" | "jpeg" => "image/jpeg".to_string(),
|
||||
"webp" => "image/webp".to_string(),
|
||||
"png" => "image/png".to_string(),
|
||||
value if !value.is_empty() => format!("image/{value}"),
|
||||
_ => "image/png".to_string(),
|
||||
};
|
||||
format!(
|
||||
"",
|
||||
frame.b64_json
|
||||
)
|
||||
}
|
||||
|
||||
fn openai_image_chat_usage_counts(usage: Option<&Value>) -> Option<(u64, u64, u64, u64)> {
|
||||
let usage = usage.and_then(Value::as_object)?;
|
||||
let mut input_tokens = usage
|
||||
.get("input_tokens")
|
||||
.or_else(|| usage.get("prompt_tokens"))
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let output_tokens = usage
|
||||
.get("output_tokens")
|
||||
.or_else(|| usage.get("completion_tokens"))
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let total_tokens = usage
|
||||
.get("total_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(input_tokens.saturating_add(output_tokens));
|
||||
if input_tokens == 0 && total_tokens > output_tokens {
|
||||
input_tokens = total_tokens.saturating_sub(output_tokens);
|
||||
}
|
||||
(total_tokens > 0).then_some((input_tokens, output_tokens, total_tokens, 0))
|
||||
}
|
||||
|
||||
fn image_failure_error(event: &Value) -> Value {
|
||||
let mut error = event
|
||||
.get("error")
|
||||
@@ -340,6 +1034,48 @@ fn image_request_operation(report_context: &Value) -> Option<&str> {
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn image_request_output_format(report_context: Option<&Value>) -> Option<String> {
|
||||
report_context
|
||||
.and_then(|value| value.get("image_request"))
|
||||
.and_then(|value| value.get("output_format"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn image_request_size(report_context: Option<&Value>) -> Option<String> {
|
||||
report_context
|
||||
.and_then(|value| value.get("image_request"))
|
||||
.and_then(|value| value.get("size"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn image_request_quality(report_context: Option<&Value>) -> Option<String> {
|
||||
report_context
|
||||
.and_then(|value| value.get("image_request"))
|
||||
.and_then(|value| value.get("quality"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn image_bridge_model(report_context: Option<&Value>) -> Option<String> {
|
||||
report_context.and_then(|context| {
|
||||
context
|
||||
.get("mapped_model")
|
||||
.or_else(|| context.get("model"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
}
|
||||
|
||||
fn find_sse_block_end(buffer: &[u8]) -> Option<usize> {
|
||||
buffer
|
||||
.windows(2)
|
||||
|
||||
@@ -92,8 +92,7 @@ pub fn canonical_usage_from_openai_usage(value: Option<&Value>) -> Option<Canoni
|
||||
let total_tokens = usage.get("total_tokens").and_then(Value::as_u64).unwrap_or(
|
||||
input_tokens
|
||||
.saturating_add(output_tokens)
|
||||
.saturating_add(cache_creation_tokens)
|
||||
.saturating_add(cache_read_tokens),
|
||||
.saturating_add(reasoning_tokens),
|
||||
);
|
||||
if input_tokens == 0 && total_tokens > output_tokens {
|
||||
input_tokens = total_tokens.saturating_sub(output_tokens);
|
||||
@@ -150,8 +149,7 @@ pub fn canonical_usage_from_claude_usage(value: Option<&Value>) -> Option<Canoni
|
||||
output_tokens,
|
||||
total_tokens: input_tokens
|
||||
.saturating_add(output_tokens)
|
||||
.saturating_add(cache_creation_tokens)
|
||||
.saturating_add(cache_read_tokens),
|
||||
.saturating_add(reasoning_tokens),
|
||||
cache_creation_tokens,
|
||||
cache_creation_ephemeral_5m_tokens,
|
||||
cache_creation_ephemeral_1h_tokens,
|
||||
@@ -235,7 +233,7 @@ pub fn canonical_usage_from_gemini_usage(value: Option<&Value>) -> Option<Canoni
|
||||
.unwrap_or(
|
||||
input_tokens
|
||||
.saturating_add(output_tokens)
|
||||
.saturating_add(cache_read_tokens),
|
||||
.saturating_add(reasoning_tokens),
|
||||
);
|
||||
Some(CanonicalUsage {
|
||||
input_tokens,
|
||||
|
||||
@@ -8,6 +8,7 @@ use crate::formats::openai::chat::stream::{
|
||||
OpenAIChatClientEmitter, OpenAIChatProviderState, OpenAIResponsesClientEmitter,
|
||||
OpenAIResponsesProviderState,
|
||||
};
|
||||
use crate::formats::openai::image::stream::OpenAiImageStreamTerminalState;
|
||||
use crate::formats::shared::error_body::{
|
||||
build_core_error_body_for_client_format, LocalCoreSyncErrorKind,
|
||||
};
|
||||
@@ -97,7 +98,7 @@ impl StreamingStandardFormatMatrix {
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct StreamingStandardTerminalObserver {
|
||||
provider: Option<ProviderStreamParser>,
|
||||
provider: Option<TerminalStreamParser>,
|
||||
latest_summary: Option<ExecutionStreamTerminalSummary>,
|
||||
}
|
||||
|
||||
@@ -111,8 +112,17 @@ impl StreamingStandardTerminalObserver {
|
||||
let Some(provider) = self.provider.as_mut() else {
|
||||
return Ok(());
|
||||
};
|
||||
let frames = provider.push_line(report_context, line)?;
|
||||
self.observe_frames(frames);
|
||||
match provider {
|
||||
TerminalStreamParser::Standard(provider) => {
|
||||
let frames = provider.push_line(report_context, line)?;
|
||||
self.observe_frames(frames);
|
||||
}
|
||||
TerminalStreamParser::OpenAIImage(provider) => {
|
||||
if let Some(summary) = provider.push_line(report_context, line)? {
|
||||
self.latest_summary = Some(summary);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -124,8 +134,17 @@ impl StreamingStandardTerminalObserver {
|
||||
let Some(provider) = self.provider.as_mut() else {
|
||||
return Ok(self.latest_summary.clone());
|
||||
};
|
||||
let frames = provider.finish(report_context)?;
|
||||
self.observe_frames(frames);
|
||||
match provider {
|
||||
TerminalStreamParser::Standard(provider) => {
|
||||
let frames = provider.finish(report_context)?;
|
||||
self.observe_frames(frames);
|
||||
}
|
||||
TerminalStreamParser::OpenAIImage(provider) => {
|
||||
if let Some(summary) = provider.finish(report_context)? {
|
||||
self.latest_summary = Some(summary);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(self.latest_summary.clone())
|
||||
}
|
||||
|
||||
@@ -153,7 +172,7 @@ impl StreamingStandardTerminalObserver {
|
||||
return;
|
||||
}
|
||||
let provider_api_format = provider_api_format_for_context(report_context);
|
||||
self.provider = ProviderStreamParser::for_api_format(provider_api_format.as_str());
|
||||
self.provider = TerminalStreamParser::for_api_format(provider_api_format.as_str());
|
||||
}
|
||||
|
||||
fn observe_frames(&mut self, frames: Vec<CanonicalStreamFrame>) {
|
||||
@@ -194,6 +213,23 @@ impl StreamingStandardTerminalObserver {
|
||||
}
|
||||
}
|
||||
|
||||
enum TerminalStreamParser {
|
||||
Standard(ProviderStreamParser),
|
||||
OpenAIImage(OpenAiImageStreamTerminalState),
|
||||
}
|
||||
|
||||
impl TerminalStreamParser {
|
||||
fn for_api_format(provider_api_format: &str) -> Option<Self> {
|
||||
if provider_api_format
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("openai:image")
|
||||
{
|
||||
return Some(Self::OpenAIImage(OpenAiImageStreamTerminalState::default()));
|
||||
}
|
||||
ProviderStreamParser::for_api_format(provider_api_format).map(Self::Standard)
|
||||
}
|
||||
}
|
||||
|
||||
enum ProviderStreamParser {
|
||||
OpenAIChat(OpenAIChatProviderState),
|
||||
OpenAIResponses(OpenAIResponsesProviderState),
|
||||
@@ -936,4 +972,86 @@ mod tests {
|
||||
assert_eq!(summary.unknown_event_count, 1);
|
||||
assert!(!summary.observed_finish);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn terminal_observer_tracks_openai_image_stream_usage() {
|
||||
let mut report_context = report_context("openai:image", "openai:chat");
|
||||
report_context["image_request"] = json!({
|
||||
"size": "1024x1024",
|
||||
"quality": "medium",
|
||||
"output_format": "png",
|
||||
});
|
||||
let mut observer = StreamingStandardTerminalObserver::default();
|
||||
|
||||
observer
|
||||
.push_line(
|
||||
&report_context,
|
||||
data_line(json!({
|
||||
"type": "response.output_item.done",
|
||||
"output_index": 0,
|
||||
"item": {
|
||||
"id": "ig_123",
|
||||
"type": "image_generation_call",
|
||||
"result": "aGVsbG8=",
|
||||
},
|
||||
})),
|
||||
)
|
||||
.expect("image output item should parse");
|
||||
observer
|
||||
.push_line(&report_context, b"\n".to_vec())
|
||||
.expect("image output event should flush");
|
||||
observer
|
||||
.push_line(
|
||||
&report_context,
|
||||
data_line(json!({
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_image_123",
|
||||
"model": "gpt-image-2",
|
||||
"output": [],
|
||||
"tool_usage": {
|
||||
"image_gen": {
|
||||
"input_tokens": 40,
|
||||
"output_tokens": 60,
|
||||
"total_tokens": 100,
|
||||
},
|
||||
},
|
||||
},
|
||||
})),
|
||||
)
|
||||
.expect("image completed should parse");
|
||||
observer
|
||||
.push_line(&report_context, b"\n".to_vec())
|
||||
.expect("image completed event should flush");
|
||||
|
||||
let summary = observer
|
||||
.finish(&report_context)
|
||||
.expect("image summary should finish")
|
||||
.expect("summary should exist");
|
||||
let usage = summary
|
||||
.standardized_usage
|
||||
.expect("standardized usage should exist");
|
||||
|
||||
assert_eq!(summary.response_id.as_deref(), Some("resp_image_123"));
|
||||
assert_eq!(summary.model.as_deref(), Some("gpt-image-2"));
|
||||
assert_eq!(summary.finish_reason.as_deref(), Some("stop"));
|
||||
assert!(summary.observed_finish);
|
||||
assert_eq!(usage.input_tokens, 40);
|
||||
assert_eq!(usage.output_tokens, 60);
|
||||
assert_eq!(usage.request_count, 1);
|
||||
assert_eq!(usage.dimensions.get("image_count"), Some(&json!(1)));
|
||||
assert_eq!(usage.dimensions.get("total_tokens"), Some(&json!(100)));
|
||||
assert_eq!(
|
||||
usage.dimensions.get("image_size"),
|
||||
Some(&json!("1024x1024"))
|
||||
);
|
||||
assert_eq!(
|
||||
usage.dimensions.get("image_output_format"),
|
||||
Some(&json!("png"))
|
||||
);
|
||||
assert_eq!(
|
||||
usage.dimensions.get("image_quality"),
|
||||
Some(&json!("medium"))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::formats::openai::image::stream::OpenAiImageStreamState;
|
||||
use crate::formats::openai::image::stream::{OpenAiImageChatStreamState, OpenAiImageStreamState};
|
||||
use crate::formats::shared::model_directives::model_directive_display_model_from_report_context;
|
||||
use crate::formats::shared::stream_core::StreamingStandardFormatMatrix;
|
||||
use crate::formats::shared::AiSurfaceFinalizeError;
|
||||
@@ -15,6 +15,7 @@ pub enum FinalizeStreamRewriteMode {
|
||||
EnvelopeUnwrap,
|
||||
ModelDirectiveDisplay,
|
||||
OpenAiImage,
|
||||
OpenAiImageToOpenAiChat,
|
||||
Standard,
|
||||
KiroToClaudeCli,
|
||||
KiroToClaudeCliThenStandard,
|
||||
@@ -57,6 +58,14 @@ pub fn resolve_finalize_stream_rewrite_mode(
|
||||
.then_some(FinalizeStreamRewriteMode::KiroToClaudeCliThenStandard);
|
||||
}
|
||||
|
||||
if provider_api_format == "openai:image" && client_api_format == "openai:chat" {
|
||||
return Some(FinalizeStreamRewriteMode::OpenAiImageToOpenAiChat);
|
||||
}
|
||||
|
||||
if provider_api_format == "openai:image" && client_api_format == "openai:image" {
|
||||
return Some(FinalizeStreamRewriteMode::OpenAiImage);
|
||||
}
|
||||
|
||||
if needs_conversion {
|
||||
// CPA strategy: when provider and client share the same wire format
|
||||
// (exact match or same family), pass through the stream verbatim.
|
||||
@@ -73,10 +82,6 @@ pub fn resolve_finalize_stream_rewrite_mode(
|
||||
.then_some(FinalizeStreamRewriteMode::Standard);
|
||||
}
|
||||
|
||||
if provider_api_format == "openai:image" && client_api_format == "openai:image" {
|
||||
return Some(FinalizeStreamRewriteMode::OpenAiImage);
|
||||
}
|
||||
|
||||
if envelope_name.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME) {
|
||||
return (provider_api_format == "claude:messages"
|
||||
&& client_api_format == "claude:messages")
|
||||
@@ -106,6 +111,7 @@ enum AiSurfaceStreamRewriteState {
|
||||
EnvelopeUnwrap,
|
||||
ModelDirectiveDisplay,
|
||||
OpenAiImage(Box<OpenAiImageStreamState>),
|
||||
OpenAiImageToOpenAiChat(Box<OpenAiImageChatStreamState>),
|
||||
Standard(Box<StreamingStandardFormatMatrix>),
|
||||
KiroToClaudeCli(Box<KiroToClaudeCliStreamState>),
|
||||
KiroToClaudeCliThenStandard {
|
||||
@@ -132,6 +138,11 @@ pub fn maybe_build_ai_surface_stream_rewriter<'a>(
|
||||
FinalizeStreamRewriteMode::OpenAiImage => {
|
||||
AiSurfaceStreamRewriteState::OpenAiImage(Box::<OpenAiImageStreamState>::default())
|
||||
}
|
||||
FinalizeStreamRewriteMode::OpenAiImageToOpenAiChat => {
|
||||
AiSurfaceStreamRewriteState::OpenAiImageToOpenAiChat(
|
||||
Box::<OpenAiImageChatStreamState>::default(),
|
||||
)
|
||||
}
|
||||
FinalizeStreamRewriteMode::Standard => {
|
||||
AiSurfaceStreamRewriteState::Standard(Box::<StreamingStandardFormatMatrix>::default())
|
||||
}
|
||||
@@ -159,6 +170,9 @@ impl AiSurfaceStreamRewriter<'_> {
|
||||
AiSurfaceStreamRewriteState::OpenAiImage(state) => {
|
||||
state.push_chunk(self.report_context, chunk)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::OpenAiImageToOpenAiChat(state) => {
|
||||
state.push_chunk(self.report_context, chunk)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::KiroToClaudeCli(state) => {
|
||||
state.push_chunk(self.report_context, chunk)
|
||||
}
|
||||
@@ -183,6 +197,9 @@ impl AiSurfaceStreamRewriter<'_> {
|
||||
pub fn finish(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
match &mut self.state {
|
||||
AiSurfaceStreamRewriteState::OpenAiImage(state) => state.finish(self.report_context),
|
||||
AiSurfaceStreamRewriteState::OpenAiImageToOpenAiChat(state) => {
|
||||
state.finish(self.report_context)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::KiroToClaudeCli(state) => {
|
||||
state.finish(self.report_context)
|
||||
}
|
||||
@@ -228,6 +245,7 @@ impl AiSurfaceStreamRewriter<'_> {
|
||||
transform_standard_line(state, self.report_context, line)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::OpenAiImage(_)
|
||||
| AiSurfaceStreamRewriteState::OpenAiImageToOpenAiChat(_)
|
||||
| AiSurfaceStreamRewriteState::KiroToClaudeCli(_)
|
||||
| AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard { .. } => Ok(Vec::new()),
|
||||
}
|
||||
@@ -686,4 +704,58 @@ data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"thinki
|
||||
Some(FinalizeStreamRewriteMode::OpenAiImage)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rewrites_openai_image_stream_to_openai_chat_final_chunk() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:chat",
|
||||
"mapped_model": "gpt-image-2",
|
||||
"request_id": "trace-image-chat-stream",
|
||||
"needs_conversion": false,
|
||||
});
|
||||
assert_eq!(
|
||||
resolve_finalize_stream_rewrite_mode(&report_context),
|
||||
Some(FinalizeStreamRewriteMode::OpenAiImageToOpenAiChat)
|
||||
);
|
||||
let mut rewriter = maybe_build_ai_surface_stream_rewriter(Some(&report_context))
|
||||
.expect("image to chat stream rewriter should exist");
|
||||
|
||||
let progress = rewriter
|
||||
.push_chunk(
|
||||
br#"event: response.image_generation_call.partial_image
|
||||
data: {"type":"response.image_generation_call.partial_image","partial_image_b64":"cGFydGlhbA=="}
|
||||
|
||||
"#,
|
||||
)
|
||||
.expect("partial image should rewrite as progress");
|
||||
let progress_text = String::from_utf8(progress).expect("progress output should be utf8");
|
||||
assert!(progress_text.contains("\"object\":\"chat.completion.chunk\""));
|
||||
assert!(!progress_text.contains("cGFydGlhbA=="));
|
||||
|
||||
let output_item = rewriter
|
||||
.push_chunk(
|
||||
br#"event: response.output_item.done
|
||||
data: {"type":"response.output_item.done","item":{"type":"image_generation_call","id":"ig_1","result":"aGVsbG8=","output_format":"png"}}
|
||||
|
||||
"#,
|
||||
)
|
||||
.expect("output item should rewrite");
|
||||
let output_item_text = String::from_utf8(output_item).expect("output item should be utf8");
|
||||
assert!(output_item_text.is_empty());
|
||||
|
||||
let final_output = rewriter
|
||||
.push_chunk(
|
||||
br#"event: response.completed
|
||||
data: {"type":"response.completed","response":{"id":"resp_123","model":"gpt-image-2","tool_usage":{"image_gen":{"total_tokens":0}},"output":[]}}
|
||||
|
||||
"#,
|
||||
)
|
||||
.expect("completed event should rewrite");
|
||||
let final_text = String::from_utf8(final_output).expect("final output should be utf8");
|
||||
assert!(final_text.contains("\"object\":\"chat.completion.chunk\""));
|
||||
assert!(final_text.contains(""));
|
||||
assert!(final_text.contains("data: [DONE]"));
|
||||
assert!(!final_text.contains("image_generation.completed"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,16 +1,21 @@
|
||||
use std::borrow::Cow;
|
||||
|
||||
use aether_ai_formats::formats::conversion::response::{
|
||||
convert_claude_response_to_openai_responses, convert_gemini_response_to_openai_responses,
|
||||
convert_openai_chat_response_to_openai_responses,
|
||||
};
|
||||
use aether_contracts::{ExecutionStreamTerminalSummary, StandardizedUsage};
|
||||
use serde_json::{json, Value};
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use crate::formats::claude::messages::stream::ClaudeClientEmitter;
|
||||
use crate::formats::gemini::generate_content::stream::GeminiClientEmitter;
|
||||
use crate::formats::openai::chat::stream::{
|
||||
OpenAIChatClientEmitter, OpenAIResponsesClientEmitter, OpenAIResponsesProviderState,
|
||||
};
|
||||
use crate::formats::shared::sse::encode_json_sse;
|
||||
use crate::formats::shared::sse::{encode_done_sse, encode_json_sse};
|
||||
use crate::formats::shared::stream_core::common::{
|
||||
build_openai_chat_chunk, build_openai_chat_finish_chunk, build_openai_chat_usage_chunk,
|
||||
};
|
||||
use crate::formats::shared::stream_core::CanonicalStreamFrame;
|
||||
use crate::formats::shared::AiSurfaceFinalizeError;
|
||||
|
||||
@@ -27,12 +32,25 @@ pub fn maybe_bridge_standard_sync_json_to_stream(
|
||||
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
|
||||
let provider_api_format = normalize_api_format(provider_api_format);
|
||||
let client_api_format = normalize_api_format(client_api_format);
|
||||
if client_api_format == "openai:image"
|
||||
&& matches!(
|
||||
provider_api_format.as_str(),
|
||||
"openai:image" | "gemini:generate_content"
|
||||
)
|
||||
{
|
||||
if provider_api_format == "openai:image" {
|
||||
return match client_api_format.as_str() {
|
||||
"openai:image" => {
|
||||
maybe_bridge_openai_image_sync_json_to_stream(provider_body_json, report_context)
|
||||
}
|
||||
"openai:chat" => maybe_bridge_openai_image_sync_json_to_chat_stream(
|
||||
provider_body_json,
|
||||
report_context,
|
||||
),
|
||||
"openai:responses" | "openai:responses:compact" => {
|
||||
maybe_bridge_openai_image_sync_json_to_responses_stream(
|
||||
provider_body_json,
|
||||
report_context,
|
||||
)
|
||||
}
|
||||
_ => Ok(None),
|
||||
};
|
||||
}
|
||||
if client_api_format == "openai:image" && provider_api_format == "gemini:generate_content" {
|
||||
return maybe_bridge_openai_image_sync_json_to_stream(provider_body_json, report_context);
|
||||
}
|
||||
if !is_standard_api_format(provider_api_format.as_str())
|
||||
@@ -72,49 +90,19 @@ fn maybe_bridge_openai_image_sync_json_to_stream(
|
||||
provider_body_json: &Value,
|
||||
report_context: Option<&Value>,
|
||||
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
|
||||
let provider_api_format = report_context
|
||||
.and_then(|value| value.get("provider_api_format"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or("openai:image");
|
||||
let owned_response;
|
||||
let provider_body_json = if provider_api_format == "gemini:generate_content" {
|
||||
let Some(converted) =
|
||||
crate::formats::shared::image_bridge::build_openai_image_response_from_gemini_response(
|
||||
provider_body_json,
|
||||
report_context,
|
||||
)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
owned_response = converted;
|
||||
&owned_response
|
||||
} else if provider_body_json.get("output").is_some() && provider_body_json.get("data").is_none()
|
||||
{
|
||||
let Some(converted) = crate::formats::shared::image_bridge::build_openai_image_response_from_response_stream_sync_body(
|
||||
provider_body_json,
|
||||
report_context,
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
owned_response = converted;
|
||||
&owned_response
|
||||
} else {
|
||||
provider_body_json
|
||||
};
|
||||
let Some(response) = provider_body_json.as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(image) = response
|
||||
.get("data")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(Value::as_object)
|
||||
.find_map(extract_openai_image_sync_b64_json)
|
||||
let Some(provider_body_json) =
|
||||
normalize_openai_image_sync_response(provider_body_json, report_context)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(response) = provider_body_json.as_ref().as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let outputs = collect_openai_image_outputs(response, report_context);
|
||||
let Some(image) = outputs.iter().find_map(OpenAiImageOutput::b64_json) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let image_count = openai_image_response_image_count(response).max(outputs.len() as u64);
|
||||
let usage = response.get("usage").cloned().unwrap_or(Value::Null);
|
||||
let event_name = openai_image_completed_event_name(report_context);
|
||||
let sse_body = encode_json_sse(
|
||||
@@ -128,27 +116,461 @@ fn maybe_bridge_openai_image_sync_json_to_stream(
|
||||
|
||||
Ok(Some(SyncToStreamBridgeOutcome {
|
||||
sse_body,
|
||||
terminal_summary: Some(ExecutionStreamTerminalSummary {
|
||||
standardized_usage: response
|
||||
.get("usage")
|
||||
.and_then(standardized_usage_from_openai_usage),
|
||||
finish_reason: Some("stop".to_string()),
|
||||
response_id: response
|
||||
.get("id")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned),
|
||||
model: response
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| image_bridge_model(report_context)),
|
||||
observed_finish: true,
|
||||
unknown_event_count: 0,
|
||||
parser_error: None,
|
||||
}),
|
||||
terminal_summary: Some(openai_image_terminal_summary(
|
||||
response,
|
||||
report_context,
|
||||
image_count,
|
||||
)),
|
||||
}))
|
||||
}
|
||||
|
||||
fn maybe_bridge_openai_image_sync_json_to_chat_stream(
|
||||
provider_body_json: &Value,
|
||||
report_context: Option<&Value>,
|
||||
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
|
||||
let Some(provider_body_json) =
|
||||
normalize_openai_image_sync_response(provider_body_json, report_context)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(response) = provider_body_json.as_ref().as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let outputs = collect_openai_image_outputs(response, report_context);
|
||||
if outputs.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let image_count = openai_image_response_image_count(response).max(outputs.len() as u64);
|
||||
let summary = openai_image_terminal_summary(response, report_context, image_count);
|
||||
let response_id = openai_image_bridge_response_id(response, report_context, "chatcmpl-image");
|
||||
let model = openai_image_bridge_response_model(response, report_context);
|
||||
let content = outputs
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, output)| output.markdown(index))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n");
|
||||
|
||||
let mut sse_body = Vec::new();
|
||||
sse_body.extend(encode_json_sse(
|
||||
None,
|
||||
&build_openai_chat_chunk(&response_id, &model, content, None, None),
|
||||
)?);
|
||||
sse_body.extend(encode_json_sse(
|
||||
None,
|
||||
&build_openai_chat_finish_chunk(&response_id, &model, Some("stop")),
|
||||
)?);
|
||||
if let Some((input_tokens, output_tokens, total_tokens, reasoning_tokens)) = summary
|
||||
.standardized_usage
|
||||
.as_ref()
|
||||
.and_then(openai_chat_usage_counts)
|
||||
{
|
||||
sse_body.extend(encode_json_sse(
|
||||
None,
|
||||
&build_openai_chat_usage_chunk(
|
||||
&response_id,
|
||||
&model,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
reasoning_tokens,
|
||||
),
|
||||
)?);
|
||||
}
|
||||
sse_body.extend(encode_done_sse());
|
||||
|
||||
Ok(Some(SyncToStreamBridgeOutcome {
|
||||
sse_body,
|
||||
terminal_summary: Some(summary),
|
||||
}))
|
||||
}
|
||||
|
||||
fn maybe_bridge_openai_image_sync_json_to_responses_stream(
|
||||
provider_body_json: &Value,
|
||||
report_context: Option<&Value>,
|
||||
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
|
||||
let Some(provider_body_json) =
|
||||
normalize_openai_image_sync_response(provider_body_json, report_context)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(response) = provider_body_json.as_ref().as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let outputs = collect_openai_image_outputs(response, report_context);
|
||||
if outputs.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let response_id = openai_image_bridge_response_id(response, report_context, "resp-image");
|
||||
let model = openai_image_bridge_response_model(response, report_context);
|
||||
let mut response_output = Vec::new();
|
||||
for (index, output) in outputs.iter().enumerate() {
|
||||
response_output.push(output.responses_image_generation_item(&response_id, index));
|
||||
}
|
||||
|
||||
let mut response_object = Map::new();
|
||||
response_object.insert("id".to_string(), Value::String(response_id.clone()));
|
||||
response_object.insert("object".to_string(), Value::String("response".to_string()));
|
||||
response_object.insert("model".to_string(), Value::String(model));
|
||||
response_object.insert("status".to_string(), Value::String("completed".to_string()));
|
||||
response_object.insert("output".to_string(), Value::Array(response_output.clone()));
|
||||
if let Some(created) = response.get("created").and_then(Value::as_i64) {
|
||||
response_object.insert("created_at".to_string(), json!(created));
|
||||
}
|
||||
if let Some(usage) = response.get("usage").filter(|value| value.is_object()) {
|
||||
response_object.insert("usage".to_string(), usage.clone());
|
||||
}
|
||||
|
||||
let mut sse_body = Vec::new();
|
||||
for (index, item) in response_output.iter().enumerate() {
|
||||
sse_body.extend(encode_json_sse(
|
||||
Some("response.output_item.done"),
|
||||
&json!({
|
||||
"type": "response.output_item.done",
|
||||
"output_index": index,
|
||||
"item": item,
|
||||
}),
|
||||
)?);
|
||||
}
|
||||
sse_body.extend(encode_json_sse(
|
||||
Some("response.completed"),
|
||||
&json!({
|
||||
"type": "response.completed",
|
||||
"response": Value::Object(response_object),
|
||||
}),
|
||||
)?);
|
||||
|
||||
let image_count = openai_image_response_image_count(response).max(outputs.len() as u64);
|
||||
Ok(Some(SyncToStreamBridgeOutcome {
|
||||
sse_body,
|
||||
terminal_summary: Some(openai_image_terminal_summary(
|
||||
response,
|
||||
report_context,
|
||||
image_count,
|
||||
)),
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct OpenAiImageOutput {
|
||||
b64_json: Option<String>,
|
||||
url: Option<String>,
|
||||
mime_type: String,
|
||||
output_format: Option<String>,
|
||||
revised_prompt: Option<String>,
|
||||
}
|
||||
|
||||
impl OpenAiImageOutput {
|
||||
fn b64_json(&self) -> Option<String> {
|
||||
self.b64_json
|
||||
.clone()
|
||||
.or_else(|| self.url.as_deref().and_then(extract_base64_from_data_url))
|
||||
}
|
||||
|
||||
fn source_url(&self) -> Option<String> {
|
||||
self.url.clone().or_else(|| {
|
||||
self.b64_json
|
||||
.as_ref()
|
||||
.map(|value| format!("data:{};base64,{value}", self.mime_type))
|
||||
})
|
||||
}
|
||||
|
||||
fn markdown(&self, index: usize) -> String {
|
||||
let alt = if index == 0 {
|
||||
"generated image".to_string()
|
||||
} else {
|
||||
format!("generated image {}", index + 1)
|
||||
};
|
||||
match self.source_url() {
|
||||
Some(url) => format!(""),
|
||||
None => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn responses_image_generation_item(&self, response_id: &str, index: usize) -> Value {
|
||||
let mut item = Map::new();
|
||||
item.insert(
|
||||
"id".to_string(),
|
||||
Value::String(format!("{response_id}_img_{index}")),
|
||||
);
|
||||
item.insert(
|
||||
"type".to_string(),
|
||||
Value::String("image_generation_call".to_string()),
|
||||
);
|
||||
item.insert("status".to_string(), Value::String("completed".to_string()));
|
||||
if let Some(result) = self.b64_json().or_else(|| self.url.clone()) {
|
||||
item.insert("result".to_string(), Value::String(result));
|
||||
}
|
||||
if let Some(output_format) = self.output_format.as_ref() {
|
||||
item.insert(
|
||||
"output_format".to_string(),
|
||||
Value::String(output_format.clone()),
|
||||
);
|
||||
}
|
||||
if let Some(revised_prompt) = self.revised_prompt.as_ref() {
|
||||
item.insert(
|
||||
"revised_prompt".to_string(),
|
||||
Value::String(revised_prompt.clone()),
|
||||
);
|
||||
}
|
||||
Value::Object(item)
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_openai_image_sync_response<'a>(
|
||||
provider_body_json: &'a Value,
|
||||
report_context: Option<&Value>,
|
||||
) -> Result<Option<Cow<'a, Value>>, AiSurfaceFinalizeError> {
|
||||
let provider_api_format = report_context
|
||||
.and_then(|value| value.get("provider_api_format"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or("openai:image");
|
||||
if provider_api_format == "gemini:generate_content" {
|
||||
let Some(converted) =
|
||||
crate::formats::shared::image_bridge::build_openai_image_response_from_gemini_response(
|
||||
provider_body_json,
|
||||
report_context,
|
||||
)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
return Ok(Some(Cow::Owned(converted)));
|
||||
}
|
||||
if provider_body_json.get("output").is_some() && provider_body_json.get("data").is_none() {
|
||||
let Some(converted) = crate::formats::shared::image_bridge::build_openai_image_response_from_response_stream_sync_body(
|
||||
provider_body_json,
|
||||
report_context,
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
return Ok(Some(Cow::Owned(converted)));
|
||||
}
|
||||
Ok(Some(Cow::Borrowed(provider_body_json)))
|
||||
}
|
||||
|
||||
fn collect_openai_image_outputs(
|
||||
response: &Map<String, Value>,
|
||||
report_context: Option<&Value>,
|
||||
) -> Vec<OpenAiImageOutput> {
|
||||
response
|
||||
.get("data")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(Value::as_object)
|
||||
.filter_map(|item| openai_image_output_from_item(item, report_context))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn openai_image_output_from_item(
|
||||
item: &Map<String, Value>,
|
||||
report_context: Option<&Value>,
|
||||
) -> Option<OpenAiImageOutput> {
|
||||
let b64_json = extract_openai_image_sync_b64_json(item);
|
||||
let url = item
|
||||
.get("url")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
if b64_json.is_none() && url.is_none() {
|
||||
return None;
|
||||
}
|
||||
let output_format = item
|
||||
.get("output_format")
|
||||
.or_else(|| item.get("format"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| image_request_output_format(report_context));
|
||||
let mime_type = url
|
||||
.as_deref()
|
||||
.and_then(extract_mime_type_from_data_url)
|
||||
.or_else(|| {
|
||||
output_format
|
||||
.as_deref()
|
||||
.map(mime_type_from_image_output_format)
|
||||
})
|
||||
.unwrap_or_else(|| "image/png".to_string());
|
||||
let revised_prompt = item
|
||||
.get("revised_prompt")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
|
||||
Some(OpenAiImageOutput {
|
||||
b64_json,
|
||||
url,
|
||||
mime_type,
|
||||
output_format,
|
||||
revised_prompt,
|
||||
})
|
||||
}
|
||||
|
||||
fn openai_image_response_image_count(response: &Map<String, Value>) -> u64 {
|
||||
response
|
||||
.get("data")
|
||||
.and_then(Value::as_array)
|
||||
.map(|items| items.len() as u64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn openai_image_terminal_summary(
|
||||
response: &Map<String, Value>,
|
||||
report_context: Option<&Value>,
|
||||
image_count: u64,
|
||||
) -> ExecutionStreamTerminalSummary {
|
||||
ExecutionStreamTerminalSummary {
|
||||
standardized_usage: openai_image_standardized_usage(
|
||||
response.get("usage"),
|
||||
report_context,
|
||||
image_count,
|
||||
),
|
||||
finish_reason: Some("stop".to_string()),
|
||||
response_id: response
|
||||
.get("id")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned),
|
||||
model: response
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| image_bridge_model(report_context)),
|
||||
observed_finish: true,
|
||||
unknown_event_count: 0,
|
||||
parser_error: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn openai_image_standardized_usage(
|
||||
usage: Option<&Value>,
|
||||
report_context: Option<&Value>,
|
||||
image_count: u64,
|
||||
) -> Option<StandardizedUsage> {
|
||||
let mut standardized_usage = usage
|
||||
.and_then(standardized_usage_from_openai_usage)
|
||||
.unwrap_or_else(StandardizedUsage::new);
|
||||
if image_count > 0 {
|
||||
standardized_usage.request_count = i64::try_from(image_count).unwrap_or(i64::MAX);
|
||||
standardized_usage
|
||||
.dimensions
|
||||
.insert("image_count".to_string(), json!(image_count));
|
||||
}
|
||||
if let Some(output_format) = image_request_output_format(report_context) {
|
||||
standardized_usage
|
||||
.dimensions
|
||||
.insert("image_output_format".to_string(), json!(output_format));
|
||||
}
|
||||
if let Some(size) = image_request_size(report_context) {
|
||||
standardized_usage
|
||||
.dimensions
|
||||
.insert("image_size".to_string(), json!(size));
|
||||
}
|
||||
if let Some(quality) = image_request_quality(report_context) {
|
||||
standardized_usage
|
||||
.dimensions
|
||||
.insert("image_quality".to_string(), json!(quality));
|
||||
}
|
||||
(standardized_usage.signal_score() > 0).then_some(standardized_usage)
|
||||
}
|
||||
|
||||
fn openai_chat_usage_counts(usage: &StandardizedUsage) -> Option<(u64, u64, u64, u64)> {
|
||||
let input_tokens = usage.input_tokens.max(0) as u64;
|
||||
let output_tokens = usage.output_tokens.max(0) as u64;
|
||||
let reasoning_tokens = usage.reasoning_tokens.max(0) as u64;
|
||||
let total_tokens = usage
|
||||
.dimensions
|
||||
.get("total_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or_else(|| {
|
||||
input_tokens
|
||||
.saturating_add(output_tokens)
|
||||
.saturating_add(reasoning_tokens)
|
||||
});
|
||||
(total_tokens > 0).then_some((input_tokens, output_tokens, total_tokens, reasoning_tokens))
|
||||
}
|
||||
|
||||
fn openai_image_bridge_response_id(
|
||||
response: &Map<String, Value>,
|
||||
report_context: Option<&Value>,
|
||||
fallback_prefix: &str,
|
||||
) -> String {
|
||||
response
|
||||
.get("id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| {
|
||||
report_context
|
||||
.and_then(|value| value.get("request_id"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(|value| format!("{fallback_prefix}-{value}"))
|
||||
})
|
||||
.unwrap_or_else(|| fallback_prefix.to_string())
|
||||
}
|
||||
|
||||
fn openai_image_bridge_response_model(
|
||||
response: &Map<String, Value>,
|
||||
report_context: Option<&Value>,
|
||||
) -> String {
|
||||
response
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| image_bridge_model(report_context))
|
||||
.unwrap_or_else(|| "gpt-image".to_string())
|
||||
}
|
||||
|
||||
fn image_request_output_format(report_context: Option<&Value>) -> Option<String> {
|
||||
report_context
|
||||
.and_then(|value| value.get("image_request"))
|
||||
.and_then(|value| value.get("output_format"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn image_request_size(report_context: Option<&Value>) -> Option<String> {
|
||||
report_context
|
||||
.and_then(|value| value.get("image_request"))
|
||||
.and_then(|value| value.get("size"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn image_request_quality(report_context: Option<&Value>) -> Option<String> {
|
||||
report_context
|
||||
.and_then(|value| value.get("image_request"))
|
||||
.and_then(|value| value.get("quality"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn mime_type_from_image_output_format(output_format: &str) -> String {
|
||||
match output_format.trim().to_ascii_lowercase().as_str() {
|
||||
"jpg" | "jpeg" => "image/jpeg".to_string(),
|
||||
"webp" => "image/webp".to_string(),
|
||||
"png" => "image/png".to_string(),
|
||||
value if !value.is_empty() => format!("image/{value}"),
|
||||
_ => "image/png".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_api_format(value: &str) -> String {
|
||||
aether_ai_formats::normalize_api_format_alias(value)
|
||||
}
|
||||
@@ -186,6 +608,14 @@ fn extract_base64_from_data_url(value: &str) -> Option<String> {
|
||||
(!payload.trim().is_empty()).then(|| payload.trim().to_string())
|
||||
}
|
||||
|
||||
fn extract_mime_type_from_data_url(value: &str) -> Option<String> {
|
||||
let trimmed = value.trim();
|
||||
let (metadata, _) = trimmed.split_once(',')?;
|
||||
let mime_type = metadata.strip_prefix("data:")?.strip_suffix(";base64")?;
|
||||
let mime_type = mime_type.trim();
|
||||
(!mime_type.is_empty()).then(|| mime_type.to_string())
|
||||
}
|
||||
|
||||
fn openai_image_completed_event_name(report_context: Option<&Value>) -> &'static str {
|
||||
if openai_image_request_operation(report_context) == Some("edit") {
|
||||
"image_edit.completed"
|
||||
@@ -565,6 +995,76 @@ mod tests {
|
||||
.cloned(),
|
||||
Some(json!(100))
|
||||
);
|
||||
assert_eq!(
|
||||
summary
|
||||
.standardized_usage
|
||||
.as_ref()
|
||||
.and_then(|usage| usage.dimensions.get("image_count"))
|
||||
.cloned(),
|
||||
Some(json!(1))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bridges_openai_image_sync_json_to_openai_chat_sse() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:chat",
|
||||
"mapped_model": "gpt-image-2",
|
||||
"image_request": {
|
||||
"operation": "generate",
|
||||
"output_format": "png",
|
||||
"size": "1024x1024",
|
||||
"quality": "medium"
|
||||
}
|
||||
});
|
||||
let outcome = maybe_bridge_standard_sync_json_to_stream(
|
||||
&json!({
|
||||
"id": "img_123",
|
||||
"created": 1776971267,
|
||||
"model": "gpt-image-2",
|
||||
"data": [
|
||||
{"b64_json": "aGVsbG8="},
|
||||
{"b64_json": "d29ybGQ="}
|
||||
],
|
||||
"usage": {
|
||||
"total_tokens": 100,
|
||||
"input_tokens": 50,
|
||||
"output_tokens": 50
|
||||
}
|
||||
}),
|
||||
"openai:image",
|
||||
"openai:chat",
|
||||
Some(&report_context),
|
||||
)
|
||||
.expect("bridge should succeed")
|
||||
.expect("bridge should produce sse");
|
||||
|
||||
let output = utf8(outcome.sse_body);
|
||||
assert!(output.contains("\"object\":\"chat.completion.chunk\""));
|
||||
assert!(output.contains(""));
|
||||
assert!(output.contains(""));
|
||||
assert!(output.contains("\"finish_reason\":\"stop\""));
|
||||
assert!(output.contains("data: [DONE]"));
|
||||
assert!(!output.contains("image_generation.completed"));
|
||||
|
||||
let summary = outcome
|
||||
.terminal_summary
|
||||
.expect("terminal summary should exist");
|
||||
let usage = summary
|
||||
.standardized_usage
|
||||
.as_ref()
|
||||
.expect("standard usage should exist");
|
||||
assert_eq!(usage.request_count, 2);
|
||||
assert_eq!(usage.dimensions.get("image_count"), Some(&json!(2)));
|
||||
assert_eq!(
|
||||
usage.dimensions.get("image_size"),
|
||||
Some(&json!("1024x1024"))
|
||||
);
|
||||
assert_eq!(
|
||||
usage.dimensions.get("image_quality"),
|
||||
Some(&json!("medium"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -23,14 +23,28 @@ impl DefaultBillingRuleGenerator {
|
||||
pricing: &BillingModelPricingSnapshot,
|
||||
task_type: &str,
|
||||
) -> Option<VirtualBillingRule> {
|
||||
let pricing_config = pricing.effective_tiered_pricing();
|
||||
let tiers = pricing
|
||||
.effective_tiered_pricing()
|
||||
.and_then(|value| value.get("tiers"))
|
||||
.and_then(Value::as_array)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let explicit_image_output_price_default =
|
||||
explicit_image_output_price_default(pricing_config);
|
||||
let image_output_price_default = explicit_image_output_price_default.unwrap_or(0.0);
|
||||
let has_image_output_matrix = explicit_image_output_price_entries(pricing_config)
|
||||
.is_some_and(|entries| !entries.is_empty());
|
||||
let has_image_output_ranges = explicit_image_output_price_ranges(pricing_config)
|
||||
.is_some_and(|ranges| !ranges.is_empty());
|
||||
let has_image_output_pricing = has_image_output_matrix
|
||||
|| has_image_output_ranges
|
||||
|| explicit_image_output_price_default.is_some();
|
||||
|
||||
if tiers.is_empty() && pricing.effective_price_per_request().is_none() {
|
||||
if tiers.is_empty()
|
||||
&& pricing.effective_price_per_request().is_none()
|
||||
&& !has_image_output_pricing
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
@@ -63,6 +77,10 @@ impl DefaultBillingRuleGenerator {
|
||||
json!(base_cache_read_price),
|
||||
);
|
||||
variables.insert("price_per_request".to_string(), json!(base_request_price));
|
||||
variables.insert(
|
||||
"image_output_price_per_image".to_string(),
|
||||
json!(image_output_price_default),
|
||||
);
|
||||
|
||||
let mut dimension_mappings = BTreeMap::new();
|
||||
for (name, key, default) in [
|
||||
@@ -86,6 +104,14 @@ impl DefaultBillingRuleGenerator {
|
||||
),
|
||||
("cache_read_tokens", "cache_read_tokens", json!(0)),
|
||||
("request_count", "request_count", json!(1)),
|
||||
("image_count", "image_count", json!(0)),
|
||||
("image_count_unmetered", "image_count_unmetered", json!(0)),
|
||||
("image_price_key", "image_price_key", json!("default")),
|
||||
(
|
||||
"image_output_price_per_image",
|
||||
"image_output_price_per_image",
|
||||
json!(image_output_price_default),
|
||||
),
|
||||
] {
|
||||
dimension_mappings.insert(
|
||||
name.to_string(),
|
||||
@@ -121,6 +147,10 @@ impl DefaultBillingRuleGenerator {
|
||||
"cache_read_cost",
|
||||
"cache_read_tokens * cache_read_price_per_1m / 1000000",
|
||||
),
|
||||
(
|
||||
"image_output_cost",
|
||||
"image_count_unmetered * image_output_price_per_image",
|
||||
),
|
||||
("request_cost", "request_count * price_per_request"),
|
||||
] {
|
||||
dimension_mappings.insert(
|
||||
@@ -209,7 +239,7 @@ impl DefaultBillingRuleGenerator {
|
||||
id: "__default__".to_string(),
|
||||
name: format!("Default rule for {}", pricing.global_model_name),
|
||||
task_type: normalize_task_type(task_type).to_string(),
|
||||
expression: "input_cost + output_cost + cache_creation_uncategorized_cost + cache_creation_ephemeral_5m_cost + cache_creation_ephemeral_1h_cost + cache_read_cost + request_cost".to_string(),
|
||||
expression: "input_cost + output_cost + cache_creation_uncategorized_cost + cache_creation_ephemeral_5m_cost + cache_creation_ephemeral_1h_cost + cache_read_cost + image_output_cost + request_cost".to_string(),
|
||||
variables,
|
||||
dimension_mappings,
|
||||
scope: "default".to_string(),
|
||||
@@ -267,3 +297,201 @@ fn build_tier_entries(
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn explicit_image_output_price_entries(
|
||||
pricing_config: Option<&Value>,
|
||||
) -> Option<BTreeMap<String, Value>> {
|
||||
let pricing_config = pricing_config?;
|
||||
let mut entries = BTreeMap::new();
|
||||
for key in [
|
||||
"image_output_prices",
|
||||
"image_output_price_per_image",
|
||||
"image_output_price_matrix",
|
||||
"image_prices",
|
||||
] {
|
||||
if let Some(value) = pricing_config.get(key) {
|
||||
collect_image_output_price_entries(value, &mut entries);
|
||||
}
|
||||
}
|
||||
Some(entries)
|
||||
}
|
||||
|
||||
pub(crate) fn explicit_image_output_price_ranges(
|
||||
pricing_config: Option<&Value>,
|
||||
) -> Option<Vec<Value>> {
|
||||
let pricing_config = pricing_config?;
|
||||
let Some(value) = pricing_config.get("image_output_price_ranges") else {
|
||||
return Some(Vec::new());
|
||||
};
|
||||
|
||||
let mut ranges = Vec::new();
|
||||
match value {
|
||||
Value::Array(items) => {
|
||||
for item in items {
|
||||
let Some(object) = item.as_object() else {
|
||||
continue;
|
||||
};
|
||||
let mut range = serde_json::Map::new();
|
||||
if let Some(up_to_pixels) = object
|
||||
.get("up_to_pixels")
|
||||
.or_else(|| object.get("up_to"))
|
||||
.or_else(|| object.get("max_pixels"))
|
||||
{
|
||||
range.insert("up_to_pixels".to_string(), up_to_pixels.clone());
|
||||
}
|
||||
if let Some(label) = object.get("label").cloned() {
|
||||
range.insert("label".to_string(), label);
|
||||
}
|
||||
if let Some(prices) = object.get("prices") {
|
||||
range.insert("prices".to_string(), prices.clone());
|
||||
} else {
|
||||
let mut prices = serde_json::Map::new();
|
||||
for quality in ["low", "medium", "high"] {
|
||||
if let Some(price) = object.get(quality).and_then(Value::as_f64) {
|
||||
prices.insert(quality.to_string(), json!(price));
|
||||
}
|
||||
}
|
||||
if prices.is_empty() {
|
||||
if let Some(price) = object
|
||||
.get("price_per_image")
|
||||
.or_else(|| object.get("price"))
|
||||
.or_else(|| object.get("value"))
|
||||
.and_then(Value::as_f64)
|
||||
{
|
||||
prices.insert("default".to_string(), json!(price));
|
||||
}
|
||||
}
|
||||
if !prices.is_empty() {
|
||||
range.insert("prices".to_string(), Value::Object(prices));
|
||||
}
|
||||
}
|
||||
if !range.is_empty() {
|
||||
ranges.push(Value::Object(range));
|
||||
}
|
||||
}
|
||||
}
|
||||
Value::Object(object) => {
|
||||
for (key, item) in object {
|
||||
let Some(entry) = item.as_object() else {
|
||||
continue;
|
||||
};
|
||||
let mut range = serde_json::Map::new();
|
||||
if let Some(up_to_pixels) = entry
|
||||
.get("up_to_pixels")
|
||||
.or_else(|| entry.get("up_to"))
|
||||
.or_else(|| entry.get("max_pixels"))
|
||||
{
|
||||
range.insert("up_to_pixels".to_string(), up_to_pixels.clone());
|
||||
} else if let Ok(parsed) = key.parse::<u64>() {
|
||||
range.insert("up_to_pixels".to_string(), json!(parsed));
|
||||
}
|
||||
if let Some(label) = entry.get("label").cloned() {
|
||||
range.insert("label".to_string(), label);
|
||||
}
|
||||
if let Some(prices) = entry.get("prices") {
|
||||
range.insert("prices".to_string(), prices.clone());
|
||||
}
|
||||
if !range.is_empty() {
|
||||
ranges.push(Value::Object(range));
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
Some(ranges)
|
||||
}
|
||||
|
||||
pub(crate) fn explicit_image_output_price_default(pricing_config: Option<&Value>) -> Option<f64> {
|
||||
let pricing_config = pricing_config?;
|
||||
pricing_config
|
||||
.get("image_output_price_default")
|
||||
.or_else(|| pricing_config.get("image_price_default"))
|
||||
.or_else(|| {
|
||||
pricing_config
|
||||
.get("image_output_prices")
|
||||
.and_then(|value| value.get("default"))
|
||||
})
|
||||
.and_then(Value::as_f64)
|
||||
}
|
||||
|
||||
fn collect_image_output_price_entries(value: &Value, entries: &mut BTreeMap<String, Value>) {
|
||||
if let Some(object) = value.as_object() {
|
||||
for (key, value) in object {
|
||||
if key.eq_ignore_ascii_case("default") {
|
||||
continue;
|
||||
}
|
||||
if let Some(price) = value.as_f64() {
|
||||
entries.insert(normalize_image_price_key(key), json!(price));
|
||||
continue;
|
||||
}
|
||||
let Some(nested) = value.as_object() else {
|
||||
continue;
|
||||
};
|
||||
let key_is_quality = matches_quality_key(key);
|
||||
for (nested_key, nested_value) in nested {
|
||||
let Some(price) = nested_value.as_f64() else {
|
||||
continue;
|
||||
};
|
||||
let (size, quality) = if key_is_quality {
|
||||
(nested_key.as_str(), key.as_str())
|
||||
} else {
|
||||
(key.as_str(), nested_key.as_str())
|
||||
};
|
||||
entries.insert(image_price_key(size, quality), json!(price));
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if let Some(items) = value.as_array() {
|
||||
for item in items.iter().filter_map(Value::as_object) {
|
||||
let Some(size) = item.get("size").and_then(Value::as_str) else {
|
||||
continue;
|
||||
};
|
||||
let quality = item
|
||||
.get("quality")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("medium");
|
||||
let Some(price) = item
|
||||
.get("price_per_image")
|
||||
.or_else(|| item.get("price"))
|
||||
.or_else(|| item.get("cost"))
|
||||
.and_then(Value::as_f64)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
entries.insert(image_price_key(size, quality), json!(price));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_image_price_key(value: &str) -> String {
|
||||
if let Some((size, quality)) = value.split_once(':').or_else(|| value.split_once('|')) {
|
||||
return image_price_key(size, quality);
|
||||
}
|
||||
value.trim().to_ascii_lowercase().replace(' ', "")
|
||||
}
|
||||
|
||||
fn image_price_key(size: &str, quality: &str) -> String {
|
||||
format!(
|
||||
"{}:{}",
|
||||
normalize_image_size(size),
|
||||
normalize_image_quality(quality)
|
||||
)
|
||||
}
|
||||
|
||||
fn normalize_image_size(value: &str) -> String {
|
||||
value.trim().to_ascii_lowercase().replace(' ', "")
|
||||
}
|
||||
|
||||
fn normalize_image_quality(value: &str) -> String {
|
||||
value.trim().to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn matches_quality_key(value: &str) -> bool {
|
||||
matches!(
|
||||
normalize_image_quality(value).as_str(),
|
||||
"low" | "medium" | "high"
|
||||
)
|
||||
}
|
||||
|
||||
@@ -35,7 +35,10 @@ pub async fn enrich_usage_event_with_billing(
|
||||
data: &dyn BillingModelContextLookup,
|
||||
event: &mut UsageEvent,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if !matches!(event.event_type, UsageEventType::Completed) {
|
||||
if !matches!(
|
||||
event.event_type,
|
||||
UsageEventType::Completed | UsageEventType::Cancelled
|
||||
) {
|
||||
event.data.total_cost_usd = Some(0.0);
|
||||
event.data.actual_total_cost_usd = Some(0.0);
|
||||
return Ok(());
|
||||
@@ -122,24 +125,37 @@ fn calculate_billing_computation(
|
||||
pricing: &BillingModelPricingSnapshot,
|
||||
event: &UsageEvent,
|
||||
) -> Result<BillingComputation, DataLayerError> {
|
||||
let failed =
|
||||
event.data.status_code.unwrap_or_default() >= 400 || event.data.error_message.is_some();
|
||||
let is_image_usage = usage_event_is_image_usage(&event.data);
|
||||
let image_count = if failed {
|
||||
0
|
||||
} else {
|
||||
usage_event_image_count(&event.data).unwrap_or(0)
|
||||
};
|
||||
let request_count = if failed {
|
||||
0
|
||||
} else if is_image_usage && image_count > 0 {
|
||||
image_count
|
||||
} else {
|
||||
1
|
||||
};
|
||||
let input = BillingUsageInput {
|
||||
task_type: event
|
||||
.data
|
||||
.request_type
|
||||
.clone()
|
||||
.unwrap_or_else(|| "chat".to_string()),
|
||||
task_type: if is_image_usage {
|
||||
"image".to_string()
|
||||
} else {
|
||||
event
|
||||
.data
|
||||
.request_type
|
||||
.clone()
|
||||
.unwrap_or_else(|| "chat".to_string())
|
||||
},
|
||||
api_format: event
|
||||
.data
|
||||
.endpoint_api_format
|
||||
.clone()
|
||||
.or_else(|| event.data.api_format.clone()),
|
||||
request_count: if event.data.status_code.unwrap_or_default() >= 400
|
||||
|| event.data.error_message.is_some()
|
||||
{
|
||||
0
|
||||
} else {
|
||||
1
|
||||
},
|
||||
request_count,
|
||||
input_tokens: event.data.input_tokens.unwrap_or_default() as i64,
|
||||
output_tokens: event.data.output_tokens.unwrap_or_default() as i64,
|
||||
cache_creation_tokens: event.data.cache_creation_input_tokens.unwrap_or_default() as i64,
|
||||
@@ -152,6 +168,10 @@ fn calculate_billing_computation(
|
||||
.cache_creation_ephemeral_1h_input_tokens
|
||||
.unwrap_or_default() as i64,
|
||||
cache_read_tokens: event.data.cache_read_input_tokens.unwrap_or_default() as i64,
|
||||
image_count,
|
||||
image_size: usage_event_dimension_string(&event.data, "image_size"),
|
||||
image_quality: usage_event_dimension_string(&event.data, "image_quality"),
|
||||
image_output_format: usage_event_dimension_string(&event.data, "image_output_format"),
|
||||
cache_ttl_minutes: pricing.provider_api_key_cache_ttl_minutes,
|
||||
};
|
||||
|
||||
@@ -162,6 +182,82 @@ fn calculate_billing_computation(
|
||||
})
|
||||
}
|
||||
|
||||
fn usage_event_is_image_usage(data: &aether_usage_runtime::UsageEventData) -> bool {
|
||||
data.request_type
|
||||
.as_deref()
|
||||
.is_some_and(|value| value.eq_ignore_ascii_case("image"))
|
||||
|| api_format_endpoint_kind(data.endpoint_api_format.as_deref()) == Some("image")
|
||||
|| api_format_endpoint_kind(data.api_format.as_deref()) == Some("image")
|
||||
|| usage_event_image_count(data).is_some_and(|value| value > 0)
|
||||
}
|
||||
|
||||
fn usage_event_image_count(data: &aether_usage_runtime::UsageEventData) -> Option<i64> {
|
||||
metadata_dimension_i64(data.request_metadata.as_ref(), "dimensions", "image_count")
|
||||
.or_else(|| {
|
||||
metadata_dimension_i64(
|
||||
data.request_metadata.as_ref(),
|
||||
"billing_dimensions",
|
||||
"image_count",
|
||||
)
|
||||
})
|
||||
.filter(|value| *value > 0)
|
||||
}
|
||||
|
||||
fn usage_event_dimension_string(
|
||||
data: &aether_usage_runtime::UsageEventData,
|
||||
dimension_key: &str,
|
||||
) -> Option<String> {
|
||||
metadata_dimension_string(data.request_metadata.as_ref(), "dimensions", dimension_key).or_else(
|
||||
|| {
|
||||
metadata_dimension_string(
|
||||
data.request_metadata.as_ref(),
|
||||
"billing_dimensions",
|
||||
dimension_key,
|
||||
)
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
fn metadata_dimension_string(
|
||||
metadata: Option<&Value>,
|
||||
bag_key: &str,
|
||||
dimension_key: &str,
|
||||
) -> Option<String> {
|
||||
metadata
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|object| object.get(bag_key))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|object| object.get(dimension_key))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn metadata_dimension_i64(
|
||||
metadata: Option<&Value>,
|
||||
bag_key: &str,
|
||||
dimension_key: &str,
|
||||
) -> Option<i64> {
|
||||
metadata
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|object| object.get(bag_key))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|object| object.get(dimension_key))
|
||||
.and_then(|value| {
|
||||
value
|
||||
.as_i64()
|
||||
.or_else(|| value.as_u64().and_then(|number| i64::try_from(number).ok()))
|
||||
})
|
||||
}
|
||||
|
||||
fn api_format_endpoint_kind(api_format: Option<&str>) -> Option<&str> {
|
||||
api_format
|
||||
.and_then(|value| value.split_once(':').map(|(_, kind)| kind))
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn apply_billing_computation(
|
||||
event: &mut UsageEvent,
|
||||
pricing: &BillingModelPricingSnapshot,
|
||||
@@ -378,6 +474,307 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn image_usage_uses_image_count_for_request_cost() {
|
||||
let lookup = TestLookup {
|
||||
name_context: Some(
|
||||
StoredBillingModelContext::new(
|
||||
"provider-1".to_string(),
|
||||
Some("pay_as_you_go".to_string()),
|
||||
Some("key-1".to_string()),
|
||||
None,
|
||||
None,
|
||||
"global-image-1".to_string(),
|
||||
"gpt-image-2".to_string(),
|
||||
None,
|
||||
Some(0.02),
|
||||
None,
|
||||
Some("model-image-1".to_string()),
|
||||
Some("gpt-image-2".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("billing context should build"),
|
||||
),
|
||||
model_id_context: None,
|
||||
};
|
||||
let mut event = UsageEvent::new(
|
||||
UsageEventType::Completed,
|
||||
"req-image-billing-1",
|
||||
UsageEventData {
|
||||
provider_name: "OpenAI Image".to_string(),
|
||||
model: "gpt-image-2".to_string(),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
provider_api_key_id: Some("key-1".to_string()),
|
||||
request_type: Some("chat".to_string()),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
endpoint_api_format: Some("openai:image".to_string()),
|
||||
request_metadata: Some(json!({
|
||||
"dimensions": {
|
||||
"image_count": 3
|
||||
}
|
||||
})),
|
||||
status_code: Some(200),
|
||||
..UsageEventData::default()
|
||||
},
|
||||
);
|
||||
|
||||
enrich_usage_event_with_billing(&lookup, &mut event)
|
||||
.await
|
||||
.expect("billing should succeed");
|
||||
|
||||
assert_eq!(event.data.total_cost_usd, Some(0.06));
|
||||
assert_eq!(event.data.actual_total_cost_usd, Some(0.06));
|
||||
assert_eq!(
|
||||
event
|
||||
.data
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("billing_dimensions"))
|
||||
.and_then(|value| value.get("request_count"))
|
||||
.and_then(Value::as_i64),
|
||||
Some(3)
|
||||
);
|
||||
assert_eq!(
|
||||
event
|
||||
.data
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("billing_dimensions"))
|
||||
.and_then(|value| value.get("image_count"))
|
||||
.and_then(Value::as_i64),
|
||||
Some(3)
|
||||
);
|
||||
assert_eq!(
|
||||
event
|
||||
.data
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("billing_dimensions"))
|
||||
.and_then(|value| value.get("effective_task_type"))
|
||||
.and_then(Value::as_str),
|
||||
Some("image")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn image_usage_uses_configured_output_price_matrix() {
|
||||
let lookup = TestLookup {
|
||||
name_context: Some(
|
||||
StoredBillingModelContext::new(
|
||||
"provider-1".to_string(),
|
||||
Some("pay_as_you_go".to_string()),
|
||||
Some("key-1".to_string()),
|
||||
None,
|
||||
None,
|
||||
"global-image-1".to_string(),
|
||||
"gpt-image-2".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(json!({
|
||||
"tiers": [{
|
||||
"up_to": null,
|
||||
"input_price_per_1m": 5.0,
|
||||
"output_price_per_1m": 30.0,
|
||||
"cache_read_price_per_1m": 1.25
|
||||
}],
|
||||
"image_output_price_default": 0.01,
|
||||
"image_output_prices": {
|
||||
"1024x1024": {"low": 0.006, "medium": 0.053, "high": 0.211},
|
||||
"1536x1024": {"low": 0.005, "medium": 0.041, "high": 0.165},
|
||||
"1024x1536": {"low": 0.005, "medium": 0.041, "high": 0.165}
|
||||
}
|
||||
})),
|
||||
Some("model-image-1".to_string()),
|
||||
Some("gpt-image-2".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("billing context should build"),
|
||||
),
|
||||
model_id_context: None,
|
||||
};
|
||||
let mut event = UsageEvent::new(
|
||||
UsageEventType::Completed,
|
||||
"req-image-billing-matrix-1",
|
||||
UsageEventData {
|
||||
provider_name: "OpenAI Image".to_string(),
|
||||
model: "gpt-image-2".to_string(),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
provider_api_key_id: Some("key-1".to_string()),
|
||||
request_type: Some("chat".to_string()),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
endpoint_api_format: Some("openai:image".to_string()),
|
||||
request_metadata: Some(json!({
|
||||
"dimensions": {
|
||||
"image_count": 2,
|
||||
"image_size": "1536x1024",
|
||||
"image_quality": "medium",
|
||||
"image_output_format": "png"
|
||||
}
|
||||
})),
|
||||
status_code: Some(200),
|
||||
..UsageEventData::default()
|
||||
},
|
||||
);
|
||||
|
||||
enrich_usage_event_with_billing(&lookup, &mut event)
|
||||
.await
|
||||
.expect("billing should succeed");
|
||||
|
||||
assert_eq!(event.data.total_cost_usd, Some(0.082));
|
||||
assert_eq!(event.data.actual_total_cost_usd, Some(0.082));
|
||||
let metadata = event.data.request_metadata.as_ref().expect("metadata");
|
||||
assert_eq!(
|
||||
metadata
|
||||
.get("billing_dimensions")
|
||||
.and_then(|value| value.get("image_price_key"))
|
||||
.and_then(Value::as_str),
|
||||
Some("1536x1024:medium")
|
||||
);
|
||||
assert_eq!(
|
||||
metadata
|
||||
.get("billing_snapshot")
|
||||
.and_then(|value| value.get("resolved_variables"))
|
||||
.and_then(|value| value.get("image_output_price_per_image"))
|
||||
.and_then(Value::as_f64),
|
||||
Some(0.041)
|
||||
);
|
||||
assert_eq!(
|
||||
metadata
|
||||
.get("billing_snapshot")
|
||||
.and_then(|value| value.get("cost_breakdown"))
|
||||
.and_then(|value| value.get("image_output_cost"))
|
||||
.and_then(Value::as_f64),
|
||||
Some(0.082)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn enriches_cancelled_usage_event_with_billing_snapshot() {
|
||||
let lookup = TestLookup {
|
||||
name_context: Some(
|
||||
StoredBillingModelContext::new(
|
||||
"provider-1".to_string(),
|
||||
Some("pay_as_you_go".to_string()),
|
||||
Some("key-1".to_string()),
|
||||
Some(json!({"openai:responses": 0.5})),
|
||||
Some(60),
|
||||
"global-model-1".to_string(),
|
||||
"gpt-5".to_string(),
|
||||
None,
|
||||
Some(0.02),
|
||||
Some(json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0,"cache_creation_price_per_1m":3.75,"cache_read_price_per_1m":0.30}]})),
|
||||
Some("model-1".to_string()),
|
||||
Some("gpt-5-upstream".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("billing context should build"),
|
||||
),
|
||||
model_id_context: None,
|
||||
};
|
||||
let mut event = UsageEvent::new(
|
||||
UsageEventType::Cancelled,
|
||||
"req-billing-cancelled-1",
|
||||
UsageEventData {
|
||||
provider_name: "OpenAI".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
provider_api_key_id: Some("key-1".to_string()),
|
||||
request_type: Some("chat".to_string()),
|
||||
api_format: Some("openai:responses".to_string()),
|
||||
endpoint_api_format: Some("openai:responses".to_string()),
|
||||
input_tokens: Some(1_000),
|
||||
output_tokens: Some(500),
|
||||
cache_read_input_tokens: Some(100),
|
||||
status_code: Some(499),
|
||||
..UsageEventData::default()
|
||||
},
|
||||
);
|
||||
|
||||
enrich_usage_event_with_billing(&lookup, &mut event)
|
||||
.await
|
||||
.expect("billing should succeed");
|
||||
|
||||
assert!(event.data.total_cost_usd.unwrap_or_default() > 0.0);
|
||||
assert!(event.data.actual_total_cost_usd.unwrap_or_default() > 0.0);
|
||||
assert_eq!(
|
||||
event
|
||||
.data
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("billing_snapshot"))
|
||||
.and_then(|value| value.get("status"))
|
||||
.and_then(Value::as_str),
|
||||
Some("complete")
|
||||
);
|
||||
assert_eq!(
|
||||
event
|
||||
.data
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("billing_dimensions"))
|
||||
.and_then(|value| value.get("request_count"))
|
||||
.and_then(Value::as_i64),
|
||||
Some(0)
|
||||
);
|
||||
assert_eq!(
|
||||
event
|
||||
.data
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("billing_dimensions"))
|
||||
.and_then(|value| value.get("input_tokens"))
|
||||
.and_then(Value::as_i64),
|
||||
Some(900)
|
||||
);
|
||||
assert_eq!(
|
||||
event
|
||||
.data
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("billing_dimensions"))
|
||||
.and_then(|value| value.get("cache_read_tokens"))
|
||||
.and_then(Value::as_i64),
|
||||
Some(100)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn failed_usage_event_remains_unbilled() {
|
||||
let lookup = TestLookup {
|
||||
name_context: None,
|
||||
model_id_context: None,
|
||||
};
|
||||
let mut event = UsageEvent::new(
|
||||
UsageEventType::Failed,
|
||||
"req-billing-failed-1",
|
||||
UsageEventData {
|
||||
provider_name: "OpenAI".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
provider_api_key_id: Some("key-1".to_string()),
|
||||
request_type: Some("chat".to_string()),
|
||||
input_tokens: Some(1_000),
|
||||
output_tokens: Some(500),
|
||||
status_code: Some(500),
|
||||
..UsageEventData::default()
|
||||
},
|
||||
);
|
||||
|
||||
enrich_usage_event_with_billing(&lookup, &mut event)
|
||||
.await
|
||||
.expect("billing should succeed");
|
||||
|
||||
assert_eq!(event.data.total_cost_usd, Some(0.0));
|
||||
assert_eq!(event.data.actual_total_cost_usd, Some(0.0));
|
||||
assert!(event.data.request_metadata.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn enriches_by_provider_model_id_before_name_fallback() {
|
||||
let blank_name_context = StoredBillingModelContext::new(
|
||||
|
||||
@@ -24,7 +24,7 @@ impl BillingModelPricingSnapshot {
|
||||
pub fn effective_tiered_pricing(&self) -> Option<&Value> {
|
||||
self.model_tiered_pricing
|
||||
.as_ref()
|
||||
.filter(|value| has_tiered_pricing_tiers(value))
|
||||
.filter(|value| has_pricing_data(value))
|
||||
.or(self.default_tiered_pricing.as_ref())
|
||||
}
|
||||
|
||||
@@ -37,7 +37,7 @@ impl BillingModelPricingSnapshot {
|
||||
if self
|
||||
.model_tiered_pricing
|
||||
.as_ref()
|
||||
.is_some_and(has_tiered_pricing_tiers)
|
||||
.is_some_and(has_pricing_data)
|
||||
|| self.model_price_per_request.is_some()
|
||||
{
|
||||
"provider_override"
|
||||
@@ -75,11 +75,29 @@ impl BillingModelPricingSnapshot {
|
||||
}
|
||||
}
|
||||
|
||||
fn has_tiered_pricing_tiers(value: &Value) -> bool {
|
||||
fn has_pricing_data(value: &Value) -> bool {
|
||||
value
|
||||
.get("tiers")
|
||||
.and_then(Value::as_array)
|
||||
.is_some_and(|tiers| !tiers.is_empty())
|
||||
|| value
|
||||
.get("image_output_price_default")
|
||||
.and_then(Value::as_f64)
|
||||
.is_some()
|
||||
|| [
|
||||
"image_output_prices",
|
||||
"image_output_price_ranges",
|
||||
"image_output_price_per_image",
|
||||
"image_output_price_matrix",
|
||||
"image_prices",
|
||||
]
|
||||
.iter()
|
||||
.any(|key| value.get(key).is_some_and(value_has_entries))
|
||||
}
|
||||
|
||||
fn value_has_entries(value: &Value) -> bool {
|
||||
value.as_object().is_some_and(|object| !object.is_empty())
|
||||
|| value.as_array().is_some_and(|items| !items.is_empty())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -147,6 +165,10 @@ pub struct BillingUsageInput {
|
||||
pub cache_creation_ephemeral_5m_tokens: i64,
|
||||
pub cache_creation_ephemeral_1h_tokens: i64,
|
||||
pub cache_read_tokens: i64,
|
||||
pub image_count: i64,
|
||||
pub image_size: Option<String>,
|
||||
pub image_quality: Option<String>,
|
||||
pub image_output_format: Option<String>,
|
||||
pub cache_ttl_minutes: Option<i64>,
|
||||
}
|
||||
|
||||
@@ -162,6 +184,10 @@ impl BillingUsageInput {
|
||||
cache_creation_ephemeral_5m_tokens: 0,
|
||||
cache_creation_ephemeral_1h_tokens: 0,
|
||||
cache_read_tokens: 0,
|
||||
image_count: 0,
|
||||
image_size: None,
|
||||
image_quality: None,
|
||||
image_output_format: None,
|
||||
cache_ttl_minutes: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,15 +3,18 @@ use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::default_rule::{normalize_task_type, DefaultBillingRuleGenerator};
|
||||
use crate::default_rule::{
|
||||
explicit_image_output_price_default, explicit_image_output_price_entries,
|
||||
explicit_image_output_price_ranges, normalize_task_type, DefaultBillingRuleGenerator,
|
||||
};
|
||||
use crate::precision::quantize_cost;
|
||||
use crate::pricing::{BillingComputation, BillingModelPricingSnapshot, BillingUsageInput};
|
||||
use crate::schema::{
|
||||
BillingSnapshot, BillingSnapshotStatus, CostResult, BILLING_SNAPSHOT_SCHEMA_VERSION,
|
||||
};
|
||||
use crate::{
|
||||
normalize_input_tokens_for_billing, ExpressionEvaluationError, FormulaEngine,
|
||||
FormulaEvaluationStatus,
|
||||
normalize_input_tokens_for_billing, normalize_total_input_context_for_cache_hit_rate,
|
||||
ExpressionEvaluationError, FormulaEngine, FormulaEvaluationStatus,
|
||||
};
|
||||
|
||||
pub struct BillingService {
|
||||
@@ -43,7 +46,7 @@ impl BillingService {
|
||||
rule_name: None,
|
||||
scope: None,
|
||||
expression: None,
|
||||
resolved_dimensions: build_dimensions(input),
|
||||
resolved_dimensions: build_dimensions(input, pricing),
|
||||
resolved_variables: BTreeMap::new(),
|
||||
cost_breakdown: BTreeMap::new(),
|
||||
total_cost: 0.0,
|
||||
@@ -62,7 +65,7 @@ impl BillingService {
|
||||
});
|
||||
};
|
||||
|
||||
let dims = build_dimensions(input);
|
||||
let dims = build_dimensions(input, pricing);
|
||||
let result = self.engine.evaluate(
|
||||
&rule.expression,
|
||||
Some(&rule.variables),
|
||||
@@ -123,7 +126,10 @@ impl Default for BillingService {
|
||||
}
|
||||
}
|
||||
|
||||
fn build_dimensions(input: &BillingUsageInput) -> BTreeMap<String, Value> {
|
||||
fn build_dimensions(
|
||||
input: &BillingUsageInput,
|
||||
pricing: &BillingModelPricingSnapshot,
|
||||
) -> BTreeMap<String, Value> {
|
||||
let normalized_input_tokens = normalize_input_tokens_for_billing(
|
||||
input.api_format.as_deref(),
|
||||
input.input_tokens,
|
||||
@@ -136,10 +142,14 @@ fn build_dimensions(input: &BillingUsageInput) -> BTreeMap<String, Value> {
|
||||
.cache_creation_tokens
|
||||
.saturating_sub(classified_cache_creation_tokens)
|
||||
.max(0);
|
||||
let total_input_context = input
|
||||
.input_tokens
|
||||
.saturating_add(input.cache_creation_tokens)
|
||||
.saturating_add(input.cache_read_tokens);
|
||||
let total_input_context = normalize_total_input_context_for_cache_hit_rate(
|
||||
input.api_format.as_deref(),
|
||||
input.input_tokens,
|
||||
input.cache_creation_tokens,
|
||||
input.cache_read_tokens,
|
||||
);
|
||||
let image_output_pricing = image_output_pricing_state(pricing);
|
||||
let image_output_resolution = resolve_image_output_price_resolution(pricing, input);
|
||||
|
||||
let mut out = BTreeMap::from([
|
||||
("input_tokens".to_string(), json!(normalized_input_tokens)),
|
||||
@@ -168,6 +178,35 @@ fn build_dimensions(input: &BillingUsageInput) -> BTreeMap<String, Value> {
|
||||
"request_count".to_string(),
|
||||
json!(input.request_count.max(0)),
|
||||
),
|
||||
("image_count".to_string(), json!(input.image_count.max(0))),
|
||||
(
|
||||
"image_count_unmetered".to_string(),
|
||||
json!(if image_output_pricing.enabled {
|
||||
input.image_count.max(0)
|
||||
} else {
|
||||
0
|
||||
}),
|
||||
),
|
||||
(
|
||||
"image_output_pricing_enabled".to_string(),
|
||||
json!(image_output_pricing.enabled),
|
||||
),
|
||||
(
|
||||
"image_output_matrix_enabled".to_string(),
|
||||
json!(image_output_pricing.matrix_enabled),
|
||||
),
|
||||
(
|
||||
"image_output_range_enabled".to_string(),
|
||||
json!(image_output_pricing.range_enabled),
|
||||
),
|
||||
(
|
||||
"image_output_pricing_mode".to_string(),
|
||||
json!(image_output_resolution.pricing_mode),
|
||||
),
|
||||
(
|
||||
"image_output_price_per_image".to_string(),
|
||||
json!(image_output_resolution.price_per_image),
|
||||
),
|
||||
(
|
||||
"total_input_context".to_string(),
|
||||
json!(total_input_context),
|
||||
@@ -193,9 +232,346 @@ fn build_dimensions(input: &BillingUsageInput) -> BTreeMap<String, Value> {
|
||||
json!(cache_ttl_minutes.max(0)),
|
||||
);
|
||||
}
|
||||
if let Some(image_pixels) = image_output_resolution.image_pixels {
|
||||
out.insert("image_pixels".to_string(), json!(image_pixels));
|
||||
}
|
||||
if let Some(price_bucket) = image_output_resolution.price_bucket.as_ref() {
|
||||
out.insert("image_output_price_bucket".to_string(), json!(price_bucket));
|
||||
}
|
||||
if input.image_count > 0 {
|
||||
let image_size = input
|
||||
.image_size
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let image_quality = input
|
||||
.image_quality
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
if let Some(image_size) = image_size.as_ref() {
|
||||
out.insert("image_size".to_string(), json!(image_size));
|
||||
}
|
||||
if let Some(image_quality) = image_quality.as_ref() {
|
||||
out.insert("image_quality".to_string(), json!(image_quality));
|
||||
}
|
||||
if let (Some(image_size), Some(image_quality)) =
|
||||
(image_size.as_ref(), image_quality.as_ref())
|
||||
{
|
||||
out.insert(
|
||||
"image_price_key".to_string(),
|
||||
json!(format!(
|
||||
"{}:{}",
|
||||
normalize_image_output_size(image_size),
|
||||
normalize_image_output_quality(image_quality)
|
||||
)),
|
||||
);
|
||||
}
|
||||
}
|
||||
if let Some(output_format) = input
|
||||
.image_output_format
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
out.insert("image_output_format".to_string(), json!(output_format));
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct ImageOutputPricingState {
|
||||
enabled: bool,
|
||||
matrix_enabled: bool,
|
||||
range_enabled: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct ImageOutputPriceResolution {
|
||||
price_per_image: f64,
|
||||
pricing_mode: &'static str,
|
||||
price_bucket: Option<String>,
|
||||
image_pixels: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct ParsedImageOutputPriceRange {
|
||||
up_to_pixels: Option<i64>,
|
||||
label: Option<String>,
|
||||
prices: BTreeMap<String, f64>,
|
||||
}
|
||||
|
||||
fn image_output_pricing_state(pricing: &BillingModelPricingSnapshot) -> ImageOutputPricingState {
|
||||
let matrix_enabled = pricing_has_image_output_matrix(pricing);
|
||||
let range_enabled = pricing_has_image_output_ranges(pricing);
|
||||
let default_enabled = pricing_has_image_output_default_price(pricing);
|
||||
ImageOutputPricingState {
|
||||
enabled: matrix_enabled || range_enabled || default_enabled,
|
||||
matrix_enabled,
|
||||
range_enabled,
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_image_output_price_resolution(
|
||||
pricing: &BillingModelPricingSnapshot,
|
||||
input: &BillingUsageInput,
|
||||
) -> ImageOutputPriceResolution {
|
||||
let pricing_config = pricing.effective_tiered_pricing();
|
||||
let default_price = explicit_image_output_price_default(pricing_config);
|
||||
let image_size = input
|
||||
.image_size
|
||||
.as_deref()
|
||||
.map(normalize_image_output_size)
|
||||
.filter(|value| !value.is_empty());
|
||||
let image_quality = input
|
||||
.image_quality
|
||||
.as_deref()
|
||||
.map(normalize_image_output_quality)
|
||||
.filter(|value| !value.is_empty());
|
||||
let image_pixels = image_size.as_deref().and_then(parse_image_size_pixels);
|
||||
|
||||
if let (Some(size), Some(entries)) = (
|
||||
image_size.as_deref(),
|
||||
explicit_image_output_price_entries(pricing_config),
|
||||
) {
|
||||
for key in image_price_lookup_keys(size, image_quality.as_deref()) {
|
||||
if let Some(price) = entries.get(&key).and_then(Value::as_f64) {
|
||||
return ImageOutputPriceResolution {
|
||||
price_per_image: price,
|
||||
pricing_mode: "matrix",
|
||||
price_bucket: None,
|
||||
image_pixels,
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(pixels) = image_pixels {
|
||||
if let Some((price, bucket)) = resolve_image_output_range_price(
|
||||
explicit_image_output_price_ranges(pricing_config).unwrap_or_default(),
|
||||
pixels,
|
||||
image_quality.as_deref(),
|
||||
default_price,
|
||||
) {
|
||||
return ImageOutputPriceResolution {
|
||||
price_per_image: price,
|
||||
pricing_mode: "pixel_tiers",
|
||||
price_bucket: Some(bucket),
|
||||
image_pixels,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(price) = default_price {
|
||||
return ImageOutputPriceResolution {
|
||||
price_per_image: price,
|
||||
pricing_mode: "per_image",
|
||||
price_bucket: Some("default".to_string()),
|
||||
image_pixels,
|
||||
};
|
||||
}
|
||||
|
||||
ImageOutputPriceResolution {
|
||||
price_per_image: 0.0,
|
||||
pricing_mode: "none",
|
||||
price_bucket: None,
|
||||
image_pixels,
|
||||
}
|
||||
}
|
||||
|
||||
fn pricing_has_image_output_matrix(pricing: &BillingModelPricingSnapshot) -> bool {
|
||||
let Some(config) = pricing.effective_tiered_pricing() else {
|
||||
return false;
|
||||
};
|
||||
[
|
||||
"image_output_prices",
|
||||
"image_output_price_per_image",
|
||||
"image_output_price_matrix",
|
||||
"image_prices",
|
||||
]
|
||||
.iter()
|
||||
.any(|key| {
|
||||
config
|
||||
.get(key)
|
||||
.is_some_and(image_price_entries_have_matrix_values)
|
||||
})
|
||||
}
|
||||
|
||||
fn pricing_has_image_output_ranges(pricing: &BillingModelPricingSnapshot) -> bool {
|
||||
explicit_image_output_price_ranges(pricing.effective_tiered_pricing())
|
||||
.is_some_and(|ranges| !ranges.is_empty())
|
||||
}
|
||||
|
||||
fn pricing_has_image_output_default_price(pricing: &BillingModelPricingSnapshot) -> bool {
|
||||
let Some(config) = pricing.effective_tiered_pricing() else {
|
||||
return false;
|
||||
};
|
||||
config
|
||||
.get("image_output_price_default")
|
||||
.or_else(|| config.get("image_price_default"))
|
||||
.or_else(|| {
|
||||
config
|
||||
.get("image_output_prices")
|
||||
.and_then(|value| value.get("default"))
|
||||
})
|
||||
.and_then(Value::as_f64)
|
||||
.is_some()
|
||||
}
|
||||
|
||||
fn image_price_entries_have_matrix_values(value: &Value) -> bool {
|
||||
match value {
|
||||
Value::Object(object) => object.iter().any(|(key, value)| {
|
||||
!key.eq_ignore_ascii_case("default")
|
||||
&& (value.as_f64().is_some() || image_price_entries_have_matrix_values(value))
|
||||
}),
|
||||
Value::Array(items) => items.iter().any(image_price_entries_have_matrix_values),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_image_output_range_price(
|
||||
ranges: Vec<Value>,
|
||||
image_pixels: i64,
|
||||
image_quality: Option<&str>,
|
||||
default_price: Option<f64>,
|
||||
) -> Option<(f64, String)> {
|
||||
let mut parsed_ranges = ranges
|
||||
.iter()
|
||||
.filter_map(parse_image_output_price_range)
|
||||
.collect::<Vec<_>>();
|
||||
parsed_ranges.sort_by(
|
||||
|left, right| match (left.up_to_pixels, right.up_to_pixels) {
|
||||
(Some(left), Some(right)) => left.cmp(&right),
|
||||
(Some(_), None) => std::cmp::Ordering::Less,
|
||||
(None, Some(_)) => std::cmp::Ordering::Greater,
|
||||
(None, None) => std::cmp::Ordering::Equal,
|
||||
},
|
||||
);
|
||||
|
||||
for range in parsed_ranges {
|
||||
if !range
|
||||
.up_to_pixels
|
||||
.map(|up_to| image_pixels <= up_to)
|
||||
.unwrap_or(true)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let Some(price) =
|
||||
image_output_price_for_quality(&range.prices, image_quality).or(default_price)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
return Some((price, image_output_range_bucket(&range)));
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn parse_image_output_price_range(value: &Value) -> Option<ParsedImageOutputPriceRange> {
|
||||
let object = value.as_object()?;
|
||||
let prices = object
|
||||
.get("prices")
|
||||
.and_then(Value::as_object)?
|
||||
.iter()
|
||||
.filter_map(|(key, value)| {
|
||||
value
|
||||
.as_f64()
|
||||
.map(|price| (key.to_ascii_lowercase(), price))
|
||||
})
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
if prices.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(ParsedImageOutputPriceRange {
|
||||
up_to_pixels: object.get("up_to_pixels").and_then(value_as_positive_i64),
|
||||
label: object
|
||||
.get("label")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
prices,
|
||||
})
|
||||
}
|
||||
|
||||
fn image_output_price_for_quality(
|
||||
prices: &BTreeMap<String, f64>,
|
||||
image_quality: Option<&str>,
|
||||
) -> Option<f64> {
|
||||
for key in image_quality_lookup_keys(image_quality) {
|
||||
if let Some(price) = prices.get(&key) {
|
||||
return Some(*price);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn image_price_lookup_keys(size: &str, image_quality: Option<&str>) -> Vec<String> {
|
||||
image_quality_lookup_keys(image_quality)
|
||||
.into_iter()
|
||||
.filter(|quality| quality != "default")
|
||||
.map(|quality| format!("{}:{}", size, quality))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn image_quality_lookup_keys(image_quality: Option<&str>) -> Vec<String> {
|
||||
let quality = image_quality
|
||||
.map(normalize_image_output_quality)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or_else(|| "medium".to_string());
|
||||
let mut keys = vec![quality.clone()];
|
||||
if quality == "auto" {
|
||||
keys.push("medium".to_string());
|
||||
}
|
||||
keys.push("default".to_string());
|
||||
keys
|
||||
}
|
||||
|
||||
fn image_output_range_bucket(range: &ParsedImageOutputPriceRange) -> String {
|
||||
range
|
||||
.label
|
||||
.clone()
|
||||
.unwrap_or_else(|| match range.up_to_pixels {
|
||||
Some(up_to_pixels) => format!("<={up_to_pixels}px"),
|
||||
None => "unbounded".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn normalize_image_output_size(value: &str) -> String {
|
||||
value
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
.replace('×', "x")
|
||||
.chars()
|
||||
.filter(|ch| !ch.is_whitespace())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn normalize_image_output_quality(value: &str) -> String {
|
||||
value.trim().to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn parse_image_size_pixels(size: &str) -> Option<i64> {
|
||||
let (width, height) = size.split_once('x')?;
|
||||
let width = width.parse::<i64>().ok()?;
|
||||
let height = height.parse::<i64>().ok()?;
|
||||
if width <= 0 || height <= 0 {
|
||||
return None;
|
||||
}
|
||||
width.checked_mul(height)
|
||||
}
|
||||
|
||||
fn value_as_positive_i64(value: &Value) -> Option<i64> {
|
||||
let parsed = value
|
||||
.as_i64()
|
||||
.or_else(|| value.as_u64().and_then(|value| i64::try_from(value).ok()))
|
||||
.or_else(|| value.as_f64().map(|value| value as i64))
|
||||
.or_else(|| value.as_str().and_then(|value| value.trim().parse().ok()))?;
|
||||
(parsed > 0).then_some(parsed)
|
||||
}
|
||||
|
||||
fn now_marker() -> String {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
@@ -254,6 +630,10 @@ mod tests {
|
||||
cache_creation_ephemeral_5m_tokens: 0,
|
||||
cache_creation_ephemeral_1h_tokens: 0,
|
||||
cache_read_tokens: 100,
|
||||
image_count: 0,
|
||||
image_size: None,
|
||||
image_quality: None,
|
||||
image_output_format: None,
|
||||
cache_ttl_minutes: Some(60),
|
||||
},
|
||||
)
|
||||
@@ -265,6 +645,414 @@ mod tests {
|
||||
assert_eq!(result.rate_multiplier, 0.5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_cache_hit_context_does_not_double_count_cache_read() {
|
||||
let result = BillingService::new()
|
||||
.calculate(
|
||||
&pricing(),
|
||||
&BillingUsageInput {
|
||||
task_type: "chat".to_string(),
|
||||
api_format: Some("openai:responses".to_string()),
|
||||
request_count: 1,
|
||||
input_tokens: 1_000,
|
||||
output_tokens: 10,
|
||||
cache_creation_tokens: 0,
|
||||
cache_creation_ephemeral_5m_tokens: 0,
|
||||
cache_creation_ephemeral_1h_tokens: 0,
|
||||
cache_read_tokens: 800,
|
||||
image_count: 0,
|
||||
image_size: None,
|
||||
image_quality: None,
|
||||
image_output_format: None,
|
||||
cache_ttl_minutes: Some(60),
|
||||
},
|
||||
)
|
||||
.expect("billing should calculate");
|
||||
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.resolved_dimensions
|
||||
.get("input_tokens"),
|
||||
Some(&json!(200))
|
||||
);
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.resolved_dimensions
|
||||
.get("total_input_context"),
|
||||
Some(&json!(1_000))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_token_usage_without_image_output_price_bills_tokens_only() {
|
||||
let pricing = BillingModelPricingSnapshot {
|
||||
default_price_per_request: None,
|
||||
default_tiered_pricing: Some(json!({
|
||||
"tiers": [{
|
||||
"up_to": null,
|
||||
"input_price_per_1m": 1.0,
|
||||
"output_price_per_1m": 2.0
|
||||
}]
|
||||
})),
|
||||
..pricing()
|
||||
};
|
||||
|
||||
let result = BillingService::new()
|
||||
.calculate(
|
||||
&pricing,
|
||||
&BillingUsageInput {
|
||||
task_type: "image".to_string(),
|
||||
api_format: Some("openai:image".to_string()),
|
||||
request_count: 1,
|
||||
input_tokens: 1_000,
|
||||
output_tokens: 20_000,
|
||||
cache_creation_tokens: 0,
|
||||
cache_creation_ephemeral_5m_tokens: 0,
|
||||
cache_creation_ephemeral_1h_tokens: 0,
|
||||
cache_read_tokens: 0,
|
||||
image_count: 1,
|
||||
image_size: Some("1024x1024".to_string()),
|
||||
image_quality: Some("medium".to_string()),
|
||||
image_output_format: Some("png".to_string()),
|
||||
cache_ttl_minutes: None,
|
||||
},
|
||||
)
|
||||
.expect("billing should calculate");
|
||||
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.resolved_dimensions
|
||||
.get("image_output_pricing_mode"),
|
||||
Some(&json!("none"))
|
||||
);
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.resolved_dimensions
|
||||
.get("image_count_unmetered"),
|
||||
Some(&json!(0))
|
||||
);
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.cost_breakdown
|
||||
.get("image_output_cost"),
|
||||
Some(&0.0)
|
||||
);
|
||||
assert_eq!(result.cost_result.cost, 0.041);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_default_output_price_adds_image_cost_even_with_token_usage() {
|
||||
let pricing = BillingModelPricingSnapshot {
|
||||
default_price_per_request: None,
|
||||
default_tiered_pricing: Some(json!({
|
||||
"tiers": [{
|
||||
"up_to": null,
|
||||
"input_price_per_1m": 1.0,
|
||||
"output_price_per_1m": 2.0
|
||||
}],
|
||||
"image_output_price_default": 0.05
|
||||
})),
|
||||
..pricing()
|
||||
};
|
||||
|
||||
let result = BillingService::new()
|
||||
.calculate(
|
||||
&pricing,
|
||||
&BillingUsageInput {
|
||||
task_type: "image".to_string(),
|
||||
api_format: Some("openai:image".to_string()),
|
||||
request_count: 1,
|
||||
input_tokens: 1_000,
|
||||
output_tokens: 20_000,
|
||||
cache_creation_tokens: 0,
|
||||
cache_creation_ephemeral_5m_tokens: 0,
|
||||
cache_creation_ephemeral_1h_tokens: 0,
|
||||
cache_read_tokens: 0,
|
||||
image_count: 1,
|
||||
image_size: Some("1024x1024".to_string()),
|
||||
image_quality: Some("medium".to_string()),
|
||||
image_output_format: Some("png".to_string()),
|
||||
cache_ttl_minutes: None,
|
||||
},
|
||||
)
|
||||
.expect("billing should calculate");
|
||||
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.resolved_dimensions
|
||||
.get("image_output_pricing_mode"),
|
||||
Some(&json!("per_image"))
|
||||
);
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.cost_breakdown
|
||||
.get("image_output_cost"),
|
||||
Some(&0.05)
|
||||
);
|
||||
assert_eq!(result.cost_result.cost, 0.091);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_default_output_price_generates_rule_without_token_tiers() {
|
||||
let pricing = BillingModelPricingSnapshot {
|
||||
default_price_per_request: None,
|
||||
default_tiered_pricing: Some(json!({
|
||||
"image_output_price_default": 0.05
|
||||
})),
|
||||
..pricing()
|
||||
};
|
||||
|
||||
let result = BillingService::new()
|
||||
.calculate(
|
||||
&pricing,
|
||||
&BillingUsageInput {
|
||||
task_type: "image".to_string(),
|
||||
api_format: Some("openai:image".to_string()),
|
||||
request_count: 1,
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
cache_creation_tokens: 0,
|
||||
cache_creation_ephemeral_5m_tokens: 0,
|
||||
cache_creation_ephemeral_1h_tokens: 0,
|
||||
cache_read_tokens: 0,
|
||||
image_count: 2,
|
||||
image_size: Some("1024x1024".to_string()),
|
||||
image_quality: Some("medium".to_string()),
|
||||
image_output_format: Some("png".to_string()),
|
||||
cache_ttl_minutes: None,
|
||||
},
|
||||
)
|
||||
.expect("billing should calculate");
|
||||
|
||||
assert_eq!(result.cost_result.status, BillingSnapshotStatus::Complete);
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.cost_breakdown
|
||||
.get("image_output_cost"),
|
||||
Some(&0.1)
|
||||
);
|
||||
assert_eq!(result.cost_result.cost, 0.1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_pixel_ranges_generate_rule_without_token_tiers() {
|
||||
let pricing = BillingModelPricingSnapshot {
|
||||
default_price_per_request: None,
|
||||
default_tiered_pricing: Some(json!({
|
||||
"image_output_price_ranges": [{
|
||||
"up_to_pixels": null,
|
||||
"prices": { "medium": 0.04 }
|
||||
}]
|
||||
})),
|
||||
..pricing()
|
||||
};
|
||||
|
||||
let result = BillingService::new()
|
||||
.calculate(
|
||||
&pricing,
|
||||
&BillingUsageInput {
|
||||
task_type: "image".to_string(),
|
||||
api_format: Some("openai:image".to_string()),
|
||||
request_count: 1,
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
cache_creation_tokens: 0,
|
||||
cache_creation_ephemeral_5m_tokens: 0,
|
||||
cache_creation_ephemeral_1h_tokens: 0,
|
||||
cache_read_tokens: 0,
|
||||
image_count: 2,
|
||||
image_size: Some("1024x1024".to_string()),
|
||||
image_quality: Some("medium".to_string()),
|
||||
image_output_format: Some("png".to_string()),
|
||||
cache_ttl_minutes: None,
|
||||
},
|
||||
)
|
||||
.expect("billing should calculate");
|
||||
|
||||
assert_eq!(result.cost_result.status, BillingSnapshotStatus::Complete);
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.resolved_dimensions
|
||||
.get("image_output_pricing_mode"),
|
||||
Some(&json!("pixel_tiers"))
|
||||
);
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.cost_breakdown
|
||||
.get("image_output_cost"),
|
||||
Some(&0.08)
|
||||
);
|
||||
assert_eq!(result.cost_result.cost, 0.08);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_token_usage_with_matrix_adds_matrix_image_cost() {
|
||||
let pricing = BillingModelPricingSnapshot {
|
||||
default_price_per_request: None,
|
||||
default_tiered_pricing: Some(json!({
|
||||
"tiers": [{
|
||||
"up_to": null,
|
||||
"input_price_per_1m": 1.0,
|
||||
"output_price_per_1m": 2.0
|
||||
}],
|
||||
"image_output_price_default": 0.01,
|
||||
"image_output_prices": {
|
||||
"1024x1024": { "medium": 0.05 }
|
||||
}
|
||||
})),
|
||||
..pricing()
|
||||
};
|
||||
|
||||
let result = BillingService::new()
|
||||
.calculate(
|
||||
&pricing,
|
||||
&BillingUsageInput {
|
||||
task_type: "image".to_string(),
|
||||
api_format: Some("openai:image".to_string()),
|
||||
request_count: 1,
|
||||
input_tokens: 1_000,
|
||||
output_tokens: 20_000,
|
||||
cache_creation_tokens: 0,
|
||||
cache_creation_ephemeral_5m_tokens: 0,
|
||||
cache_creation_ephemeral_1h_tokens: 0,
|
||||
cache_read_tokens: 0,
|
||||
image_count: 1,
|
||||
image_size: Some("1024x1024".to_string()),
|
||||
image_quality: Some("medium".to_string()),
|
||||
image_output_format: Some("png".to_string()),
|
||||
cache_ttl_minutes: None,
|
||||
},
|
||||
)
|
||||
.expect("billing should calculate");
|
||||
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.resolved_dimensions
|
||||
.get("image_output_pricing_mode"),
|
||||
Some(&json!("matrix"))
|
||||
);
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.cost_breakdown
|
||||
.get("image_output_cost"),
|
||||
Some(&0.05)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_token_usage_with_pixel_ranges_adds_range_image_cost() {
|
||||
let pricing = BillingModelPricingSnapshot {
|
||||
default_price_per_request: None,
|
||||
default_tiered_pricing: Some(json!({
|
||||
"tiers": [{
|
||||
"up_to": null,
|
||||
"input_price_per_1m": 1.0,
|
||||
"output_price_per_1m": 2.0
|
||||
}],
|
||||
"image_output_price_default": 0.01,
|
||||
"image_output_price_ranges": [
|
||||
{
|
||||
"up_to_pixels": 1_048_576,
|
||||
"prices": { "medium": 0.04 }
|
||||
},
|
||||
{
|
||||
"up_to_pixels": 2_097_152,
|
||||
"prices": { "medium": 0.08 }
|
||||
}
|
||||
]
|
||||
})),
|
||||
..pricing()
|
||||
};
|
||||
|
||||
let result = BillingService::new()
|
||||
.calculate(
|
||||
&pricing,
|
||||
&BillingUsageInput {
|
||||
task_type: "image".to_string(),
|
||||
api_format: Some("openai:image".to_string()),
|
||||
request_count: 1,
|
||||
input_tokens: 1_000,
|
||||
output_tokens: 20_000,
|
||||
cache_creation_tokens: 0,
|
||||
cache_creation_ephemeral_5m_tokens: 0,
|
||||
cache_creation_ephemeral_1h_tokens: 0,
|
||||
cache_read_tokens: 0,
|
||||
image_count: 1,
|
||||
image_size: Some("1536 x 1024".to_string()),
|
||||
image_quality: Some("medium".to_string()),
|
||||
image_output_format: Some("png".to_string()),
|
||||
cache_ttl_minutes: None,
|
||||
},
|
||||
)
|
||||
.expect("billing should calculate");
|
||||
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.resolved_dimensions
|
||||
.get("image_output_pricing_mode"),
|
||||
Some(&json!("pixel_tiers"))
|
||||
);
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.resolved_dimensions
|
||||
.get("image_pixels"),
|
||||
Some(&json!(1_572_864))
|
||||
);
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.resolved_dimensions
|
||||
.get("image_output_price_bucket"),
|
||||
Some(&json!("<=2097152px"))
|
||||
);
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.resolved_variables
|
||||
.get("image_output_price_per_image"),
|
||||
Some(&json!(0.08))
|
||||
);
|
||||
assert_eq!(
|
||||
result
|
||||
.cost_result
|
||||
.snapshot
|
||||
.cost_breakdown
|
||||
.get("image_output_cost"),
|
||||
Some(&0.08)
|
||||
);
|
||||
assert_eq!(result.cost_result.cost, 0.121);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn five_minute_cache_ttl_uses_base_cache_prices() {
|
||||
let pricing = BillingModelPricingSnapshot {
|
||||
@@ -311,6 +1099,10 @@ mod tests {
|
||||
cache_creation_ephemeral_5m_tokens: 0,
|
||||
cache_creation_ephemeral_1h_tokens: 0,
|
||||
cache_read_tokens: 100,
|
||||
image_count: 0,
|
||||
image_size: None,
|
||||
image_quality: None,
|
||||
image_output_format: None,
|
||||
cache_ttl_minutes: Some(5),
|
||||
},
|
||||
)
|
||||
@@ -380,6 +1172,10 @@ mod tests {
|
||||
cache_creation_ephemeral_5m_tokens: 0,
|
||||
cache_creation_ephemeral_1h_tokens: 0,
|
||||
cache_read_tokens: 100,
|
||||
image_count: 0,
|
||||
image_size: None,
|
||||
image_quality: None,
|
||||
image_output_format: None,
|
||||
cache_ttl_minutes: Some(60),
|
||||
},
|
||||
)
|
||||
|
||||
@@ -4,8 +4,8 @@ use std::sync::{Arc, RwLock};
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::{
|
||||
plan_finite_wallet_debit, SettlementWriteRepository, StoredUsageSettlement,
|
||||
UsageSettlementInput, SETTLEMENT_EPSILON_USD,
|
||||
plan_finite_wallet_debit, settlement_billing_status_for_usage_status,
|
||||
SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
|
||||
};
|
||||
use crate::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
|
||||
use crate::DataLayerError;
|
||||
@@ -102,11 +102,8 @@ impl SettlementWriteRepository for InMemorySettlementRepository {
|
||||
})));
|
||||
}
|
||||
|
||||
let mut final_billing_status = if input.status == "completed" {
|
||||
"settled".to_string()
|
||||
} else {
|
||||
"void".to_string()
|
||||
};
|
||||
let mut final_billing_status =
|
||||
settlement_billing_status_for_usage_status(&input.status).to_string();
|
||||
let mut settlement = self.wallets.with_mut(|wallets| {
|
||||
let wallet_id = input
|
||||
.api_key_id
|
||||
@@ -306,6 +303,32 @@ mod tests {
|
||||
assert_eq!(settlement.wallet_balance_after, Some(9.0));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn settles_cancelled_usage_against_wallet_and_provider_quota() {
|
||||
let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]);
|
||||
let settlement = repository
|
||||
.settle_usage(UsageSettlementInput {
|
||||
request_id: "req-cancelled".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("key-1".to_string()),
|
||||
api_key_is_standalone: false,
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
status: "cancelled".to_string(),
|
||||
billing_status: "pending".to_string(),
|
||||
total_cost_usd: 3.0,
|
||||
actual_total_cost_usd: 1.5,
|
||||
finalized_at_unix_secs: Some(200),
|
||||
})
|
||||
.await
|
||||
.expect("settlement should succeed")
|
||||
.expect("settlement should exist");
|
||||
|
||||
assert_eq!(settlement.billing_status, "settled");
|
||||
assert_eq!(settlement.wallet_balance_before, Some(12.0));
|
||||
assert_eq!(settlement.wallet_balance_after, Some(9.0));
|
||||
assert_eq!(settlement.provider_monthly_used_usd, Some(1.5));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn standalone_key_settlement_never_falls_back_to_owner_wallet() {
|
||||
let repository = InMemorySettlementRepository::seed(vec![sample_user_wallet(
|
||||
|
||||
@@ -36,6 +36,31 @@ fn plan_finite_wallet_debit(
|
||||
}
|
||||
}
|
||||
|
||||
fn settlement_billing_status_for_usage_status(status: &str) -> &'static str {
|
||||
match status {
|
||||
"completed" | "cancelled" => "settled",
|
||||
_ => "void",
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::settlement_billing_status_for_usage_status;
|
||||
|
||||
#[test]
|
||||
fn cancelled_usage_status_is_billable() {
|
||||
assert_eq!(
|
||||
settlement_billing_status_for_usage_status("completed"),
|
||||
"settled"
|
||||
);
|
||||
assert_eq!(
|
||||
settlement_billing_status_for_usage_status("cancelled"),
|
||||
"settled"
|
||||
);
|
||||
assert_eq!(settlement_billing_status_for_usage_status("failed"), "void");
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(unused_imports)]
|
||||
pub(crate) use aether_data_contracts::repository::settlement::{
|
||||
SettlementRepository, SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput,
|
||||
|
||||
@@ -2,8 +2,9 @@ use async_trait::async_trait;
|
||||
use sqlx::{mysql::MySqlRow, Row};
|
||||
|
||||
use super::{
|
||||
finite_wallet_available_usd, plan_finite_wallet_debit, SettlementWriteRepository,
|
||||
StoredUsageSettlement, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
|
||||
finite_wallet_available_usd, plan_finite_wallet_debit,
|
||||
settlement_billing_status_for_usage_status, SettlementWriteRepository, StoredUsageSettlement,
|
||||
UsageSettlementInput, SETTLEMENT_EPSILON_USD,
|
||||
};
|
||||
use crate::driver::mysql::MysqlPool;
|
||||
use crate::error::SqlResultExt;
|
||||
@@ -364,11 +365,8 @@ impl SettlementWriteRepository for MysqlSettlementRepository {
|
||||
return Ok(Some(settlement));
|
||||
}
|
||||
|
||||
let mut final_billing_status = if input.status == "completed" {
|
||||
"settled".to_string()
|
||||
} else {
|
||||
"void".to_string()
|
||||
};
|
||||
let mut final_billing_status =
|
||||
settlement_billing_status_for_usage_status(&input.status).to_string();
|
||||
let mut settlement = StoredUsageSettlement {
|
||||
request_id: input.request_id.clone(),
|
||||
wallet_id: None,
|
||||
|
||||
@@ -2,8 +2,9 @@ use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::{
|
||||
finite_wallet_available_usd, plan_finite_wallet_debit, SettlementWriteRepository,
|
||||
StoredUsageSettlement, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
|
||||
finite_wallet_available_usd, plan_finite_wallet_debit,
|
||||
settlement_billing_status_for_usage_status, SettlementWriteRepository, StoredUsageSettlement,
|
||||
UsageSettlementInput, SETTLEMENT_EPSILON_USD,
|
||||
};
|
||||
use crate::driver::postgres::PostgresTransactionRunner;
|
||||
use crate::error::SqlxResultExt;
|
||||
@@ -455,11 +456,8 @@ impl SettlementWriteRepository for SqlxSettlementRepository {
|
||||
return settlement_from_row(&usage_row).map(Some);
|
||||
}
|
||||
|
||||
let mut final_billing_status = if input.status == "completed" {
|
||||
"settled".to_string()
|
||||
} else {
|
||||
"void".to_string()
|
||||
};
|
||||
let mut final_billing_status =
|
||||
settlement_billing_status_for_usage_status(&input.status).to_string();
|
||||
let finalized_at =
|
||||
i64::try_from(input.finalized_at_unix_secs.unwrap_or_else(|| {
|
||||
std::time::SystemTime::now()
|
||||
|
||||
@@ -2,8 +2,9 @@ use async_trait::async_trait;
|
||||
use sqlx::{sqlite::SqliteRow, Row};
|
||||
|
||||
use super::{
|
||||
finite_wallet_available_usd, plan_finite_wallet_debit, SettlementWriteRepository,
|
||||
StoredUsageSettlement, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
|
||||
finite_wallet_available_usd, plan_finite_wallet_debit,
|
||||
settlement_billing_status_for_usage_status, SettlementWriteRepository, StoredUsageSettlement,
|
||||
UsageSettlementInput, SETTLEMENT_EPSILON_USD,
|
||||
};
|
||||
use crate::driver::sqlite::{sqlite_optional_real, sqlite_real, SqlitePool};
|
||||
use crate::error::SqlResultExt;
|
||||
@@ -377,11 +378,8 @@ impl SettlementWriteRepository for SqliteSettlementRepository {
|
||||
return Ok(Some(settlement));
|
||||
}
|
||||
|
||||
let mut final_billing_status = if input.status == "completed" {
|
||||
"settled".to_string()
|
||||
} else {
|
||||
"void".to_string()
|
||||
};
|
||||
let mut final_billing_status =
|
||||
settlement_billing_status_for_usage_status(&input.status).to_string();
|
||||
let mut settlement = StoredUsageSettlement {
|
||||
request_id: input.request_id.clone(),
|
||||
wallet_id: None,
|
||||
|
||||
@@ -136,7 +136,7 @@ fn lifecycle_status_and_billing(event_type: UsageEventType) -> (&'static str, &'
|
||||
UsageEventType::Streaming => ("streaming", "pending"),
|
||||
UsageEventType::Completed => ("completed", "pending"),
|
||||
UsageEventType::Failed => ("failed", "void"),
|
||||
UsageEventType::Cancelled => ("cancelled", "void"),
|
||||
UsageEventType::Cancelled => ("cancelled", "pending"),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -181,6 +181,38 @@ mod tests {
|
||||
assert_eq!(record.finalized_at_unix_secs, Some(1_700_000_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancelled_terminal_record_stays_pending_for_settlement() {
|
||||
let record = build_upsert_usage_record_from_event(&UsageEvent {
|
||||
event_type: UsageEventType::Cancelled,
|
||||
request_id: "req-cancelled".to_string(),
|
||||
timestamp_ms: 1_700_000_000_000,
|
||||
data: UsageEventData {
|
||||
provider_name: "OpenAI".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
input_tokens: Some(10),
|
||||
output_tokens: Some(20),
|
||||
total_tokens: Some(30),
|
||||
total_cost_usd: Some(0.03),
|
||||
actual_total_cost_usd: Some(0.02),
|
||||
status_code: Some(499),
|
||||
response_time_ms: Some(200),
|
||||
first_byte_time_ms: Some(50),
|
||||
..UsageEventData::default()
|
||||
},
|
||||
})
|
||||
.expect("record should build");
|
||||
|
||||
assert_eq!(record.status, "cancelled");
|
||||
assert_eq!(record.billing_status, "pending");
|
||||
assert_eq!(record.total_tokens, Some(30));
|
||||
assert_eq!(record.total_cost_usd, Some(0.03));
|
||||
assert_eq!(record.actual_total_cost_usd, Some(0.02));
|
||||
assert_eq!(record.status_code, Some(499));
|
||||
assert_eq!(record.response_time_ms, Some(200));
|
||||
assert_eq!(record.first_byte_time_ms, Some(50));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitizes_request_metadata_before_building_upsert_record() {
|
||||
let record = build_upsert_usage_record_from_event(&UsageEvent {
|
||||
|
||||
@@ -164,6 +164,29 @@ mod tests {
|
||||
assert!(!inputs[0].api_key_is_standalone);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn settles_pending_cancelled_usage() {
|
||||
let writer = TestSettlementWriter {
|
||||
has_writer: true,
|
||||
..Default::default()
|
||||
};
|
||||
let mut usage = sample_usage();
|
||||
usage.status = "cancelled".to_string();
|
||||
usage.status_code = Some(499);
|
||||
|
||||
settle_usage_if_needed(&writer, &usage)
|
||||
.await
|
||||
.expect("settlement should succeed");
|
||||
|
||||
let inputs = writer.inputs.lock().expect("settlement inputs lock");
|
||||
assert_eq!(inputs.len(), 1);
|
||||
assert_eq!(inputs[0].request_id, "req-1");
|
||||
assert_eq!(inputs[0].status, "cancelled");
|
||||
assert_eq!(inputs[0].billing_status, "pending");
|
||||
assert_eq!(inputs[0].total_cost_usd, 1.25);
|
||||
assert_eq!(inputs[0].actual_total_cost_usd, 0.75);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn propagates_standalone_key_flag_from_usage_metadata() {
|
||||
let writer = TestSettlementWriter {
|
||||
|
||||
@@ -27,15 +27,21 @@ impl UsageMapper {
|
||||
}
|
||||
|
||||
derive_missing_input_tokens(raw_usage, api_format, &mut usage);
|
||||
copy_explicit_total_tokens(raw_usage, api_format, &mut usage);
|
||||
usage.normalize_cache_creation_breakdown()
|
||||
}
|
||||
|
||||
pub fn map_from_response(response: &serde_json::Value, api_format: &str) -> StandardizedUsage {
|
||||
let family = api_family(api_format);
|
||||
let Some(usage_value) = resolve_usage_value(response, family.as_str()) else {
|
||||
return StandardizedUsage::new();
|
||||
let mut usage = if let Some(usage_value) = resolve_usage_value(response, family.as_str()) {
|
||||
Self::map(usage_value, api_format, None)
|
||||
} else {
|
||||
StandardizedUsage::new()
|
||||
};
|
||||
Self::map(usage_value, api_format, None)
|
||||
if is_openai_image_api(api_format) {
|
||||
apply_openai_image_response_dimensions(response, &mut usage);
|
||||
}
|
||||
usage
|
||||
}
|
||||
}
|
||||
|
||||
@@ -59,6 +65,53 @@ fn api_family(api_format: &str) -> String {
|
||||
.to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn api_kind(api_format: &str) -> String {
|
||||
api_format
|
||||
.split(':')
|
||||
.nth(1)
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn is_openai_image_api(api_format: &str) -> bool {
|
||||
api_family(api_format).as_str() == "openai" && api_kind(api_format).as_str() == "image"
|
||||
}
|
||||
|
||||
fn apply_openai_image_response_dimensions(
|
||||
response: &serde_json::Value,
|
||||
usage: &mut StandardizedUsage,
|
||||
) {
|
||||
let image_count = openai_image_response_image_count(response);
|
||||
if image_count <= 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
usage.request_count = image_count;
|
||||
usage
|
||||
.dimensions
|
||||
.insert("image_count".to_string(), serde_json::json!(image_count));
|
||||
}
|
||||
|
||||
fn openai_image_response_image_count(response: &serde_json::Value) -> i64 {
|
||||
response
|
||||
.get("data")
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.map(|items| items.len() as i64)
|
||||
.filter(|value| *value > 0)
|
||||
.or_else(|| image_result_count(response.get("result")))
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn image_result_count(value: Option<&serde_json::Value>) -> Option<i64> {
|
||||
match value? {
|
||||
serde_json::Value::Array(items) => Some(items.len() as i64).filter(|count| *count > 0),
|
||||
serde_json::Value::Object(object) if !object.is_empty() => Some(1),
|
||||
serde_json::Value::String(text) if !text.trim().is_empty() => Some(1),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn base_mapping(api_format: &str) -> BTreeMap<String, String> {
|
||||
let mut mapping = BTreeMap::new();
|
||||
match api_family(api_format).as_str() {
|
||||
@@ -189,6 +242,22 @@ fn derive_missing_input_tokens(
|
||||
}
|
||||
}
|
||||
|
||||
fn copy_explicit_total_tokens(
|
||||
raw_usage: &serde_json::Value,
|
||||
api_format: &str,
|
||||
usage: &mut StandardizedUsage,
|
||||
) {
|
||||
let total_tokens = match api_family(api_format).as_str() {
|
||||
"gemini" => numeric_i64(raw_usage.get("totalTokenCount")),
|
||||
_ => numeric_i64(raw_usage.get("total_tokens")),
|
||||
};
|
||||
if let Some(total_tokens) = total_tokens.filter(|value| *value > 0) {
|
||||
usage
|
||||
.dimensions
|
||||
.insert("total_tokens".to_string(), serde_json::json!(total_tokens));
|
||||
}
|
||||
}
|
||||
|
||||
fn numeric_i64(value: Option<&serde_json::Value>) -> Option<i64> {
|
||||
value.and_then(|value| {
|
||||
value
|
||||
@@ -603,4 +672,47 @@ mod tests {
|
||||
assert_eq!(usage.output_tokens, 6);
|
||||
assert_eq!(usage.cache_read_tokens, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_openai_image_response_dimensions_without_usage() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"created": 1_700_000_000,
|
||||
"data": [
|
||||
{ "b64_json": "abc" },
|
||||
{ "url": "https://example.test/image.png" }
|
||||
]
|
||||
}),
|
||||
"openai:image",
|
||||
);
|
||||
|
||||
assert_eq!(usage.request_count, 2);
|
||||
assert_eq!(
|
||||
usage.dimensions.get("image_count"),
|
||||
Some(&serde_json::json!(2))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_openai_image_response_dimensions_with_native_usage() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"usage": {
|
||||
"input_tokens": 11,
|
||||
"output_tokens": 22,
|
||||
"total_tokens": 33
|
||||
},
|
||||
"data": [{ "b64_json": "abc" }]
|
||||
}),
|
||||
"openai:image",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 11);
|
||||
assert_eq!(usage.output_tokens, 22);
|
||||
assert_eq!(usage.request_count, 1);
|
||||
assert_eq!(
|
||||
usage.dimensions.get("image_count"),
|
||||
Some(&serde_json::json!(1))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -196,6 +196,7 @@ export interface RequestDetail {
|
||||
total_cost?: number
|
||||
cache_creation_cost?: number
|
||||
cache_read_cost?: number
|
||||
image_output_cost?: number
|
||||
request_cost?: number // 按次计费费用
|
||||
// Historical pricing fields (per 1M tokens)
|
||||
input_price_per_1m?: number
|
||||
|
||||
@@ -18,9 +18,20 @@ export interface PricingTier {
|
||||
cache_ttl_pricing?: CacheTTLPricing[]
|
||||
}
|
||||
|
||||
export type ImageOutputQuality = 'low' | 'medium' | 'high'
|
||||
|
||||
export interface ImageOutputPriceRange {
|
||||
up_to_pixels: number | null
|
||||
prices: Partial<Record<ImageOutputQuality, number>>
|
||||
label?: string | null
|
||||
}
|
||||
|
||||
/** 阶梯计费配置 */
|
||||
export interface TieredPricingConfig {
|
||||
tiers: PricingTier[]
|
||||
image_output_prices?: Record<string, Record<string, number>> | null
|
||||
image_output_price_default?: number | null
|
||||
image_output_price_ranges?: ImageOutputPriceRange[] | null
|
||||
}
|
||||
|
||||
export interface Model {
|
||||
|
||||
@@ -2,8 +2,8 @@
|
||||
<input
|
||||
type="checkbox"
|
||||
:class="checkboxClass"
|
||||
:checked="isChecked"
|
||||
v-bind="$attrs"
|
||||
:checked="isChecked"
|
||||
@change="handleChange"
|
||||
>
|
||||
</template>
|
||||
|
||||
@@ -230,6 +230,21 @@
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="flex items-start gap-2 border-t border-border/60 pt-3">
|
||||
<Checkbox
|
||||
:checked="isImageGenerationEnabled"
|
||||
class="mt-0.5"
|
||||
@update:checked="setImageGenerationEnabled"
|
||||
/>
|
||||
<div class="space-y-1">
|
||||
<div class="text-sm font-medium">
|
||||
图片模型
|
||||
</div>
|
||||
<p class="text-xs text-muted-foreground">
|
||||
启用图片输出计费,并展开尺寸 × 质量矩阵价格。
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
@@ -242,6 +257,7 @@
|
||||
ref="tieredPricingEditorRef"
|
||||
v-model="tieredPricing"
|
||||
:show-cache1h="true"
|
||||
:show-image-pricing="isImageGenerationEnabled"
|
||||
/>
|
||||
<div class="flex items-center gap-3 pt-2 border-t">
|
||||
<Label class="text-xs whitespace-nowrap">按次计费</Label>
|
||||
@@ -575,6 +591,7 @@ const defaultForm = (): FormData => ({
|
||||
})
|
||||
|
||||
const form = ref<FormData>(defaultForm())
|
||||
const imageGenerationExplicitOverride = ref<boolean | null>(null)
|
||||
|
||||
const isEmbeddingEnabled = computed(() => {
|
||||
return form.value.supported_capabilities?.includes('embedding') === true
|
||||
@@ -582,6 +599,18 @@ const isEmbeddingEnabled = computed(() => {
|
||||
|| form.value.config?.model_type === 'embedding'
|
||||
})
|
||||
|
||||
const isImageGenerationEnabled = computed(() => {
|
||||
if (imageGenerationExplicitOverride.value !== null) {
|
||||
return imageGenerationExplicitOverride.value
|
||||
}
|
||||
return form.value.supported_capabilities?.includes('image_generation') === true
|
||||
|| form.value.config?.image_generation === true
|
||||
|| form.value.config?.model_type === 'image'
|
||||
|| (Array.isArray(form.value.config?.api_formats)
|
||||
&& form.value.config.api_formats.some((format) => String(format).endsWith(':image')))
|
||||
|| tieredPricingHasImageOutputPricing(tieredPricing.value)
|
||||
})
|
||||
|
||||
const KEEP_FALSE_CONFIG_KEYS = new Set(['streaming'])
|
||||
|
||||
// 设置 config 字段
|
||||
@@ -624,6 +653,21 @@ function setEmbeddingEnabled(enabled: boolean) {
|
||||
form.value.supported_capabilities = [...caps]
|
||||
}
|
||||
|
||||
function setImageGenerationEnabled(value: boolean | 'indeterminate') {
|
||||
const enabled = value === true
|
||||
imageGenerationExplicitOverride.value = enabled
|
||||
const caps = new Set(form.value.supported_capabilities || [])
|
||||
if (enabled) {
|
||||
caps.add('image_generation')
|
||||
setConfigField('image_generation', true)
|
||||
} else {
|
||||
caps.delete('image_generation')
|
||||
setConfigField('image_generation', undefined)
|
||||
if (form.value.config?.model_type === 'image') setConfigField('model_type', undefined)
|
||||
}
|
||||
form.value.supported_capabilities = [...caps]
|
||||
}
|
||||
|
||||
function getNested(obj: unknown, path: string): unknown {
|
||||
if (!obj || typeof obj !== 'object') return undefined
|
||||
const parts = path.split('.').filter(Boolean)
|
||||
@@ -781,6 +825,7 @@ watch(() => props.open, (isOpen) => {
|
||||
|
||||
// 选择模型并填充表单
|
||||
function selectModel(model: ModelsDevModelItem) {
|
||||
imageGenerationExplicitOverride.value = null
|
||||
manualModelMode.value = false
|
||||
selectedModel.value = model
|
||||
expandedProvider.value = model.providerId
|
||||
@@ -806,7 +851,10 @@ function selectModel(model: ModelsDevModelItem) {
|
||||
if (model.inputModalities?.length) config.input_modalities = model.inputModalities
|
||||
if (model.outputModalities?.length) config.output_modalities = model.outputModalities
|
||||
form.value.config = config
|
||||
form.value.supported_capabilities = model.supportsEmbedding ? ['embedding'] : []
|
||||
const supportedCapabilities = new Set<string>()
|
||||
if (model.supportsEmbedding) supportedCapabilities.add('embedding')
|
||||
if (model.outputModalities?.includes('image')) supportedCapabilities.add('image_generation')
|
||||
form.value.supported_capabilities = [...supportedCapabilities]
|
||||
if (model.supportsEmbedding) {
|
||||
setEmbeddingEnabled(true)
|
||||
}
|
||||
@@ -827,6 +875,7 @@ function selectModel(model: ModelsDevModelItem) {
|
||||
|
||||
// 清除选择(手动填写)
|
||||
function clearSelection() {
|
||||
imageGenerationExplicitOverride.value = null
|
||||
manualModelMode.value = false
|
||||
selectedModel.value = null
|
||||
form.value = defaultForm()
|
||||
@@ -841,6 +890,7 @@ function handleLogoError(event: Event) {
|
||||
|
||||
// 重置表单
|
||||
function resetForm() {
|
||||
imageGenerationExplicitOverride.value = null
|
||||
form.value = defaultForm()
|
||||
tieredPricing.value = null
|
||||
videoResolutionPrices.value = []
|
||||
@@ -854,23 +904,30 @@ function resetForm() {
|
||||
// 加载模型数据(编辑模式)
|
||||
function loadModelData() {
|
||||
if (!props.model) return
|
||||
imageGenerationExplicitOverride.value = null
|
||||
// 先重置创建模式的残留状态
|
||||
selectedModel.value = null
|
||||
searchQuery.value = ''
|
||||
expandedProvider.value = null
|
||||
|
||||
const modelTieredPricing = props.model.default_tiered_pricing
|
||||
? JSON.parse(JSON.stringify(props.model.default_tiered_pricing))
|
||||
: null
|
||||
const supportedCapabilities = new Set(props.model.supported_capabilities || [])
|
||||
if (tieredPricingHasImageOutputPricing(modelTieredPricing)) {
|
||||
supportedCapabilities.add('image_generation')
|
||||
}
|
||||
|
||||
form.value = {
|
||||
name: props.model.name,
|
||||
display_name: props.model.display_name,
|
||||
default_price_per_request: props.model.default_price_per_request,
|
||||
supported_capabilities: [...(props.model.supported_capabilities || [])],
|
||||
supported_capabilities: [...supportedCapabilities],
|
||||
config: props.model.config ? { ...props.model.config } : { streaming: true },
|
||||
is_active: props.model.is_active,
|
||||
}
|
||||
// 确保 tieredPricing 也被正确设置或重置
|
||||
tieredPricing.value = props.model.default_tiered_pricing
|
||||
? JSON.parse(JSON.stringify(props.model.default_tiered_pricing))
|
||||
: null
|
||||
tieredPricing.value = modelTieredPricing
|
||||
loadVideoPricingFromConfig()
|
||||
}
|
||||
|
||||
@@ -898,8 +955,7 @@ async function handleSubmit() {
|
||||
return
|
||||
}
|
||||
|
||||
const finalTiers = tieredPricingEditorRef.value?.getFinalTiers()
|
||||
const finalTieredPricing = finalTiers ? { tiers: finalTiers } : tieredPricing.value
|
||||
const finalTieredPricing = tieredPricingEditorRef.value?.getFinalPricing() ?? tieredPricing.value
|
||||
|
||||
if (!finalTieredPricing?.tiers?.length) {
|
||||
showError('请配置至少一个价格阶梯')
|
||||
@@ -920,6 +976,9 @@ async function handleSubmit() {
|
||||
} else {
|
||||
caps.delete('cache_1h')
|
||||
}
|
||||
if (tieredPricingHasImageOutputPricing(finalTieredPricing)) {
|
||||
caps.add('image_generation')
|
||||
}
|
||||
form.value.supported_capabilities = caps.size > 0 ? [...caps] : []
|
||||
|
||||
// 清理空的 config
|
||||
@@ -949,4 +1008,29 @@ async function handleSubmit() {
|
||||
submitting.value = false
|
||||
}
|
||||
}
|
||||
|
||||
function tieredPricingHasImageOutputPricing(pricing: TieredPricingConfig | null | undefined): boolean {
|
||||
if (!pricing) return false
|
||||
if (toFinitePrice(pricing.image_output_price_default) !== null) return true
|
||||
if (Object.values(pricing.image_output_prices || {}).some((prices) => {
|
||||
if (!prices || typeof prices !== 'object') return false
|
||||
return Object.values(prices).some((price) => toFinitePrice(price) !== null)
|
||||
})) return true
|
||||
return (pricing.image_output_price_ranges || []).some((range) => {
|
||||
if (!range || typeof range !== 'object') return false
|
||||
const prices = range.prices && typeof range.prices === 'object'
|
||||
? range.prices
|
||||
: range as Record<string, unknown>
|
||||
return Object.values(prices).some((price) => toFinitePrice(price) !== null)
|
||||
})
|
||||
}
|
||||
|
||||
function toFinitePrice(value: unknown): number | null {
|
||||
if (typeof value === 'number' && Number.isFinite(value)) return value
|
||||
if (typeof value === 'string' && value.trim()) {
|
||||
const parsed = Number(value)
|
||||
return Number.isFinite(parsed) ? parsed : null
|
||||
}
|
||||
return null
|
||||
}
|
||||
</script>
|
||||
|
||||
@@ -137,6 +137,7 @@
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
</div>
|
||||
|
||||
<!-- 默认定价 -->
|
||||
@@ -145,6 +146,114 @@
|
||||
默认定价
|
||||
</h4>
|
||||
|
||||
<!-- 图片输出计费 -->
|
||||
<div
|
||||
v-if="hasImagePricing"
|
||||
class="space-y-2"
|
||||
>
|
||||
<div class="flex items-center justify-between gap-3 text-sm text-muted-foreground">
|
||||
<div class="flex items-center gap-2">
|
||||
<span>图片输出计费</span>
|
||||
<Badge
|
||||
v-if="imagePricingEntries.length > 0"
|
||||
variant="outline"
|
||||
class="text-[10px] h-5 px-1.5"
|
||||
>
|
||||
矩阵
|
||||
</Badge>
|
||||
<Badge
|
||||
v-if="imagePriceRangeEntries.length > 0"
|
||||
variant="outline"
|
||||
class="text-[10px] h-5 px-1.5"
|
||||
>
|
||||
区间
|
||||
</Badge>
|
||||
</div>
|
||||
<span
|
||||
v-if="imageOutputDefaultPrice !== null"
|
||||
class="text-xs font-mono"
|
||||
>默认 ${{ imageOutputDefaultPrice.toFixed(6) }}/张</span>
|
||||
</div>
|
||||
<div
|
||||
v-if="imagePricingEntries.length > 0"
|
||||
class="border rounded-lg overflow-hidden"
|
||||
>
|
||||
<Table>
|
||||
<TableHeader>
|
||||
<TableRow class="bg-muted/30">
|
||||
<TableHead class="text-xs h-9">
|
||||
分辨率
|
||||
</TableHead>
|
||||
<TableHead
|
||||
v-for="quality in IMAGE_OUTPUT_QUALITIES"
|
||||
:key="quality"
|
||||
class="text-xs h-9 text-right"
|
||||
>
|
||||
{{ quality }}
|
||||
</TableHead>
|
||||
</TableRow>
|
||||
</TableHeader>
|
||||
<TableBody>
|
||||
<TableRow
|
||||
v-for="entry in imagePricingEntries"
|
||||
:key="entry.size"
|
||||
class="text-xs"
|
||||
>
|
||||
<TableCell class="py-2 font-mono">
|
||||
{{ formatImageSize(entry.size) }}
|
||||
</TableCell>
|
||||
<TableCell
|
||||
v-for="quality in IMAGE_OUTPUT_QUALITIES"
|
||||
:key="`${entry.size}-${quality}`"
|
||||
class="py-2 text-right font-mono"
|
||||
>
|
||||
{{ formatImagePrice(entry.prices[quality]) }}
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
</TableBody>
|
||||
</Table>
|
||||
</div>
|
||||
<div
|
||||
v-if="imagePriceRangeEntries.length > 0"
|
||||
class="border rounded-lg overflow-hidden"
|
||||
>
|
||||
<Table>
|
||||
<TableHeader>
|
||||
<TableRow class="bg-muted/30">
|
||||
<TableHead class="text-xs h-9">
|
||||
上限像素
|
||||
</TableHead>
|
||||
<TableHead
|
||||
v-for="quality in IMAGE_OUTPUT_QUALITIES"
|
||||
:key="quality"
|
||||
class="text-xs h-9 text-right"
|
||||
>
|
||||
{{ quality }}
|
||||
</TableHead>
|
||||
</TableRow>
|
||||
</TableHeader>
|
||||
<TableBody>
|
||||
<TableRow
|
||||
v-for="entry in imagePriceRangeEntries"
|
||||
:key="entry.key"
|
||||
class="text-xs"
|
||||
>
|
||||
<TableCell class="py-2 font-mono">
|
||||
{{ formatPixelLimit(entry.upToPixels) }}
|
||||
</TableCell>
|
||||
<TableCell
|
||||
v-for="quality in IMAGE_OUTPUT_QUALITIES"
|
||||
:key="`${entry.key}-${quality}`"
|
||||
class="py-2 text-right font-mono"
|
||||
>
|
||||
{{ formatImagePrice(entry.prices[quality]) }}
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
</TableBody>
|
||||
</Table>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- 单阶梯(固定价格)展示 -->
|
||||
<div
|
||||
v-if="getTierCount(model.default_tiered_pricing) <= 1"
|
||||
@@ -561,6 +670,79 @@ const videoPricingEntries = computed(() => {
|
||||
return sortResolutionEntries(Object.entries(priceByResolution))
|
||||
})
|
||||
|
||||
const IMAGE_OUTPUT_QUALITIES = ['low', 'medium', 'high'] as const
|
||||
|
||||
const imageOutputDefaultPrice = computed(() => {
|
||||
const value = props.model?.default_tiered_pricing?.image_output_price_default
|
||||
return typeof value === 'number' && Number.isFinite(value) ? value : null
|
||||
})
|
||||
|
||||
const imagePricingEntries = computed(() => {
|
||||
const prices = props.model?.default_tiered_pricing?.image_output_prices
|
||||
if (!prices || typeof prices !== 'object') return []
|
||||
return sortResolutionEntries(Object.entries(prices)).map(([size, qualityPrices]) => ({
|
||||
size,
|
||||
prices: normalizeImageQualityPrices(qualityPrices),
|
||||
})).filter(entry => Object.values(entry.prices).some(price => price !== null))
|
||||
})
|
||||
|
||||
const imagePriceRangeEntries = computed(() => {
|
||||
const ranges = props.model?.default_tiered_pricing?.image_output_price_ranges
|
||||
if (!Array.isArray(ranges)) return []
|
||||
return ranges.map((range, index) => {
|
||||
const object = range && typeof range === 'object' ? range as Record<string, unknown> : {}
|
||||
const rawPrices = object.prices && typeof object.prices === 'object'
|
||||
? object.prices
|
||||
: object
|
||||
return {
|
||||
key: `${object.up_to_pixels ?? 'unbounded'}-${index}`,
|
||||
upToPixels: toFiniteNumber(object.up_to_pixels),
|
||||
prices: normalizeImageQualityPrices(rawPrices),
|
||||
}
|
||||
}).filter(entry => Object.values(entry.prices).some(price => price !== null))
|
||||
})
|
||||
|
||||
const hasImagePricing = computed(() =>
|
||||
imageOutputDefaultPrice.value !== null
|
||||
|| imagePricingEntries.value.length > 0
|
||||
|| imagePriceRangeEntries.value.length > 0,
|
||||
)
|
||||
|
||||
function normalizeImageQualityPrices(value: unknown): Record<typeof IMAGE_OUTPUT_QUALITIES[number], number | null> {
|
||||
const object = value && typeof value === 'object' ? value as Record<string, unknown> : {}
|
||||
return {
|
||||
low: toFiniteNumber(object.low),
|
||||
medium: toFiniteNumber(object.medium),
|
||||
high: toFiniteNumber(object.high),
|
||||
}
|
||||
}
|
||||
|
||||
function toFiniteNumber(value: unknown): number | null {
|
||||
return typeof value === 'number' && Number.isFinite(value) ? value : null
|
||||
}
|
||||
|
||||
function formatImagePrice(value: number | null): string {
|
||||
return value === null ? '-' : `$${value.toFixed(6)}`
|
||||
}
|
||||
|
||||
function formatImageSize(value: string): string {
|
||||
return value.replace(/\s*[xX×]\s*/g, ' x ')
|
||||
}
|
||||
|
||||
function formatPixelLimit(value: number | null): string {
|
||||
return value === null ? '无上限' : `<= ${formatPixels(value)}`
|
||||
}
|
||||
|
||||
function formatPixels(value: number): string {
|
||||
if (value >= 1_000_000) {
|
||||
return `${(value / 1_000_000).toFixed(value % 1_000_000 === 0 ? 0 : 2)}M px`
|
||||
}
|
||||
if (value >= 1_000) {
|
||||
return `${(value / 1_000).toFixed(0)}K px`
|
||||
}
|
||||
return `${value} px`
|
||||
}
|
||||
|
||||
const detailTab = ref('basic')
|
||||
|
||||
// 处理背景点击
|
||||
|
||||
@@ -784,7 +784,9 @@ function targetFormatsForEndpoint(
|
||||
provider: RoutingProviderInfo,
|
||||
endpoint: RoutingEndpointInfo
|
||||
): string[] {
|
||||
return STANDARD_ROUTING_API_FORMATS.filter(format =>
|
||||
const endpointFormat = normalizeLegacyOpenAIFormatAlias(endpoint.api_format)
|
||||
const candidateFormats = Array.from(new Set([...STANDARD_ROUTING_API_FORMATS, endpointFormat]))
|
||||
return candidateFormats.filter(format =>
|
||||
endpointSupportsClientFormat(provider, endpoint, format, endpoint.api_format)
|
||||
)
|
||||
}
|
||||
|
||||
@@ -137,6 +137,141 @@
|
||||
添加价格阶梯
|
||||
</Button>
|
||||
|
||||
<div
|
||||
v-if="showImagePricing"
|
||||
class="rounded-lg border bg-muted/10 p-3 space-y-3"
|
||||
>
|
||||
<div class="flex flex-wrap items-end justify-between gap-3">
|
||||
<Label class="text-xs font-medium">图像输出计费 ($/张)</Label>
|
||||
<div class="flex items-center gap-2">
|
||||
<Label class="text-xs text-muted-foreground">默认价</Label>
|
||||
<Input
|
||||
:model-value="imageOutputPriceDefault"
|
||||
type="number"
|
||||
step="0.001"
|
||||
min="0"
|
||||
class="h-8 w-24"
|
||||
placeholder="0"
|
||||
@update:model-value="updateImageOutputPriceDefault"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="space-y-2">
|
||||
<div class="flex items-center justify-between gap-2">
|
||||
<Label class="text-xs text-muted-foreground">精确分辨率覆盖</Label>
|
||||
<span class="text-[11px] text-muted-foreground">优先匹配 size + quality</span>
|
||||
</div>
|
||||
<div class="grid grid-cols-[minmax(120px,1.1fr)_repeat(3,minmax(0,1fr))_32px] gap-2 text-xs text-muted-foreground">
|
||||
<span>分辨率</span>
|
||||
<span>low</span>
|
||||
<span>medium</span>
|
||||
<span>high</span>
|
||||
<span />
|
||||
</div>
|
||||
<div
|
||||
v-for="row in imageOutputPriceRows"
|
||||
:key="row.id"
|
||||
class="grid grid-cols-[minmax(120px,1.1fr)_repeat(3,minmax(0,1fr))_32px] gap-2 items-center"
|
||||
>
|
||||
<Input
|
||||
:model-value="row.size"
|
||||
class="h-8 font-mono text-xs"
|
||||
placeholder="1024x1024"
|
||||
@update:model-value="(v) => updateImageOutputSize(row.id, v)"
|
||||
/>
|
||||
<Input
|
||||
v-for="quality in IMAGE_OUTPUT_QUALITIES"
|
||||
:key="`${row.id}-${quality}`"
|
||||
:model-value="getImageOutputPrice(row, quality)"
|
||||
type="number"
|
||||
step="0.001"
|
||||
min="0"
|
||||
class="h-8"
|
||||
placeholder="0"
|
||||
@update:model-value="(v) => updateImageOutputPrice(row.id, quality, v)"
|
||||
/>
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
size="sm"
|
||||
class="h-8 w-8 p-0"
|
||||
@click="removeImageOutputSizeRow(row.id)"
|
||||
>
|
||||
<X class="w-4 h-4 text-muted-foreground hover:text-destructive" />
|
||||
</Button>
|
||||
</div>
|
||||
<Button
|
||||
type="button"
|
||||
variant="outline"
|
||||
size="sm"
|
||||
class="w-full"
|
||||
@click="addImageOutputSizeRow"
|
||||
>
|
||||
<Plus class="w-4 h-4 mr-2" />
|
||||
添加分辨率
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
<div class="space-y-2 border-t pt-3">
|
||||
<div class="flex items-center justify-between gap-2">
|
||||
<Label class="text-xs text-muted-foreground">像素区间</Label>
|
||||
<span class="text-[11px] text-muted-foreground">矩阵未命中时按宽×高落档</span>
|
||||
</div>
|
||||
<div class="grid grid-cols-[minmax(120px,1.1fr)_repeat(3,minmax(0,1fr))_32px] gap-2 text-xs text-muted-foreground">
|
||||
<span>上限像素</span>
|
||||
<span>low</span>
|
||||
<span>medium</span>
|
||||
<span>high</span>
|
||||
<span />
|
||||
</div>
|
||||
<div
|
||||
v-for="row in imageOutputPriceRangeRows"
|
||||
:key="row.id"
|
||||
class="grid grid-cols-[minmax(120px,1.1fr)_repeat(3,minmax(0,1fr))_32px] gap-2 items-center"
|
||||
>
|
||||
<Input
|
||||
:model-value="row.upToPixels"
|
||||
type="number"
|
||||
min="1"
|
||||
class="h-8 font-mono text-xs"
|
||||
placeholder="空=无上限"
|
||||
@update:model-value="(v) => updateImageOutputRangeLimit(row.id, v)"
|
||||
/>
|
||||
<Input
|
||||
v-for="quality in IMAGE_OUTPUT_QUALITIES"
|
||||
:key="`${row.id}-${quality}`"
|
||||
:model-value="getImageOutputRangePrice(row, quality)"
|
||||
type="number"
|
||||
step="0.001"
|
||||
min="0"
|
||||
class="h-8"
|
||||
placeholder="0"
|
||||
@update:model-value="(v) => updateImageOutputRangePrice(row.id, quality, v)"
|
||||
/>
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
size="sm"
|
||||
class="h-8 w-8 p-0"
|
||||
@click="removeImageOutputRangeRow(row.id)"
|
||||
>
|
||||
<X class="w-4 h-4 text-muted-foreground hover:text-destructive" />
|
||||
</Button>
|
||||
</div>
|
||||
<Button
|
||||
type="button"
|
||||
variant="outline"
|
||||
size="sm"
|
||||
class="w-full"
|
||||
@click="addImageOutputRangeRow"
|
||||
>
|
||||
<Plus class="w-4 h-4 mr-2" />
|
||||
添加像素区间
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- 验证提示 -->
|
||||
<p
|
||||
v-if="validationError"
|
||||
@@ -151,11 +286,28 @@
|
||||
import { ref, computed, watch, reactive } from 'vue'
|
||||
import { Plus, X } from 'lucide-vue-next'
|
||||
import { Button, Input, Label } from '@/components/ui'
|
||||
import type { TieredPricingConfig, PricingTier } from '@/api/endpoints/types'
|
||||
import type { TieredPricingConfig, PricingTier, ImageOutputPriceRange } from '@/api/endpoints/types'
|
||||
|
||||
type ImageOutputQuality = 'low' | 'medium' | 'high'
|
||||
type ImageOutputPriceRow = {
|
||||
id: string
|
||||
size: string
|
||||
prices: Partial<Record<ImageOutputQuality, number>>
|
||||
}
|
||||
type ImageOutputPriceRangeRow = {
|
||||
id: string
|
||||
upToPixels: string
|
||||
prices: Partial<Record<ImageOutputQuality, number>>
|
||||
}
|
||||
|
||||
const DEFAULT_IMAGE_OUTPUT_SIZES = ['1024x1024', '1536x1024', '1024x1536']
|
||||
const DEFAULT_IMAGE_OUTPUT_PIXEL_LIMITS = [1_048_576, 1_572_864, 2_097_152]
|
||||
const IMAGE_OUTPUT_QUALITIES: ImageOutputQuality[] = ['low', 'medium', 'high']
|
||||
|
||||
const props = defineProps<{
|
||||
modelValue?: TieredPricingConfig | null
|
||||
showCache1h?: boolean
|
||||
showImagePricing?: boolean
|
||||
}>()
|
||||
|
||||
const emit = defineEmits<{
|
||||
@@ -164,6 +316,12 @@ const emit = defineEmits<{
|
||||
|
||||
// 本地状态
|
||||
const localTiers = ref<PricingTier[]>([])
|
||||
const imageOutputPriceRows = ref<ImageOutputPriceRow[]>([])
|
||||
const imageOutputPriceRangeRows = ref<ImageOutputPriceRangeRow[]>([])
|
||||
const imageOutputPriceDefault = ref<string>('')
|
||||
const lastEmittedPricingJson = ref<string>('')
|
||||
let imageOutputPriceRowId = 0
|
||||
let imageOutputPriceRangeRowId = 0
|
||||
|
||||
// 跟踪每个阶梯的缓存价格是否被手动设置
|
||||
const cacheManuallySet = reactive<Record<number, { creation: boolean; read: boolean; cache1h: boolean }>>({})
|
||||
@@ -186,8 +344,16 @@ const customInputValue = reactive<Record<number, string>>({})
|
||||
watch(
|
||||
() => props.modelValue,
|
||||
(newValue) => {
|
||||
if (lastEmittedPricingJson.value && JSON.stringify(newValue ?? null) === lastEmittedPricingJson.value) {
|
||||
return
|
||||
}
|
||||
if (newValue?.tiers) {
|
||||
localTiers.value = newValue.tiers.map(t => ({ ...t }))
|
||||
imageOutputPriceRows.value = createImageOutputPriceRows(newValue.image_output_prices)
|
||||
imageOutputPriceRangeRows.value = createImageOutputPriceRangeRows(newValue.image_output_price_ranges)
|
||||
imageOutputPriceDefault.value = newValue.image_output_price_default != null
|
||||
? String(newValue.image_output_price_default)
|
||||
: ''
|
||||
// 如果已有缓存价格,标记为手动设置
|
||||
newValue.tiers.forEach((t, i) => {
|
||||
const has1hCache = t.cache_ttl_pricing?.some(c => c.ttl_minutes === 60) ?? false
|
||||
@@ -203,6 +369,9 @@ watch(
|
||||
input_price_per_1m: 0,
|
||||
output_price_per_1m: 0,
|
||||
}]
|
||||
imageOutputPriceRows.value = createImageOutputPriceRows(null)
|
||||
imageOutputPriceRangeRows.value = createImageOutputPriceRangeRows(null)
|
||||
imageOutputPriceDefault.value = ''
|
||||
cacheManuallySet[0] = { creation: false, read: false, cache1h: false }
|
||||
}
|
||||
},
|
||||
@@ -367,7 +536,9 @@ function syncToParent() {
|
||||
return tier
|
||||
})
|
||||
|
||||
emit('update:modelValue', { tiers })
|
||||
const value = buildPricingConfig(tiers)
|
||||
lastEmittedPricingJson.value = JSON.stringify(value ?? null)
|
||||
emit('update:modelValue', value)
|
||||
}
|
||||
|
||||
// 获取最终提交的数据(包含自动计算的缓存价格)
|
||||
@@ -406,11 +577,239 @@ function getFinalTiers(): PricingTier[] {
|
||||
})
|
||||
}
|
||||
|
||||
function getFinalPricing(): TieredPricingConfig {
|
||||
return buildPricingConfig(getFinalTiers())
|
||||
}
|
||||
|
||||
// 暴露给父组件调用
|
||||
defineExpose({
|
||||
getFinalTiers,
|
||||
getFinalPricing,
|
||||
})
|
||||
|
||||
function buildPricingConfig(tiers: PricingTier[]): TieredPricingConfig {
|
||||
const config: TieredPricingConfig = { tiers }
|
||||
if (!props.showImagePricing) {
|
||||
return config
|
||||
}
|
||||
const matrix = normalizedImageOutputPrices()
|
||||
if (Object.keys(matrix).length > 0) {
|
||||
config.image_output_prices = matrix
|
||||
}
|
||||
const ranges = normalizedImageOutputPriceRanges()
|
||||
if (ranges.length > 0) {
|
||||
config.image_output_price_ranges = ranges
|
||||
}
|
||||
const defaultPrice = parseOptionalFloat(imageOutputPriceDefault.value)
|
||||
if (defaultPrice != null) {
|
||||
config.image_output_price_default = defaultPrice
|
||||
}
|
||||
return config
|
||||
}
|
||||
|
||||
function createImageOutputPriceRows(value: TieredPricingConfig['image_output_prices']): ImageOutputPriceRow[] {
|
||||
const rows: ImageOutputPriceRow[] = []
|
||||
if (!value || typeof value !== 'object') {
|
||||
return DEFAULT_IMAGE_OUTPUT_SIZES.map(size => createImageOutputPriceRow(size))
|
||||
}
|
||||
for (const [size, prices] of Object.entries(value)) {
|
||||
if (!prices || typeof prices !== 'object') continue
|
||||
const rowPrices: Partial<Record<ImageOutputQuality, number>> = {}
|
||||
for (const quality of IMAGE_OUTPUT_QUALITIES) {
|
||||
const price = (prices as Record<string, unknown>)[quality]
|
||||
if (typeof price === 'number' && Number.isFinite(price)) {
|
||||
rowPrices[quality] = price
|
||||
}
|
||||
}
|
||||
rows.push(createImageOutputPriceRow(size, rowPrices))
|
||||
}
|
||||
if (rows.length > 0) return rows
|
||||
return DEFAULT_IMAGE_OUTPUT_SIZES.map(size => createImageOutputPriceRow(size))
|
||||
}
|
||||
|
||||
function createImageOutputPriceRangeRows(value: TieredPricingConfig['image_output_price_ranges']): ImageOutputPriceRangeRow[] {
|
||||
const rows: ImageOutputPriceRangeRow[] = []
|
||||
if (!Array.isArray(value)) {
|
||||
return rows
|
||||
}
|
||||
for (const range of value) {
|
||||
if (!range || typeof range !== 'object') continue
|
||||
const rowPrices: Partial<Record<ImageOutputQuality, number>> = {}
|
||||
const rawPrices = 'prices' in range && range.prices && typeof range.prices === 'object'
|
||||
? range.prices as Record<string, unknown>
|
||||
: range as Record<string, unknown>
|
||||
for (const quality of IMAGE_OUTPUT_QUALITIES) {
|
||||
const price = rawPrices[quality]
|
||||
if (typeof price === 'number' && Number.isFinite(price)) {
|
||||
rowPrices[quality] = price
|
||||
}
|
||||
}
|
||||
const upToPixels = 'up_to_pixels' in range && range.up_to_pixels != null
|
||||
? String(range.up_to_pixels)
|
||||
: ''
|
||||
rows.push(createImageOutputPriceRangeRow(upToPixels, rowPrices))
|
||||
}
|
||||
return rows
|
||||
}
|
||||
|
||||
function createImageOutputPriceRow(
|
||||
size = '',
|
||||
prices: Partial<Record<ImageOutputQuality, number>> = {},
|
||||
): ImageOutputPriceRow {
|
||||
imageOutputPriceRowId += 1
|
||||
return {
|
||||
id: `image-output-size-${imageOutputPriceRowId}`,
|
||||
size,
|
||||
prices: { ...prices },
|
||||
}
|
||||
}
|
||||
|
||||
function createImageOutputPriceRangeRow(
|
||||
upToPixels = '',
|
||||
prices: Partial<Record<ImageOutputQuality, number>> = {},
|
||||
): ImageOutputPriceRangeRow {
|
||||
imageOutputPriceRangeRowId += 1
|
||||
return {
|
||||
id: `image-output-range-${imageOutputPriceRangeRowId}`,
|
||||
upToPixels,
|
||||
prices: { ...prices },
|
||||
}
|
||||
}
|
||||
|
||||
function normalizedImageOutputPrices(): Record<string, Record<string, number>> {
|
||||
const out: Record<string, Record<string, number>> = {}
|
||||
for (const row of imageOutputPriceRows.value) {
|
||||
const size = normalizeImageOutputSize(row.size)
|
||||
if (!size) continue
|
||||
for (const quality of IMAGE_OUTPUT_QUALITIES) {
|
||||
const price = row.prices[quality]
|
||||
if (price != null && Number.isFinite(price)) {
|
||||
out[size] = { ...(out[size] || {}), [quality]: price }
|
||||
}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
function normalizedImageOutputPriceRanges(): ImageOutputPriceRange[] {
|
||||
const ranges: ImageOutputPriceRange[] = []
|
||||
for (const row of imageOutputPriceRangeRows.value) {
|
||||
const prices: Partial<Record<ImageOutputQuality, number>> = {}
|
||||
for (const quality of IMAGE_OUTPUT_QUALITIES) {
|
||||
const price = row.prices[quality]
|
||||
if (price != null && Number.isFinite(price)) {
|
||||
prices[quality] = price
|
||||
}
|
||||
}
|
||||
if (Object.keys(prices).length === 0) continue
|
||||
ranges.push({
|
||||
up_to_pixels: parseOptionalInteger(row.upToPixels),
|
||||
prices,
|
||||
})
|
||||
}
|
||||
return ranges.sort((a, b) => {
|
||||
if (a.up_to_pixels == null && b.up_to_pixels == null) return 0
|
||||
if (a.up_to_pixels == null) return 1
|
||||
if (b.up_to_pixels == null) return -1
|
||||
return a.up_to_pixels - b.up_to_pixels
|
||||
})
|
||||
}
|
||||
|
||||
function parseOptionalFloat(value: string | number): number | null {
|
||||
if (value === '' || value === null || value === undefined) return null
|
||||
const number = typeof value === 'string' ? parseFloat(value) : value
|
||||
return Number.isFinite(number) ? number : null
|
||||
}
|
||||
|
||||
function parseOptionalInteger(value: string | number): number | null {
|
||||
if (value === '' || value === null || value === undefined) return null
|
||||
const number = typeof value === 'string' ? parseInt(value, 10) : value
|
||||
return Number.isFinite(number) && number > 0 ? Math.trunc(number) : null
|
||||
}
|
||||
|
||||
function normalizeImageOutputSize(size: string): string {
|
||||
return String(size || '').trim().replace(/\s*[xX×]\s*/g, 'x')
|
||||
}
|
||||
|
||||
function getImageOutputPrice(row: ImageOutputPriceRow, quality: ImageOutputQuality): string | number {
|
||||
return row.prices[quality] ?? ''
|
||||
}
|
||||
|
||||
function getImageOutputRangePrice(row: ImageOutputPriceRangeRow, quality: ImageOutputQuality): string | number {
|
||||
return row.prices[quality] ?? ''
|
||||
}
|
||||
|
||||
function updateImageOutputSize(rowId: string, value: string | number) {
|
||||
const row = imageOutputPriceRows.value.find(item => item.id === rowId)
|
||||
if (!row) return
|
||||
row.size = normalizeImageOutputSize(String(value ?? ''))
|
||||
imageOutputPriceRows.value = [...imageOutputPriceRows.value]
|
||||
syncToParent()
|
||||
}
|
||||
|
||||
function updateImageOutputPrice(rowId: string, quality: ImageOutputQuality, value: string | number) {
|
||||
const row = imageOutputPriceRows.value.find(item => item.id === rowId)
|
||||
if (!row) return
|
||||
const price = parseOptionalFloat(value)
|
||||
if (price == null) {
|
||||
delete row.prices[quality]
|
||||
} else {
|
||||
row.prices[quality] = price
|
||||
}
|
||||
imageOutputPriceRows.value = [...imageOutputPriceRows.value]
|
||||
syncToParent()
|
||||
}
|
||||
|
||||
function addImageOutputSizeRow() {
|
||||
const usedSizes = new Set(imageOutputPriceRows.value.map(row => normalizeImageOutputSize(row.size)).filter(Boolean))
|
||||
const suggestedSize = DEFAULT_IMAGE_OUTPUT_SIZES.find(size => !usedSizes.has(size)) || ''
|
||||
imageOutputPriceRows.value = [...imageOutputPriceRows.value, createImageOutputPriceRow(suggestedSize)]
|
||||
syncToParent()
|
||||
}
|
||||
|
||||
function removeImageOutputSizeRow(rowId: string) {
|
||||
imageOutputPriceRows.value = imageOutputPriceRows.value.filter(row => row.id !== rowId)
|
||||
syncToParent()
|
||||
}
|
||||
|
||||
function updateImageOutputRangeLimit(rowId: string, value: string | number) {
|
||||
const row = imageOutputPriceRangeRows.value.find(item => item.id === rowId)
|
||||
if (!row) return
|
||||
row.upToPixels = String(value ?? '')
|
||||
imageOutputPriceRangeRows.value = [...imageOutputPriceRangeRows.value]
|
||||
syncToParent()
|
||||
}
|
||||
|
||||
function updateImageOutputRangePrice(rowId: string, quality: ImageOutputQuality, value: string | number) {
|
||||
const row = imageOutputPriceRangeRows.value.find(item => item.id === rowId)
|
||||
if (!row) return
|
||||
const price = parseOptionalFloat(value)
|
||||
if (price == null) {
|
||||
delete row.prices[quality]
|
||||
} else {
|
||||
row.prices[quality] = price
|
||||
}
|
||||
imageOutputPriceRangeRows.value = [...imageOutputPriceRangeRows.value]
|
||||
syncToParent()
|
||||
}
|
||||
|
||||
function addImageOutputRangeRow() {
|
||||
const usedLimits = new Set(imageOutputPriceRangeRows.value.map(row => parseOptionalInteger(row.upToPixels)).filter((value): value is number => value !== null))
|
||||
const suggestedLimit = DEFAULT_IMAGE_OUTPUT_PIXEL_LIMITS.find(limit => !usedLimits.has(limit))
|
||||
imageOutputPriceRangeRows.value = [...imageOutputPriceRangeRows.value, createImageOutputPriceRangeRow(suggestedLimit ? String(suggestedLimit) : '')]
|
||||
syncToParent()
|
||||
}
|
||||
|
||||
function removeImageOutputRangeRow(rowId: string) {
|
||||
imageOutputPriceRangeRows.value = imageOutputPriceRangeRows.value.filter(row => row.id !== rowId)
|
||||
syncToParent()
|
||||
}
|
||||
|
||||
function updateImageOutputPriceDefault(value: string | number) {
|
||||
imageOutputPriceDefault.value = String(value ?? '')
|
||||
syncToParent()
|
||||
}
|
||||
|
||||
function parseFloatInput(value: string | number): number {
|
||||
const num = typeof value === 'string' ? parseFloat(value) : value
|
||||
return isNaN(num) ? 0 : num
|
||||
|
||||
@@ -94,6 +94,24 @@
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="rounded-lg border border-border/60 bg-muted/20 px-3 py-2">
|
||||
<div class="flex items-start gap-2">
|
||||
<Checkbox
|
||||
:checked="isImageGenerationEnabled"
|
||||
class="mt-0.5"
|
||||
@update:checked="setImageGenerationEnabled"
|
||||
/>
|
||||
<div class="space-y-1">
|
||||
<div class="text-sm font-medium">
|
||||
图片模型
|
||||
</div>
|
||||
<p class="text-xs text-muted-foreground">
|
||||
启用图片输出计费,并展开尺寸 × 质量矩阵价格。
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- 价格配置 -->
|
||||
<div class="space-y-4">
|
||||
<h4 class="font-semibold text-sm border-b pb-2">
|
||||
@@ -103,6 +121,7 @@
|
||||
ref="tieredPricingEditorRef"
|
||||
v-model="tieredPricing"
|
||||
:show-cache1h="showCache1h"
|
||||
:show-image-pricing="isImageGenerationEnabled"
|
||||
/>
|
||||
|
||||
<!-- 按次计费 -->
|
||||
@@ -249,6 +268,7 @@ import {
|
||||
SelectContent,
|
||||
SelectItem,
|
||||
Badge,
|
||||
Checkbox,
|
||||
} from '@/components/ui'
|
||||
import { useToast } from '@/composables/useToast'
|
||||
import { parseNumberInput, sortResolutionEntries } from '@/utils/form'
|
||||
@@ -290,10 +310,28 @@ const selectedGlobalModel = computed(() => {
|
||||
})
|
||||
|
||||
const selectedGlobalModelSupportsEmbedding = computed(() => modelSupportsEmbedding(selectedGlobalModel.value))
|
||||
const selectedGlobalModelSupportsImageGeneration = computed(() => modelSupportsImageGeneration(selectedGlobalModel.value))
|
||||
const editingModelSupportsEmbedding = computed(() => {
|
||||
return props.editingModel?.effective_supports_embedding === true
|
||||
|| modelSupportsEmbedding(props.editingModel)
|
||||
})
|
||||
const editingModelSupportsImageGeneration = computed(() => {
|
||||
return props.editingModel?.effective_supports_image_generation === true
|
||||
|| modelSupportsImageGeneration(props.editingModel)
|
||||
})
|
||||
|
||||
const isImageGenerationEnabled = computed(() => {
|
||||
if (imageGenerationExplicitOverride.value !== null) {
|
||||
return imageGenerationExplicitOverride.value
|
||||
}
|
||||
if (form.value.supports_image_generation !== undefined) {
|
||||
return form.value.supports_image_generation === true
|
||||
}
|
||||
const supportsImageGeneration = isEditing.value
|
||||
? editingModelSupportsImageGeneration.value
|
||||
: selectedGlobalModelSupportsImageGeneration.value
|
||||
return supportsImageGeneration || tieredPricingHasImageOutputPricing(tieredPricing.value)
|
||||
})
|
||||
|
||||
// 1h 缓存定价始终显示
|
||||
const showCache1h = true
|
||||
@@ -349,6 +387,7 @@ const form = ref({
|
||||
supports_image_generation: undefined as boolean | undefined,
|
||||
is_active: true
|
||||
})
|
||||
const imageGenerationExplicitOverride = ref<boolean | null>(null)
|
||||
|
||||
const canSubmitCreate = computed(() => {
|
||||
if (isEditing.value) return true
|
||||
@@ -364,6 +403,7 @@ watch(() => props.open, async (newOpen) => {
|
||||
// 编辑模式:填充表单
|
||||
// 使用有效配置(合并全局模型的默认值)供用户查看和编辑
|
||||
const effectiveConfig = props.editingModel.effective_config || props.editingModel.config || {}
|
||||
const supportsImageGeneration = modelSupportsImageGeneration(props.editingModel)
|
||||
form.value = {
|
||||
global_model_id: props.editingModel.global_model_id || '',
|
||||
provider_model_name: props.editingModel.provider_model_name || '',
|
||||
@@ -374,7 +414,7 @@ watch(() => props.open, async (newOpen) => {
|
||||
supports_function_calling: props.editingModel.supports_function_calling ?? undefined,
|
||||
supports_streaming: props.editingModel.supports_streaming ?? undefined,
|
||||
supports_extended_thinking: props.editingModel.supports_extended_thinking ?? undefined,
|
||||
supports_image_generation: props.editingModel.supports_image_generation ?? undefined,
|
||||
supports_image_generation: supportsImageGeneration ? true : props.editingModel.supports_image_generation ?? undefined,
|
||||
is_active: props.editingModel.is_active
|
||||
}
|
||||
// 从有效配置中加载视频费用
|
||||
@@ -425,6 +465,7 @@ watch(tieredPricing, (newValue) => {
|
||||
|
||||
// 重置表单
|
||||
function resetForm() {
|
||||
imageGenerationExplicitOverride.value = null
|
||||
form.value = {
|
||||
global_model_id: '',
|
||||
provider_model_name: '',
|
||||
@@ -446,11 +487,66 @@ function resetForm() {
|
||||
}
|
||||
|
||||
function handleGlobalModelSelect(value: string) {
|
||||
imageGenerationExplicitOverride.value = null
|
||||
form.value.supports_image_generation = undefined
|
||||
form.value.global_model_id = value
|
||||
const selectedModel = availableGlobalModels.value.find(model => model.id === value)
|
||||
form.value.provider_model_name = selectedModel?.name || form.value.provider_model_name
|
||||
}
|
||||
|
||||
function modelSupportsImageGeneration(model: {
|
||||
supported_capabilities?: string[] | null
|
||||
supports_image_generation?: boolean | null
|
||||
effective_supports_image_generation?: boolean | null
|
||||
default_tiered_pricing?: TieredPricingConfig | null
|
||||
tiered_pricing?: TieredPricingConfig | null
|
||||
effective_tiered_pricing?: TieredPricingConfig | null
|
||||
config?: Record<string, unknown> | null
|
||||
} | null | undefined): boolean {
|
||||
if (!model) return false
|
||||
if (model.effective_supports_image_generation === true) return true
|
||||
if (model.supports_image_generation === true) return true
|
||||
const config = model.config || {}
|
||||
return model.supported_capabilities?.includes('image_generation') === true
|
||||
|| config.image_generation === true
|
||||
|| config.model_type === 'image'
|
||||
|| (Array.isArray(config.api_formats) && config.api_formats.some((format) => String(format).endsWith(':image')))
|
||||
|| tieredPricingHasImageOutputPricing(model.default_tiered_pricing)
|
||||
|| tieredPricingHasImageOutputPricing(model.tiered_pricing)
|
||||
|| tieredPricingHasImageOutputPricing(model.effective_tiered_pricing)
|
||||
}
|
||||
|
||||
function tieredPricingHasImageOutputPricing(pricing: TieredPricingConfig | null | undefined): boolean {
|
||||
if (!pricing) return false
|
||||
if (toFinitePrice(pricing.image_output_price_default) !== null) return true
|
||||
if (Object.values(pricing.image_output_prices || {}).some((prices) => {
|
||||
if (!prices || typeof prices !== 'object') return false
|
||||
return Object.values(prices).some((price) => toFinitePrice(price) !== null)
|
||||
})) return true
|
||||
return (pricing.image_output_price_ranges || []).some((range) => {
|
||||
if (!range || typeof range !== 'object') return false
|
||||
const prices = range.prices && typeof range.prices === 'object'
|
||||
? range.prices
|
||||
: range as Record<string, unknown>
|
||||
return Object.values(prices).some((price) => toFinitePrice(price) !== null)
|
||||
})
|
||||
}
|
||||
|
||||
function toFinitePrice(value: unknown): number | null {
|
||||
if (typeof value === 'number' && Number.isFinite(value)) return value
|
||||
if (typeof value === 'string' && value.trim()) {
|
||||
const parsed = Number(value)
|
||||
return Number.isFinite(parsed) ? parsed : null
|
||||
}
|
||||
return null
|
||||
}
|
||||
|
||||
function setImageGenerationEnabled(value: boolean | 'indeterminate') {
|
||||
const enabled = value === true
|
||||
imageGenerationExplicitOverride.value = enabled
|
||||
form.value.supports_image_generation = enabled
|
||||
}
|
||||
|
||||
function getNested(obj: Record<string, unknown>, path: string): unknown {
|
||||
if (!obj || typeof obj !== 'object') return undefined
|
||||
const parts = path.split('.').filter(Boolean)
|
||||
@@ -631,8 +727,9 @@ async function handleSubmit() {
|
||||
submitting.value = true
|
||||
try {
|
||||
// 获取包含自动计算缓存价格的最终数据
|
||||
const finalTiers = tieredPricingEditorRef.value?.getFinalTiers()
|
||||
const finalTieredPricing = finalTiers ? { tiers: finalTiers } : tieredPricing.value
|
||||
const finalTieredPricing = tieredPricingEditorRef.value?.getFinalPricing() ?? tieredPricing.value
|
||||
const supportsImageGeneration = isImageGenerationEnabled.value
|
||||
|| tieredPricingHasImageOutputPricing(finalTieredPricing)
|
||||
|
||||
// Apply billing (video) pricing into config.
|
||||
applyVideoPricingToConfig(form.value.config)
|
||||
@@ -651,7 +748,7 @@ async function handleSubmit() {
|
||||
supportsFunctionCalling: form.value.supports_function_calling,
|
||||
supportsStreaming: form.value.supports_streaming,
|
||||
supportsExtendedThinking: form.value.supports_extended_thinking,
|
||||
supportsImageGeneration: form.value.supports_image_generation,
|
||||
supportsImageGeneration,
|
||||
isActive: form.value.is_active
|
||||
}))
|
||||
showSuccess('模型配置已更新')
|
||||
@@ -674,7 +771,7 @@ async function handleSubmit() {
|
||||
supportsFunctionCalling: form.value.supports_function_calling,
|
||||
supportsStreaming: form.value.supports_streaming,
|
||||
supportsExtendedThinking: form.value.supports_extended_thinking,
|
||||
supportsImageGeneration: form.value.supports_image_generation,
|
||||
supportsImageGeneration,
|
||||
isActive: form.value.is_active
|
||||
}))
|
||||
showSuccess('模型已添加')
|
||||
|
||||
@@ -443,9 +443,10 @@ const activeEndpoints = computed(() => (props.endpoints ?? [])
|
||||
if (typeof endpoint.active_keys === 'number') {
|
||||
return endpoint.is_active !== false
|
||||
&& isModelTestableApiFormat(endpoint.api_format)
|
||||
&& endpoint.active_keys > 0
|
||||
&& (endpoint.active_keys > 0
|
||||
|| isModelTestableEndpoint(endpoint, providerKeysState.value, props.provider.provider_type))
|
||||
}
|
||||
return isModelTestableEndpoint(endpoint, providerKeysState.value)
|
||||
return isModelTestableEndpoint(endpoint, providerKeysState.value, props.provider.provider_type)
|
||||
}))
|
||||
const selectableTestEndpoints = computed(() => mappingTestEndpoints.value ?? activeEndpoints.value)
|
||||
const parsedTestRequestHeaders = computed(() => parseModelTestRequestHeadersDraft(testRequestHeadersDraft.value))
|
||||
|
||||
@@ -313,9 +313,10 @@ const activeEndpoints = computed(() => (props.endpoints ?? [])
|
||||
if (typeof endpoint.active_keys === 'number') {
|
||||
return endpoint.is_active !== false
|
||||
&& isModelTestableApiFormat(endpoint.api_format)
|
||||
&& endpoint.active_keys > 0
|
||||
&& (endpoint.active_keys > 0
|
||||
|| isModelTestableEndpoint(endpoint, props.providerKeys ?? [], props.provider.provider_type))
|
||||
}
|
||||
return isModelTestableEndpoint(endpoint, props.providerKeys ?? [])
|
||||
return isModelTestableEndpoint(endpoint, props.providerKeys ?? [], props.provider.provider_type)
|
||||
}))
|
||||
const parsedTestRequestHeaders = computed(() => parseModelTestRequestHeadersDraft(testRequestHeadersDraft.value))
|
||||
const testRequestHeadersError = computed(() => parsedTestRequestHeaders.value.error)
|
||||
|
||||
@@ -58,7 +58,7 @@ describe('buildDefaultModelTestRequestBody', () => {
|
||||
expect(body.input).toBeUndefined()
|
||||
})
|
||||
|
||||
it('uses prompt payloads for openai image api formats', () => {
|
||||
it('uses image prompt payloads for OpenAI image test requests', () => {
|
||||
const body = JSON.parse(buildDefaultModelTestRequestBody('gpt-image-2', 'openai:image'))
|
||||
|
||||
expect(body).toEqual({
|
||||
@@ -240,6 +240,8 @@ describe('isModelTestableApiFormat', () => {
|
||||
it.each([
|
||||
'openai:chat',
|
||||
'openai:responses',
|
||||
'openai:responses:compact',
|
||||
'openai:image',
|
||||
'claude:messages',
|
||||
'gemini:generate_content',
|
||||
'openai:image',
|
||||
@@ -350,6 +352,23 @@ describe('isModelTestableEndpoint', () => {
|
||||
is_active: true,
|
||||
}, keys)).toBe(true)
|
||||
})
|
||||
|
||||
it('lets fixed provider OAuth keys inherit testable endpoint formats', () => {
|
||||
const keys = [{
|
||||
api_formats: ['legacy:mismatch'],
|
||||
auth_type: 'oauth',
|
||||
is_active: true,
|
||||
}]
|
||||
|
||||
expect(isModelTestableEndpoint({
|
||||
api_format: 'openai:image',
|
||||
is_active: true,
|
||||
}, keys, 'chatgpt_web')).toBe(true)
|
||||
expect(isModelTestableEndpoint({
|
||||
api_format: 'openai:image',
|
||||
is_active: true,
|
||||
}, keys, 'custom')).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe('formatModelTestDiagnostic', () => {
|
||||
|
||||
@@ -15,6 +15,9 @@ export type ModelTestImageSource = {
|
||||
export type ModelTestKeySource = {
|
||||
api_formats?: string[] | null
|
||||
is_active?: boolean | null
|
||||
auth_type?: string | null
|
||||
credential_kind?: string | null
|
||||
oauth_managed?: boolean | null
|
||||
}
|
||||
|
||||
const MODEL_TEST_UNSUPPORTED_API_FORMATS = new Set([
|
||||
@@ -23,6 +26,20 @@ const MODEL_TEST_UNSUPPORTED_API_FORMATS = new Set([
|
||||
'gemini:files',
|
||||
])
|
||||
|
||||
const MODEL_TEST_OAUTH_INHERITS_PROVIDER_FORMATS = new Set([
|
||||
'claude_code',
|
||||
'codex',
|
||||
'chatgpt_web',
|
||||
'gemini_cli',
|
||||
'vertex_ai',
|
||||
'antigravity',
|
||||
'kiro',
|
||||
])
|
||||
|
||||
const MODEL_TEST_BEARER_INHERITS_PROVIDER_FORMATS = new Set([
|
||||
'chatgpt_web',
|
||||
])
|
||||
|
||||
const MODEL_TEST_DIAGNOSTIC_LABELS: Record<string, string> = {
|
||||
pool_account_blocked: '账号已失效,需重新授权',
|
||||
}
|
||||
@@ -41,12 +58,15 @@ export function isModelTestableApiFormat(apiFormat: string | null | undefined):
|
||||
export function modelTestKeySupportsEndpoint(
|
||||
key: ModelTestKeySource,
|
||||
endpoint: ModelTestEndpointSource,
|
||||
providerType?: string | null,
|
||||
): boolean {
|
||||
if (key.is_active === false) return false
|
||||
|
||||
const endpointFormat = normalizeApiFormatAlias(endpoint.api_format)
|
||||
if (!isModelTestableApiFormat(endpointFormat)) return false
|
||||
|
||||
if (modelTestKeyInheritsProviderFormats(key, providerType)) return true
|
||||
|
||||
const keyFormats = normalizeModelTestStringList(key.api_formats)
|
||||
if (keyFormats.length === 0) return true
|
||||
|
||||
@@ -56,10 +76,32 @@ export function modelTestKeySupportsEndpoint(
|
||||
export function isModelTestableEndpoint(
|
||||
endpoint: ModelTestEndpointSource,
|
||||
keys: ModelTestKeySource[],
|
||||
providerType?: string | null,
|
||||
): boolean {
|
||||
return endpoint.is_active !== false
|
||||
&& isModelTestableApiFormat(endpoint.api_format)
|
||||
&& keys.some(key => modelTestKeySupportsEndpoint(key, endpoint))
|
||||
&& keys.some(key => modelTestKeySupportsEndpoint(key, endpoint, providerType))
|
||||
}
|
||||
|
||||
function modelTestKeyInheritsProviderFormats(
|
||||
key: ModelTestKeySource,
|
||||
providerType: string | null | undefined,
|
||||
): boolean {
|
||||
const normalizedProviderType = providerType?.trim().toLowerCase()
|
||||
if (!normalizedProviderType) return false
|
||||
|
||||
const authType = key.auth_type?.trim().toLowerCase()
|
||||
const credentialKind = key.credential_kind?.trim().toLowerCase()
|
||||
const oauthManaged = key.oauth_managed === true
|
||||
|| credentialKind === 'oauth_session'
|
||||
|| authType === 'oauth'
|
||||
|
||||
if (oauthManaged && MODEL_TEST_OAUTH_INHERITS_PROVIDER_FORMATS.has(normalizedProviderType)) {
|
||||
return true
|
||||
}
|
||||
|
||||
return authType === 'bearer'
|
||||
&& MODEL_TEST_BEARER_INHERITS_PROVIDER_FORMATS.has(normalizedProviderType)
|
||||
}
|
||||
|
||||
export function selectPreferredModelTestEndpoint<T extends ModelTestEndpointSource>(
|
||||
|
||||
@@ -149,14 +149,16 @@ export function buildDefaultModelTestRequestBody(
|
||||
apiFormat?: string | null,
|
||||
model?: ModelTestImageSource | null,
|
||||
): string {
|
||||
if (apiFormat?.trim().toLowerCase().endsWith(':embedding')) {
|
||||
const normalizedApiFormat = normalizeApiFormatAlias(apiFormat ?? '')
|
||||
|
||||
if (normalizedApiFormat.endsWith(':embedding')) {
|
||||
return JSON.stringify({
|
||||
model: modelName,
|
||||
input: 'This is a test embedding input.',
|
||||
}, null, 2)
|
||||
}
|
||||
|
||||
if (apiFormat?.trim().toLowerCase().endsWith(':rerank')) {
|
||||
if (normalizedApiFormat.endsWith(':rerank')) {
|
||||
return JSON.stringify({
|
||||
model: modelName,
|
||||
query: 'Apple',
|
||||
@@ -171,7 +173,7 @@ export function buildDefaultModelTestRequestBody(
|
||||
}, null, 2)
|
||||
}
|
||||
|
||||
if (normalizeApiFormatAlias(apiFormat ?? '') === 'openai:image') {
|
||||
if (normalizedApiFormat === 'openai:image') {
|
||||
return JSON.stringify({
|
||||
model: modelName,
|
||||
prompt: DEFAULT_MODEL_TEST_MESSAGE,
|
||||
@@ -181,7 +183,7 @@ export function buildDefaultModelTestRequestBody(
|
||||
}, null, 2)
|
||||
}
|
||||
|
||||
if (normalizeApiFormatAlias(apiFormat ?? '') === 'openai:responses' && modelSupportsImageGeneration(model)) {
|
||||
if (normalizedApiFormat === 'openai:responses' && modelSupportsImageGeneration(model)) {
|
||||
return JSON.stringify({
|
||||
model: modelName,
|
||||
input: DEFAULT_MODEL_TEST_MESSAGE,
|
||||
@@ -271,4 +273,4 @@ export function parseModelTestRequestHeadersDraft(
|
||||
emptyError: null,
|
||||
invalidTypeError: '测试请求头必须是 JSON 对象',
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -275,6 +275,9 @@
|
||||
<template v-if="perRequestCost > 0">
|
||||
+ 按次费用 <span class="font-medium">${{ perRequestCost.toFixed(6) }}</span>
|
||||
</template>
|
||||
<template v-if="imageOutputCostTotal > 0">
|
||||
+ 图片输出费用 <span class="font-medium">${{ imageOutputCostTotal.toFixed(6) }}</span>
|
||||
</template>
|
||||
<template v-if="videoCostTotal > 0">
|
||||
+ {{ detail.video_billing?.task_type === 'image' ? '图像' : detail.video_billing?.task_type === 'audio' ? '音频' : '视频' }}费用 <span class="font-medium">${{ videoCostTotal.toFixed(6) }}</span>
|
||||
</template>
|
||||
@@ -426,7 +429,52 @@
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- ========== 4. 视频/图像/音频计费(独立隔离,与Token计费风格一致) ========== -->
|
||||
<!-- ========== 4. 图片输出计费 ========== -->
|
||||
<div
|
||||
v-if="hasImageBillingDetail"
|
||||
class="rounded-lg p-3 space-y-2 bg-primary/5 border border-primary/30 mb-3"
|
||||
>
|
||||
<div class="flex flex-col sm:flex-row sm:items-center sm:justify-between gap-1 sm:gap-2 text-xs">
|
||||
<div class="flex items-center gap-2 flex-wrap">
|
||||
<span class="font-medium text-primary">图片输出</span>
|
||||
<Badge
|
||||
variant="outline"
|
||||
class="text-[10px] px-1.5 py-0 h-4"
|
||||
>
|
||||
{{ imageOutputBillingLabel }}
|
||||
</Badge>
|
||||
<span
|
||||
v-if="imageOutputPricingDescriptor"
|
||||
class="text-muted-foreground font-mono"
|
||||
>{{ imageOutputPricingDescriptor }}</span>
|
||||
</div>
|
||||
<div class="text-muted-foreground flex items-center gap-2 flex-wrap">
|
||||
<span
|
||||
v-if="imageOutputPricePerImage !== null"
|
||||
class="font-mono"
|
||||
>{{ formatNumber(imageOutputCount) }} 张 × ${{ imageOutputPricePerImage.toFixed(6) }}/张 = ${{ imageOutputCostTotal.toFixed(6) }}</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="flex items-center">
|
||||
<div class="flex items-center flex-1">
|
||||
<span class="text-xs text-muted-foreground w-[56px]">数量</span>
|
||||
<span class="text-sm font-semibold font-mono flex-1 text-center">{{ formatNumber(imageOutputCount) }}</span>
|
||||
<span class="text-xs font-mono">${{ imageOutputCostTotal.toFixed(6) }}</span>
|
||||
</div>
|
||||
<Separator
|
||||
orientation="vertical"
|
||||
class="h-4 mx-4"
|
||||
/>
|
||||
<div class="flex items-center flex-1">
|
||||
<span class="text-xs text-muted-foreground w-[56px]">格式</span>
|
||||
<span class="text-sm font-semibold font-mono flex-1 text-center">{{ imageOutputFormat || '-' }}</span>
|
||||
<span class="text-xs font-mono text-muted-foreground">{{ imageOutputBillingLabel }}</span>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- ========== 5. 视频/图像/音频计费(独立隔离,与Token计费风格一致) ========== -->
|
||||
<div
|
||||
v-if="detail.video_billing"
|
||||
class="rounded-lg p-3 space-y-2 bg-primary/5 border border-primary/30"
|
||||
@@ -943,6 +991,11 @@ function getNestedNumber(record: JsonRecord | null, ...path: string[]): number |
|
||||
return toNumber(getNestedValue(record, ...path))
|
||||
}
|
||||
|
||||
function getNestedString(record: JsonRecord | null, ...path: string[]): string | null {
|
||||
const value = getNestedValue(record, ...path)
|
||||
return typeof value === 'string' && value.trim() ? value.trim() : null
|
||||
}
|
||||
|
||||
function normalizeCacheTtlPricing(value: unknown): CacheTTLPriceEntry[] {
|
||||
if (!Array.isArray(value)) return []
|
||||
return value
|
||||
@@ -1080,6 +1133,10 @@ const billingResolvedVariables = computed<JsonRecord | null>(() =>
|
||||
asRecord(billingSnapshot.value?.resolved_variables),
|
||||
)
|
||||
|
||||
const billingResolvedDimensions = computed<JsonRecord | null>(() =>
|
||||
asRecord(billingSnapshot.value?.resolved_dimensions),
|
||||
)
|
||||
|
||||
const billingCostBreakdown = computed<JsonRecord | null>(() =>
|
||||
asRecord(billingSnapshot.value?.cost_breakdown),
|
||||
)
|
||||
@@ -1421,6 +1478,108 @@ const effectiveRequestCost = computed(() => {
|
||||
return 0
|
||||
})
|
||||
|
||||
const effectiveImageOutputCost = computed(() =>
|
||||
getNestedNumber(billingCostBreakdown.value, 'image_output_cost')
|
||||
?? toNumber(detail.value?.image_output_cost)
|
||||
?? 0,
|
||||
)
|
||||
|
||||
const imageOutputCostTotal = computed(() => effectiveImageOutputCost.value)
|
||||
|
||||
const imageOutputPricePerImage = computed(() =>
|
||||
getNestedNumber(billingResolvedVariables.value, 'image_output_price_per_image'),
|
||||
)
|
||||
|
||||
const imageOutputCount = computed(() =>
|
||||
getNestedNumber(billingResolvedDimensions.value, 'image_count')
|
||||
?? getNestedNumber(traceRequestMetadata.value, 'billing_dimensions', 'image_count')
|
||||
?? getNestedNumber(traceRequestMetadata.value, 'dimensions', 'image_count')
|
||||
?? 0,
|
||||
)
|
||||
|
||||
const imageOutputSize = computed(() =>
|
||||
getNestedString(billingResolvedDimensions.value, 'image_size')
|
||||
?? getNestedString(traceRequestMetadata.value, 'billing_dimensions', 'image_size')
|
||||
?? getNestedString(traceRequestMetadata.value, 'dimensions', 'image_size'),
|
||||
)
|
||||
|
||||
const imageOutputQuality = computed(() =>
|
||||
getNestedString(billingResolvedDimensions.value, 'image_quality')
|
||||
?? getNestedString(traceRequestMetadata.value, 'billing_dimensions', 'image_quality')
|
||||
?? getNestedString(traceRequestMetadata.value, 'dimensions', 'image_quality'),
|
||||
)
|
||||
|
||||
const imageOutputFormat = computed(() =>
|
||||
getNestedString(billingResolvedDimensions.value, 'image_output_format')
|
||||
?? getNestedString(traceRequestMetadata.value, 'billing_dimensions', 'image_output_format')
|
||||
?? getNestedString(traceRequestMetadata.value, 'dimensions', 'image_output_format'),
|
||||
)
|
||||
|
||||
const imagePriceKey = computed(() => {
|
||||
const snapshotKey = getNestedString(billingResolvedDimensions.value, 'image_price_key')
|
||||
if (snapshotKey) return snapshotKey
|
||||
const fallbackKey = [imageOutputSize.value, imageOutputQuality.value].filter(Boolean).join(':')
|
||||
return fallbackKey || null
|
||||
})
|
||||
|
||||
const imageOutputPriceBucket = computed(() =>
|
||||
getNestedString(billingResolvedDimensions.value, 'image_output_price_bucket'),
|
||||
)
|
||||
|
||||
const imageOutputPixels = computed(() =>
|
||||
getNestedNumber(billingResolvedDimensions.value, 'image_pixels')
|
||||
?? parseImageSizePixels(imageOutputSize.value),
|
||||
)
|
||||
|
||||
const imageOutputPricingMode = computed(() =>
|
||||
getNestedString(billingResolvedDimensions.value, 'image_output_pricing_mode'),
|
||||
)
|
||||
|
||||
const imageOutputPricingEnabled = computed(() =>
|
||||
getNestedValue(billingResolvedDimensions.value, 'image_output_pricing_enabled') === true
|
||||
|| imageOutputPricingMode.value === 'matrix'
|
||||
|| imageOutputPricingMode.value === 'pixel_tiers'
|
||||
|| imageOutputPricingMode.value === 'per_image'
|
||||
|| imageOutputCostTotal.value > 0,
|
||||
)
|
||||
|
||||
const imageOutputMatrixEnabled = computed(() => {
|
||||
if (imageOutputPricingMode.value) return imageOutputPricingMode.value === 'matrix'
|
||||
return getNestedValue(billingResolvedDimensions.value, 'image_output_matrix_enabled') === true
|
||||
})
|
||||
|
||||
const imageOutputRangeEnabled = computed(() => {
|
||||
if (imageOutputPricingMode.value) return imageOutputPricingMode.value === 'pixel_tiers'
|
||||
return getNestedValue(billingResolvedDimensions.value, 'image_output_range_enabled') === true
|
||||
})
|
||||
|
||||
const imageOutputBillingLabel = computed(() => {
|
||||
if (imageOutputMatrixEnabled.value) return '矩阵计费'
|
||||
if (imageOutputRangeEnabled.value) return '像素区间'
|
||||
return '默认计费'
|
||||
})
|
||||
|
||||
const imageOutputPricingDescriptor = computed(() => {
|
||||
if (imageOutputMatrixEnabled.value && imagePriceKey.value) return imagePriceKey.value
|
||||
|
||||
const parts: string[] = []
|
||||
if (imageOutputPriceBucket.value && imageOutputPriceBucket.value !== 'default') {
|
||||
parts.push(formatImagePriceBucket(imageOutputPriceBucket.value))
|
||||
}
|
||||
const sizeQuality = [imageOutputSize.value, imageOutputQuality.value].filter(Boolean).join(' / ')
|
||||
if (sizeQuality) parts.push(sizeQuality)
|
||||
if (imageOutputRangeEnabled.value && imageOutputPixels.value !== null) {
|
||||
parts.push(formatPixels(imageOutputPixels.value))
|
||||
}
|
||||
if (parts.length > 0) return parts.join(' · ')
|
||||
if (imageOutputPriceBucket.value === 'default') return '默认价'
|
||||
return null
|
||||
})
|
||||
|
||||
const hasImageBillingDetail = computed(() =>
|
||||
imageOutputPricingEnabled.value && (imageOutputCount.value > 0 || imageOutputCostTotal.value > 0),
|
||||
)
|
||||
|
||||
const fallbackCacheTtlPricing = computed<CacheTTLPriceEntry[]>(() => {
|
||||
const tierPricing = normalizeCacheTtlPricing(billingTierInfo.value?.cache_ttl_pricing)
|
||||
if (tierPricing.length > 0) return tierPricing
|
||||
@@ -2123,6 +2282,34 @@ function formatNumber(num: number): string {
|
||||
return num.toLocaleString()
|
||||
}
|
||||
|
||||
function parseImageSizePixels(size: string | null): number | null {
|
||||
if (!size) return null
|
||||
const normalized = size.trim().toLowerCase().replace(/\s+/g, '').replace(/×/g, 'x')
|
||||
const [widthText, heightText] = normalized.split('x')
|
||||
const width = Number(widthText)
|
||||
const height = Number(heightText)
|
||||
if (!Number.isFinite(width) || !Number.isFinite(height) || width <= 0 || height <= 0) return null
|
||||
return Math.trunc(width * height)
|
||||
}
|
||||
|
||||
function formatImagePriceBucket(bucket: string): string {
|
||||
if (bucket === 'default') return '默认价'
|
||||
if (bucket === 'unbounded') return '无上限'
|
||||
const match = bucket.match(/^<=([0-9]+)px$/)
|
||||
if (match) return `<= ${formatPixels(Number(match[1]))}`
|
||||
return bucket
|
||||
}
|
||||
|
||||
function formatPixels(value: number): string {
|
||||
if (value >= 1_000_000) {
|
||||
return `${(value / 1_000_000).toFixed(value % 1_000_000 === 0 ? 0 : 2)}M px`
|
||||
}
|
||||
if (value >= 1_000) {
|
||||
return `${(value / 1_000).toFixed(0)}K px`
|
||||
}
|
||||
return `${value} px`
|
||||
}
|
||||
|
||||
// 格式化响应时间,自动选择合适的单位
|
||||
function formatResponseTime(ms: number): { value: string; unit: string } {
|
||||
if (ms >= 1_000) {
|
||||
|
||||
Reference in New Issue
Block a user