mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
refactor ai serving modules and crates
This commit is contained in:
21
apps/aether-gateway/src/ai_serving/adaptation/mod.rs
Normal file
21
apps/aether-gateway/src/ai_serving/adaptation/mod.rs
Normal file
@@ -0,0 +1,21 @@
|
||||
pub(crate) mod private_envelope;
|
||||
|
||||
pub(crate) mod kiro {
|
||||
pub(crate) use crate::ai_serving::pure::KiroToClaudeCliStreamState;
|
||||
}
|
||||
|
||||
pub(crate) use crate::ai_serving::{
|
||||
provider_adaptation_allows_sync_finalize_envelope, provider_adaptation_anchor_api_format,
|
||||
provider_adaptation_descriptor_for_envelope, provider_adaptation_descriptor_for_provider_type,
|
||||
provider_adaptation_requires_eventstream_accept,
|
||||
provider_adaptation_should_unwrap_stream_envelope, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME,
|
||||
GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME,
|
||||
};
|
||||
pub(crate) use kiro::KiroToClaudeCliStreamState;
|
||||
pub(crate) use private_envelope::{
|
||||
maybe_build_provider_private_stream_normalizer,
|
||||
maybe_normalize_provider_private_sync_report_payload,
|
||||
normalize_provider_private_report_context, normalize_provider_private_response_value,
|
||||
provider_private_response_allows_sync_finalize, transform_provider_private_stream_line,
|
||||
ProviderPrivateStreamNormalizer,
|
||||
};
|
||||
@@ -0,0 +1,10 @@
|
||||
#[path = "private_envelope/sync.rs"]
|
||||
mod sync;
|
||||
|
||||
pub(crate) use self::sync::maybe_normalize_provider_private_sync_report_payload;
|
||||
pub(crate) use crate::ai_serving::{
|
||||
maybe_build_provider_private_stream_normalizer, normalize_provider_private_report_context,
|
||||
normalize_provider_private_response_value, provider_private_response_allows_sync_finalize,
|
||||
stream_body_contains_error_event, transform_provider_private_stream_line,
|
||||
ProviderPrivateStreamNormalizer,
|
||||
};
|
||||
@@ -0,0 +1,70 @@
|
||||
use base64::Engine as _;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
|
||||
use super::{
|
||||
maybe_build_provider_private_stream_normalizer, normalize_provider_private_report_context,
|
||||
normalize_provider_private_response_value, provider_private_response_allows_sync_finalize,
|
||||
stream_body_contains_error_event, ProviderPrivateStreamNormalizer,
|
||||
};
|
||||
|
||||
pub(crate) fn maybe_normalize_provider_private_sync_report_payload(
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<GatewaySyncReportRequest>, GatewayError> {
|
||||
let Some(report_context) = payload.report_context.as_ref() else {
|
||||
return Ok(Some(payload.clone()));
|
||||
};
|
||||
if !report_context
|
||||
.get("has_envelope")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return Ok(Some(payload.clone()));
|
||||
}
|
||||
if !provider_private_response_allows_sync_finalize(report_context) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let mut normalized = payload.clone();
|
||||
normalized.report_context = normalize_provider_private_report_context(Some(report_context));
|
||||
|
||||
if let Some(body_json) = payload.body_json.clone() {
|
||||
normalized.body_json = normalize_provider_private_response_value(body_json, report_context);
|
||||
if normalized.body_json.is_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(body_base64) = payload.body_base64.as_deref() {
|
||||
let body_bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(body_base64)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let Some(normalized_bytes) =
|
||||
normalize_provider_private_stream_bytes(report_context, &body_bytes)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
if stream_body_contains_error_event(&normalized_bytes) {
|
||||
return Ok(None);
|
||||
}
|
||||
normalized.body_base64 = (!normalized_bytes.is_empty())
|
||||
.then(|| base64::engine::general_purpose::STANDARD.encode(normalized_bytes));
|
||||
}
|
||||
|
||||
Ok(Some(normalized))
|
||||
}
|
||||
|
||||
fn normalize_provider_private_stream_bytes(
|
||||
report_context: &Value,
|
||||
body: &[u8],
|
||||
) -> Result<Option<Vec<u8>>, GatewayError> {
|
||||
let Some(mut normalizer): Option<ProviderPrivateStreamNormalizer<'_>> =
|
||||
maybe_build_provider_private_stream_normalizer(Some(report_context))
|
||||
else {
|
||||
return Ok(Some(body.to_vec()));
|
||||
};
|
||||
let mut normalized = normalizer.push_chunk(body).map_err(GatewayError::from)?;
|
||||
normalized.extend(normalizer.finish().map_err(GatewayError::from)?);
|
||||
Ok(Some(normalized))
|
||||
}
|
||||
120
apps/aether-gateway/src/ai_serving/api.rs
Normal file
120
apps/aether-gateway/src/ai_serving/api.rs
Normal file
@@ -0,0 +1,120 @@
|
||||
use crate::ai_serving::{is_json_request, GatewayControlDecision};
|
||||
|
||||
pub(crate) use crate::ai_serving::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_local_gemini_files_stream_plan_and_reports_for_kind,
|
||||
build_local_gemini_files_sync_plan_and_reports_for_kind,
|
||||
build_local_image_stream_plan_and_reports_for_kind,
|
||||
build_local_image_sync_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_sync_plan_and_reports_for_kind,
|
||||
build_local_openai_responses_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_responses_sync_plan_and_reports_for_kind,
|
||||
build_local_same_format_stream_plan_and_reports, build_local_same_format_sync_plan_and_reports,
|
||||
build_local_video_sync_plan_and_reports_for_kind,
|
||||
build_openai_responses_stream_plan_from_decision,
|
||||
build_openai_responses_sync_plan_from_decision, build_passthrough_sync_plan_from_decision,
|
||||
build_standard_family_stream_plan_and_reports, build_standard_family_sync_plan_and_reports,
|
||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
||||
maybe_build_stream_decision_payload, maybe_build_stream_plan_payload,
|
||||
maybe_build_sync_decision_payload, maybe_build_sync_plan_payload,
|
||||
set_local_openai_chat_execution_exhausted_diagnostic,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::{
|
||||
maybe_bridge_standard_sync_json_to_stream, maybe_build_provider_private_stream_normalizer,
|
||||
maybe_build_stream_response_rewriter, maybe_build_sync_finalize_outcome,
|
||||
maybe_compile_sync_finalize_response, LocalCoreSyncFinalizeOutcome,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::{
|
||||
AiExecutionDecision, AiExecutionPlanPayload, AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
pub(crate) use aether_ai_surfaces::api::{
|
||||
build_core_error_body_for_client_format, core_error_background_report_kind,
|
||||
core_error_default_client_api_format, core_success_background_report_kind,
|
||||
encode_kiro_sse_events, implicit_sync_finalize_report_kind, is_core_error_finalize_kind,
|
||||
normalize_provider_private_report_context, normalize_provider_private_response_value,
|
||||
provider_private_response_allows_sync_finalize, resolve_claude_stream_spec,
|
||||
resolve_claude_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
|
||||
resolve_local_image_stream_spec, resolve_local_image_sync_spec,
|
||||
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
|
||||
ExecutionRuntimeAuthContext, GatewayControlPlanRequest, LocalCoreSyncErrorKind,
|
||||
LocalOpenAiImageSpec, LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
|
||||
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
|
||||
StreamingStandardTerminalObserver, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_FILES_DOWNLOAD_PLAN_KIND,
|
||||
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
|
||||
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
|
||||
pub(crate) fn parse_direct_request_body(
|
||||
parts: &http::request::Parts,
|
||||
body_bytes: &axum::body::Bytes,
|
||||
) -> Option<(serde_json::Value, Option<String>)> {
|
||||
aether_ai_surfaces::api::parse_direct_request_body(
|
||||
is_json_request(&parts.headers),
|
||||
body_bytes.as_ref(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_execution_runtime_stream_plan_kind(
|
||||
parts: &http::request::Parts,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<&'static str> {
|
||||
aether_ai_surfaces::api::resolve_execution_runtime_stream_plan_kind(
|
||||
decision.route_class.as_deref(),
|
||||
decision.route_family.as_deref(),
|
||||
decision.route_kind.as_deref(),
|
||||
&parts.method,
|
||||
parts.uri.path(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_execution_runtime_sync_plan_kind(
|
||||
parts: &http::request::Parts,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<&'static str> {
|
||||
aether_ai_surfaces::api::resolve_execution_runtime_sync_plan_kind(
|
||||
decision.route_class.as_deref(),
|
||||
decision.route_family.as_deref(),
|
||||
decision.route_kind.as_deref(),
|
||||
&parts.method,
|
||||
parts.uri.path(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn is_matching_stream_request(
|
||||
plan_kind: &str,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> bool {
|
||||
crate::ai_serving::planner_is_matching_stream_request(plan_kind, parts, body_json, body_base64)
|
||||
}
|
||||
|
||||
pub(crate) fn supports_sync_scheduler_decision_kind(plan_kind: &str) -> bool {
|
||||
aether_ai_surfaces::api::supports_sync_scheduler_decision_kind(plan_kind)
|
||||
}
|
||||
|
||||
pub(crate) fn supports_stream_scheduler_decision_kind(plan_kind: &str) -> bool {
|
||||
aether_ai_surfaces::api::supports_stream_scheduler_decision_kind(plan_kind)
|
||||
}
|
||||
|
||||
pub(crate) fn aggregate_openai_chat_stream_sync_response(body: &[u8]) -> Option<serde_json::Value> {
|
||||
aether_ai_surfaces::api::aggregate_openai_chat_stream_sync_response(body)
|
||||
}
|
||||
|
||||
pub(crate) fn aggregate_openai_responses_stream_sync_response(
|
||||
body: &[u8],
|
||||
) -> Option<serde_json::Value> {
|
||||
aether_ai_surfaces::api::aggregate_openai_responses_stream_sync_response(body)
|
||||
}
|
||||
|
||||
pub(crate) fn aggregate_claude_stream_sync_response(body: &[u8]) -> Option<serde_json::Value> {
|
||||
aether_ai_surfaces::api::aggregate_claude_stream_sync_response(body)
|
||||
}
|
||||
|
||||
pub(crate) fn aggregate_gemini_stream_sync_response(body: &[u8]) -> Option<serde_json::Value> {
|
||||
aether_ai_surfaces::api::aggregate_gemini_stream_sync_response(body)
|
||||
}
|
||||
139
apps/aether-gateway/src/ai_serving/finalize/common.rs
Normal file
139
apps/aether-gateway/src/ai_serving/finalize/common.rs
Normal file
@@ -0,0 +1,139 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use serde_json::Value;
|
||||
|
||||
pub(crate) use crate::ai_serving::api::{
|
||||
normalize_provider_private_response_value as unwrap_local_finalize_response_value,
|
||||
provider_private_response_allows_sync_finalize as local_finalize_allows_envelope,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
build_generated_tool_call_id,
|
||||
build_local_success_background_report as build_local_success_background_report_impl,
|
||||
build_local_success_conversion_background_report as build_local_success_conversion_background_report_impl,
|
||||
canonicalize_tool_arguments,
|
||||
prepare_local_success_response_parts as prepare_local_success_response_parts_impl,
|
||||
GatewayControlDecision, LocalSyncReportParts,
|
||||
};
|
||||
use crate::api::response::build_client_response_from_parts;
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
|
||||
pub(crate) struct LocalCoreSyncFinalizeOutcome {
|
||||
pub(crate) response: Response<Body>,
|
||||
pub(crate) background_report: Option<GatewaySyncReportRequest>,
|
||||
}
|
||||
|
||||
fn build_local_success_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
status_code: u16,
|
||||
body_bytes: Vec<u8>,
|
||||
headers: BTreeMap<String, String>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
build_client_response_from_parts(
|
||||
status_code,
|
||||
&headers,
|
||||
Body::from(body_bytes),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)
|
||||
}
|
||||
|
||||
fn surface_report_parts_from_gateway(payload: &GatewaySyncReportRequest) -> LocalSyncReportParts {
|
||||
LocalSyncReportParts {
|
||||
trace_id: payload.trace_id.clone(),
|
||||
report_kind: payload.report_kind.clone(),
|
||||
report_context: payload.report_context.clone(),
|
||||
status_code: payload.status_code,
|
||||
headers: payload.headers.clone(),
|
||||
body_json: payload.body_json.clone(),
|
||||
client_body_json: payload.client_body_json.clone(),
|
||||
body_base64: payload.body_base64.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
fn gateway_report_from_surface(
|
||||
source: &GatewaySyncReportRequest,
|
||||
report: LocalSyncReportParts,
|
||||
) -> GatewaySyncReportRequest {
|
||||
GatewaySyncReportRequest {
|
||||
trace_id: report.trace_id,
|
||||
report_kind: report.report_kind,
|
||||
report_context: report.report_context,
|
||||
status_code: report.status_code,
|
||||
headers: report.headers,
|
||||
body_json: report.body_json,
|
||||
client_body_json: report.client_body_json,
|
||||
body_base64: report.body_base64,
|
||||
telemetry: source.telemetry.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_success_outcome(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
body_json: Value,
|
||||
) -> Result<LocalCoreSyncFinalizeOutcome, GatewayError> {
|
||||
let report_headers = payload.headers.clone();
|
||||
let (body_bytes, response_headers) =
|
||||
prepare_local_success_response_parts_impl(&payload.headers, &body_json)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let surface_payload = surface_report_parts_from_gateway(payload);
|
||||
let background_report =
|
||||
build_local_success_background_report_impl(&surface_payload, body_json, report_headers)
|
||||
.map(|report| gateway_report_from_surface(payload, report));
|
||||
build_local_success_outcome_with_report(
|
||||
trace_id,
|
||||
decision,
|
||||
payload.status_code,
|
||||
body_bytes,
|
||||
response_headers,
|
||||
background_report,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_success_outcome_with_report(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
status_code: u16,
|
||||
body_bytes: Vec<u8>,
|
||||
headers: BTreeMap<String, String>,
|
||||
background_report: Option<GatewaySyncReportRequest>,
|
||||
) -> Result<LocalCoreSyncFinalizeOutcome, GatewayError> {
|
||||
let response =
|
||||
build_local_success_response(trace_id, decision, status_code, body_bytes, headers)?;
|
||||
Ok(LocalCoreSyncFinalizeOutcome {
|
||||
response,
|
||||
background_report,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_success_outcome_with_conversion_report(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
client_body_json: Value,
|
||||
provider_body_json: Value,
|
||||
) -> Result<LocalCoreSyncFinalizeOutcome, GatewayError> {
|
||||
let (body_bytes, response_headers) =
|
||||
prepare_local_success_response_parts_impl(&payload.headers, &client_body_json)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let surface_payload = surface_report_parts_from_gateway(payload);
|
||||
let report_payload = build_local_success_conversion_background_report_impl(
|
||||
&surface_payload,
|
||||
client_body_json,
|
||||
provider_body_json,
|
||||
)
|
||||
.map(|report| gateway_report_from_surface(payload, report));
|
||||
|
||||
build_local_success_outcome_with_report(
|
||||
trace_id,
|
||||
decision,
|
||||
payload.status_code,
|
||||
body_bytes,
|
||||
response_headers,
|
||||
report_payload,
|
||||
)
|
||||
}
|
||||
44
apps/aether-gateway/src/ai_serving/finalize/internal/mod.rs
Normal file
44
apps/aether-gateway/src/ai_serving/finalize/internal/mod.rs
Normal file
@@ -0,0 +1,44 @@
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
|
||||
#[path = "stream_rewrite.rs"]
|
||||
pub(crate) mod stream;
|
||||
#[path = "sync_finalize.rs"]
|
||||
pub(crate) mod sync;
|
||||
#[path = "sync_to_stream.rs"]
|
||||
pub(crate) mod sync_to_stream;
|
||||
|
||||
pub(crate) use stream::LocalStreamRewriter;
|
||||
pub(crate) use sync::LocalCoreSyncFinalizeOutcome;
|
||||
pub(crate) use sync_to_stream::{
|
||||
maybe_bridge_standard_sync_json_to_stream, SyncToStreamBridgeOutcome,
|
||||
};
|
||||
|
||||
pub(crate) fn maybe_build_sync_finalize_outcome(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<LocalCoreSyncFinalizeOutcome>, GatewayError> {
|
||||
sync::maybe_build_local_core_sync_finalize_response(trace_id, decision, payload)
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_compile_sync_finalize_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
Ok(
|
||||
maybe_build_sync_finalize_outcome(trace_id, decision, payload)?
|
||||
.map(|outcome| outcome.response),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_build_stream_response_rewriter(
|
||||
report_context: Option<&Value>,
|
||||
) -> Option<LocalStreamRewriter<'_>> {
|
||||
stream::maybe_build_local_stream_rewriter(report_context)
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::{
|
||||
maybe_build_ai_surface_stream_rewriter, AiSurfaceFinalizeError, AiSurfaceStreamRewriter,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
|
||||
pub(crate) struct LocalStreamRewriter<'a> {
|
||||
inner: AiSurfaceStreamRewriter<'a>,
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_build_local_stream_rewriter<'a>(
|
||||
report_context: Option<&'a Value>,
|
||||
) -> Option<LocalStreamRewriter<'a>> {
|
||||
maybe_build_ai_surface_stream_rewriter(report_context)
|
||||
.map(|inner| LocalStreamRewriter { inner })
|
||||
}
|
||||
|
||||
impl LocalStreamRewriter<'_> {
|
||||
pub(crate) fn push_chunk(&mut self, chunk: &[u8]) -> Result<Vec<u8>, GatewayError> {
|
||||
self.inner.push_chunk(chunk).map_err(map_surface_error)
|
||||
}
|
||||
|
||||
pub(crate) fn finish(&mut self) -> Result<Vec<u8>, GatewayError> {
|
||||
self.inner.finish().map_err(map_surface_error)
|
||||
}
|
||||
}
|
||||
|
||||
fn map_surface_error(error: AiSurfaceFinalizeError) -> GatewayError {
|
||||
error.into()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "../tests_stream.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,113 @@
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::ai_serving::{build_generated_tool_call_id, canonicalize_tool_arguments};
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
|
||||
pub(crate) use crate::ai_serving::finalize::common::{
|
||||
build_local_success_outcome, build_local_success_outcome_with_conversion_report,
|
||||
local_finalize_allows_envelope, unwrap_local_finalize_response_value,
|
||||
LocalCoreSyncFinalizeOutcome,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::finalize::standard::{
|
||||
maybe_build_standard_sync_finalize_product_from_normalized_payload,
|
||||
StandardSyncFinalizeNormalizedProduct,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::{
|
||||
aggregate_claude_stream_sync_response, aggregate_gemini_stream_sync_response,
|
||||
aggregate_openai_chat_stream_sync_response, aggregate_openai_responses_stream_sync_response,
|
||||
maybe_build_openai_image_sync_finalize_product,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::{
|
||||
convert_claude_chat_response_to_openai_chat, convert_claude_response_to_openai_responses,
|
||||
convert_gemini_chat_response_to_openai_chat, convert_gemini_response_to_openai_responses,
|
||||
};
|
||||
|
||||
pub(crate) fn maybe_build_local_core_sync_finalize_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<LocalCoreSyncFinalizeOutcome>, GatewayError> {
|
||||
if let Some(outcome) =
|
||||
maybe_build_local_openai_image_sync_finalize_response(trace_id, decision, payload)?
|
||||
{
|
||||
return Ok(Some(outcome));
|
||||
}
|
||||
|
||||
let Some(normalized_payload) =
|
||||
crate::ai_serving::adaptation::private_envelope::maybe_normalize_provider_private_sync_report_payload(payload)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let payload = &normalized_payload;
|
||||
let Some(report_context) = payload.report_context.as_ref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !local_finalize_allows_envelope(report_context) {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(product) = maybe_build_standard_sync_finalize_product_from_normalized_payload(
|
||||
payload.report_kind.as_str(),
|
||||
payload.status_code,
|
||||
Some(report_context),
|
||||
payload.body_json.as_ref(),
|
||||
payload.body_base64.as_deref(),
|
||||
)
|
||||
.map_err(GatewayError::from)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
match product {
|
||||
StandardSyncFinalizeNormalizedProduct::SuccessBody(body_json) => {
|
||||
let Some(body_json) = unwrap_local_finalize_response_value(body_json, report_context)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some(build_local_success_outcome(
|
||||
trace_id, decision, payload, body_json,
|
||||
)?))
|
||||
}
|
||||
StandardSyncFinalizeNormalizedProduct::CrossFormat(product) => {
|
||||
let Some(provider_body_json) =
|
||||
unwrap_local_finalize_response_value(product.provider_body_json, report_context)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some(build_local_success_outcome_with_conversion_report(
|
||||
trace_id,
|
||||
decision,
|
||||
payload,
|
||||
product.client_body_json,
|
||||
provider_body_json,
|
||||
)?))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn maybe_build_local_openai_image_sync_finalize_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<LocalCoreSyncFinalizeOutcome>, GatewayError> {
|
||||
let Some(product) = maybe_build_openai_image_sync_finalize_product(
|
||||
payload.report_kind.as_str(),
|
||||
payload.status_code,
|
||||
payload.report_context.as_ref(),
|
||||
payload.body_base64.as_deref(),
|
||||
)
|
||||
.map_err(GatewayError::from)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
Ok(Some(build_local_success_outcome_with_conversion_report(
|
||||
trace_id,
|
||||
decision,
|
||||
payload,
|
||||
product.client_body_json,
|
||||
product.provider_body_json,
|
||||
)?))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "../tests_sync.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,20 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::GatewayError;
|
||||
|
||||
pub(crate) use crate::ai_serving::pure::SyncToStreamBridgeOutcome;
|
||||
|
||||
pub(crate) fn maybe_bridge_standard_sync_json_to_stream(
|
||||
provider_body_json: &Value,
|
||||
provider_api_format: &str,
|
||||
client_api_format: &str,
|
||||
report_context: Option<&Value>,
|
||||
) -> Result<Option<SyncToStreamBridgeOutcome>, GatewayError> {
|
||||
crate::ai_serving::pure::maybe_bridge_standard_sync_json_to_stream(
|
||||
provider_body_json,
|
||||
provider_api_format,
|
||||
client_api_format,
|
||||
report_context,
|
||||
)
|
||||
.map_err(GatewayError::from)
|
||||
}
|
||||
4
apps/aether-gateway/src/ai_serving/finalize/mod.rs
Normal file
4
apps/aether-gateway/src/ai_serving/finalize/mod.rs
Normal file
@@ -0,0 +1,4 @@
|
||||
pub(crate) mod common;
|
||||
pub(crate) mod internal;
|
||||
pub(crate) mod sse;
|
||||
pub(crate) mod standard;
|
||||
15
apps/aether-gateway/src/ai_serving/finalize/sse.rs
Normal file
15
apps/aether-gateway/src/ai_serving/finalize/sse.rs
Normal file
@@ -0,0 +1,15 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::{
|
||||
encode_done_sse, encode_json_sse as encode_json_sse_impl, map_claude_stop_reason,
|
||||
AiSurfaceFinalizeError,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
|
||||
fn map_error(err: AiSurfaceFinalizeError) -> GatewayError {
|
||||
err.into()
|
||||
}
|
||||
|
||||
pub(crate) fn encode_json_sse(event: Option<&str>, value: &Value) -> Result<Vec<u8>, GatewayError> {
|
||||
encode_json_sse_impl(event, value).map_err(map_error)
|
||||
}
|
||||
19
apps/aether-gateway/src/ai_serving/finalize/standard/mod.rs
Normal file
19
apps/aether-gateway/src/ai_serving/finalize/standard/mod.rs
Normal file
@@ -0,0 +1,19 @@
|
||||
//! Standard finalize surface for standard contract sync/stream compilation.
|
||||
|
||||
pub(crate) use crate::ai_serving::{
|
||||
aggregate_standard_chat_stream_sync_response, aggregate_standard_cli_stream_sync_response,
|
||||
build_openai_responses_response, convert_claude_chat_response_to_openai_chat,
|
||||
convert_claude_response_to_openai_responses, convert_gemini_chat_response_to_openai_chat,
|
||||
convert_gemini_response_to_openai_responses, convert_openai_chat_response_to_claude_chat,
|
||||
convert_openai_chat_response_to_gemini_chat, convert_openai_chat_response_to_openai_responses,
|
||||
convert_openai_responses_response_to_openai_chat, convert_standard_chat_response,
|
||||
convert_standard_cli_response,
|
||||
maybe_build_openai_chat_cross_format_sync_product_from_normalized_payload,
|
||||
maybe_build_openai_responses_cross_format_sync_product_from_normalized_payload,
|
||||
maybe_build_openai_responses_same_family_sync_body_from_normalized_payload,
|
||||
maybe_build_standard_cross_format_sync_product,
|
||||
maybe_build_standard_cross_format_sync_product_from_normalized_payload,
|
||||
maybe_build_standard_same_format_sync_body_from_normalized_payload,
|
||||
maybe_build_standard_sync_finalize_product_from_normalized_payload,
|
||||
StandardCrossFormatSyncProduct, StandardSyncFinalizeNormalizedProduct,
|
||||
};
|
||||
1099
apps/aether-gateway/src/ai_serving/finalize/tests_stream.rs
Normal file
1099
apps/aether-gateway/src/ai_serving/finalize/tests_stream.rs
Normal file
File diff suppressed because it is too large
Load Diff
2165
apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs
Normal file
2165
apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs
Normal file
File diff suppressed because it is too large
Load Diff
166
apps/aether-gateway/src/ai_serving/mod.rs
Normal file
166
apps/aether-gateway/src/ai_serving/mod.rs
Normal file
@@ -0,0 +1,166 @@
|
||||
mod adaptation;
|
||||
pub(crate) mod api;
|
||||
mod finalize;
|
||||
mod planner;
|
||||
mod pure;
|
||||
pub(crate) mod transport;
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::http::{Response, Uri};
|
||||
|
||||
use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError};
|
||||
|
||||
pub(crate) use self::adaptation::{
|
||||
maybe_build_provider_private_stream_normalizer, ProviderPrivateStreamNormalizer,
|
||||
};
|
||||
pub(crate) use self::finalize::common::LocalCoreSyncFinalizeOutcome;
|
||||
pub(crate) use self::finalize::internal::{
|
||||
maybe_bridge_standard_sync_json_to_stream, maybe_build_stream_response_rewriter,
|
||||
maybe_build_sync_finalize_outcome, maybe_compile_sync_finalize_response,
|
||||
SyncToStreamBridgeOutcome,
|
||||
};
|
||||
pub(crate) use self::planner::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_local_gemini_files_stream_plan_and_reports_for_kind,
|
||||
build_local_gemini_files_sync_plan_and_reports_for_kind,
|
||||
build_local_image_stream_plan_and_reports_for_kind,
|
||||
build_local_image_sync_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_sync_plan_and_reports_for_kind,
|
||||
build_local_openai_responses_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_responses_sync_plan_and_reports_for_kind,
|
||||
build_local_same_format_stream_plan_and_reports, build_local_same_format_sync_plan_and_reports,
|
||||
build_local_video_sync_plan_and_reports_for_kind,
|
||||
build_openai_responses_stream_plan_from_decision,
|
||||
build_openai_responses_sync_plan_from_decision, build_passthrough_sync_plan_from_decision,
|
||||
build_standard_family_stream_plan_and_reports, build_standard_family_sync_plan_and_reports,
|
||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
||||
extract_pool_sticky_session_token, maybe_build_stream_decision_payload,
|
||||
maybe_build_stream_plan_payload, maybe_build_sync_decision_payload,
|
||||
maybe_build_sync_plan_payload, planner_is_matching_stream_request,
|
||||
set_local_openai_chat_execution_exhausted_diagnostic, CandidateFailureDiagnostic,
|
||||
CandidateFailureDiagnosticKind, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot,
|
||||
LocalResolvedOAuthRequestAuth, PlannerAppState,
|
||||
};
|
||||
pub(crate) use self::pure::*;
|
||||
pub(crate) use self::transport::{
|
||||
append_transport_diagnostics_to_value, build_request_trace_proxy_value,
|
||||
candidate_common_transport_skip_reason, candidate_transport_pair_skip_reason,
|
||||
request_conversion_direct_auth, request_conversion_enabled_for_transport,
|
||||
request_conversion_transport_supported, request_conversion_transport_unsupported_reason,
|
||||
request_pair_allowed_for_transport, CandidateTransportPolicyFacts,
|
||||
};
|
||||
pub(crate) use crate::control::GatewayControlDecision;
|
||||
pub(crate) use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
|
||||
pub(crate) use crate::headers::RequestOrigin;
|
||||
pub(crate) use aether_ai_serving::{
|
||||
ai_local_execution_contract_for_formats, augment_sync_report_context,
|
||||
build_ai_report_context_original_request_echo as build_report_context_original_request_echo,
|
||||
extract_ai_gemini_model_from_path as extract_gemini_model_from_path,
|
||||
generic_decision_missing_exact_provider_request as generic_decision_missing_exact_provider_request_impl,
|
||||
AiExecutionDecision, AiExecutionPlanPayload, AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
|
||||
pub(crate) fn build_provider_transport_request_url(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
provider_api_format: &str,
|
||||
mapped_model: Option<&str>,
|
||||
upstream_is_stream: bool,
|
||||
request_query: Option<&str>,
|
||||
kiro_api_region: Option<&str>,
|
||||
) -> Option<String> {
|
||||
self::transport::build_transport_request_url(
|
||||
transport,
|
||||
self::transport::TransportRequestUrlParams {
|
||||
provider_api_format,
|
||||
mapped_model,
|
||||
upstream_is_stream,
|
||||
request_query,
|
||||
kiro_api_region,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_execution_runtime_auth_context(
|
||||
state: &AppState,
|
||||
decision: &GatewayControlDecision,
|
||||
headers: &http::HeaderMap,
|
||||
uri: &Uri,
|
||||
trace_id: &str,
|
||||
) -> Result<Option<crate::control::GatewayControlAuthContext>, GatewayError> {
|
||||
crate::control::resolve_execution_runtime_auth_context(state, decision, headers, uri, trace_id)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) fn collect_control_headers(
|
||||
headers: &http::HeaderMap,
|
||||
) -> std::collections::BTreeMap<String, String> {
|
||||
crate::headers::collect_control_headers(headers)
|
||||
}
|
||||
|
||||
pub(crate) fn request_origin_from_headers(headers: &http::HeaderMap) -> RequestOrigin {
|
||||
crate::headers::request_origin_from_headers(headers)
|
||||
}
|
||||
|
||||
pub(crate) fn request_origin_from_parts(parts: &http::request::Parts) -> RequestOrigin {
|
||||
crate::headers::request_origin_from_parts(parts)
|
||||
}
|
||||
|
||||
pub(crate) fn is_json_request(headers: &http::HeaderMap) -> bool {
|
||||
crate::headers::is_json_request(headers)
|
||||
}
|
||||
|
||||
pub(crate) fn build_execution_runtime_auth_context(
|
||||
auth_context: &crate::control::GatewayControlAuthContext,
|
||||
) -> ExecutionRuntimeAuthContext {
|
||||
ExecutionRuntimeAuthContext {
|
||||
user_id: auth_context.user_id.clone(),
|
||||
api_key_id: auth_context.api_key_id.clone(),
|
||||
username: auth_context.username.clone(),
|
||||
api_key_name: auth_context.api_key_name.clone(),
|
||||
balance_remaining: auth_context.balance_remaining,
|
||||
access_allowed: auth_context.access_allowed,
|
||||
api_key_is_standalone: auth_context.api_key_is_standalone,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_decision_execution_runtime_auth_context(
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<ExecutionRuntimeAuthContext> {
|
||||
decision
|
||||
.auth_context
|
||||
.as_ref()
|
||||
.map(build_execution_runtime_auth_context)
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_local_decision_execution_runtime_auth_context(
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<ExecutionRuntimeAuthContext> {
|
||||
resolve_decision_execution_runtime_auth_context(decision).filter(|auth_context| {
|
||||
!auth_context.user_id.trim().is_empty() && !auth_context.api_key_id.trim().is_empty()
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn generic_decision_missing_exact_provider_request(
|
||||
payload: &AiExecutionDecision,
|
||||
) -> bool {
|
||||
if !generic_decision_missing_exact_provider_request_impl(payload) {
|
||||
return false;
|
||||
}
|
||||
|
||||
tracing::warn!(
|
||||
decision_kind = payload.decision_kind.as_deref().unwrap_or_default(),
|
||||
provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default(),
|
||||
client_api_format = payload.client_api_format.as_deref().unwrap_or_default(),
|
||||
"gateway generic decision missing exact provider request; falling back to plan"
|
||||
);
|
||||
true
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_build_local_sync_finalize_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
crate::execution_runtime::maybe_build_local_sync_finalize_response(trace_id, decision, payload)
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
use aether_scheduler_core::{
|
||||
build_scheduler_affinity_cache_key_for_api_key_id, SchedulerAffinityTarget,
|
||||
SchedulerMinimalCandidateSelectionCandidate,
|
||||
};
|
||||
|
||||
use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
|
||||
|
||||
const PLANNER_SCHEDULER_AFFINITY_MAX_ENTRIES: usize = 10_000;
|
||||
|
||||
pub(crate) fn read_cached_scheduler_affinity_target(
|
||||
state: PlannerAppState<'_>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_api_format: &str,
|
||||
requested_model: Option<&str>,
|
||||
) -> Option<SchedulerAffinityTarget> {
|
||||
let requested_model = requested_model
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())?;
|
||||
let api_key_id = auth_snapshot
|
||||
.map(|snapshot| snapshot.api_key_id.trim())
|
||||
.filter(|value| !value.is_empty())?;
|
||||
let cache_key = build_scheduler_affinity_cache_key_for_api_key_id(
|
||||
api_key_id,
|
||||
client_api_format,
|
||||
requested_model,
|
||||
)?;
|
||||
|
||||
state
|
||||
.app()
|
||||
.read_scheduler_affinity_target(&cache_key, SCHEDULER_AFFINITY_TTL)
|
||||
}
|
||||
|
||||
pub(crate) fn remember_scheduler_affinity_for_candidate(
|
||||
state: PlannerAppState<'_>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_api_format: &str,
|
||||
requested_model: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) {
|
||||
let Some(api_key_id) = auth_snapshot
|
||||
.map(|snapshot| snapshot.api_key_id.trim())
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let Some(cache_key) = build_scheduler_affinity_cache_key_for_api_key_id(
|
||||
api_key_id,
|
||||
client_api_format,
|
||||
requested_model,
|
||||
) else {
|
||||
return;
|
||||
};
|
||||
|
||||
state.app().remember_scheduler_affinity_target(
|
||||
&cache_key,
|
||||
SchedulerAffinityTarget {
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
},
|
||||
SCHEDULER_AFFINITY_TTL,
|
||||
PLANNER_SCHEDULER_AFFINITY_MAX_ENTRIES,
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,887 @@
|
||||
use aether_ai_serving::{
|
||||
ai_candidate_extra_data_with_ranking, ai_should_persist_available_candidate_for_pool_key,
|
||||
ai_should_persist_skipped_candidate_for_pool_membership,
|
||||
run_ai_available_candidate_persistence, run_ai_candidate_materialization,
|
||||
run_ai_skipped_candidate_persistence, AiAvailableCandidatePersistencePort,
|
||||
AiCandidateMaterializationOutcome, AiCandidateMaterializationPort, AiCandidateResolutionMode,
|
||||
AiSkippedCandidatePersistencePort,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
use std::convert::Infallible;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_serving::planner::candidate_affinity_cache::remember_scheduler_affinity_for_candidate;
|
||||
use crate::ai_serving::planner::candidate_resolution::{
|
||||
resolve_and_rank_local_execution_candidates,
|
||||
resolve_and_rank_local_execution_candidates_without_transport_pair_gate,
|
||||
EligibleLocalExecutionCandidate, SkippedLocalExecutionCandidate,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::LocalCandidatePersistencePolicy;
|
||||
use crate::ai_serving::planner::runtime_miss::record_local_runtime_candidate_skip_reason;
|
||||
use crate::ai_serving::planner::CandidateFailureDiagnostic;
|
||||
use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::clock::current_unix_ms;
|
||||
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
|
||||
use crate::orchestration::{local_attempt_slot_count, ExecutionAttemptIdentity};
|
||||
use crate::AppState;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalExecutionCandidateAttempt {
|
||||
pub(crate) eligible: EligibleLocalExecutionCandidate,
|
||||
pub(crate) candidate_index: u32,
|
||||
pub(crate) retry_index: u32,
|
||||
pub(crate) candidate_id: String,
|
||||
}
|
||||
|
||||
impl LocalExecutionCandidateAttempt {
|
||||
pub(crate) fn attempt_identity(&self) -> ExecutionAttemptIdentity {
|
||||
ExecutionAttemptIdentity::new(self.candidate_index, self.retry_index)
|
||||
.with_pool_key_index(self.eligible.orchestration.pool_key_index)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct LocalAvailableCandidatePersistenceContext<'a> {
|
||||
pub(crate) user_id: &'a str,
|
||||
pub(crate) api_key_id: &'a str,
|
||||
pub(crate) required_capabilities: Option<&'a Value>,
|
||||
pub(crate) error_context: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct LocalSkippedCandidatePersistenceContext<'a> {
|
||||
pub(crate) user_id: &'a str,
|
||||
pub(crate) api_key_id: &'a str,
|
||||
pub(crate) required_capabilities: Option<&'a Value>,
|
||||
pub(crate) error_context: &'static str,
|
||||
pub(crate) record_runtime_miss_diagnostic: bool,
|
||||
}
|
||||
|
||||
pub(crate) use aether_ai_serving::AiCandidateResolutionMode as LocalCandidateResolutionMode;
|
||||
|
||||
struct GatewayLocalCandidateMaterializationPort<'a, F, G> {
|
||||
state: PlannerAppState<'a>,
|
||||
trace_id: &'a str,
|
||||
client_api_format: &'a str,
|
||||
requested_model: Option<&'a str>,
|
||||
auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&'a Value>,
|
||||
sticky_session_token: Option<&'a str>,
|
||||
persistence_policy: LocalCandidatePersistencePolicy<'a>,
|
||||
resolution_mode: LocalCandidateResolutionMode,
|
||||
build_available_extra_data: F,
|
||||
decorate_skipped_candidate: G,
|
||||
}
|
||||
|
||||
struct GatewayAvailableCandidatePersistencePort<'a, F> {
|
||||
state: PlannerAppState<'a>,
|
||||
trace_id: &'a str,
|
||||
user_id: &'a str,
|
||||
api_key_id: &'a str,
|
||||
required_capabilities: Option<&'a Value>,
|
||||
error_context: &'static str,
|
||||
created_at_unix_ms: u64,
|
||||
build_extra_data: F,
|
||||
}
|
||||
|
||||
struct GatewaySkippedCandidatePersistencePort<'a> {
|
||||
state: &'a AppState,
|
||||
trace_id: &'a str,
|
||||
user_id: &'a str,
|
||||
api_key_id: &'a str,
|
||||
required_capabilities: Option<&'a Value>,
|
||||
error_context: &'static str,
|
||||
record_runtime_miss_diagnostic: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<F, G> AiCandidateMaterializationPort for GatewayLocalCandidateMaterializationPort<'_, F, G>
|
||||
where
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
|
||||
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync,
|
||||
{
|
||||
type Candidate = SchedulerMinimalCandidateSelectionCandidate;
|
||||
type Eligible = EligibleLocalExecutionCandidate;
|
||||
type Skipped = SkippedLocalExecutionCandidate;
|
||||
type Attempt = LocalExecutionCandidateAttempt;
|
||||
type Error = Infallible;
|
||||
|
||||
async fn resolve_and_rank_candidates(
|
||||
&self,
|
||||
candidates: Vec<Self::Candidate>,
|
||||
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error> {
|
||||
let requested_model = self.requested_model.map(str::to_string);
|
||||
let resolved = match self.resolution_mode {
|
||||
AiCandidateResolutionMode::Standard => {
|
||||
resolve_and_rank_local_execution_candidates(
|
||||
self.state,
|
||||
candidates,
|
||||
self.client_api_format,
|
||||
requested_model.as_deref().unwrap_or_default(),
|
||||
self.auth_snapshot,
|
||||
self.required_capabilities,
|
||||
self.sticky_session_token,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiCandidateResolutionMode::WithoutTransportPairGate => {
|
||||
resolve_and_rank_local_execution_candidates_without_transport_pair_gate(
|
||||
self.state,
|
||||
candidates,
|
||||
self.client_api_format,
|
||||
requested_model.as_deref(),
|
||||
self.auth_snapshot,
|
||||
self.required_capabilities,
|
||||
self.sticky_session_token,
|
||||
)
|
||||
.await
|
||||
}
|
||||
};
|
||||
Ok(resolved)
|
||||
}
|
||||
|
||||
fn decorate_skipped_candidate(&self, skipped: Self::Skipped) -> Self::Skipped {
|
||||
(self.decorate_skipped_candidate)(skipped)
|
||||
}
|
||||
|
||||
fn remember_first_candidate_affinity(&self, candidates: &[Self::Eligible]) {
|
||||
remember_first_local_candidate_affinity(
|
||||
self.state,
|
||||
self.auth_snapshot,
|
||||
self.client_api_format,
|
||||
self.requested_model,
|
||||
candidates,
|
||||
);
|
||||
}
|
||||
|
||||
async fn persist_available_candidates(
|
||||
&self,
|
||||
candidates: Vec<Self::Eligible>,
|
||||
) -> Result<Vec<Self::Attempt>, Self::Error> {
|
||||
Ok(persist_available_local_execution_candidates_with_context(
|
||||
self.state,
|
||||
self.trace_id,
|
||||
self.persistence_policy.available,
|
||||
candidates,
|
||||
&self.build_available_extra_data,
|
||||
)
|
||||
.await)
|
||||
}
|
||||
|
||||
async fn persist_skipped_candidates(
|
||||
&self,
|
||||
starting_candidate_index: u32,
|
||||
skipped_candidates: Vec<Self::Skipped>,
|
||||
) -> Result<(), Self::Error> {
|
||||
persist_skipped_local_execution_candidates_with_context(
|
||||
self.state.app(),
|
||||
self.trace_id,
|
||||
self.persistence_policy.skipped,
|
||||
starting_candidate_index,
|
||||
skipped_candidates,
|
||||
)
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<F> AiAvailableCandidatePersistencePort for GatewayAvailableCandidatePersistencePort<'_, F>
|
||||
where
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
|
||||
{
|
||||
type Candidate = EligibleLocalExecutionCandidate;
|
||||
type Attempt = LocalExecutionCandidateAttempt;
|
||||
type ExtraData = Value;
|
||||
type Error = Infallible;
|
||||
|
||||
fn attempt_slot_count(&self, candidate: &Self::Candidate) -> u32 {
|
||||
local_attempt_slot_count(&candidate.transport)
|
||||
}
|
||||
|
||||
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData> {
|
||||
ai_candidate_extra_data_with_ranking(
|
||||
(self.build_extra_data)(candidate),
|
||||
candidate.ranking.as_ref(),
|
||||
)
|
||||
}
|
||||
|
||||
fn generate_candidate_id(&self) -> String {
|
||||
Uuid::new_v4().to_string()
|
||||
}
|
||||
|
||||
fn should_persist_available_candidate(&self, candidate: &Self::Candidate) -> bool {
|
||||
should_persist_available_local_candidate(candidate)
|
||||
}
|
||||
|
||||
async fn persist_available_candidate(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
candidate_index: u32,
|
||||
retry_index: u32,
|
||||
generated_candidate_id: &str,
|
||||
extra_data: Option<Self::ExtraData>,
|
||||
) -> Result<String, Self::Error> {
|
||||
Ok(self
|
||||
.state
|
||||
.persist_available_local_candidate(
|
||||
self.trace_id,
|
||||
self.user_id,
|
||||
self.api_key_id,
|
||||
&candidate.candidate,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
generated_candidate_id,
|
||||
self.required_capabilities,
|
||||
extra_data,
|
||||
self.created_at_unix_ms,
|
||||
self.error_context,
|
||||
)
|
||||
.await)
|
||||
}
|
||||
|
||||
fn build_attempt(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
candidate_index: u32,
|
||||
retry_index: u32,
|
||||
candidate_id: String,
|
||||
) -> Self::Attempt {
|
||||
LocalExecutionCandidateAttempt {
|
||||
eligible: candidate,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
candidate_id,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiSkippedCandidatePersistencePort for GatewaySkippedCandidatePersistencePort<'_> {
|
||||
type Skipped = SkippedLocalExecutionCandidate;
|
||||
type ExtraData = Value;
|
||||
type Error = Infallible;
|
||||
|
||||
fn should_persist_skipped_candidate(&self, candidate: &Self::Skipped) -> bool {
|
||||
should_persist_skipped_local_candidate(candidate)
|
||||
}
|
||||
|
||||
fn build_extra_data(&self, candidate: &Self::Skipped) -> Option<Self::ExtraData> {
|
||||
ai_candidate_extra_data_with_ranking(
|
||||
candidate.extra_data.clone(),
|
||||
candidate.ranking.as_ref(),
|
||||
)
|
||||
}
|
||||
|
||||
fn generate_candidate_id(&self) -> String {
|
||||
Uuid::new_v4().to_string()
|
||||
}
|
||||
|
||||
async fn persist_skipped_candidate(
|
||||
&self,
|
||||
candidate: &Self::Skipped,
|
||||
candidate_index: u32,
|
||||
generated_candidate_id: &str,
|
||||
extra_data: Option<Self::ExtraData>,
|
||||
) -> Result<(), Self::Error> {
|
||||
persist_skipped_local_execution_candidate(
|
||||
self.state,
|
||||
self.trace_id,
|
||||
self.user_id,
|
||||
self.api_key_id,
|
||||
&candidate.candidate,
|
||||
candidate_index,
|
||||
generated_candidate_id,
|
||||
self.required_capabilities,
|
||||
candidate.skip_reason,
|
||||
extra_data,
|
||||
self.error_context,
|
||||
self.record_runtime_miss_diagnostic,
|
||||
)
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn materialize_local_execution_candidates_with_serving<F, G>(
|
||||
state: PlannerAppState<'_>,
|
||||
trace_id: &str,
|
||||
client_api_format: &str,
|
||||
requested_model: Option<&str>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&Value>,
|
||||
sticky_session_token: Option<&str>,
|
||||
persistence_policy: LocalCandidatePersistencePolicy<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
preselection_skipped: Vec<SkippedLocalExecutionCandidate>,
|
||||
resolution_mode: LocalCandidateResolutionMode,
|
||||
build_available_extra_data: F,
|
||||
decorate_skipped_candidate: G,
|
||||
) -> AiCandidateMaterializationOutcome<LocalExecutionCandidateAttempt>
|
||||
where
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
|
||||
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync,
|
||||
{
|
||||
let port = GatewayLocalCandidateMaterializationPort {
|
||||
state,
|
||||
trace_id,
|
||||
client_api_format,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
sticky_session_token,
|
||||
persistence_policy,
|
||||
resolution_mode,
|
||||
build_available_extra_data,
|
||||
decorate_skipped_candidate,
|
||||
};
|
||||
|
||||
match run_ai_candidate_materialization(&port, candidates, preselection_skipped).await {
|
||||
Ok(outcome) => outcome,
|
||||
Err(error) => match error {},
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn remember_first_local_candidate_affinity(
|
||||
state: PlannerAppState<'_>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_api_format: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidates: &[EligibleLocalExecutionCandidate],
|
||||
) {
|
||||
let Some(first_candidate) = candidates.first() else {
|
||||
return;
|
||||
};
|
||||
let affinity_requested_model = requested_model
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(first_candidate.candidate.global_model_name.as_str());
|
||||
remember_scheduler_affinity_for_candidate(
|
||||
state,
|
||||
auth_snapshot,
|
||||
client_api_format,
|
||||
affinity_requested_model,
|
||||
&first_candidate.candidate,
|
||||
);
|
||||
}
|
||||
|
||||
fn should_persist_available_local_candidate(eligible: &EligibleLocalExecutionCandidate) -> bool {
|
||||
ai_should_persist_available_candidate_for_pool_key(eligible.orchestration.pool_key_index)
|
||||
}
|
||||
|
||||
fn should_persist_skipped_local_candidate(candidate: &SkippedLocalExecutionCandidate) -> bool {
|
||||
let is_pool_candidate = candidate.transport.as_ref().is_some_and(|transport| {
|
||||
admin_provider_pool_config_from_config_value(transport.provider.config.as_ref()).is_some()
|
||||
});
|
||||
ai_should_persist_skipped_candidate_for_pool_membership(is_pool_candidate)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn persist_available_local_execution_candidates<F>(
|
||||
state: PlannerAppState<'_>,
|
||||
trace_id: &str,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
required_capabilities: Option<&Value>,
|
||||
candidates: Vec<EligibleLocalExecutionCandidate>,
|
||||
error_context: &'static str,
|
||||
build_extra_data: F,
|
||||
) -> Vec<LocalExecutionCandidateAttempt>
|
||||
where
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
|
||||
{
|
||||
let port = GatewayAvailableCandidatePersistencePort {
|
||||
state,
|
||||
trace_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
required_capabilities,
|
||||
error_context,
|
||||
created_at_unix_ms: current_unix_ms(),
|
||||
build_extra_data,
|
||||
};
|
||||
|
||||
match run_ai_available_candidate_persistence(&port, candidates).await {
|
||||
Ok(attempts) => attempts,
|
||||
Err(error) => match error {},
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn persist_available_local_execution_candidates_with_context<F>(
|
||||
state: PlannerAppState<'_>,
|
||||
trace_id: &str,
|
||||
context: LocalAvailableCandidatePersistenceContext<'_>,
|
||||
candidates: Vec<EligibleLocalExecutionCandidate>,
|
||||
build_extra_data: F,
|
||||
) -> Vec<LocalExecutionCandidateAttempt>
|
||||
where
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
|
||||
{
|
||||
persist_available_local_execution_candidates(
|
||||
state,
|
||||
trace_id,
|
||||
context.user_id,
|
||||
context.api_key_id,
|
||||
context.required_capabilities,
|
||||
candidates,
|
||||
context.error_context,
|
||||
build_extra_data,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn persist_skipped_local_execution_candidate(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
required_capabilities: Option<&Value>,
|
||||
skip_reason: &'static str,
|
||||
extra_data: Option<Value>,
|
||||
error_context: &'static str,
|
||||
record_runtime_miss_diagnostic: bool,
|
||||
) {
|
||||
if record_runtime_miss_diagnostic {
|
||||
record_local_runtime_candidate_skip_reason(state, trace_id, skip_reason);
|
||||
}
|
||||
|
||||
PlannerAppState::new(state)
|
||||
.persist_skipped_local_candidate(
|
||||
trace_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
0,
|
||||
candidate_id,
|
||||
required_capabilities,
|
||||
skip_reason,
|
||||
extra_data,
|
||||
current_unix_ms(),
|
||||
error_context,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(crate) async fn mark_skipped_local_execution_candidate(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
context: LocalSkippedCandidatePersistenceContext<'_>,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
persist_skipped_local_execution_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
context.user_id,
|
||||
context.api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
context.required_capabilities,
|
||||
skip_reason,
|
||||
None,
|
||||
context.error_context,
|
||||
context.record_runtime_miss_diagnostic,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(crate) async fn mark_skipped_local_execution_candidate_with_extra_data(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
context: LocalSkippedCandidatePersistenceContext<'_>,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
extra_data: Option<Value>,
|
||||
) {
|
||||
persist_skipped_local_execution_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
context.user_id,
|
||||
context.api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
context.required_capabilities,
|
||||
skip_reason,
|
||||
extra_data,
|
||||
context.error_context,
|
||||
context.record_runtime_miss_diagnostic,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(crate) async fn mark_skipped_local_execution_candidate_with_failure_diagnostic(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
context: LocalSkippedCandidatePersistenceContext<'_>,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
diagnostic: CandidateFailureDiagnostic,
|
||||
) {
|
||||
mark_skipped_local_execution_candidate_with_extra_data(
|
||||
state,
|
||||
trace_id,
|
||||
context,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
Some(diagnostic.to_extra_data()),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn persist_skipped_local_execution_candidates(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
required_capabilities: Option<&Value>,
|
||||
starting_candidate_index: u32,
|
||||
skipped_candidates: Vec<SkippedLocalExecutionCandidate>,
|
||||
error_context: &'static str,
|
||||
record_runtime_miss_diagnostic: bool,
|
||||
) {
|
||||
let port = GatewaySkippedCandidatePersistencePort {
|
||||
state,
|
||||
trace_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
required_capabilities,
|
||||
error_context,
|
||||
record_runtime_miss_diagnostic,
|
||||
};
|
||||
|
||||
match run_ai_skipped_candidate_persistence(&port, starting_candidate_index, skipped_candidates)
|
||||
.await
|
||||
{
|
||||
Ok(()) => {}
|
||||
Err(error) => match error {},
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn persist_skipped_local_execution_candidates_with_context(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
context: LocalSkippedCandidatePersistenceContext<'_>,
|
||||
starting_candidate_index: u32,
|
||||
skipped_candidates: Vec<SkippedLocalExecutionCandidate>,
|
||||
) {
|
||||
persist_skipped_local_execution_candidates(
|
||||
state,
|
||||
trace_id,
|
||||
context.user_id,
|
||||
context.api_key_id,
|
||||
context.required_capabilities,
|
||||
starting_candidate_index,
|
||||
skipped_candidates,
|
||||
context.error_context,
|
||||
context.record_runtime_miss_diagnostic,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
|
||||
use aether_provider_transport::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider,
|
||||
};
|
||||
use aether_scheduler_core::{
|
||||
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode, SchedulerRankingMode,
|
||||
SchedulerRankingOutcome,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||
|
||||
fn sample_candidate(key_id: &str) -> SchedulerMinimalCandidateSelectionCandidate {
|
||||
SchedulerMinimalCandidateSelectionCandidate {
|
||||
provider_id: "provider-1".to_string(),
|
||||
provider_name: "provider-1".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
provider_priority: 10,
|
||||
endpoint_id: "endpoint-1".to_string(),
|
||||
endpoint_api_format: "openai:chat".to_string(),
|
||||
key_id: key_id.to_string(),
|
||||
key_name: key_id.to_string(),
|
||||
key_auth_type: "api_key".to_string(),
|
||||
key_internal_priority: 10,
|
||||
key_global_priority_for_format: Some(10),
|
||||
key_capabilities: None,
|
||||
model_id: "model-1".to_string(),
|
||||
global_model_id: "global-model-1".to_string(),
|
||||
global_model_name: "gpt-5".to_string(),
|
||||
selected_provider_model_name: "gpt-5".to_string(),
|
||||
mapping_matched_model: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_transport(
|
||||
key_id: &str,
|
||||
provider_config: Option<serde_json::Value>,
|
||||
) -> Arc<crate::ai_serving::GatewayProviderTransportSnapshot> {
|
||||
Arc::new(crate::ai_serving::GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-1".to_string(),
|
||||
name: "provider-1".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: false,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: provider_config,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
api_format: "openai:chat".to_string(),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
is_active: true,
|
||||
base_url: "https://example.com".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: key_id.to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
name: key_id.to_string(),
|
||||
auth_type: "api_key".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["openai:chat".to_string()]),
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
decrypted_api_key: "secret".to_string(),
|
||||
decrypted_auth_config: None,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
fn sample_eligible(
|
||||
key_id: &str,
|
||||
pool_key_index: Option<u32>,
|
||||
) -> EligibleLocalExecutionCandidate {
|
||||
EligibleLocalExecutionCandidate {
|
||||
candidate: sample_candidate(key_id),
|
||||
transport: sample_transport(
|
||||
key_id,
|
||||
pool_key_index.map(|_| json!({ "pool_advanced": {} })),
|
||||
),
|
||||
provider_api_format: "openai:chat".to_string(),
|
||||
orchestration: LocalExecutionCandidateMetadata {
|
||||
candidate_group_id: pool_key_index.map(|_| "pool-group".to_string()),
|
||||
pool_key_index,
|
||||
},
|
||||
ranking: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_group_representatives_are_persisted_as_available_before_attempt() {
|
||||
let repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let app = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_request_candidate_repository_for_tests(Arc::clone(
|
||||
&repository,
|
||||
)),
|
||||
);
|
||||
|
||||
let attempts = persist_available_local_execution_candidates(
|
||||
PlannerAppState::new(&app),
|
||||
"trace-pool-lazy",
|
||||
"user-1",
|
||||
"api-key-1",
|
||||
None,
|
||||
vec![
|
||||
sample_eligible("pool-key", Some(0)),
|
||||
sample_eligible("pool-key-internal", Some(1)),
|
||||
sample_eligible("normal-key", None),
|
||||
],
|
||||
"persist should not fail",
|
||||
|_| None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(attempts.len(), 3);
|
||||
let stored = app
|
||||
.read_request_candidates_by_request_id("trace-pool-lazy")
|
||||
.await
|
||||
.expect("request candidates should read");
|
||||
assert_eq!(stored.len(), 2);
|
||||
assert_eq!(stored[0].key_id.as_deref(), Some("pool-key"));
|
||||
assert_eq!(stored[0].candidate_index, 0);
|
||||
assert_eq!(stored[1].key_id.as_deref(), Some("normal-key"));
|
||||
assert_eq!(stored[1].candidate_index, 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn available_candidates_persist_ranking_metadata_in_extra_data() {
|
||||
let repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let app = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_request_candidate_repository_for_tests(Arc::clone(
|
||||
&repository,
|
||||
)),
|
||||
);
|
||||
let mut eligible = sample_eligible("ranked-key", None);
|
||||
eligible.ranking = Some(SchedulerRankingOutcome {
|
||||
original_index: 1,
|
||||
ranking_index: 0,
|
||||
priority_mode: SchedulerPriorityMode::Provider,
|
||||
ranking_mode: SchedulerRankingMode::CacheAffinity,
|
||||
priority_slot: 7,
|
||||
promoted_by: Some("cached_affinity"),
|
||||
demoted_by: Some("cross_format"),
|
||||
});
|
||||
|
||||
persist_available_local_execution_candidates(
|
||||
PlannerAppState::new(&app),
|
||||
"trace-ranking-extra-data",
|
||||
"user-1",
|
||||
"api-key-1",
|
||||
None,
|
||||
vec![eligible],
|
||||
"persist should not fail",
|
||||
|_| Some(json!({ "existing": "value" })),
|
||||
)
|
||||
.await;
|
||||
|
||||
let stored = app
|
||||
.read_request_candidates_by_request_id("trace-ranking-extra-data")
|
||||
.await
|
||||
.expect("request candidates should read");
|
||||
assert_eq!(stored.len(), 1);
|
||||
let extra_data = stored[0]
|
||||
.extra_data
|
||||
.as_ref()
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.expect("ranking metadata should persist as object extra data");
|
||||
assert_eq!(extra_data.get("existing"), Some(&json!("value")));
|
||||
assert_eq!(
|
||||
extra_data.get("ranking_mode"),
|
||||
Some(&json!("CacheAffinity"))
|
||||
);
|
||||
assert_eq!(extra_data.get("priority_mode"), Some(&json!("Provider")));
|
||||
assert_eq!(extra_data.get("ranking_index"), Some(&json!(0)));
|
||||
assert_eq!(extra_data.get("priority_slot"), Some(&json!(7)));
|
||||
assert_eq!(
|
||||
extra_data.get("promoted_by"),
|
||||
Some(&json!("cached_affinity"))
|
||||
);
|
||||
assert_eq!(extra_data.get("demoted_by"), Some(&json!("cross_format")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_internal_skipped_candidates_are_not_persisted() {
|
||||
let repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let app = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_request_candidate_repository_for_tests(Arc::clone(
|
||||
&repository,
|
||||
)),
|
||||
);
|
||||
|
||||
persist_skipped_local_execution_candidates(
|
||||
&app,
|
||||
"trace-pool-skipped",
|
||||
"user-1",
|
||||
"api-key-1",
|
||||
None,
|
||||
0,
|
||||
vec![
|
||||
SkippedLocalExecutionCandidate {
|
||||
candidate: sample_candidate("pool-skipped"),
|
||||
skip_reason: "pool_cooldown",
|
||||
transport: Some(sample_transport(
|
||||
"pool-skipped",
|
||||
Some(json!({ "pool_advanced": {} })),
|
||||
)),
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
},
|
||||
SkippedLocalExecutionCandidate {
|
||||
candidate: sample_candidate("normal-skipped"),
|
||||
skip_reason: "key_inactive",
|
||||
transport: None,
|
||||
ranking: Some(SchedulerRankingOutcome {
|
||||
original_index: 2,
|
||||
ranking_index: 1,
|
||||
priority_mode: SchedulerPriorityMode::Provider,
|
||||
ranking_mode: SchedulerRankingMode::CacheAffinity,
|
||||
priority_slot: 9,
|
||||
promoted_by: None,
|
||||
demoted_by: Some("cross_format"),
|
||||
}),
|
||||
extra_data: Some(json!({ "existing": "value" })),
|
||||
},
|
||||
],
|
||||
"persist skipped should not fail",
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
|
||||
let stored = app
|
||||
.read_request_candidates_by_request_id("trace-pool-skipped")
|
||||
.await
|
||||
.expect("request candidates should read");
|
||||
assert_eq!(stored.len(), 1);
|
||||
assert_eq!(stored[0].key_id.as_deref(), Some("normal-skipped"));
|
||||
assert_eq!(stored[0].candidate_index, 0);
|
||||
let extra_data = stored[0]
|
||||
.extra_data
|
||||
.as_ref()
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.expect("skipped ranking metadata should persist");
|
||||
assert_eq!(extra_data.get("existing"), Some(&json!("value")));
|
||||
assert_eq!(
|
||||
extra_data.get("ranking_mode"),
|
||||
Some(&json!("CacheAffinity"))
|
||||
);
|
||||
assert_eq!(extra_data.get("priority_mode"), Some(&json!("Provider")));
|
||||
assert_eq!(extra_data.get("ranking_index"), Some(&json!(1)));
|
||||
assert_eq!(extra_data.get("priority_slot"), Some(&json!(9)));
|
||||
assert_eq!(extra_data.get("demoted_by"), Some(&json!("cross_format")));
|
||||
}
|
||||
}
|
||||
320
apps/aether-gateway/src/ai_serving/planner/candidate_metadata.rs
Normal file
320
apps/aether-gateway/src/ai_serving/planner/candidate_metadata.rs
Normal file
@@ -0,0 +1,320 @@
|
||||
use aether_ai_serving::{
|
||||
append_ai_execution_contract_fields_to_value, append_ai_ranking_metadata_to_object,
|
||||
build_ai_candidate_metadata_from_candidate,
|
||||
};
|
||||
use aether_scheduler_core::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRankingOutcome};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use crate::ai_serving::transport::append_transport_diagnostics_to_value;
|
||||
use crate::ai_serving::GatewayProviderTransportSnapshot;
|
||||
use crate::ai_serving::{ConversionMode, ExecutionStrategy};
|
||||
|
||||
pub(crate) struct LocalExecutionCandidateMetadataParts<'a> {
|
||||
pub(crate) eligible: &'a EligibleLocalExecutionCandidate,
|
||||
pub(crate) provider_api_format: &'a str,
|
||||
pub(crate) client_api_format: &'a str,
|
||||
pub(crate) extra_fields: Map<String, Value>,
|
||||
}
|
||||
|
||||
pub(crate) fn append_ranking_metadata_to_object(
|
||||
object: &mut Map<String, Value>,
|
||||
ranking: &SchedulerRankingOutcome,
|
||||
) {
|
||||
append_ai_ranking_metadata_to_object(object, ranking);
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_candidate_metadata(
|
||||
parts: LocalExecutionCandidateMetadataParts<'_>,
|
||||
) -> Value {
|
||||
build_local_execution_candidate_metadata_for_candidate(
|
||||
&parts.eligible.candidate,
|
||||
Some(parts.eligible.transport.as_ref()),
|
||||
parts.provider_api_format,
|
||||
parts.client_api_format,
|
||||
parts.extra_fields,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_candidate_metadata_for_candidate(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
transport: Option<&GatewayProviderTransportSnapshot>,
|
||||
provider_api_format: &str,
|
||||
client_api_format: &str,
|
||||
extra_fields: Map<String, Value>,
|
||||
) -> Value {
|
||||
append_transport_diagnostics_to_value(
|
||||
build_ai_candidate_metadata_from_candidate(
|
||||
candidate,
|
||||
provider_api_format,
|
||||
client_api_format,
|
||||
extra_fields,
|
||||
),
|
||||
transport,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_candidate_contract_metadata(
|
||||
parts: LocalExecutionCandidateMetadataParts<'_>,
|
||||
execution_strategy: ExecutionStrategy,
|
||||
conversion_mode: ConversionMode,
|
||||
provider_contract: &str,
|
||||
) -> Value {
|
||||
append_ai_execution_contract_fields_to_value(
|
||||
build_local_execution_candidate_metadata_for_candidate(
|
||||
&parts.eligible.candidate,
|
||||
Some(parts.eligible.transport.as_ref()),
|
||||
parts.provider_api_format,
|
||||
parts.client_api_format,
|
||||
parts.extra_fields,
|
||||
),
|
||||
execution_strategy.as_str(),
|
||||
conversion_mode.as_str(),
|
||||
parts.client_api_format,
|
||||
provider_contract,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_candidate_contract_metadata_for_candidate(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
transport: Option<&GatewayProviderTransportSnapshot>,
|
||||
provider_api_format: &str,
|
||||
client_api_format: &str,
|
||||
extra_fields: Map<String, Value>,
|
||||
execution_strategy: ExecutionStrategy,
|
||||
conversion_mode: ConversionMode,
|
||||
provider_contract: &str,
|
||||
) -> Value {
|
||||
append_ai_execution_contract_fields_to_value(
|
||||
build_local_execution_candidate_metadata_for_candidate(
|
||||
candidate,
|
||||
transport,
|
||||
provider_api_format,
|
||||
client_api_format,
|
||||
extra_fields,
|
||||
),
|
||||
execution_strategy.as_str(),
|
||||
conversion_mode.as_str(),
|
||||
client_api_format,
|
||||
provider_contract,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
build_local_execution_candidate_contract_metadata_for_candidate,
|
||||
build_local_execution_candidate_metadata_for_candidate,
|
||||
};
|
||||
use crate::ai_serving::transport::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider,
|
||||
};
|
||||
use crate::ai_serving::{ConversionMode, ExecutionStrategy, GatewayProviderTransportSnapshot};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
fn sample_candidate() -> SchedulerMinimalCandidateSelectionCandidate {
|
||||
SchedulerMinimalCandidateSelectionCandidate {
|
||||
provider_id: "provider-1".to_string(),
|
||||
provider_name: "RightCode".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
provider_priority: 22,
|
||||
endpoint_id: "endpoint-1".to_string(),
|
||||
endpoint_api_format: "openai:responses".to_string(),
|
||||
key_id: "key-1".to_string(),
|
||||
key_name: "codex".to_string(),
|
||||
key_auth_type: "oauth".to_string(),
|
||||
key_internal_priority: 10,
|
||||
key_global_priority_for_format: None,
|
||||
key_capabilities: None,
|
||||
model_id: "model-1".to_string(),
|
||||
global_model_id: "global-1".to_string(),
|
||||
global_model_name: "gpt-5.4".to_string(),
|
||||
selected_provider_model_name: "gpt-5.4".to_string(),
|
||||
mapping_matched_model: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_transport() -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-1".to_string(),
|
||||
name: "RightCode".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: true,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: Some(json!({"enabled": true, "mode": "node", "node_id": "proxy-node-1"})),
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
api_format: "openai:responses".to_string(),
|
||||
api_family: None,
|
||||
endpoint_kind: None,
|
||||
is_active: true,
|
||||
base_url: "https://example.com".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: Some("/v1/responses".to_string()),
|
||||
config: None,
|
||||
format_acceptance_config: Some(json!({
|
||||
"enabled": true,
|
||||
"accept_formats": ["claude:messages"]
|
||||
})),
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
name: "codex".to_string(),
|
||||
auth_type: "oauth".to_string(),
|
||||
is_active: true,
|
||||
api_formats: None,
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: Some(json!({
|
||||
"tls_profile": "chrome_136",
|
||||
"user_agent": "Mozilla/5.0"
|
||||
})),
|
||||
decrypted_api_key: "sk-test".to_string(),
|
||||
decrypted_auth_config: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_claude_code_transport_without_auth() -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-cc-1".to_string(),
|
||||
name: "NekoCode".to_string(),
|
||||
provider_type: "claude_code".to_string(),
|
||||
website: Some("https://nekocode.ai".to_string()),
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: true,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-cc-1".to_string(),
|
||||
provider_id: "provider-cc-1".to_string(),
|
||||
api_format: "claude:messages".to_string(),
|
||||
api_family: Some("claude".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
is_active: true,
|
||||
base_url: "https://api.anthropic.com".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-cc-1".to_string(),
|
||||
provider_id: "provider-cc-1".to_string(),
|
||||
name: "CC-特价-0.4".to_string(),
|
||||
auth_type: "api_key".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["claude:messages".to_string()]),
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
decrypted_api_key: "__placeholder__".to_string(),
|
||||
decrypted_auth_config: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_contract_metadata_includes_transport_diagnostics() {
|
||||
let metadata = build_local_execution_candidate_contract_metadata_for_candidate(
|
||||
&sample_candidate(),
|
||||
Some(&sample_transport()),
|
||||
"openai:responses",
|
||||
"claude:messages",
|
||||
serde_json::Map::new(),
|
||||
ExecutionStrategy::LocalCrossFormat,
|
||||
ConversionMode::Bidirectional,
|
||||
"openai:responses",
|
||||
);
|
||||
|
||||
assert_eq!(metadata["transport_diagnostics"]["provider_type"], "codex");
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["fingerprint"]["tls_profile"],
|
||||
"chrome_136"
|
||||
);
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["resolved_tls_profile"],
|
||||
"chrome_136"
|
||||
);
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["request_pair"]["conversion_enabled"],
|
||||
Value::Bool(true)
|
||||
);
|
||||
assert!(
|
||||
metadata["transport_diagnostics"]["request_pair"]["transport_unsupported_reason"]
|
||||
.is_null()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_metadata_marks_missing_transport_snapshot() {
|
||||
let metadata = build_local_execution_candidate_metadata_for_candidate(
|
||||
&sample_candidate(),
|
||||
None,
|
||||
"openai:responses",
|
||||
"openai:responses",
|
||||
serde_json::Map::new(),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["transport_snapshot_available"],
|
||||
Value::Bool(false)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_metadata_uses_same_format_provider_specific_transport_reason() {
|
||||
let metadata = build_local_execution_candidate_metadata_for_candidate(
|
||||
&sample_candidate(),
|
||||
Some(&sample_claude_code_transport_without_auth()),
|
||||
"claude:messages",
|
||||
"claude:messages",
|
||||
serde_json::Map::new(),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["request_pair"]["transport_unsupported_reason"],
|
||||
Value::String("transport_auth_unavailable".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,198 @@
|
||||
use aether_ai_serving::{
|
||||
prepare_ai_header_authenticated_candidate, resolve_ai_candidate_mapped_model,
|
||||
AiPreparedHeaderAuthenticatedCandidate,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::{
|
||||
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
|
||||
};
|
||||
|
||||
pub(crate) type PreparedHeaderAuthenticatedCandidate = AiPreparedHeaderAuthenticatedCandidate;
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct OauthPreparationContext<'a> {
|
||||
pub(crate) trace_id: &'a str,
|
||||
pub(crate) api_format: &'a str,
|
||||
pub(crate) operation: &'a str,
|
||||
}
|
||||
|
||||
pub(crate) async fn prepare_header_authenticated_candidate(
|
||||
state: PlannerAppState<'_>,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
direct_auth: Option<(String, String)>,
|
||||
context: OauthPreparationContext<'_>,
|
||||
) -> Result<PreparedHeaderAuthenticatedCandidate, &'static str> {
|
||||
let oauth_auth = if direct_auth.is_none() {
|
||||
match resolve_candidate_oauth_auth(state, transport, context).await {
|
||||
Some(LocalResolvedOAuthRequestAuth::Header { name, value }) => Some((name, value)),
|
||||
Some(LocalResolvedOAuthRequestAuth::Kiro(_)) => None,
|
||||
None => None,
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
prepare_ai_header_authenticated_candidate(
|
||||
direct_auth,
|
||||
oauth_auth,
|
||||
candidate.selected_provider_model_name.as_str(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn prepare_header_authenticated_candidate_from_auth(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
auth_header: String,
|
||||
auth_value: String,
|
||||
) -> Result<PreparedHeaderAuthenticatedCandidate, &'static str> {
|
||||
prepare_ai_header_authenticated_candidate(
|
||||
Some((auth_header, auth_value)),
|
||||
None,
|
||||
candidate.selected_provider_model_name.as_str(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_candidate_mapped_model(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> Result<String, &'static str> {
|
||||
resolve_ai_candidate_mapped_model(candidate.selected_provider_model_name.as_str())
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_candidate_oauth_auth(
|
||||
state: PlannerAppState<'_>,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
context: OauthPreparationContext<'_>,
|
||||
) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||
match state.resolve_local_oauth_request_auth(transport).await {
|
||||
Ok(Some(auth)) => Some(auth),
|
||||
Ok(None) => None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "candidate_preparation_oauth_auth_resolution_failed",
|
||||
log_type = "event",
|
||||
trace_id = %context.trace_id,
|
||||
api_format = %context.api_format,
|
||||
operation = %context.operation,
|
||||
provider_type = %transport.provider.provider_type,
|
||||
error = ?err,
|
||||
"failed to resolve oauth auth while preparing local candidate"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use aether_provider_transport::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
use super::{prepare_header_authenticated_candidate, OauthPreparationContext};
|
||||
use crate::ai_serving::PlannerAppState;
|
||||
|
||||
fn sample_transport() -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-1".to_string(),
|
||||
name: "provider".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: false,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
api_format: "openai:chat".to_string(),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
is_active: true,
|
||||
base_url: "https://example.test".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
name: "key".to_string(),
|
||||
auth_type: "api_key".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["openai:chat".to_string()]),
|
||||
auth_type_by_format: None,
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
decrypted_api_key: String::new(),
|
||||
decrypted_auth_config: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_candidate() -> SchedulerMinimalCandidateSelectionCandidate {
|
||||
SchedulerMinimalCandidateSelectionCandidate {
|
||||
provider_id: "provider-1".to_string(),
|
||||
provider_name: "provider".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 1,
|
||||
endpoint_id: "endpoint-1".to_string(),
|
||||
endpoint_api_format: "openai:chat".to_string(),
|
||||
key_id: "key-1".to_string(),
|
||||
key_name: "key".to_string(),
|
||||
key_auth_type: "api_key".to_string(),
|
||||
key_internal_priority: 1,
|
||||
key_global_priority_for_format: None,
|
||||
key_capabilities: None,
|
||||
model_id: "model-1".to_string(),
|
||||
global_model_id: "global-model-1".to_string(),
|
||||
global_model_name: "gpt-test".to_string(),
|
||||
selected_provider_model_name: "gpt-test-upstream".to_string(),
|
||||
mapping_matched_model: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn header_auth_preparation_allows_empty_auth_value() {
|
||||
let state = crate::AppState::new().expect("state should build");
|
||||
let transport = sample_transport();
|
||||
let candidate = sample_candidate();
|
||||
|
||||
let prepared = prepare_header_authenticated_candidate(
|
||||
PlannerAppState::new(&state),
|
||||
&transport,
|
||||
&candidate,
|
||||
Some(("authorization".to_string(), String::new())),
|
||||
OauthPreparationContext {
|
||||
trace_id: "trace-empty-auth",
|
||||
api_format: "openai:chat",
|
||||
operation: "test",
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("empty auth value should still prepare the candidate");
|
||||
|
||||
assert_eq!(prepared.auth_header, "authorization");
|
||||
assert_eq!(prepared.auth_value, "");
|
||||
assert_eq!(prepared.mapped_model, "gpt-test-upstream");
|
||||
}
|
||||
}
|
||||
1929
apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs
Normal file
1929
apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,287 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_ai_serving::{
|
||||
run_ai_candidate_resolution, AiCandidateResolutionMode, AiCandidateResolutionPort,
|
||||
AiCandidateResolutionRequest,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use std::convert::Infallible;
|
||||
use tracing::warn;
|
||||
|
||||
use aether_scheduler_core::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRankingOutcome};
|
||||
|
||||
use crate::ai_serving::{
|
||||
candidate_common_transport_skip_reason, candidate_transport_pair_skip_reason,
|
||||
CandidateTransportPolicyFacts, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot,
|
||||
PlannerAppState,
|
||||
};
|
||||
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||
|
||||
use super::candidate_ranking::rank_eligible_local_execution_candidates;
|
||||
use super::pool_scheduler::apply_local_execution_pool_scheduler;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub(crate) struct EligibleLocalExecutionCandidate {
|
||||
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub(crate) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
pub(crate) provider_api_format: String,
|
||||
pub(crate) orchestration: LocalExecutionCandidateMetadata,
|
||||
pub(crate) ranking: Option<SchedulerRankingOutcome>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub(crate) struct SkippedLocalExecutionCandidate {
|
||||
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub(crate) skip_reason: &'static str,
|
||||
pub(crate) transport: Option<Arc<GatewayProviderTransportSnapshot>>,
|
||||
pub(crate) ranking: Option<SchedulerRankingOutcome>,
|
||||
pub(crate) extra_data: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
impl SkippedLocalExecutionCandidate {
|
||||
pub(crate) fn transport_ref(&self) -> Option<&GatewayProviderTransportSnapshot> {
|
||||
self.transport.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
struct GatewayLocalCandidateResolutionPort<'a> {
|
||||
state: PlannerAppState<'a>,
|
||||
requested_model: Option<&'a str>,
|
||||
auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&'a serde_json::Value>,
|
||||
sticky_session_token: Option<&'a str>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> {
|
||||
type Candidate = SchedulerMinimalCandidateSelectionCandidate;
|
||||
type Transport = GatewayProviderTransportSnapshot;
|
||||
type Eligible = EligibleLocalExecutionCandidate;
|
||||
type Skipped = SkippedLocalExecutionCandidate;
|
||||
type Error = Infallible;
|
||||
|
||||
async fn read_candidate_transport(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
) -> Result<Option<Self::Transport>, Self::Error> {
|
||||
Ok(read_candidate_transport_snapshot(self.state, candidate).await)
|
||||
}
|
||||
|
||||
fn build_missing_transport_skipped_candidate(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
) -> Self::Skipped {
|
||||
SkippedLocalExecutionCandidate {
|
||||
candidate,
|
||||
skip_reason: "transport_snapshot_missing",
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn candidate_common_skip_reason(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
transport: &Self::Transport,
|
||||
requested_model: Option<&str>,
|
||||
) -> Option<&'static str> {
|
||||
candidate_common_transport_skip_reason(
|
||||
transport,
|
||||
candidate_transport_policy_facts(candidate),
|
||||
requested_model,
|
||||
)
|
||||
}
|
||||
|
||||
fn candidate_transport_pair_skip_reason(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
transport: &Self::Transport,
|
||||
normalized_client_api_format: &str,
|
||||
requested_model: &str,
|
||||
) -> Option<&'static str> {
|
||||
let _ = (candidate, requested_model);
|
||||
candidate_transport_pair_skip_reason(transport, normalized_client_api_format)
|
||||
}
|
||||
|
||||
fn build_skipped_candidate(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
transport: Self::Transport,
|
||||
skip_reason: &'static str,
|
||||
) -> Self::Skipped {
|
||||
SkippedLocalExecutionCandidate {
|
||||
candidate,
|
||||
skip_reason,
|
||||
transport: Some(Arc::new(transport)),
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn build_eligible_candidate(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
transport: Self::Transport,
|
||||
) -> Self::Eligible {
|
||||
let provider_api_format = transport.endpoint.api_format.trim().to_ascii_lowercase();
|
||||
EligibleLocalExecutionCandidate {
|
||||
candidate,
|
||||
transport: Arc::new(transport),
|
||||
provider_api_format,
|
||||
orchestration: LocalExecutionCandidateMetadata::default(),
|
||||
ranking: None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn rank_eligible_candidates(
|
||||
&self,
|
||||
candidates: Vec<Self::Eligible>,
|
||||
normalized_client_api_format: &str,
|
||||
) -> Result<Vec<Self::Eligible>, Self::Error> {
|
||||
Ok(rank_eligible_local_execution_candidates(
|
||||
self.state,
|
||||
candidates,
|
||||
normalized_client_api_format,
|
||||
self.requested_model,
|
||||
self.auth_snapshot,
|
||||
self.required_capabilities,
|
||||
)
|
||||
.await)
|
||||
}
|
||||
|
||||
async fn apply_pool_scheduler(
|
||||
&self,
|
||||
candidates: Vec<Self::Eligible>,
|
||||
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error> {
|
||||
Ok(
|
||||
apply_local_execution_pool_scheduler(self.state, candidates, self.sticky_session_token)
|
||||
.await,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_and_rank_local_execution_candidates(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
client_api_format: &str,
|
||||
requested_model: &str,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
sticky_session_token: Option<&str>,
|
||||
) -> (
|
||||
Vec<EligibleLocalExecutionCandidate>,
|
||||
Vec<SkippedLocalExecutionCandidate>,
|
||||
) {
|
||||
let requested_model = requested_model.trim();
|
||||
resolve_and_rank_local_execution_candidates_with_mode(
|
||||
state,
|
||||
candidates,
|
||||
client_api_format,
|
||||
Some(requested_model),
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
sticky_session_token,
|
||||
AiCandidateResolutionMode::Standard,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_and_rank_local_execution_candidates_without_transport_pair_gate(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
client_api_format: &str,
|
||||
requested_model: Option<&str>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
sticky_session_token: Option<&str>,
|
||||
) -> (
|
||||
Vec<EligibleLocalExecutionCandidate>,
|
||||
Vec<SkippedLocalExecutionCandidate>,
|
||||
) {
|
||||
let requested_model = requested_model.map(str::trim);
|
||||
resolve_and_rank_local_execution_candidates_with_mode(
|
||||
state,
|
||||
candidates,
|
||||
client_api_format,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
sticky_session_token,
|
||||
AiCandidateResolutionMode::WithoutTransportPairGate,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn resolve_and_rank_local_execution_candidates_with_mode(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
client_api_format: &str,
|
||||
requested_model: Option<&str>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
sticky_session_token: Option<&str>,
|
||||
mode: AiCandidateResolutionMode,
|
||||
) -> (
|
||||
Vec<EligibleLocalExecutionCandidate>,
|
||||
Vec<SkippedLocalExecutionCandidate>,
|
||||
) {
|
||||
let port = GatewayLocalCandidateResolutionPort {
|
||||
state,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
sticky_session_token,
|
||||
};
|
||||
|
||||
let request = AiCandidateResolutionRequest {
|
||||
client_api_format,
|
||||
requested_model,
|
||||
mode,
|
||||
};
|
||||
|
||||
match run_ai_candidate_resolution(&port, candidates, request).await {
|
||||
Ok(outcome) => (outcome.eligible_candidates, outcome.skipped_candidates),
|
||||
Err(error) => match error {},
|
||||
}
|
||||
}
|
||||
|
||||
fn candidate_transport_policy_facts(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> CandidateTransportPolicyFacts<'_> {
|
||||
CandidateTransportPolicyFacts {
|
||||
endpoint_api_format: candidate.endpoint_api_format.as_str(),
|
||||
global_model_name: candidate.global_model_name.as_str(),
|
||||
selected_provider_model_name: candidate.selected_provider_model_name.as_str(),
|
||||
mapping_matched_model: candidate.mapping_matched_model.as_deref(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn read_candidate_transport_snapshot(
|
||||
state: PlannerAppState<'_>,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> Option<GatewayProviderTransportSnapshot> {
|
||||
match state
|
||||
.read_provider_transport_snapshot(
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(transport)) => Some(transport),
|
||||
Ok(None) => None,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
event_name = "candidate_resolution_transport_load_failed",
|
||||
log_type = "event",
|
||||
provider_id = %candidate.provider_id,
|
||||
endpoint_id = %candidate.endpoint_id,
|
||||
key_id = %candidate.key_id,
|
||||
error = ?error,
|
||||
"failed to load provider transport while evaluating local candidate eligibility"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
218
apps/aether-gateway/src/ai_serving/planner/candidate_source.rs
Normal file
218
apps/aether-gateway/src/ai_serving/planner/candidate_source.rs
Normal file
@@ -0,0 +1,218 @@
|
||||
use aether_ai_serving::{
|
||||
run_ai_candidate_preselection, AiCandidatePreselectionOutcome, AiCandidatePreselectionPort,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::scheduler::candidate::SchedulerSkippedCandidate;
|
||||
use crate::GatewayError;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum LocalCandidatePreselectionKeyMode {
|
||||
ProviderEndpointKeyModel,
|
||||
ProviderEndpointKeyModelAndApiFormat,
|
||||
}
|
||||
|
||||
struct GatewayLocalCandidatePreselectionPort<'a> {
|
||||
state: PlannerAppState<'a>,
|
||||
client_api_format: &'a str,
|
||||
requested_model: &'a str,
|
||||
require_streaming: bool,
|
||||
required_capabilities: Option<&'a serde_json::Value>,
|
||||
auth_snapshot: &'a GatewayAuthApiKeySnapshot,
|
||||
use_api_format_alias_match: bool,
|
||||
key_mode: LocalCandidatePreselectionKeyMode,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
|
||||
type Candidate = SchedulerMinimalCandidateSelectionCandidate;
|
||||
type Skipped = SkippedLocalExecutionCandidate;
|
||||
type Error = GatewayError;
|
||||
|
||||
fn candidate_api_formats(&self) -> Vec<String> {
|
||||
crate::ai_serving::request_candidate_api_formats(
|
||||
self.client_api_format,
|
||||
self.require_streaming,
|
||||
)
|
||||
.into_iter()
|
||||
.map(str::to_string)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn candidate_api_format_matches_client(&self, candidate_api_format: &str) -> bool {
|
||||
if self.use_api_format_alias_match {
|
||||
crate::ai_serving::api_format_alias_matches(
|
||||
candidate_api_format,
|
||||
self.client_api_format,
|
||||
)
|
||||
} else {
|
||||
candidate_api_format == self.client_api_format
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_candidates_for_api_format(
|
||||
&self,
|
||||
candidate_api_format: &str,
|
||||
matches_client_format: bool,
|
||||
) -> Result<(Vec<Self::Candidate>, Vec<Self::Skipped>), Self::Error> {
|
||||
let auth_snapshot = matches_client_format.then_some(self.auth_snapshot);
|
||||
let (candidates, skipped_candidates) = self
|
||||
.state
|
||||
.list_selectable_candidates_with_skip_reasons(
|
||||
candidate_api_format,
|
||||
self.requested_model,
|
||||
self.require_streaming,
|
||||
self.required_capabilities,
|
||||
auth_snapshot,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok((
|
||||
candidates,
|
||||
skipped_candidates
|
||||
.into_iter()
|
||||
.map(skipped_local_execution_candidate_from_scheduler_skip)
|
||||
.collect(),
|
||||
))
|
||||
}
|
||||
|
||||
fn candidate_allowed(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
_candidate_api_format: &str,
|
||||
matches_client_format: bool,
|
||||
) -> bool {
|
||||
matches_client_format
|
||||
|| auth_snapshot_allows_cross_format_candidate(
|
||||
self.auth_snapshot,
|
||||
self.requested_model,
|
||||
candidate,
|
||||
)
|
||||
}
|
||||
|
||||
fn skipped_candidate_allowed(
|
||||
&self,
|
||||
skipped_candidate: &Self::Skipped,
|
||||
_candidate_api_format: &str,
|
||||
matches_client_format: bool,
|
||||
) -> bool {
|
||||
matches_client_format
|
||||
|| auth_snapshot_allows_cross_format_candidate(
|
||||
self.auth_snapshot,
|
||||
self.requested_model,
|
||||
&skipped_candidate.candidate,
|
||||
)
|
||||
}
|
||||
|
||||
fn candidate_key(&self, candidate: &Self::Candidate) -> String {
|
||||
local_candidate_preselection_key(candidate, self.key_mode)
|
||||
}
|
||||
|
||||
fn skipped_candidate_key(&self, skipped_candidate: &Self::Skipped) -> String {
|
||||
local_candidate_preselection_key(&skipped_candidate.candidate, self.key_mode)
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn preselect_local_execution_candidates_with_serving(
|
||||
state: PlannerAppState<'_>,
|
||||
client_api_format: &str,
|
||||
requested_model: &str,
|
||||
require_streaming: bool,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
auth_snapshot: &GatewayAuthApiKeySnapshot,
|
||||
use_api_format_alias_match: bool,
|
||||
key_mode: LocalCandidatePreselectionKeyMode,
|
||||
) -> Result<
|
||||
AiCandidatePreselectionOutcome<
|
||||
SchedulerMinimalCandidateSelectionCandidate,
|
||||
SkippedLocalExecutionCandidate,
|
||||
>,
|
||||
GatewayError,
|
||||
> {
|
||||
let port = GatewayLocalCandidatePreselectionPort {
|
||||
state,
|
||||
client_api_format,
|
||||
requested_model,
|
||||
require_streaming,
|
||||
required_capabilities,
|
||||
auth_snapshot,
|
||||
use_api_format_alias_match,
|
||||
key_mode,
|
||||
};
|
||||
|
||||
run_ai_candidate_preselection(&port).await
|
||||
}
|
||||
|
||||
fn skipped_local_execution_candidate_from_scheduler_skip(
|
||||
skipped_candidate: SchedulerSkippedCandidate,
|
||||
) -> SkippedLocalExecutionCandidate {
|
||||
SkippedLocalExecutionCandidate {
|
||||
candidate: skipped_candidate.candidate,
|
||||
skip_reason: skipped_candidate.skip_reason,
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn local_candidate_preselection_key(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
mode: LocalCandidatePreselectionKeyMode,
|
||||
) -> String {
|
||||
match mode {
|
||||
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModel => format!(
|
||||
"{}:{}:{}:{}:{}",
|
||||
candidate.provider_id,
|
||||
candidate.endpoint_id,
|
||||
candidate.key_id,
|
||||
candidate.model_id,
|
||||
candidate.selected_provider_model_name,
|
||||
),
|
||||
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat => format!(
|
||||
"{}:{}:{}:{}:{}:{}",
|
||||
candidate.provider_id,
|
||||
candidate.endpoint_id,
|
||||
candidate.key_id,
|
||||
candidate.model_id,
|
||||
candidate.selected_provider_model_name,
|
||||
candidate.endpoint_api_format,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn auth_snapshot_allows_cross_format_candidate(
|
||||
auth_snapshot: &GatewayAuthApiKeySnapshot,
|
||||
requested_model: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> bool {
|
||||
if let Some(allowed_providers) = auth_snapshot.effective_allowed_providers() {
|
||||
let provider_allowed = allowed_providers.iter().any(|value| {
|
||||
aether_scheduler_core::provider_matches_allowed_value(
|
||||
value,
|
||||
&candidate.provider_id,
|
||||
&candidate.provider_name,
|
||||
&candidate.provider_type,
|
||||
)
|
||||
});
|
||||
if !provider_allowed {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(allowed_models) = auth_snapshot.effective_allowed_models() {
|
||||
let model_allowed = allowed_models
|
||||
.iter()
|
||||
.any(|value| value == requested_model || value == &candidate.global_model_name);
|
||||
if !model_allowed {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_scheduler_core::{
|
||||
SchedulerMinimalCandidateSelectionCandidate, SchedulerTunnelAffinityBucket,
|
||||
};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::{GatewayProviderTransportSnapshot, PlannerAppState};
|
||||
use crate::scheduler::config::SchedulerOrderingConfig;
|
||||
|
||||
use super::candidate_resolution::read_candidate_transport_snapshot;
|
||||
|
||||
pub(super) type CandidateTransportIdentity<'a> = (&'a str, &'a str, &'a str);
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) struct CandidateTransportRankingFacts {
|
||||
pub(super) tunnel_bucket: SchedulerTunnelAffinityBucket,
|
||||
pub(super) keep_priority_on_conversion: bool,
|
||||
}
|
||||
|
||||
pub(super) async fn resolve_cached_candidate_transport_ranking_facts<'a>(
|
||||
state: PlannerAppState<'_>,
|
||||
cache: &mut BTreeMap<CandidateTransportIdentity<'a>, CandidateTransportRankingFacts>,
|
||||
candidate: &'a SchedulerMinimalCandidateSelectionCandidate,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> CandidateTransportRankingFacts {
|
||||
let identity = candidate_transport_identity(candidate);
|
||||
if let Some(facts) = cache.get(&identity).copied() {
|
||||
return facts;
|
||||
}
|
||||
|
||||
let facts = resolve_candidate_transport_ranking_facts(state, candidate, ordering_config).await;
|
||||
cache.insert(identity, facts);
|
||||
facts
|
||||
}
|
||||
|
||||
pub(super) async fn resolve_cached_transport_ranking_facts<'a>(
|
||||
state: PlannerAppState<'_>,
|
||||
cache: &mut BTreeMap<CandidateTransportIdentity<'a>, CandidateTransportRankingFacts>,
|
||||
candidate: &'a SchedulerMinimalCandidateSelectionCandidate,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> CandidateTransportRankingFacts {
|
||||
let identity = candidate_transport_identity(candidate);
|
||||
if let Some(facts) = cache.get(&identity).copied() {
|
||||
return facts;
|
||||
}
|
||||
|
||||
let facts =
|
||||
resolve_candidate_transport_ranking_facts_from_transport(state, transport, ordering_config)
|
||||
.await;
|
||||
cache.insert(identity, facts);
|
||||
facts
|
||||
}
|
||||
|
||||
async fn resolve_candidate_transport_ranking_facts(
|
||||
state: PlannerAppState<'_>,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> CandidateTransportRankingFacts {
|
||||
let Some(transport) = read_candidate_transport_snapshot(state, candidate).await else {
|
||||
return CandidateTransportRankingFacts {
|
||||
tunnel_bucket: SchedulerTunnelAffinityBucket::Neutral,
|
||||
keep_priority_on_conversion: ordering_config.keep_priority_on_conversion,
|
||||
};
|
||||
};
|
||||
|
||||
resolve_candidate_transport_ranking_facts_from_transport(state, &transport, ordering_config)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn resolve_candidate_transport_ranking_facts_from_transport(
|
||||
state: PlannerAppState<'_>,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> CandidateTransportRankingFacts {
|
||||
CandidateTransportRankingFacts {
|
||||
tunnel_bucket: resolve_tunnel_owner_affinity_from_transport(state, transport).await,
|
||||
keep_priority_on_conversion: ordering_config.keep_priority_on_conversion
|
||||
|| transport.provider.keep_priority_on_conversion,
|
||||
}
|
||||
}
|
||||
|
||||
async fn resolve_tunnel_owner_affinity_from_transport(
|
||||
state: PlannerAppState<'_>,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> SchedulerTunnelAffinityBucket {
|
||||
let Some(proxy) = state
|
||||
.app()
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
|
||||
.await
|
||||
else {
|
||||
return SchedulerTunnelAffinityBucket::Neutral;
|
||||
};
|
||||
if proxy.enabled == Some(false) {
|
||||
return SchedulerTunnelAffinityBucket::Neutral;
|
||||
}
|
||||
let Some(node_id) = proxy
|
||||
.node_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return SchedulerTunnelAffinityBucket::Neutral;
|
||||
};
|
||||
|
||||
if state.app().tunnel.has_local_proxy(node_id) {
|
||||
return SchedulerTunnelAffinityBucket::LocalTunnel;
|
||||
}
|
||||
|
||||
match state
|
||||
.app()
|
||||
.tunnel
|
||||
.lookup_attachment_owner(state.app().data.as_ref(), node_id)
|
||||
.await
|
||||
{
|
||||
Ok(Some(owner)) if owner.gateway_instance_id == state.app().tunnel.local_instance_id() => {
|
||||
SchedulerTunnelAffinityBucket::LocalTunnel
|
||||
}
|
||||
Ok(Some(_)) => SchedulerTunnelAffinityBucket::RemoteTunnel,
|
||||
Ok(None) => SchedulerTunnelAffinityBucket::Neutral,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
event_name = "candidate_transport_ranking_facts_tunnel_owner_lookup_failed",
|
||||
log_type = "event",
|
||||
node_id = node_id,
|
||||
error = %error,
|
||||
"failed to load tunnel attachment owner while evaluating candidate transport ranking facts"
|
||||
);
|
||||
SchedulerTunnelAffinityBucket::Neutral
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn candidate_transport_identity(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> CandidateTransportIdentity<'_> {
|
||||
(
|
||||
candidate.provider_id.as_str(),
|
||||
candidate.endpoint_id.as_str(),
|
||||
candidate.key_id.as_str(),
|
||||
)
|
||||
}
|
||||
131
apps/aether-gateway/src/ai_serving/planner/common.rs
Normal file
131
apps/aether-gateway/src/ai_serving/planner/common.rs
Normal file
@@ -0,0 +1,131 @@
|
||||
use axum::body::Bytes;
|
||||
|
||||
use crate::ai_serving::{
|
||||
force_upstream_streaming_for_provider as force_upstream_streaming_for_provider_impl,
|
||||
is_json_request, parse_direct_request_body as parse_direct_request_body_impl,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::{
|
||||
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
|
||||
CLAUDE_CLI_SYNC_PLAN_KIND, EXECUTION_RUNTIME_STREAM_ACTION,
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION,
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND,
|
||||
GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
|
||||
GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND,
|
||||
GEMINI_FILES_LIST_PLAN_KIND, GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
|
||||
pub(crate) use aether_ai_serving::AiRequestedModelFamily as RequestedModelFamily;
|
||||
|
||||
pub(crate) fn parse_direct_request_body(
|
||||
parts: &http::request::Parts,
|
||||
body_bytes: &Bytes,
|
||||
) -> Option<(serde_json::Value, Option<String>)> {
|
||||
parse_direct_request_body_impl(is_json_request(&parts.headers), body_bytes.as_ref())
|
||||
}
|
||||
|
||||
pub(crate) fn force_upstream_streaming_for_provider(
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
) -> bool {
|
||||
force_upstream_streaming_for_provider_impl(provider_type, provider_api_format)
|
||||
}
|
||||
|
||||
pub(crate) fn extract_standard_requested_model(body_json: &serde_json::Value) -> Option<String> {
|
||||
aether_ai_serving::extract_ai_standard_requested_model(body_json)
|
||||
}
|
||||
|
||||
pub(crate) fn extract_requested_model_from_request(
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
family: RequestedModelFamily,
|
||||
) -> Option<String> {
|
||||
aether_ai_serving::extract_ai_requested_model_from_request_path(
|
||||
parts.uri.path(),
|
||||
body_json,
|
||||
family,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
extract_requested_model_from_request, extract_standard_requested_model,
|
||||
force_upstream_streaming_for_provider, RequestedModelFamily,
|
||||
};
|
||||
use axum::http::Request;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn forces_streaming_for_codex_openai_responses() {
|
||||
assert!(force_upstream_streaming_for_provider(
|
||||
"codex",
|
||||
"openai:responses"
|
||||
));
|
||||
assert!(!force_upstream_streaming_for_provider(
|
||||
"codex",
|
||||
"openai:responses:compact"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn does_not_force_streaming_for_compact_or_other_provider_types() {
|
||||
assert!(!force_upstream_streaming_for_provider(
|
||||
"codex",
|
||||
"openai:responses:compact"
|
||||
));
|
||||
assert!(!force_upstream_streaming_for_provider(
|
||||
"codex",
|
||||
"openai:responses:compact"
|
||||
));
|
||||
assert!(!force_upstream_streaming_for_provider(
|
||||
"openai",
|
||||
"openai:responses"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_standard_requested_model_from_request_body() {
|
||||
let requested_model =
|
||||
extract_standard_requested_model(&json!({ "model": " claude-sonnet-4 " }));
|
||||
|
||||
assert_eq!(requested_model.as_deref(), Some("claude-sonnet-4"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_family_helper_delegates_standard_model_extraction() {
|
||||
let request = Request::builder()
|
||||
.uri("https://example.test/v1/chat/completions")
|
||||
.body(())
|
||||
.expect("request should build");
|
||||
let (parts, _) = request.into_parts();
|
||||
|
||||
let requested_model = extract_requested_model_from_request(
|
||||
&parts,
|
||||
&json!({ "model": " claude-sonnet-4 " }),
|
||||
RequestedModelFamily::Standard,
|
||||
);
|
||||
|
||||
assert_eq!(requested_model.as_deref(), Some("claude-sonnet-4"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_gemini_requested_model_from_request_path() {
|
||||
let request = Request::builder()
|
||||
.uri("https://example.test/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse")
|
||||
.body(())
|
||||
.expect("request should build");
|
||||
let (parts, _) = request.into_parts();
|
||||
|
||||
let requested_model =
|
||||
extract_requested_model_from_request(&parts, &json!({}), RequestedModelFamily::Gemini);
|
||||
|
||||
assert_eq!(requested_model.as_deref(), Some("gemini-2.5-pro"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
use crate::ai_serving::planner::common::{
|
||||
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
|
||||
CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
|
||||
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_FILES_DELETE_PLAN_KIND,
|
||||
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND,
|
||||
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND,
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
|
||||
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND,
|
||||
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
|
||||
OPENAI_RESPONSES_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_serving::planner::plan_builders::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_openai_chat_stream_plan_from_decision, build_openai_chat_sync_plan_from_decision,
|
||||
build_openai_responses_stream_plan_from_decision,
|
||||
build_openai_responses_sync_plan_from_decision, build_passthrough_stream_plan_from_decision,
|
||||
build_passthrough_sync_plan_from_decision, build_standard_stream_plan_from_decision,
|
||||
build_standard_sync_plan_from_decision,
|
||||
};
|
||||
use crate::ai_serving::planner::route::{
|
||||
resolve_execution_runtime_stream_plan_kind as resolve_stream_plan_kind,
|
||||
resolve_execution_runtime_sync_plan_kind as resolve_sync_plan_kind,
|
||||
};
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{AiExecutionDecision, AiExecutionPlanPayload, AppState, GatewayError};
|
||||
use aether_ai_serving::{
|
||||
build_ai_stream_execution_plan_payload, build_ai_sync_execution_plan_payload,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_build_sync_plan_payload_impl(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
) -> Result<Option<AiExecutionPlanPayload>, GatewayError> {
|
||||
let Some(plan_kind) = resolve_sync_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(payload) = super::maybe_build_sync_decision_payload(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
body_base64,
|
||||
body_is_empty,
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
build_sync_plan_payload_from_decision(parts, body_json, plan_kind, payload)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_stream_plan_payload_impl(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<AiExecutionPlanPayload>, GatewayError> {
|
||||
let Some(plan_kind) = resolve_stream_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(payload) = super::maybe_build_stream_decision_payload(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
body_base64,
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
build_stream_plan_payload_from_decision(parts, body_json, plan_kind, payload)
|
||||
}
|
||||
|
||||
fn build_sync_plan_payload_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
mut payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiExecutionPlanPayload>, GatewayError> {
|
||||
let auth_context = payload.auth_context.take();
|
||||
let plan_and_report = match plan_kind {
|
||||
OPENAI_CHAT_SYNC_PLAN_KIND => {
|
||||
build_openai_chat_sync_plan_from_decision(parts, body_json, payload)?
|
||||
}
|
||||
OPENAI_RESPONSES_SYNC_PLAN_KIND => {
|
||||
build_openai_responses_sync_plan_from_decision(parts, body_json, payload, false)?
|
||||
}
|
||||
OPENAI_IMAGE_SYNC_PLAN_KIND => build_passthrough_sync_plan_from_decision(parts, payload)?,
|
||||
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND => {
|
||||
build_openai_responses_sync_plan_from_decision(parts, body_json, payload, true)?
|
||||
}
|
||||
CLAUDE_CHAT_SYNC_PLAN_KIND | CLAUDE_CLI_SYNC_PLAN_KIND => {
|
||||
build_standard_sync_plan_from_decision(parts, body_json, payload)?
|
||||
}
|
||||
GEMINI_CHAT_SYNC_PLAN_KIND | GEMINI_CLI_SYNC_PLAN_KIND => {
|
||||
build_gemini_sync_plan_from_decision(parts, body_json, payload)?
|
||||
}
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND
|
||||
| OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND
|
||||
| OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND
|
||||
| OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND
|
||||
| GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND
|
||||
| GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND
|
||||
| GEMINI_FILES_LIST_PLAN_KIND
|
||||
| GEMINI_FILES_GET_PLAN_KIND
|
||||
| GEMINI_FILES_DELETE_PLAN_KIND => {
|
||||
build_passthrough_sync_plan_from_decision(parts, payload)?
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
|
||||
Ok(plan_and_report
|
||||
.map(|value| build_ai_sync_execution_plan_payload(plan_kind, value, auth_context)))
|
||||
}
|
||||
|
||||
fn build_stream_plan_payload_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
mut payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiExecutionPlanPayload>, GatewayError> {
|
||||
let auth_context = payload.auth_context.take();
|
||||
let plan_and_report = match plan_kind {
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND => {
|
||||
build_openai_chat_stream_plan_from_decision(parts, body_json, payload)?
|
||||
}
|
||||
OPENAI_RESPONSES_STREAM_PLAN_KIND => {
|
||||
build_openai_responses_stream_plan_from_decision(parts, body_json, payload, false)?
|
||||
}
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND => {
|
||||
build_standard_stream_plan_from_decision(parts, body_json, payload, false)?
|
||||
}
|
||||
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND => {
|
||||
build_openai_responses_stream_plan_from_decision(parts, body_json, payload, true)?
|
||||
}
|
||||
CLAUDE_CHAT_STREAM_PLAN_KIND | CLAUDE_CLI_STREAM_PLAN_KIND => {
|
||||
build_standard_stream_plan_from_decision(parts, body_json, payload, true)?
|
||||
}
|
||||
GEMINI_CHAT_STREAM_PLAN_KIND | GEMINI_CLI_STREAM_PLAN_KIND => {
|
||||
build_gemini_stream_plan_from_decision(parts, body_json, payload)?
|
||||
}
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND | GEMINI_FILES_DOWNLOAD_PLAN_KIND => {
|
||||
build_passthrough_stream_plan_from_decision(parts, payload)?
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
|
||||
Ok(plan_and_report
|
||||
.map(|value| build_ai_stream_execution_plan_payload(plan_kind, value, auth_context)))
|
||||
}
|
||||
26
apps/aether-gateway/src/ai_serving/planner/decision/mod.rs
Normal file
26
apps/aether-gateway/src/ai_serving/planner/decision/mod.rs
Normal file
@@ -0,0 +1,26 @@
|
||||
mod control_plan;
|
||||
mod stream;
|
||||
mod sync;
|
||||
|
||||
pub(crate) use self::control_plan::{
|
||||
maybe_build_stream_plan_payload_impl, maybe_build_sync_plan_payload_impl,
|
||||
};
|
||||
pub(crate) use self::stream::maybe_build_stream_decision_payload;
|
||||
pub(crate) use self::sync::maybe_build_sync_decision_payload;
|
||||
pub(crate) use super::passthrough::{
|
||||
maybe_build_stream_local_same_format_provider_decision_payload,
|
||||
maybe_build_sync_local_same_format_provider_decision_payload,
|
||||
};
|
||||
pub(crate) use super::specialized::{
|
||||
maybe_build_stream_local_gemini_files_decision_payload,
|
||||
maybe_build_stream_local_image_decision_payload,
|
||||
maybe_build_sync_local_gemini_files_decision_payload,
|
||||
maybe_build_sync_local_image_decision_payload, maybe_build_sync_local_video_decision_payload,
|
||||
};
|
||||
pub(crate) use super::standard::{
|
||||
maybe_build_stream_local_decision_payload,
|
||||
maybe_build_stream_local_openai_responses_decision_payload,
|
||||
maybe_build_stream_local_standard_decision_payload, maybe_build_sync_local_decision_payload,
|
||||
maybe_build_sync_local_openai_responses_decision_payload,
|
||||
maybe_build_sync_local_standard_decision_payload,
|
||||
};
|
||||
188
apps/aether-gateway/src/ai_serving/planner/decision/stream.rs
Normal file
188
apps/aether-gateway/src/ai_serving/planner/decision/stream.rs
Normal file
@@ -0,0 +1,188 @@
|
||||
use aether_ai_serving::{
|
||||
build_ai_execution_decision_from_plan, run_ai_stream_decision_path,
|
||||
AiExecutionDecisionFromPlanParts, AiStreamDecisionPathPort, AiStreamDecisionStep,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::ai_serving::planner::common::{
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_serving::planner::route::{
|
||||
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
|
||||
};
|
||||
use crate::ai_serving::{resolve_decision_execution_runtime_auth_context, GatewayControlDecision};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
pub(crate) async fn maybe_build_stream_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(plan_kind) = resolve_execution_runtime_stream_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if !is_matching_stream_request(plan_kind, parts, body_json, body_base64) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let port = GatewayStreamDecisionPathPort {
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
body_base64,
|
||||
plan_kind,
|
||||
};
|
||||
|
||||
run_ai_stream_decision_path(&port).await
|
||||
}
|
||||
|
||||
struct GatewayStreamDecisionPathPort<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
trace_id: &'a str,
|
||||
decision: &'a GatewayControlDecision,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_base64: Option<&'a str>,
|
||||
plan_kind: &'a str,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiStreamDecisionPathPort for GatewayStreamDecisionPathPort<'_> {
|
||||
type Decision = AiExecutionDecision;
|
||||
type Error = GatewayError;
|
||||
|
||||
async fn build_stream_decision_step(
|
||||
&self,
|
||||
step: AiStreamDecisionStep,
|
||||
) -> Result<Option<Self::Decision>, Self::Error> {
|
||||
match step {
|
||||
AiStreamDecisionStep::LocalVideoContent => {
|
||||
maybe_build_local_video_task_content_stream_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiStreamDecisionStep::LocalImage => {
|
||||
super::maybe_build_stream_local_image_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
self.body_base64,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiStreamDecisionStep::LocalOpenAiChat => {
|
||||
super::maybe_build_stream_local_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiStreamDecisionStep::LocalOpenAiResponses => {
|
||||
super::maybe_build_stream_local_openai_responses_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiStreamDecisionStep::LocalStandardFamily => {
|
||||
super::maybe_build_stream_local_standard_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiStreamDecisionStep::LocalSameFormatProvider => {
|
||||
super::maybe_build_stream_local_same_format_provider_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiStreamDecisionStep::LocalGeminiFiles => {
|
||||
super::maybe_build_stream_local_gemini_files_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn maybe_build_local_video_task_content_stream_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
if plan_kind != OPENAI_VIDEO_CONTENT_PLAN_KIND
|
||||
|| decision.route_family.as_deref() != Some("openai")
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let _ = state
|
||||
.hydrate_video_task_for_route(decision.route_family.as_deref(), parts.uri.path())
|
||||
.await?;
|
||||
|
||||
let Some(action) = state.video_tasks.prepare_openai_content_stream_action(
|
||||
parts.uri.path(),
|
||||
parts.uri.query(),
|
||||
trace_id,
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let crate::video_tasks::LocalVideoTaskContentAction::StreamPlan(plan) = action else {
|
||||
return Ok(None);
|
||||
};
|
||||
let plan = *plan;
|
||||
|
||||
Ok(Some(build_ai_execution_decision_from_plan(
|
||||
AiExecutionDecisionFromPlanParts {
|
||||
action: EXECUTION_RUNTIME_STREAM_DECISION_ACTION.to_string(),
|
||||
decision_kind: Some(plan_kind.to_string()),
|
||||
request_id: None,
|
||||
upstream_base_url: None,
|
||||
include_auth_pair: false,
|
||||
plan,
|
||||
report_kind: None,
|
||||
report_context: None,
|
||||
auth_context: resolve_decision_execution_runtime_auth_context(decision),
|
||||
},
|
||||
)))
|
||||
}
|
||||
257
apps/aether-gateway/src/ai_serving/planner/decision/sync.rs
Normal file
257
apps/aether-gateway/src/ai_serving/planner/decision/sync.rs
Normal file
@@ -0,0 +1,257 @@
|
||||
use aether_ai_serving::{
|
||||
build_ai_execution_decision_from_plan, infer_ai_upstream_base_url, run_ai_sync_decision_path,
|
||||
AiExecutionDecisionFromPlanParts, AiSyncDecisionPathPort, AiSyncDecisionStep,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use tracing::debug;
|
||||
|
||||
use crate::ai_serving::planner::common::{
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_FILES_DELETE_PLAN_KIND,
|
||||
GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_serving::planner::route::resolve_execution_runtime_sync_plan_kind;
|
||||
use crate::ai_serving::{
|
||||
build_execution_runtime_auth_context, resolve_execution_runtime_auth_context,
|
||||
GatewayControlDecision,
|
||||
};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
pub(crate) async fn maybe_build_sync_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(plan_kind) = resolve_execution_runtime_sync_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let port = GatewaySyncDecisionPathPort {
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
body_base64,
|
||||
body_is_empty,
|
||||
plan_kind,
|
||||
};
|
||||
|
||||
run_ai_sync_decision_path(&port).await
|
||||
}
|
||||
|
||||
struct GatewaySyncDecisionPathPort<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
trace_id: &'a str,
|
||||
decision: &'a GatewayControlDecision,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_base64: Option<&'a str>,
|
||||
body_is_empty: bool,
|
||||
plan_kind: &'a str,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiSyncDecisionPathPort for GatewaySyncDecisionPathPort<'_> {
|
||||
type Decision = AiExecutionDecision;
|
||||
type Error = GatewayError;
|
||||
|
||||
fn sync_decision_step_enabled(&self, step: AiSyncDecisionStep) -> bool {
|
||||
if step == AiSyncDecisionStep::LocalGeminiFiles {
|
||||
return matches!(
|
||||
self.plan_kind,
|
||||
GEMINI_FILES_LIST_PLAN_KIND
|
||||
| GEMINI_FILES_GET_PLAN_KIND
|
||||
| GEMINI_FILES_DELETE_PLAN_KIND
|
||||
);
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
async fn build_sync_decision_step(
|
||||
&self,
|
||||
step: AiSyncDecisionStep,
|
||||
) -> Result<Option<Self::Decision>, Self::Error> {
|
||||
match step {
|
||||
AiSyncDecisionStep::VideoTaskFollowUp => {
|
||||
maybe_build_local_video_task_follow_up_sync_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiSyncDecisionStep::LocalVideo => {
|
||||
super::maybe_build_sync_local_video_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiSyncDecisionStep::LocalImage => {
|
||||
super::maybe_build_sync_local_image_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
self.body_base64,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiSyncDecisionStep::LocalOpenAiChat => {
|
||||
super::maybe_build_sync_local_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiSyncDecisionStep::LocalOpenAiResponses => {
|
||||
super::maybe_build_sync_local_openai_responses_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiSyncDecisionStep::LocalStandardFamily => {
|
||||
super::maybe_build_sync_local_standard_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiSyncDecisionStep::LocalSameFormatProvider => {
|
||||
super::maybe_build_sync_local_same_format_provider_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiSyncDecisionStep::LocalGeminiFiles => {
|
||||
super::maybe_build_sync_local_gemini_files_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
self.body_base64,
|
||||
self.body_is_empty,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn maybe_build_local_video_task_follow_up_sync_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
if !matches!(
|
||||
plan_kind,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND
|
||||
| OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND
|
||||
| OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND
|
||||
| GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND
|
||||
) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let _ = state
|
||||
.hydrate_video_task_for_route(decision.route_family.as_deref(), parts.uri.path())
|
||||
.await?;
|
||||
|
||||
let auth_context = resolve_execution_runtime_auth_context(
|
||||
state,
|
||||
decision,
|
||||
&parts.headers,
|
||||
&parts.uri,
|
||||
trace_id,
|
||||
)
|
||||
.await?;
|
||||
let Some(auth_context) = auth_context else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(follow_up) = state.video_tasks.prepare_follow_up_sync_plan(
|
||||
plan_kind,
|
||||
parts.uri.path(),
|
||||
Some(body_json),
|
||||
Some(&auth_context),
|
||||
trace_id,
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let aether_video_tasks_core::LocalVideoTaskFollowUpPlan {
|
||||
plan,
|
||||
report_kind,
|
||||
report_context,
|
||||
} = follow_up;
|
||||
let upstream_base_url = infer_ai_upstream_base_url(&plan.url);
|
||||
|
||||
debug!(
|
||||
event_name = "local_video_follow_up_sync_decision_payload_built",
|
||||
log_type = "debug",
|
||||
trace_id = %trace_id,
|
||||
request_id = %trace_id,
|
||||
candidate_id = ?plan.candidate_id,
|
||||
provider_id = %plan.provider_id,
|
||||
endpoint_id = %plan.endpoint_id,
|
||||
key_id = %plan.key_id,
|
||||
plan_kind,
|
||||
downstream_path = %parts.uri.path(),
|
||||
provider_api_format = %plan.provider_api_format,
|
||||
client_api_format = %plan.client_api_format,
|
||||
upstream_base_url = ?upstream_base_url,
|
||||
upstream_url = %plan.url,
|
||||
"gateway built local video follow-up sync decision payload"
|
||||
);
|
||||
|
||||
Ok(Some(build_ai_execution_decision_from_plan(
|
||||
AiExecutionDecisionFromPlanParts {
|
||||
action: EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string(),
|
||||
decision_kind: Some(plan_kind.to_string()),
|
||||
request_id: Some(trace_id.to_string()),
|
||||
upstream_base_url,
|
||||
include_auth_pair: true,
|
||||
plan,
|
||||
report_kind,
|
||||
report_context,
|
||||
auth_context: Some(build_execution_runtime_auth_context(&auth_context)),
|
||||
},
|
||||
)))
|
||||
}
|
||||
127
apps/aether-gateway/src/ai_serving/planner/decision_input.rs
Normal file
127
apps/aether-gateway/src/ai_serving/planner/decision_input.rs
Normal file
@@ -0,0 +1,127 @@
|
||||
use aether_ai_serving::{run_ai_authenticated_decision_input, AiAuthenticatedDecisionInputPort};
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::ai_serving::{ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct ResolvedLocalDecisionAuthInput {
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(crate) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalRequestedModelDecisionInput {
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(crate) requested_model: String,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(crate) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalAuthenticatedDecisionInput {
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(crate) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
struct GatewayAuthenticatedDecisionInputPort<'a> {
|
||||
state: PlannerAppState<'a>,
|
||||
now_unix_secs: u64,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiAuthenticatedDecisionInputPort for GatewayAuthenticatedDecisionInputPort<'_> {
|
||||
type AuthContext = ExecutionRuntimeAuthContext;
|
||||
type AuthSnapshot = GatewayAuthApiKeySnapshot;
|
||||
type RequiredCapabilities = serde_json::Value;
|
||||
type ResolvedInput = ResolvedLocalDecisionAuthInput;
|
||||
type Error = GatewayError;
|
||||
|
||||
async fn read_auth_snapshot(
|
||||
&self,
|
||||
auth_context: &Self::AuthContext,
|
||||
) -> Result<Option<Self::AuthSnapshot>, Self::Error> {
|
||||
self.state
|
||||
.read_auth_api_key_snapshot(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
self.now_unix_secs,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn resolve_required_capabilities(
|
||||
&self,
|
||||
auth_context: &Self::AuthContext,
|
||||
requested_model: Option<&str>,
|
||||
explicit_required_capabilities: Option<&Self::RequiredCapabilities>,
|
||||
) -> Result<Option<Self::RequiredCapabilities>, Self::Error> {
|
||||
Ok(self
|
||||
.state
|
||||
.resolve_request_candidate_required_capabilities(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
requested_model,
|
||||
explicit_required_capabilities,
|
||||
)
|
||||
.await)
|
||||
}
|
||||
|
||||
fn build_resolved_input(
|
||||
&self,
|
||||
auth_context: Self::AuthContext,
|
||||
auth_snapshot: Self::AuthSnapshot,
|
||||
required_capabilities: Option<Self::RequiredCapabilities>,
|
||||
) -> Self::ResolvedInput {
|
||||
ResolvedLocalDecisionAuthInput {
|
||||
auth_context,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_requested_model_decision_input(
|
||||
resolved_input: ResolvedLocalDecisionAuthInput,
|
||||
requested_model: String,
|
||||
) -> LocalRequestedModelDecisionInput {
|
||||
LocalRequestedModelDecisionInput {
|
||||
auth_context: resolved_input.auth_context,
|
||||
requested_model,
|
||||
auth_snapshot: resolved_input.auth_snapshot,
|
||||
required_capabilities: resolved_input.required_capabilities,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_authenticated_decision_input(
|
||||
resolved_input: ResolvedLocalDecisionAuthInput,
|
||||
) -> LocalAuthenticatedDecisionInput {
|
||||
LocalAuthenticatedDecisionInput {
|
||||
auth_context: resolved_input.auth_context,
|
||||
auth_snapshot: resolved_input.auth_snapshot,
|
||||
required_capabilities: resolved_input.required_capabilities,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_local_authenticated_decision_input(
|
||||
state: &AppState,
|
||||
auth_context: ExecutionRuntimeAuthContext,
|
||||
requested_model: Option<&str>,
|
||||
explicit_required_capabilities: Option<&serde_json::Value>,
|
||||
) -> Result<Option<ResolvedLocalDecisionAuthInput>, GatewayError> {
|
||||
let port = GatewayAuthenticatedDecisionInputPort {
|
||||
state: PlannerAppState::new(state),
|
||||
now_unix_secs: current_unix_secs(),
|
||||
};
|
||||
|
||||
run_ai_authenticated_decision_input(
|
||||
&port,
|
||||
auth_context,
|
||||
requested_model,
|
||||
explicit_required_capabilities,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
use aether_ai_serving::ai_candidate_persistence_policy_spec;
|
||||
pub(crate) use aether_ai_serving::AiCandidatePersistencePolicyKind as LocalCandidatePersistencePolicyKind;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
LocalAvailableCandidatePersistenceContext, LocalSkippedCandidatePersistenceContext,
|
||||
};
|
||||
use crate::ai_serving::ExecutionRuntimeAuthContext;
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct LocalCandidatePersistencePolicy<'a> {
|
||||
pub(crate) available: LocalAvailableCandidatePersistenceContext<'a>,
|
||||
pub(crate) skipped: LocalSkippedCandidatePersistenceContext<'a>,
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_candidate_persistence_policy<'a>(
|
||||
auth_context: &'a ExecutionRuntimeAuthContext,
|
||||
required_capabilities: Option<&'a Value>,
|
||||
kind: LocalCandidatePersistencePolicyKind,
|
||||
) -> LocalCandidatePersistencePolicy<'a> {
|
||||
let spec = ai_candidate_persistence_policy_spec(kind);
|
||||
|
||||
LocalCandidatePersistencePolicy {
|
||||
available: LocalAvailableCandidatePersistenceContext {
|
||||
user_id: &auth_context.user_id,
|
||||
api_key_id: &auth_context.api_key_id,
|
||||
required_capabilities,
|
||||
error_context: spec.available_error_context,
|
||||
},
|
||||
skipped: LocalSkippedCandidatePersistenceContext {
|
||||
user_id: &auth_context.user_id,
|
||||
api_key_id: &auth_context.api_key_id,
|
||||
required_capabilities,
|
||||
error_context: spec.skipped_error_context,
|
||||
record_runtime_miss_diagnostic: spec.record_runtime_miss_diagnostic,
|
||||
},
|
||||
}
|
||||
}
|
||||
142
apps/aether-gateway/src/ai_serving/planner/mod.rs
Normal file
142
apps/aether-gateway/src/ai_serving/planner/mod.rs
Normal file
@@ -0,0 +1,142 @@
|
||||
use crate::ai_serving::{AiExecutionDecision, AiExecutionPlanPayload, GatewayControlDecision};
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
mod candidate_affinity_cache;
|
||||
mod candidate_materialization;
|
||||
mod candidate_metadata;
|
||||
mod candidate_preparation;
|
||||
mod candidate_ranking;
|
||||
mod candidate_resolution;
|
||||
mod candidate_source;
|
||||
mod candidate_transport_ranking_facts;
|
||||
mod common;
|
||||
mod decision;
|
||||
mod decision_input;
|
||||
mod materialization_policy;
|
||||
mod passthrough;
|
||||
mod plan_builders;
|
||||
mod pool_scheduler;
|
||||
mod report_context;
|
||||
mod route;
|
||||
mod runtime_miss;
|
||||
mod spec_metadata;
|
||||
mod specialized;
|
||||
mod standard;
|
||||
mod state;
|
||||
|
||||
pub(crate) use self::passthrough::{
|
||||
build_local_same_format_stream_plan_and_reports, build_local_same_format_sync_plan_and_reports,
|
||||
};
|
||||
pub(crate) use self::plan_builders::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_openai_responses_stream_plan_from_decision,
|
||||
build_openai_responses_sync_plan_from_decision, build_passthrough_sync_plan_from_decision,
|
||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
||||
AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
pub(crate) use self::route::is_matching_stream_request as planner_is_matching_stream_request;
|
||||
pub(crate) use self::specialized::{
|
||||
build_local_gemini_files_stream_plan_and_reports_for_kind,
|
||||
build_local_gemini_files_sync_plan_and_reports_for_kind,
|
||||
build_local_image_stream_plan_and_reports_for_kind,
|
||||
build_local_image_sync_plan_and_reports_for_kind,
|
||||
build_local_video_sync_plan_and_reports_for_kind,
|
||||
};
|
||||
pub(crate) use self::standard::{
|
||||
build_local_openai_chat_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_sync_plan_and_reports_for_kind,
|
||||
build_local_openai_responses_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_responses_sync_plan_and_reports_for_kind,
|
||||
build_local_stream_plan_and_reports as build_standard_family_stream_plan_and_reports,
|
||||
build_local_sync_plan_and_reports as build_standard_family_sync_plan_and_reports,
|
||||
set_local_openai_chat_execution_exhausted_diagnostic,
|
||||
};
|
||||
pub(crate) use self::state::{
|
||||
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
|
||||
PlannerAppState,
|
||||
};
|
||||
pub(crate) use aether_ai_serving::extract_ai_pool_sticky_session_token as extract_pool_sticky_session_token;
|
||||
pub(crate) use aether_ai_serving::{
|
||||
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
|
||||
CandidateFailureDiagnostic, CandidateFailureDiagnosticKind,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_build_sync_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
decision::maybe_build_sync_decision_payload(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
body_base64,
|
||||
body_is_empty,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_stream_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
decision::maybe_build_stream_decision_payload(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
body_base64,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_sync_plan_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
) -> Result<Option<AiExecutionPlanPayload>, GatewayError> {
|
||||
decision::maybe_build_sync_plan_payload_impl(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
body_base64,
|
||||
body_is_empty,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_stream_plan_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<AiExecutionPlanPayload>, GatewayError> {
|
||||
decision::maybe_build_stream_plan_payload_impl(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
body_base64,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
//! Requests that can stay in the same public/provider contract family.
|
||||
|
||||
mod provider;
|
||||
|
||||
pub(crate) use self::provider::{
|
||||
build_local_stream_plan_and_reports as build_local_same_format_stream_plan_and_reports,
|
||||
build_local_sync_plan_and_reports as build_local_same_format_sync_plan_and_reports,
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate,
|
||||
maybe_build_stream_local_same_format_provider_decision_payload,
|
||||
maybe_build_sync_local_same_format_provider_decision_payload,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::transport::provider_types::provider_type_supports_local_same_format_transport;
|
||||
@@ -0,0 +1,96 @@
|
||||
use aether_contracts::RequestBody;
|
||||
|
||||
use super::{
|
||||
augment_sync_report_context, build_ai_execution_plan_from_decision,
|
||||
resolve_ai_passthrough_sync_request_body, take_ai_decision_plan_core, take_non_empty_string,
|
||||
AiExecutionPlanFromDecisionParts, AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
use crate::{AiExecutionDecision, GatewayError};
|
||||
|
||||
pub(crate) fn build_passthrough_sync_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
let Some(core) = take_ai_decision_plan_core(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(upstream_url) = take_non_empty_string(&mut payload.upstream_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let provider_request_headers = std::mem::take(&mut payload.provider_request_headers);
|
||||
let ignored_provider_request_body = serde_json::Value::Null;
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context.take(),
|
||||
&provider_request_headers,
|
||||
&ignored_provider_request_body,
|
||||
)?;
|
||||
let request_body = resolve_ai_passthrough_sync_request_body(
|
||||
payload.provider_request_body.take(),
|
||||
payload.provider_request_body_base64.take(),
|
||||
);
|
||||
let provider_request_method = take_non_empty_string(&mut payload.provider_request_method);
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| provider_request_headers.get("content-type").cloned());
|
||||
|
||||
let plan = build_ai_execution_plan_from_decision(
|
||||
&mut payload,
|
||||
AiExecutionPlanFromDecisionParts {
|
||||
core,
|
||||
method: provider_request_method.unwrap_or_else(|| parts.method.to_string()),
|
||||
url: upstream_url,
|
||||
headers: provider_request_headers,
|
||||
content_type,
|
||||
body: request_body,
|
||||
stream: false,
|
||||
},
|
||||
);
|
||||
|
||||
Ok(Some(AiSyncAttempt {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn build_passthrough_stream_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
let Some(core) = take_ai_decision_plan_core(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(upstream_url) = take_non_empty_string(&mut payload.upstream_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let provider_request_headers = std::mem::take(&mut payload.provider_request_headers);
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| provider_request_headers.get("content-type").cloned());
|
||||
let plan = build_ai_execution_plan_from_decision(
|
||||
&mut payload,
|
||||
AiExecutionPlanFromDecisionParts {
|
||||
core,
|
||||
method: parts.method.to_string(),
|
||||
url: upstream_url,
|
||||
headers: provider_request_headers,
|
||||
content_type,
|
||||
body: RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
},
|
||||
stream: true,
|
||||
},
|
||||
);
|
||||
|
||||
Ok(Some(AiStreamAttempt {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context: payload.report_context,
|
||||
}))
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use std::collections::BTreeMap;
|
||||
use url::form_urlencoded;
|
||||
|
||||
use aether_data_contracts::repository::candidates::{
|
||||
RequestCandidateStatus, UpsertRequestCandidateRecord,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::{json, Value};
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_serving::planner::common::{
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
|
||||
};
|
||||
use crate::ai_serving::planner::plan_builders::{AiStreamAttempt, AiSyncAttempt};
|
||||
use crate::ai_serving::transport::antigravity::{
|
||||
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
|
||||
build_antigravity_v1internal_url, classify_local_antigravity_request_support,
|
||||
AntigravityEnvelopeRequestType, AntigravityRequestEnvelopeSupport,
|
||||
AntigravityRequestSideSupport, AntigravityRequestUrlAction,
|
||||
};
|
||||
use crate::ai_serving::transport::auth::{
|
||||
build_openai_passthrough_headers, resolve_local_gemini_auth, resolve_local_standard_auth,
|
||||
};
|
||||
use crate::ai_serving::transport::claude_code::{
|
||||
build_claude_code_messages_url, build_claude_code_passthrough_headers,
|
||||
sanitize_claude_code_request_body, supports_local_claude_code_transport_with_network,
|
||||
};
|
||||
use crate::ai_serving::transport::kiro::{
|
||||
build_kiro_generate_assistant_response_url, build_kiro_provider_headers,
|
||||
build_kiro_provider_request_body, supports_local_kiro_request_transport_with_network,
|
||||
KIRO_ENVELOPE_NAME,
|
||||
};
|
||||
use crate::ai_serving::transport::policy::{
|
||||
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
|
||||
};
|
||||
use crate::ai_serving::transport::url::{
|
||||
build_claude_messages_url, build_gemini_content_url, build_passthrough_path_url,
|
||||
};
|
||||
use crate::ai_serving::transport::vertex::{
|
||||
build_vertex_api_key_gemini_content_url, resolve_local_vertex_api_key_query_auth,
|
||||
supports_local_vertex_api_key_gemini_transport_with_network,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
apply_local_body_rules, apply_local_header_rules, build_passthrough_headers,
|
||||
ensure_upstream_auth_header, resolve_transport_execution_timeouts,
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
|
||||
LocalResolvedOAuthRequestAuth,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
collect_control_headers, ConversionMode, ExecutionStrategy, GatewayControlDecision,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, AiExecutionDecision, AppState, GatewayError,
|
||||
};
|
||||
|
||||
mod family;
|
||||
mod plans;
|
||||
mod request;
|
||||
|
||||
pub(crate) use self::family::{
|
||||
materialize_local_same_format_provider_candidate_attempts,
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate,
|
||||
resolve_local_same_format_provider_decision_input, LocalSameFormatProviderFamily,
|
||||
LocalSameFormatProviderSpec,
|
||||
};
|
||||
pub(crate) use self::family::{
|
||||
maybe_build_stream_local_same_format_provider_decision_payload,
|
||||
maybe_build_sync_local_same_format_provider_decision_payload,
|
||||
};
|
||||
pub(crate) use self::plans::{
|
||||
build_local_stream_plan_and_reports, build_local_sync_plan_and_reports,
|
||||
};
|
||||
|
||||
const ANTIGRAVITY_ENVELOPE_NAME: &str = "antigravity:v1internal";
|
||||
@@ -0,0 +1,149 @@
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::runtime_miss::{
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal,
|
||||
apply_local_runtime_candidate_terminal_reason, set_local_runtime_miss_diagnostic_reason,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
use super::super::plans::{resolve_stream_spec, resolve_sync_spec};
|
||||
use super::candidates::{
|
||||
materialize_local_same_format_provider_candidate_attempts,
|
||||
resolve_local_same_format_provider_decision_input,
|
||||
};
|
||||
use super::payload::maybe_build_local_same_format_provider_decision_payload_for_candidate;
|
||||
|
||||
pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let requested_model_family = spec_metadata
|
||||
.requested_model_family
|
||||
.expect("same-format provider spec metadata should include requested-model family");
|
||||
|
||||
let Some(input) = resolve_local_same_format_provider_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
extract_requested_model_from_request(parts, body_json, requested_model_family)
|
||||
.as_deref(),
|
||||
"decision_input_unavailable",
|
||||
);
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
state,
|
||||
trace_id,
|
||||
candidate_count,
|
||||
);
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) =
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_sync_plans");
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let requested_model_family = spec_metadata
|
||||
.requested_model_family
|
||||
.expect("same-format provider spec metadata should include requested-model family");
|
||||
|
||||
let Some(input) = resolve_local_same_format_provider_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
extract_requested_model_from_request(parts, body_json, requested_model_family)
|
||||
.as_deref(),
|
||||
"decision_input_unavailable",
|
||||
);
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
state,
|
||||
trace_id,
|
||||
candidate_count,
|
||||
);
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) =
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_stream_plans");
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
@@ -0,0 +1,171 @@
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
materialize_local_execution_candidates_with_serving, LocalCandidateResolutionMode,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_metadata::{
|
||||
build_local_execution_candidate_contract_metadata,
|
||||
build_local_execution_candidate_contract_metadata_for_candidate,
|
||||
LocalExecutionCandidateMetadataParts,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_serving::{
|
||||
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
|
||||
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision, PlannerAppState,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
use super::{
|
||||
LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput,
|
||||
LocalSameFormatProviderSpec,
|
||||
};
|
||||
|
||||
pub(crate) async fn resolve_local_same_format_provider_decision_input(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<LocalSameFormatProviderDecisionInput> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let requested_model = extract_requested_model_from_request(
|
||||
parts,
|
||||
body_json,
|
||||
spec_metadata
|
||||
.requested_model_family
|
||||
.expect("same-format provider specs should declare requested-model family"),
|
||||
)?;
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
auth_context,
|
||||
Some(requested_model.as_str()),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => return None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec_metadata.api_format,
|
||||
error = ?err,
|
||||
"gateway local same-format decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(build_local_requested_model_decision_input(
|
||||
resolved_input,
|
||||
requested_model,
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Result<(Vec<LocalSameFormatProviderCandidateAttempt>, usize), GatewayError> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let sticky_session_token = extract_pool_sticky_session_token(body_json);
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision,
|
||||
);
|
||||
let (candidates, preselection_skipped) = planner_state
|
||||
.list_selectable_candidates_with_skip_reasons(
|
||||
spec_metadata.api_format,
|
||||
&input.requested_model,
|
||||
spec_metadata.require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
Some(&input.auth_snapshot),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
let outcome = materialize_local_execution_candidates_with_serving(
|
||||
planner_state,
|
||||
trace_id,
|
||||
spec_metadata.api_format,
|
||||
Some(&input.requested_model),
|
||||
Some(&input.auth_snapshot),
|
||||
input.required_capabilities.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
persistence_policy,
|
||||
candidates,
|
||||
preselection_skipped
|
||||
.into_iter()
|
||||
.map(|item| SkippedLocalExecutionCandidate {
|
||||
candidate: item.candidate,
|
||||
skip_reason: item.skip_reason,
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
})
|
||||
.collect(),
|
||||
LocalCandidateResolutionMode::Standard,
|
||||
|eligible| {
|
||||
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.api_format,
|
||||
);
|
||||
Some(build_local_execution_candidate_contract_metadata(
|
||||
LocalExecutionCandidateMetadataParts {
|
||||
eligible,
|
||||
provider_api_format: spec_metadata.api_format,
|
||||
client_api_format: spec_metadata.api_format,
|
||||
extra_fields: serde_json::Map::new(),
|
||||
},
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
spec_metadata.api_format,
|
||||
))
|
||||
},
|
||||
|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(|| spec_metadata.api_format.to_string());
|
||||
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
|
||||
spec_metadata.api_format,
|
||||
provider_api_format.as_str(),
|
||||
);
|
||||
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;
|
||||
|
||||
Ok((outcome.attempts, outcome.candidate_count))
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
mod build;
|
||||
mod candidates;
|
||||
mod payload;
|
||||
mod request;
|
||||
|
||||
pub(crate) use self::build::{
|
||||
maybe_build_stream_local_same_format_provider_decision_payload,
|
||||
maybe_build_sync_local_same_format_provider_decision_payload,
|
||||
};
|
||||
pub(crate) use self::candidates::{
|
||||
materialize_local_same_format_provider_candidate_attempts,
|
||||
resolve_local_same_format_provider_decision_input,
|
||||
};
|
||||
pub(crate) use self::payload::maybe_build_local_same_format_provider_decision_payload_for_candidate;
|
||||
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalSameFormatProviderCandidateAttempt;
|
||||
pub(crate) use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput as LocalSameFormatProviderDecisionInput;
|
||||
pub(crate) use crate::ai_serving::{LocalSameFormatProviderFamily, LocalSameFormatProviderSpec};
|
||||
@@ -0,0 +1,266 @@
|
||||
use serde_json::json;
|
||||
|
||||
use crate::ai_serving::ai_local_execution_contract_for_formats;
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, LocalExecutionReportContextParts,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_serving::planner::CandidateFailureDiagnostic;
|
||||
use crate::ai_serving::planner::{
|
||||
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, append_local_failover_policy_to_value,
|
||||
AiExecutionDecision, AppState,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
use super::request::resolve_local_same_format_provider_candidate_payload_parts;
|
||||
use super::{
|
||||
LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput,
|
||||
LocalSameFormatProviderSpec,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
body_json: &serde_json::Value,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
attempt: LocalSameFormatProviderCandidateAttempt,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let LocalSameFormatProviderCandidateAttempt {
|
||||
eligible,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
..
|
||||
} = &attempt;
|
||||
let candidate = &eligible.candidate;
|
||||
let (execution_strategy, conversion_mode) =
|
||||
ai_local_execution_contract_for_formats(spec_metadata.api_format, spec_metadata.api_format);
|
||||
let resolved = resolve_local_same_format_provider_candidate_payload_parts(
|
||||
state, parts, trace_id, body_json, input, &attempt, spec,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let prompt_cache_key = resolved
|
||||
.provider_request_body
|
||||
.get("prompt_cache_key")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let proxy = state
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&resolved.transport)
|
||||
.await;
|
||||
let tls_profile = resolve_transport_tls_profile(&resolved.transport);
|
||||
let mut extra_fields = serde_json::Map::new();
|
||||
if let Some(proxy_value) =
|
||||
build_request_trace_proxy_value(Some(&resolved.transport), proxy.as_ref())
|
||||
{
|
||||
extra_fields.insert("proxy".to_string(), proxy_value);
|
||||
}
|
||||
if resolved.is_kiro {
|
||||
extra_fields.insert(
|
||||
"envelope_name".to_string(),
|
||||
json!(crate::ai_serving::transport::kiro::KIRO_ENVELOPE_NAME),
|
||||
);
|
||||
} else if resolved.is_antigravity {
|
||||
extra_fields.insert(
|
||||
"envelope_name".to_string(),
|
||||
json!(super::super::ANTIGRAVITY_ENVELOPE_NAME),
|
||||
);
|
||||
}
|
||||
let report_context = append_local_failover_policy_to_value(
|
||||
append_execution_contract_fields_to_value(
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &input.auth_context,
|
||||
request_id: trace_id,
|
||||
candidate_id,
|
||||
attempt_identity: attempt.attempt_identity(),
|
||||
model: &input.requested_model,
|
||||
provider_name: &resolved.transport.provider.name,
|
||||
provider_id: &candidate.provider_id,
|
||||
endpoint_id: &candidate.endpoint_id,
|
||||
key_id: &candidate.key_id,
|
||||
key_name: Some(&candidate.key_name),
|
||||
model_id: Some(&candidate.model_id),
|
||||
global_model_id: Some(&candidate.global_model_id),
|
||||
global_model_name: Some(&candidate.global_model_name),
|
||||
provider_api_format: spec_metadata.api_format,
|
||||
client_api_format: spec_metadata.api_format,
|
||||
mapped_model: Some(&resolved.mapped_model),
|
||||
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
|
||||
ranking: eligible.ranking.as_ref(),
|
||||
upstream_url: Some(&resolved.upstream_url),
|
||||
header_rules: resolved.transport.endpoint.header_rules.as_ref(),
|
||||
body_rules: resolved.transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: Some(serde_json::Value::Null),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: body_json
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
upstream_is_stream: resolved.upstream_is_stream,
|
||||
has_envelope: resolved.is_kiro || resolved.is_antigravity,
|
||||
needs_conversion: false,
|
||||
extra_fields,
|
||||
}),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.api_format,
|
||||
),
|
||||
&resolved.transport,
|
||||
);
|
||||
let super::request::LocalSameFormatProviderCandidatePayloadParts {
|
||||
transport,
|
||||
is_antigravity: _,
|
||||
is_kiro: _,
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
report_kind,
|
||||
upstream_is_stream,
|
||||
upstream_url,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
} = resolved;
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.to_string(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: None,
|
||||
auth_header,
|
||||
auth_value,
|
||||
provider_api_format: spec_metadata.api_format.to_string(),
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream,
|
||||
report_kind: Some(report_kind.to_string()),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_same_format_provider_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn mark_skipped_local_same_format_provider_candidate_with_extra_data(
|
||||
state: &AppState,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
extra_data: Option<serde_json::Value>,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate_with_extra_data(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
extra_data,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn mark_skipped_local_same_format_provider_candidate_with_failure_diagnostic(
|
||||
state: &AppState,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
diagnostic: CandidateFailureDiagnostic,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
diagnostic,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -0,0 +1,237 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::transport::antigravity::{
|
||||
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
|
||||
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
|
||||
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
build_same_format_provider_headers, SameFormatProviderHeadersInput,
|
||||
};
|
||||
use crate::ai_serving::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
|
||||
use crate::AppState;
|
||||
|
||||
mod policy;
|
||||
mod prepare;
|
||||
|
||||
use self::prepare::prepare_local_same_format_provider_candidate;
|
||||
use super::payload::{
|
||||
mark_skipped_local_same_format_provider_candidate,
|
||||
mark_skipped_local_same_format_provider_candidate_with_extra_data,
|
||||
mark_skipped_local_same_format_provider_candidate_with_failure_diagnostic,
|
||||
};
|
||||
use super::{
|
||||
LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput,
|
||||
LocalSameFormatProviderSpec,
|
||||
};
|
||||
use crate::ai_serving::planner::standard::same_format_provider_request_body_failure_extra_data;
|
||||
|
||||
pub(crate) struct LocalSameFormatProviderCandidatePayloadParts {
|
||||
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
pub(super) is_antigravity: bool,
|
||||
pub(super) is_kiro: bool,
|
||||
pub(super) auth_header: Option<String>,
|
||||
pub(super) auth_value: Option<String>,
|
||||
pub(super) mapped_model: String,
|
||||
pub(super) report_kind: &'static str,
|
||||
pub(super) upstream_is_stream: bool,
|
||||
pub(super) upstream_url: String,
|
||||
pub(super) provider_request_headers: BTreeMap<String, String>,
|
||||
pub(super) provider_request_body: Value,
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
body_json: &serde_json::Value,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
attempt: &LocalSameFormatProviderCandidateAttempt,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<LocalSameFormatProviderCandidatePayloadParts> {
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let prepared = prepare_local_same_format_provider_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
input,
|
||||
&attempt.eligible,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let Some(base_provider_request_body) =
|
||||
super::super::request::build_same_format_provider_request_body(
|
||||
body_json,
|
||||
&prepared.mapped_model,
|
||||
spec,
|
||||
prepared.transport.endpoint.body_rules.as_ref(),
|
||||
prepared.upstream_is_stream,
|
||||
prepared.kiro_auth.as_ref(),
|
||||
prepared.is_claude_code,
|
||||
)
|
||||
else {
|
||||
mark_skipped_local_same_format_provider_candidate_with_extra_data(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"provider_request_body_missing",
|
||||
same_format_provider_request_body_failure_extra_data(
|
||||
body_json,
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
prepared.transport.endpoint.body_rules.as_ref(),
|
||||
if prepared.kiro_auth.is_some() {
|
||||
"kiro_envelope"
|
||||
} else {
|
||||
"same_format"
|
||||
},
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let antigravity_auth = if prepared.is_antigravity {
|
||||
match classify_local_antigravity_request_support(
|
||||
&prepared.transport,
|
||||
&base_provider_request_body,
|
||||
AntigravityEnvelopeRequestType::Agent,
|
||||
) {
|
||||
AntigravityRequestSideSupport::Supported(spec) => Some(spec.auth),
|
||||
AntigravityRequestSideSupport::Unsupported(_) => {
|
||||
mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_unsupported",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let provider_request_body = if let Some(antigravity_auth) = antigravity_auth.as_ref() {
|
||||
match build_antigravity_safe_v1internal_request(
|
||||
antigravity_auth,
|
||||
trace_id,
|
||||
&prepared.mapped_model,
|
||||
&base_provider_request_body,
|
||||
AntigravityEnvelopeRequestType::Agent,
|
||||
) {
|
||||
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
|
||||
AntigravityRequestEnvelopeSupport::Unsupported(_) => {
|
||||
mark_skipped_local_same_format_provider_candidate_with_extra_data(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"provider_request_body_missing",
|
||||
same_format_provider_request_body_failure_extra_data(
|
||||
body_json,
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
prepared.transport.endpoint.body_rules.as_ref(),
|
||||
"antigravity_envelope",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
base_provider_request_body
|
||||
};
|
||||
|
||||
let Some(upstream_url) = super::super::request::build_same_format_upstream_url(
|
||||
parts,
|
||||
&prepared.transport,
|
||||
&prepared.mapped_model,
|
||||
spec,
|
||||
prepared.upstream_is_stream,
|
||||
prepared.kiro_auth.as_ref(),
|
||||
) else {
|
||||
mark_skipped_local_same_format_provider_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"upstream_url_missing",
|
||||
CandidateFailureDiagnostic::upstream_url_missing(
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
"same_format_provider_url",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let extra_headers = antigravity_auth
|
||||
.as_ref()
|
||||
.map(build_antigravity_static_identity_headers)
|
||||
.unwrap_or_default();
|
||||
let Some(provider_request_headers) =
|
||||
build_same_format_provider_headers(SameFormatProviderHeadersInput {
|
||||
headers: &parts.headers,
|
||||
provider_request_body: &provider_request_body,
|
||||
original_request_body: body_json,
|
||||
header_rules: prepared.transport.endpoint.header_rules.as_ref(),
|
||||
behavior: prepared.behavior,
|
||||
auth_header: prepared.auth_header.as_deref(),
|
||||
auth_value: prepared.auth_value.as_deref(),
|
||||
extra_headers: &extra_headers,
|
||||
key_fingerprint: prepared.transport.key.fingerprint.as_ref(),
|
||||
kiro_auth_config: prepared.kiro_auth.as_ref().map(|auth| &auth.auth_config),
|
||||
kiro_machine_id: prepared
|
||||
.kiro_auth
|
||||
.as_ref()
|
||||
.map(|auth| auth.machine_id.as_str()),
|
||||
})
|
||||
else {
|
||||
mark_skipped_local_same_format_provider_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
CandidateFailureDiagnostic::header_rules_apply_failed(
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
"same_format_provider_headers",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
Some(LocalSameFormatProviderCandidatePayloadParts {
|
||||
transport: prepared.transport,
|
||||
is_antigravity: prepared.is_antigravity,
|
||||
is_kiro: prepared.is_kiro,
|
||||
auth_header: prepared.auth_header,
|
||||
auth_value: prepared.auth_value,
|
||||
mapped_model: prepared.mapped_model,
|
||||
report_kind: prepared.report_kind,
|
||||
upstream_is_stream: prepared.upstream_is_stream,
|
||||
upstream_url,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
use crate::ai_serving::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata;
|
||||
use crate::ai_serving::transport::{
|
||||
classify_same_format_provider_request_behavior as classify_same_format_provider_request_behavior_impl,
|
||||
resolve_same_format_provider_direct_auth as resolve_same_format_provider_direct_auth_impl,
|
||||
same_format_provider_transport_supported as same_format_provider_transport_supported_impl,
|
||||
same_format_provider_transport_unsupported_reason as same_format_provider_transport_unsupported_reason_impl,
|
||||
should_try_same_format_provider_oauth_auth as should_try_same_format_provider_oauth_auth_impl,
|
||||
GatewayProviderTransportSnapshot, SameFormatProviderFamily, SameFormatProviderRequestBehavior,
|
||||
SameFormatProviderRequestBehaviorParams,
|
||||
};
|
||||
|
||||
use super::super::LocalSameFormatProviderFamily;
|
||||
|
||||
pub(super) fn classify_same_format_provider_request_behavior(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
spec_metadata: LocalExecutionSurfaceSpecMetadata,
|
||||
) -> SameFormatProviderRequestBehavior {
|
||||
classify_same_format_provider_request_behavior_impl(
|
||||
transport,
|
||||
SameFormatProviderRequestBehaviorParams {
|
||||
require_streaming: spec_metadata.require_streaming,
|
||||
report_kind: spec_metadata
|
||||
.report_kind
|
||||
.expect("same-format provider specs should declare report kind"),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn same_format_provider_transport_supported(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
same_format_provider_transport_supported_impl(
|
||||
behavior,
|
||||
transport,
|
||||
same_format_provider_family(family),
|
||||
api_format,
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn same_format_provider_transport_unsupported_reason(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
api_format: &str,
|
||||
) -> Option<&'static str> {
|
||||
same_format_provider_transport_unsupported_reason_impl(
|
||||
behavior,
|
||||
transport,
|
||||
same_format_provider_family(family),
|
||||
api_format,
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn should_try_same_format_provider_oauth_auth(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
) -> bool {
|
||||
should_try_same_format_provider_oauth_auth_impl(
|
||||
behavior,
|
||||
transport,
|
||||
same_format_provider_family(family),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn resolve_same_format_provider_direct_auth(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
) -> Option<(String, String)> {
|
||||
resolve_same_format_provider_direct_auth_impl(
|
||||
behavior,
|
||||
transport,
|
||||
same_format_provider_family(family),
|
||||
)
|
||||
}
|
||||
|
||||
fn same_format_provider_family(family: LocalSameFormatProviderFamily) -> SameFormatProviderFamily {
|
||||
match family {
|
||||
LocalSameFormatProviderFamily::Standard => SameFormatProviderFamily::Standard,
|
||||
LocalSameFormatProviderFamily::Gemini => SameFormatProviderFamily::Gemini,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,175 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::ai_serving::planner::candidate_preparation::{
|
||||
resolve_candidate_mapped_model, resolve_candidate_oauth_auth, OauthPreparationContext,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_serving::transport::kiro::KiroRequestAuth;
|
||||
use crate::ai_serving::transport::vertex::resolve_local_vertex_api_key_query_auth;
|
||||
use crate::ai_serving::transport::SameFormatProviderRequestBehavior;
|
||||
use crate::ai_serving::{
|
||||
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
|
||||
};
|
||||
use crate::AppState;
|
||||
|
||||
use super::super::LocalSameFormatProviderDecisionInput;
|
||||
use super::super::LocalSameFormatProviderSpec;
|
||||
use super::policy::{
|
||||
classify_same_format_provider_request_behavior, resolve_same_format_provider_direct_auth,
|
||||
same_format_provider_transport_supported, same_format_provider_transport_unsupported_reason,
|
||||
should_try_same_format_provider_oauth_auth,
|
||||
};
|
||||
|
||||
pub(super) struct PreparedSameFormatProviderCandidate {
|
||||
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
pub(super) behavior: SameFormatProviderRequestBehavior,
|
||||
pub(super) is_antigravity: bool,
|
||||
pub(super) is_claude_code: bool,
|
||||
pub(super) is_vertex: bool,
|
||||
pub(super) is_kiro: bool,
|
||||
pub(super) kiro_auth: Option<KiroRequestAuth>,
|
||||
pub(super) auth_header: Option<String>,
|
||||
pub(super) auth_value: Option<String>,
|
||||
pub(super) mapped_model: String,
|
||||
pub(super) report_kind: &'static str,
|
||||
pub(super) upstream_is_stream: bool,
|
||||
}
|
||||
|
||||
pub(super) async fn prepare_local_same_format_provider_candidate(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
eligible: &EligibleLocalExecutionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<PreparedSameFormatProviderCandidate> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let candidate = &eligible.candidate;
|
||||
let transport = Arc::clone(&eligible.transport);
|
||||
let behavior = classify_same_format_provider_request_behavior(&transport, spec_metadata);
|
||||
|
||||
if !same_format_provider_transport_supported(
|
||||
&behavior,
|
||||
&transport,
|
||||
spec.family,
|
||||
spec_metadata.api_format,
|
||||
) {
|
||||
let skip_reason = same_format_provider_transport_unsupported_reason(
|
||||
&behavior,
|
||||
&transport,
|
||||
spec.family,
|
||||
spec_metadata.api_format,
|
||||
)
|
||||
.unwrap_or("transport_unsupported");
|
||||
super::super::payload::mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let vertex_query_auth = if behavior.is_vertex {
|
||||
resolve_local_vertex_api_key_query_auth(&transport)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let should_try_oauth_auth =
|
||||
should_try_same_format_provider_oauth_auth(&behavior, &transport, spec.family);
|
||||
let oauth_auth = if should_try_oauth_auth {
|
||||
resolve_candidate_oauth_auth(
|
||||
planner_state,
|
||||
&transport,
|
||||
OauthPreparationContext {
|
||||
trace_id,
|
||||
api_format: spec_metadata.api_format,
|
||||
operation: "same_format_provider_prepare",
|
||||
},
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let kiro_auth = match oauth_auth.as_ref() {
|
||||
Some(LocalResolvedOAuthRequestAuth::Kiro(auth)) => Some(auth.clone()),
|
||||
_ => None,
|
||||
};
|
||||
let auth = if let Some(kiro_auth) = kiro_auth.as_ref() {
|
||||
Some((kiro_auth.name.to_string(), kiro_auth.value.clone()))
|
||||
} else if let Some(LocalResolvedOAuthRequestAuth::Header { name, value }) = oauth_auth.as_ref()
|
||||
{
|
||||
Some((name.clone(), value.clone()))
|
||||
} else {
|
||||
resolve_same_format_provider_direct_auth(&behavior, &transport, spec.family)
|
||||
};
|
||||
let (auth_header, auth_value) = match auth {
|
||||
Some((name, value)) => (Some(name), Some(value)),
|
||||
None if behavior.is_vertex && vertex_query_auth.is_some() => (None, None),
|
||||
None => {
|
||||
super::super::payload::mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
if behavior.is_vertex && vertex_query_auth.is_none() {
|
||||
super::super::payload::mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let mapped_model = match resolve_candidate_mapped_model(candidate) {
|
||||
Ok(mapped_model) => mapped_model,
|
||||
Err(skip_reason) => {
|
||||
super::super::payload::mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(PreparedSameFormatProviderCandidate {
|
||||
transport,
|
||||
behavior,
|
||||
is_antigravity: behavior.is_antigravity,
|
||||
is_claude_code: behavior.is_claude_code,
|
||||
is_vertex: behavior.is_vertex,
|
||||
is_kiro: behavior.is_kiro,
|
||||
kiro_auth,
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
report_kind: behavior.report_kind,
|
||||
upstream_is_stream: behavior.upstream_is_stream,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::runtime_miss::{
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal,
|
||||
apply_local_runtime_candidate_terminal_reason, set_local_runtime_miss_diagnostic_reason,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::{
|
||||
build_stream_plan_from_requested_model_family, build_sync_plan_from_requested_model_family,
|
||||
local_same_format_provider_spec_metadata,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::{
|
||||
resolve_local_same_format_stream_spec as resolve_stream_spec,
|
||||
resolve_local_same_format_sync_spec as resolve_sync_spec,
|
||||
};
|
||||
|
||||
use super::{
|
||||
materialize_local_same_format_provider_candidate_attempts,
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate,
|
||||
resolve_local_same_format_provider_decision_input, AiStreamAttempt, AiSyncAttempt, AppState,
|
||||
GatewayControlDecision, GatewayError, LocalSameFormatProviderSpec,
|
||||
};
|
||||
|
||||
pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let requested_model_family = spec_metadata
|
||||
.requested_model_family
|
||||
.expect("same-format provider spec metadata should include requested-model family");
|
||||
let Some(input) = resolve_local_same_format_provider_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
extract_requested_model_from_request(parts, body_json, requested_model_family)
|
||||
.as_deref(),
|
||||
"decision_input_unavailable",
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
state,
|
||||
trace_id,
|
||||
candidate_count,
|
||||
);
|
||||
if candidate_count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let built = build_sync_plan_from_requested_model_family(
|
||||
requested_model_family,
|
||||
parts,
|
||||
body_json,
|
||||
payload,
|
||||
);
|
||||
|
||||
match built {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec_metadata.api_format,
|
||||
error = ?err,
|
||||
"gateway local same-format sync decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_sync_plans");
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let requested_model_family = spec_metadata
|
||||
.requested_model_family
|
||||
.expect("same-format provider spec metadata should include requested-model family");
|
||||
let Some(input) = resolve_local_same_format_provider_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
extract_requested_model_from_request(parts, body_json, requested_model_family)
|
||||
.as_deref(),
|
||||
"decision_input_unavailable",
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
state,
|
||||
trace_id,
|
||||
candidate_count,
|
||||
);
|
||||
if candidate_count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let built = build_stream_plan_from_requested_model_family(
|
||||
requested_model_family,
|
||||
parts,
|
||||
body_json,
|
||||
payload,
|
||||
);
|
||||
|
||||
match built {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec_metadata.api_format,
|
||||
error = ?err,
|
||||
"gateway local same-format stream decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_stream_plans");
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
mod body;
|
||||
mod url;
|
||||
|
||||
pub(super) use self::body::build_same_format_provider_request_body;
|
||||
pub(super) use self::url::build_same_format_upstream_url;
|
||||
@@ -0,0 +1,36 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use super::super::LocalSameFormatProviderSpec;
|
||||
use crate::ai_serving::transport::{
|
||||
build_same_format_provider_request_body as build_same_format_provider_request_body_impl,
|
||||
SameFormatProviderFamily, SameFormatProviderRequestBodyInput,
|
||||
};
|
||||
|
||||
pub(crate) fn build_same_format_provider_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
body_rules: Option<&Value>,
|
||||
upstream_is_stream: bool,
|
||||
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
|
||||
is_claude_code: bool,
|
||||
) -> Option<Value> {
|
||||
build_same_format_provider_request_body_impl(SameFormatProviderRequestBodyInput {
|
||||
body_json,
|
||||
mapped_model,
|
||||
family: same_format_provider_family(spec.family),
|
||||
body_rules,
|
||||
upstream_is_stream,
|
||||
kiro_auth_config: kiro_auth.map(|auth| &auth.auth_config),
|
||||
is_claude_code,
|
||||
})
|
||||
}
|
||||
|
||||
fn same_format_provider_family(
|
||||
family: super::super::LocalSameFormatProviderFamily,
|
||||
) -> SameFormatProviderFamily {
|
||||
match family {
|
||||
super::super::LocalSameFormatProviderFamily::Standard => SameFormatProviderFamily::Standard,
|
||||
super::super::LocalSameFormatProviderFamily::Gemini => SameFormatProviderFamily::Gemini,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
use crate::ai_serving::transport::{
|
||||
build_same_format_provider_upstream_url as build_same_format_provider_upstream_url_impl,
|
||||
SameFormatProviderUpstreamUrlParams,
|
||||
};
|
||||
use crate::ai_serving::GatewayProviderTransportSnapshot;
|
||||
|
||||
use super::super::LocalSameFormatProviderSpec;
|
||||
|
||||
pub(crate) fn build_same_format_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
mapped_model: &str,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
upstream_is_stream: bool,
|
||||
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
|
||||
) -> Option<String> {
|
||||
build_same_format_provider_upstream_url_impl(
|
||||
transport,
|
||||
SameFormatProviderUpstreamUrlParams {
|
||||
provider_api_format: spec.api_format,
|
||||
mapped_model,
|
||||
upstream_is_stream,
|
||||
request_query: parts.uri.query(),
|
||||
kiro_api_region: kiro_auth.map(|auth| auth.auth_config.effective_api_region()),
|
||||
},
|
||||
)
|
||||
}
|
||||
50
apps/aether-gateway/src/ai_serving/planner/plan_builders.rs
Normal file
50
apps/aether-gateway/src/ai_serving/planner/plan_builders.rs
Normal file
@@ -0,0 +1,50 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
pub(crate) use aether_ai_serving::{
|
||||
build_ai_execution_plan_from_decision, resolve_ai_passthrough_sync_request_body,
|
||||
take_ai_decision_plan_core, take_ai_non_empty_string as take_non_empty_string,
|
||||
take_ai_upstream_auth_pair, AiExecutionPlanFromDecisionParts,
|
||||
};
|
||||
|
||||
use crate::ai_serving::augment_sync_report_context as augment_sync_report_context_impl;
|
||||
pub(crate) use crate::ai_serving::{
|
||||
generic_decision_missing_exact_provider_request, AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
use crate::{AiExecutionDecision, GatewayError};
|
||||
|
||||
#[path = "standard/gemini/plan_builders.rs"]
|
||||
mod gemini_builders;
|
||||
#[path = "standard/openai/plan_builders.rs"]
|
||||
mod openai_builders;
|
||||
#[path = "passthrough/plan_builders.rs"]
|
||||
mod passthrough_builders;
|
||||
#[path = "standard/plan_builders.rs"]
|
||||
mod standard_builders;
|
||||
|
||||
pub(crate) use gemini_builders::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
};
|
||||
pub(crate) use openai_builders::{
|
||||
build_openai_chat_stream_plan_from_decision, build_openai_chat_sync_plan_from_decision,
|
||||
build_openai_responses_stream_plan_from_decision,
|
||||
build_openai_responses_sync_plan_from_decision,
|
||||
};
|
||||
pub(crate) use passthrough_builders::{
|
||||
build_passthrough_stream_plan_from_decision, build_passthrough_sync_plan_from_decision,
|
||||
};
|
||||
pub(crate) use standard_builders::{
|
||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
||||
};
|
||||
|
||||
pub(super) fn augment_sync_report_context(
|
||||
report_context: Option<serde_json::Value>,
|
||||
provider_request_headers: &BTreeMap<String, String>,
|
||||
provider_request_body: &serde_json::Value,
|
||||
) -> Result<Option<serde_json::Value>, GatewayError> {
|
||||
augment_sync_report_context_impl(
|
||||
report_context,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
1249
apps/aether-gateway/src/ai_serving/planner/pool_scheduler.rs
Normal file
1249
apps/aether-gateway/src/ai_serving/planner/pool_scheduler.rs
Normal file
File diff suppressed because it is too large
Load Diff
218
apps/aether-gateway/src/ai_serving/planner/report_context.rs
Normal file
218
apps/aether-gateway/src/ai_serving/planner/report_context.rs
Normal file
@@ -0,0 +1,218 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_ai_serving::{
|
||||
build_ai_execution_report_context,
|
||||
insert_provider_stream_event_api_format as insert_ai_provider_stream_event_api_format,
|
||||
provider_stream_event_api_format_for_provider_type as ai_provider_stream_event_api_format_for_provider_type,
|
||||
AiExecutionReportContextParts, AiRequestOrigin,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerRankingOutcome;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::ai_serving::{request_origin_from_headers, ExecutionRuntimeAuthContext, RequestOrigin};
|
||||
use crate::orchestration::ExecutionAttemptIdentity;
|
||||
|
||||
pub(crate) struct LocalExecutionReportContextParts<'a> {
|
||||
pub(crate) auth_context: &'a ExecutionRuntimeAuthContext,
|
||||
pub(crate) request_id: &'a str,
|
||||
pub(crate) candidate_id: &'a str,
|
||||
pub(crate) attempt_identity: ExecutionAttemptIdentity,
|
||||
pub(crate) model: &'a str,
|
||||
pub(crate) provider_name: &'a str,
|
||||
pub(crate) provider_id: &'a str,
|
||||
pub(crate) endpoint_id: &'a str,
|
||||
pub(crate) key_id: &'a str,
|
||||
pub(crate) key_name: Option<&'a str>,
|
||||
pub(crate) model_id: Option<&'a str>,
|
||||
pub(crate) global_model_id: Option<&'a str>,
|
||||
pub(crate) global_model_name: Option<&'a str>,
|
||||
pub(crate) provider_api_format: &'a str,
|
||||
pub(crate) client_api_format: &'a str,
|
||||
pub(crate) mapped_model: Option<&'a str>,
|
||||
pub(crate) candidate_group_id: Option<&'a str>,
|
||||
pub(crate) ranking: Option<&'a SchedulerRankingOutcome>,
|
||||
pub(crate) upstream_url: Option<&'a str>,
|
||||
pub(crate) header_rules: Option<&'a Value>,
|
||||
pub(crate) body_rules: Option<&'a Value>,
|
||||
pub(crate) provider_request_method: Option<Value>,
|
||||
pub(crate) provider_request_headers: Option<&'a BTreeMap<String, String>>,
|
||||
pub(crate) original_headers: &'a http::HeaderMap,
|
||||
pub(crate) request_origin: Option<RequestOrigin>,
|
||||
pub(crate) original_request_body_json: Option<&'a Value>,
|
||||
pub(crate) original_request_body_base64: Option<&'a str>,
|
||||
pub(crate) client_requested_stream: bool,
|
||||
pub(crate) upstream_is_stream: bool,
|
||||
pub(crate) has_envelope: bool,
|
||||
pub(crate) needs_conversion: bool,
|
||||
pub(crate) extra_fields: Map<String, Value>,
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_report_context(
|
||||
parts: LocalExecutionReportContextParts<'_>,
|
||||
) -> Value {
|
||||
let RequestOrigin {
|
||||
client_ip,
|
||||
user_agent,
|
||||
} = parts
|
||||
.request_origin
|
||||
.unwrap_or_else(|| request_origin_from_headers(parts.original_headers));
|
||||
let original_headers = crate::ai_serving::collect_control_headers(parts.original_headers);
|
||||
let original_request_body = crate::ai_serving::build_report_context_original_request_echo(
|
||||
parts.original_request_body_json,
|
||||
parts.original_request_body_base64,
|
||||
);
|
||||
|
||||
build_ai_execution_report_context(AiExecutionReportContextParts {
|
||||
auth_context: parts.auth_context,
|
||||
request_id: parts.request_id,
|
||||
candidate_id: parts.candidate_id,
|
||||
candidate_index: parts.attempt_identity.candidate_index,
|
||||
retry_index: parts.attempt_identity.retry_index,
|
||||
pool_key_index: parts.attempt_identity.pool_key_index,
|
||||
model: parts.model,
|
||||
provider_name: parts.provider_name,
|
||||
provider_id: parts.provider_id,
|
||||
endpoint_id: parts.endpoint_id,
|
||||
key_id: parts.key_id,
|
||||
key_name: parts.key_name,
|
||||
model_id: parts.model_id,
|
||||
global_model_id: parts.global_model_id,
|
||||
global_model_name: parts.global_model_name,
|
||||
provider_api_format: parts.provider_api_format,
|
||||
client_api_format: parts.client_api_format,
|
||||
mapped_model: parts.mapped_model,
|
||||
candidate_group_id: parts.candidate_group_id,
|
||||
ranking: parts.ranking,
|
||||
upstream_url: parts.upstream_url,
|
||||
header_rules: parts.header_rules,
|
||||
body_rules: parts.body_rules,
|
||||
provider_request_method: parts.provider_request_method,
|
||||
provider_request_headers: parts.provider_request_headers,
|
||||
original_headers: &original_headers,
|
||||
original_request_body,
|
||||
request_origin: AiRequestOrigin {
|
||||
client_ip,
|
||||
user_agent,
|
||||
},
|
||||
client_requested_stream: parts.client_requested_stream,
|
||||
upstream_is_stream: parts.upstream_is_stream,
|
||||
has_envelope: parts.has_envelope,
|
||||
needs_conversion: parts.needs_conversion,
|
||||
extra_fields: parts.extra_fields,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn provider_stream_event_api_format_for_provider_type(
|
||||
provider_type: &str,
|
||||
) -> Option<&'static str> {
|
||||
ai_provider_stream_event_api_format_for_provider_type(provider_type)
|
||||
}
|
||||
|
||||
pub(crate) fn insert_provider_stream_event_api_format(
|
||||
extra_fields: &mut Map<String, Value>,
|
||||
provider_type: &str,
|
||||
) {
|
||||
insert_ai_provider_stream_event_api_format(extra_fields, provider_type);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use super::{
|
||||
build_local_execution_report_context, provider_stream_event_api_format_for_provider_type,
|
||||
LocalExecutionReportContextParts,
|
||||
};
|
||||
use crate::ai_serving::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_serving::RequestOrigin;
|
||||
use crate::orchestration::ExecutionAttemptIdentity;
|
||||
|
||||
#[test]
|
||||
fn codex_provider_uses_openai_responses_stream_event_format() {
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type("codex"),
|
||||
Some("openai:responses")
|
||||
);
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type("CODEX"),
|
||||
Some("openai:responses")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ordinary_providers_do_not_override_stream_event_format() {
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type("openai"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type("anthropic"),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_execution_report_context_records_request_origin() {
|
||||
let auth_context = ExecutionRuntimeAuthContext {
|
||||
user_id: "user-1".to_string(),
|
||||
api_key_id: "api-key-1".to_string(),
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
balance_remaining: None,
|
||||
access_allowed: true,
|
||||
api_key_is_standalone: false,
|
||||
};
|
||||
let original_headers = http::HeaderMap::new();
|
||||
let provider_request_headers = BTreeMap::new();
|
||||
|
||||
let report_context =
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &auth_context,
|
||||
request_id: "trace-1",
|
||||
candidate_id: "candidate-1",
|
||||
attempt_identity: ExecutionAttemptIdentity::new(0, 0),
|
||||
model: "gpt-5",
|
||||
provider_name: "OpenAI",
|
||||
provider_id: "provider-1",
|
||||
endpoint_id: "endpoint-1",
|
||||
key_id: "key-1",
|
||||
key_name: None,
|
||||
model_id: None,
|
||||
global_model_id: None,
|
||||
global_model_name: None,
|
||||
provider_api_format: "openai:chat",
|
||||
client_api_format: "openai:chat",
|
||||
mapped_model: None,
|
||||
candidate_group_id: None,
|
||||
ranking: None,
|
||||
upstream_url: None,
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
provider_request_method: None,
|
||||
provider_request_headers: Some(&provider_request_headers),
|
||||
original_headers: &original_headers,
|
||||
request_origin: Some(RequestOrigin {
|
||||
client_ip: Some("203.0.113.8".to_string()),
|
||||
user_agent: Some("Claude-Code/1.0".to_string()),
|
||||
}),
|
||||
original_request_body_json: Some(&json!({"model": "gpt-5"})),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: false,
|
||||
upstream_is_stream: false,
|
||||
has_envelope: false,
|
||||
needs_conversion: false,
|
||||
extra_fields: Map::new(),
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
report_context["client_ip"],
|
||||
Value::String("203.0.113.8".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
report_context["user_agent"],
|
||||
Value::String("Claude-Code/1.0".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
154
apps/aether-gateway/src/ai_serving/planner/route.rs
Normal file
154
apps/aether-gateway/src/ai_serving/planner/route.rs
Normal file
@@ -0,0 +1,154 @@
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::ai_serving::{
|
||||
is_matching_stream_http_request as is_matching_stream_http_request_impl,
|
||||
resolve_execution_runtime_stream_plan_kind as resolve_execution_runtime_stream_plan_kind_impl,
|
||||
resolve_execution_runtime_sync_plan_kind as resolve_execution_runtime_sync_plan_kind_impl,
|
||||
supports_stream_scheduler_decision_kind as supports_stream_scheduler_decision_kind_impl,
|
||||
supports_sync_scheduler_decision_kind as supports_sync_scheduler_decision_kind_impl,
|
||||
};
|
||||
|
||||
pub(crate) fn resolve_execution_runtime_stream_plan_kind(
|
||||
parts: &http::request::Parts,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<&'static str> {
|
||||
resolve_execution_runtime_stream_plan_kind_impl(
|
||||
decision.route_class.as_deref(),
|
||||
decision.route_family.as_deref(),
|
||||
decision.route_kind.as_deref(),
|
||||
&parts.method,
|
||||
parts.uri.path(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_execution_runtime_sync_plan_kind(
|
||||
parts: &http::request::Parts,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<&'static str> {
|
||||
resolve_execution_runtime_sync_plan_kind_impl(
|
||||
decision.route_class.as_deref(),
|
||||
decision.route_family.as_deref(),
|
||||
decision.route_kind.as_deref(),
|
||||
&parts.method,
|
||||
parts.uri.path(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn is_matching_stream_request(
|
||||
plan_kind: &str,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> bool {
|
||||
is_matching_stream_http_request_impl(plan_kind, parts, body_json, body_base64)
|
||||
}
|
||||
|
||||
pub(crate) fn supports_sync_scheduler_decision_kind(plan_kind: &str) -> bool {
|
||||
supports_sync_scheduler_decision_kind_impl(plan_kind)
|
||||
}
|
||||
|
||||
pub(crate) fn supports_stream_scheduler_decision_kind(plan_kind: &str) -> bool {
|
||||
supports_stream_scheduler_decision_kind_impl(plan_kind)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use axum::http::{Method, Request};
|
||||
use base64::Engine as _;
|
||||
|
||||
use super::{
|
||||
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
|
||||
resolve_execution_runtime_sync_plan_kind, supports_stream_scheduler_decision_kind,
|
||||
supports_sync_scheduler_decision_kind,
|
||||
};
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
|
||||
fn sample_decision(route_family: &str, route_kind: &str) -> GatewayControlDecision {
|
||||
GatewayControlDecision {
|
||||
public_path: "/".to_string(),
|
||||
public_query_string: None,
|
||||
route_class: Some("ai_public".to_string()),
|
||||
route_family: Some(route_family.to_string()),
|
||||
route_kind: Some(route_kind.to_string()),
|
||||
auth_context: None,
|
||||
admin_principal: None,
|
||||
auth_endpoint_signature: None,
|
||||
execution_runtime_candidate: true,
|
||||
local_auth_rejection: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_openai_chat_plan_kinds_via_surface_crate() {
|
||||
let request = Request::builder()
|
||||
.method(Method::POST)
|
||||
.uri("/v1/chat/completions")
|
||||
.body(())
|
||||
.expect("request should build");
|
||||
let (parts, _) = request.into_parts();
|
||||
let decision = sample_decision("openai", "chat");
|
||||
|
||||
assert_eq!(
|
||||
resolve_execution_runtime_sync_plan_kind(&parts, &decision),
|
||||
Some("openai_chat_sync")
|
||||
);
|
||||
assert_eq!(
|
||||
resolve_execution_runtime_stream_plan_kind(&parts, &decision),
|
||||
Some("openai_chat_stream")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stream_matching_uses_surface_route_logic() {
|
||||
let request = Request::builder()
|
||||
.method(Method::POST)
|
||||
.uri("/v1/chat/completions")
|
||||
.body(())
|
||||
.expect("request should build");
|
||||
let (parts, _) = request.into_parts();
|
||||
|
||||
assert!(!is_matching_stream_request(
|
||||
"openai_chat_stream",
|
||||
&parts,
|
||||
&serde_json::json!({"stream": false}),
|
||||
None,
|
||||
));
|
||||
assert!(is_matching_stream_request(
|
||||
"openai_chat_stream",
|
||||
&parts,
|
||||
&serde_json::json!({"stream": true}),
|
||||
None,
|
||||
));
|
||||
assert!(supports_sync_scheduler_decision_kind("openai_chat_sync"));
|
||||
assert!(supports_stream_scheduler_decision_kind(
|
||||
"openai_chat_stream"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_stream_matching_parses_multipart_stream_flag() {
|
||||
let request = Request::builder()
|
||||
.method(Method::POST)
|
||||
.uri("/v1/images/edits")
|
||||
.header(
|
||||
http::header::CONTENT_TYPE,
|
||||
"multipart/form-data; boundary=image-stream-boundary",
|
||||
)
|
||||
.body(())
|
||||
.expect("request should build");
|
||||
let (parts, _) = request.into_parts();
|
||||
let body = concat!(
|
||||
"--image-stream-boundary\r\n",
|
||||
"Content-Disposition: form-data; name=\"stream\"\r\n\r\n",
|
||||
"true\r\n",
|
||||
"--image-stream-boundary--\r\n"
|
||||
);
|
||||
let body_base64 = base64::engine::general_purpose::STANDARD.encode(body.as_bytes());
|
||||
|
||||
assert!(is_matching_stream_request(
|
||||
"openai_image_stream",
|
||||
&parts,
|
||||
&serde_json::json!({}),
|
||||
Some(body_base64.as_str()),
|
||||
));
|
||||
}
|
||||
}
|
||||
256
apps/aether-gateway/src/ai_serving/planner/runtime_miss.rs
Normal file
256
apps/aether-gateway/src/ai_serving/planner/runtime_miss.rs
Normal file
@@ -0,0 +1,256 @@
|
||||
use aether_ai_serving::{
|
||||
apply_ai_runtime_candidate_evaluation_progress,
|
||||
apply_ai_runtime_candidate_evaluation_progress_preserving_candidate_signal,
|
||||
apply_ai_runtime_candidate_evaluation_progress_to_diagnostic,
|
||||
apply_ai_runtime_candidate_terminal_plan_reason_to_diagnostic,
|
||||
apply_ai_runtime_candidate_terminal_reason, build_ai_runtime_candidate_evaluation_diagnostic,
|
||||
build_ai_runtime_execution_exhausted_diagnostic, record_ai_runtime_candidate_skip_reason,
|
||||
record_ai_runtime_candidate_skip_reason_on_diagnostic,
|
||||
set_ai_runtime_candidate_evaluation_diagnostic, set_ai_runtime_execution_exhausted_diagnostic,
|
||||
set_ai_runtime_miss_diagnostic_reason, AiRuntimeMissDiagnosticFields,
|
||||
AiRuntimeMissDiagnosticPort,
|
||||
};
|
||||
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{AppState, LocalExecutionRuntimeMissDiagnostic};
|
||||
|
||||
struct GatewayRuntimeMissDiagnosticPort<'a> {
|
||||
state: Option<&'a AppState>,
|
||||
}
|
||||
|
||||
impl AiRuntimeMissDiagnosticFields for LocalExecutionRuntimeMissDiagnostic {
|
||||
fn set_reason(&mut self, reason: String) {
|
||||
self.reason = reason;
|
||||
}
|
||||
|
||||
fn set_candidate_count(&mut self, candidate_count: usize) {
|
||||
self.candidate_count = Some(candidate_count);
|
||||
}
|
||||
|
||||
fn candidate_count(&self) -> Option<usize> {
|
||||
self.candidate_count
|
||||
}
|
||||
|
||||
fn skipped_candidate_count(&self) -> Option<usize> {
|
||||
self.skipped_candidate_count
|
||||
}
|
||||
|
||||
fn skip_reason_count(&self, skip_reason: &str) -> usize {
|
||||
self.skip_reasons.get(skip_reason).copied().unwrap_or(0)
|
||||
}
|
||||
|
||||
fn skip_reason_len(&self) -> usize {
|
||||
self.skip_reasons.len()
|
||||
}
|
||||
|
||||
fn record_skip_reason(&mut self, skip_reason: &'static str) {
|
||||
*self
|
||||
.skip_reasons
|
||||
.entry(skip_reason.to_string())
|
||||
.or_insert(0) += 1;
|
||||
*self.skipped_candidate_count.get_or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
|
||||
impl AiRuntimeMissDiagnosticPort for GatewayRuntimeMissDiagnosticPort<'_> {
|
||||
type Decision = GatewayControlDecision;
|
||||
type Diagnostic = LocalExecutionRuntimeMissDiagnostic;
|
||||
|
||||
fn build_runtime_miss_diagnostic(
|
||||
&self,
|
||||
decision: &Self::Decision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
reason: &str,
|
||||
) -> Self::Diagnostic {
|
||||
LocalExecutionRuntimeMissDiagnostic {
|
||||
reason: reason.to_string(),
|
||||
route_family: decision.route_family.clone(),
|
||||
route_kind: decision.route_kind.clone(),
|
||||
public_path: Some(decision.public_path.clone()),
|
||||
plan_kind: Some(plan_kind.to_string()),
|
||||
requested_model: requested_model.map(ToOwned::to_owned),
|
||||
candidate_count: None,
|
||||
skipped_candidate_count: None,
|
||||
skip_reasons: std::collections::BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn set_candidate_count(&self, diagnostic: &mut Self::Diagnostic, candidate_count: usize) {
|
||||
AiRuntimeMissDiagnosticFields::set_candidate_count(diagnostic, candidate_count);
|
||||
}
|
||||
|
||||
fn apply_candidate_evaluation_progress(
|
||||
&self,
|
||||
diagnostic: &mut Self::Diagnostic,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
apply_ai_runtime_candidate_evaluation_progress_to_diagnostic(diagnostic, candidate_count);
|
||||
}
|
||||
|
||||
fn apply_candidate_terminal_plan_reason(
|
||||
&self,
|
||||
diagnostic: &mut Self::Diagnostic,
|
||||
no_plan_reason: &'static str,
|
||||
) {
|
||||
apply_ai_runtime_candidate_terminal_plan_reason_to_diagnostic(diagnostic, no_plan_reason);
|
||||
}
|
||||
|
||||
fn record_candidate_skip_reason(
|
||||
&self,
|
||||
diagnostic: &mut Self::Diagnostic,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
record_ai_runtime_candidate_skip_reason_on_diagnostic(diagnostic, skip_reason);
|
||||
}
|
||||
|
||||
fn set_runtime_miss_diagnostic(&self, trace_id: &str, diagnostic: Self::Diagnostic) {
|
||||
self.state
|
||||
.expect("runtime miss diagnostic setter requires gateway state")
|
||||
.set_local_execution_runtime_miss_diagnostic(trace_id, diagnostic);
|
||||
}
|
||||
|
||||
fn mutate_runtime_miss_diagnostic<F>(&self, trace_id: &str, apply: F)
|
||||
where
|
||||
F: FnOnce(&mut Self::Diagnostic) + Send,
|
||||
{
|
||||
self.state
|
||||
.expect("runtime miss diagnostic mutator requires gateway state")
|
||||
.mutate_local_execution_runtime_miss_diagnostic(trace_id, apply);
|
||||
}
|
||||
|
||||
fn runtime_miss_diagnostic_has_candidate_signal(&self, trace_id: &str) -> bool {
|
||||
self.state
|
||||
.expect("runtime miss diagnostic signal check requires gateway state")
|
||||
.local_execution_runtime_miss_diagnostic_has_candidate_signal(trace_id)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn set_local_runtime_miss_diagnostic_reason(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
reason: &str,
|
||||
) {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: Some(state) };
|
||||
set_ai_runtime_miss_diagnostic_reason(
|
||||
&port,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
reason,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_runtime_execution_exhausted_diagnostic(
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) -> LocalExecutionRuntimeMissDiagnostic {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: None };
|
||||
build_ai_runtime_execution_exhausted_diagnostic(
|
||||
&port,
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
candidate_count,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn set_local_runtime_execution_exhausted_diagnostic(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: Some(state) };
|
||||
set_ai_runtime_execution_exhausted_diagnostic(
|
||||
&port,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
candidate_count,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_runtime_candidate_evaluation_diagnostic(
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) -> LocalExecutionRuntimeMissDiagnostic {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: None };
|
||||
build_ai_runtime_candidate_evaluation_diagnostic(
|
||||
&port,
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
candidate_count,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn set_local_runtime_candidate_evaluation_diagnostic(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: Some(state) };
|
||||
set_ai_runtime_candidate_evaluation_diagnostic(
|
||||
&port,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
candidate_count,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn apply_local_runtime_candidate_evaluation_progress(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: Some(state) };
|
||||
apply_ai_runtime_candidate_evaluation_progress(&port, trace_id, candidate_count);
|
||||
}
|
||||
|
||||
pub(crate) fn apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: Some(state) };
|
||||
apply_ai_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
&port,
|
||||
trace_id,
|
||||
candidate_count,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn apply_local_runtime_candidate_terminal_reason(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
no_plan_reason: &'static str,
|
||||
) {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: Some(state) };
|
||||
apply_ai_runtime_candidate_terminal_reason(&port, trace_id, no_plan_reason);
|
||||
}
|
||||
|
||||
pub(crate) fn record_local_runtime_candidate_skip_reason(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: Some(state) };
|
||||
record_ai_runtime_candidate_skip_reason(&port, trace_id, skip_reason);
|
||||
}
|
||||
53
apps/aether-gateway/src/ai_serving/planner/spec_metadata.rs
Normal file
53
apps/aether-gateway/src/ai_serving/planner/spec_metadata.rs
Normal file
@@ -0,0 +1,53 @@
|
||||
use crate::ai_serving::planner::plan_builders::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
||||
AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
use crate::ai_serving::AiExecutionDecision;
|
||||
use crate::GatewayError;
|
||||
|
||||
pub(crate) use aether_ai_serving::{
|
||||
ai_gemini_files_spec_metadata as local_gemini_files_spec_metadata,
|
||||
ai_openai_image_spec_metadata as local_openai_image_spec_metadata,
|
||||
ai_openai_responses_spec_metadata as local_openai_responses_spec_metadata,
|
||||
ai_requested_model_family_for_same_format_provider as requested_model_family_for_same_format_provider,
|
||||
ai_requested_model_family_for_standard_source as requested_model_family_for_standard_source,
|
||||
ai_requested_model_family_for_video_create as requested_model_family_for_video_create,
|
||||
ai_same_format_provider_spec_metadata as local_same_format_provider_spec_metadata,
|
||||
ai_standard_spec_metadata as local_standard_spec_metadata,
|
||||
ai_video_create_spec_metadata as local_video_create_spec_metadata,
|
||||
AiExecutionSurfaceSpecMetadata as LocalExecutionSurfaceSpecMetadata,
|
||||
AiRequestedModelFamily as RequestedModelFamily,
|
||||
};
|
||||
|
||||
pub(crate) fn build_sync_plan_from_requested_model_family(
|
||||
family: RequestedModelFamily,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
match family {
|
||||
RequestedModelFamily::Standard => {
|
||||
build_standard_sync_plan_from_decision(parts, body_json, payload)
|
||||
}
|
||||
RequestedModelFamily::Gemini => {
|
||||
build_gemini_sync_plan_from_decision(parts, body_json, payload)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_stream_plan_from_requested_model_family(
|
||||
family: RequestedModelFamily,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
match family {
|
||||
RequestedModelFamily::Standard => {
|
||||
build_standard_stream_plan_from_decision(parts, body_json, payload, false)
|
||||
}
|
||||
RequestedModelFamily::Gemini => {
|
||||
build_gemini_stream_plan_from_decision(parts, body_json, payload)
|
||||
}
|
||||
}
|
||||
}
|
||||
253
apps/aether-gateway/src/ai_serving/planner/specialized/files.rs
Normal file
253
apps/aether-gateway/src/ai_serving/planner/specialized/files.rs
Normal file
@@ -0,0 +1,253 @@
|
||||
mod decision;
|
||||
mod request;
|
||||
mod support;
|
||||
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::plan_builders::{
|
||||
build_passthrough_stream_plan_from_decision, build_passthrough_sync_plan_from_decision,
|
||||
AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_gemini_files_spec_metadata;
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::ai_serving::{
|
||||
resolve_gemini_files_stream_spec as resolve_stream_spec,
|
||||
resolve_gemini_files_sync_spec as resolve_sync_spec, LocalGeminiFilesSpec,
|
||||
};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
use self::decision::maybe_build_local_gemini_files_decision_payload_for_candidate;
|
||||
use self::support::{
|
||||
materialize_local_gemini_files_candidate_attempts, resolve_local_gemini_files_decision_input,
|
||||
};
|
||||
|
||||
pub(crate) async fn build_local_gemini_files_sync_plan_and_reports_for_kind(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
build_local_sync_plan_and_reports(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
body_is_empty,
|
||||
trace_id,
|
||||
decision,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_gemini_files_stream_plan_and_reports_for_kind(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
build_local_stream_plan_and_reports(state, parts, trace_id, decision, spec).await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_sync_local_gemini_files_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let attempts =
|
||||
materialize_local_gemini_files_candidate_attempts(state, trace_id, &input).await?;
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
body_is_empty,
|
||||
trace_id,
|
||||
&input,
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_stream_local_gemini_files_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let attempts =
|
||||
materialize_local_gemini_files_candidate_attempts(state, trace_id, &input).await?;
|
||||
|
||||
let empty_body_json = serde_json::Value::Null;
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
&empty_body_json,
|
||||
None,
|
||||
true,
|
||||
trace_id,
|
||||
&input,
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
async fn build_local_sync_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
spec: LocalGeminiFilesSpec,
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
let spec_metadata = local_gemini_files_spec_metadata(spec);
|
||||
let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let attempts =
|
||||
materialize_local_gemini_files_candidate_attempts(state, trace_id, &input).await?;
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
body_is_empty,
|
||||
trace_id,
|
||||
&input,
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_passthrough_sync_plan_from_decision(parts, payload) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = spec_metadata.decision_kind,
|
||||
error = ?err,
|
||||
"gateway local gemini files sync decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
async fn build_local_stream_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
spec: LocalGeminiFilesSpec,
|
||||
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
|
||||
let spec_metadata = local_gemini_files_spec_metadata(spec);
|
||||
let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let attempts =
|
||||
materialize_local_gemini_files_candidate_attempts(state, trace_id, &input).await?;
|
||||
|
||||
let mut plans = Vec::new();
|
||||
let empty_body_json = serde_json::Value::Null;
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
&empty_body_json,
|
||||
None,
|
||||
true,
|
||||
trace_id,
|
||||
&input,
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_passthrough_stream_plan_from_decision(parts, payload) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = spec_metadata.decision_kind,
|
||||
error = ?err,
|
||||
"gateway local gemini files stream decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
use serde_json::json;
|
||||
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, LocalExecutionReportContextParts,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_gemini_files_spec_metadata;
|
||||
use crate::ai_serving::planner::{
|
||||
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::ai_serving::{ai_local_execution_contract_for_formats, PlannerAppState};
|
||||
use crate::{AiExecutionDecision, AppState};
|
||||
|
||||
use super::request::resolve_local_gemini_files_candidate_payload_parts;
|
||||
use super::support::{
|
||||
LocalGeminiFilesCandidateAttempt, LocalGeminiFilesDecisionInput, GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
};
|
||||
use super::LocalGeminiFilesSpec;
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
trace_id: &str,
|
||||
input: &LocalGeminiFilesDecisionInput,
|
||||
attempt: LocalGeminiFilesCandidateAttempt,
|
||||
spec: LocalGeminiFilesSpec,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
let spec_metadata = local_gemini_files_spec_metadata(spec);
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let attempt_identity = attempt.attempt_identity();
|
||||
let resolved = resolve_local_gemini_files_candidate_payload_parts(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
body_is_empty,
|
||||
trace_id,
|
||||
input,
|
||||
&attempt,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
let LocalGeminiFilesCandidateAttempt {
|
||||
eligible,
|
||||
candidate_id,
|
||||
..
|
||||
} = attempt;
|
||||
let candidate = eligible.candidate;
|
||||
let transport = resolved.transport;
|
||||
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
);
|
||||
let proxy = planner_state
|
||||
.app()
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
|
||||
.await;
|
||||
let tls_profile = resolve_transport_tls_profile(&transport);
|
||||
let mut extra_fields = serde_json::Map::new();
|
||||
if let Some(proxy_value) = build_request_trace_proxy_value(Some(&transport), proxy.as_ref()) {
|
||||
extra_fields.insert("proxy".to_string(), proxy_value);
|
||||
}
|
||||
extra_fields.insert("file_key_id".to_string(), json!(candidate.key_id));
|
||||
extra_fields.insert("file_name".to_string(), json!(resolved.file_name));
|
||||
let report_context = build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &input.auth_context,
|
||||
request_id: trace_id,
|
||||
candidate_id: &candidate_id,
|
||||
attempt_identity,
|
||||
model: "gemini-files",
|
||||
provider_name: &transport.provider.name,
|
||||
provider_id: &candidate.provider_id,
|
||||
endpoint_id: &candidate.endpoint_id,
|
||||
key_id: &candidate.key_id,
|
||||
key_name: None,
|
||||
model_id: Some(&candidate.model_id),
|
||||
global_model_id: Some(&candidate.global_model_id),
|
||||
global_model_name: Some(&candidate.global_model_name),
|
||||
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
mapped_model: None,
|
||||
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
|
||||
ranking: eligible.ranking.as_ref(),
|
||||
upstream_url: None,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
body_rules: transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: None,
|
||||
provider_request_headers: None,
|
||||
original_headers: &parts.headers,
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: resolved.provider_request_body_base64.as_deref(),
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
upstream_is_stream: spec_metadata.require_streaming,
|
||||
has_envelope: false,
|
||||
needs_conversion: false,
|
||||
extra_fields,
|
||||
});
|
||||
let super::request::LocalGeminiFilesCandidatePayloadParts {
|
||||
transport: _,
|
||||
auth_header,
|
||||
auth_value,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
provider_request_body_base64,
|
||||
upstream_url,
|
||||
file_name: _,
|
||||
} = resolved;
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT.to_string(),
|
||||
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT.to_string(),
|
||||
model_name: "gemini-files".to_string(),
|
||||
mapped_model: candidate.selected_provider_model_name.clone(),
|
||||
prompt_cache_key: None,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
provider_request_body_base64,
|
||||
content_type: parts
|
||||
.headers
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream: spec_metadata.require_streaming,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
}
|
||||
@@ -0,0 +1,191 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::ai_serving::planner::spec_metadata::local_gemini_files_spec_metadata;
|
||||
use crate::ai_serving::transport::{
|
||||
build_gemini_files_headers, build_gemini_files_request_body, build_gemini_files_upstream_url,
|
||||
gemini_files_transport_unsupported_reason, resolve_gemini_files_auth, GeminiFilesHeadersInput,
|
||||
GeminiFilesRequestBodyError,
|
||||
};
|
||||
use crate::ai_serving::GEMINI_FILES_UPLOAD_PLAN_KIND;
|
||||
use crate::ai_serving::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
|
||||
use crate::AppState;
|
||||
|
||||
use super::support::{
|
||||
mark_skipped_local_gemini_files_candidate,
|
||||
mark_skipped_local_gemini_files_candidate_with_failure_diagnostic,
|
||||
LocalGeminiFilesCandidateAttempt, LocalGeminiFilesDecisionInput,
|
||||
GEMINI_FILES_CANDIDATE_API_FORMAT, GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
};
|
||||
use super::LocalGeminiFilesSpec;
|
||||
|
||||
pub(super) struct LocalGeminiFilesCandidatePayloadParts {
|
||||
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
pub(super) auth_header: String,
|
||||
pub(super) auth_value: String,
|
||||
pub(super) provider_request_headers: BTreeMap<String, String>,
|
||||
pub(super) provider_request_body: Option<serde_json::Value>,
|
||||
pub(super) provider_request_body_base64: Option<String>,
|
||||
pub(super) upstream_url: String,
|
||||
pub(super) file_name: String,
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
trace_id: &str,
|
||||
input: &LocalGeminiFilesDecisionInput,
|
||||
attempt: &LocalGeminiFilesCandidateAttempt,
|
||||
spec: LocalGeminiFilesSpec,
|
||||
) -> Option<LocalGeminiFilesCandidatePayloadParts> {
|
||||
let spec_metadata = local_gemini_files_spec_metadata(spec);
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let transport = &attempt.eligible.transport;
|
||||
|
||||
if let Some(skip_reason) =
|
||||
gemini_files_transport_unsupported_reason(transport, GEMINI_FILES_CANDIDATE_API_FORMAT)
|
||||
{
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let Some((auth_header, auth_value)) = resolve_gemini_files_auth(transport) else {
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(upstream_url) =
|
||||
build_gemini_files_upstream_url(transport, parts.uri.path(), parts.uri.query())
|
||||
else {
|
||||
mark_skipped_local_gemini_files_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"upstream_url_missing",
|
||||
CandidateFailureDiagnostic::upstream_url_missing(
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
GEMINI_FILES_CANDIDATE_API_FORMAT,
|
||||
"gemini_files_passthrough_url",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let body_parts = match build_gemini_files_request_body(
|
||||
body_json,
|
||||
body_base64,
|
||||
body_is_empty,
|
||||
spec_metadata.decision_kind == GEMINI_FILES_UPLOAD_PLAN_KIND,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
) {
|
||||
Ok(parts) => parts,
|
||||
Err(GeminiFilesRequestBodyError::BodyRulesUnsupportedForBinaryUpload) => {
|
||||
mark_skipped_local_gemini_files_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_body_rules_unsupported_for_binary_upload",
|
||||
CandidateFailureDiagnostic::body_rules_unsupported_for_binary_upload(
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
GEMINI_FILES_CANDIDATE_API_FORMAT,
|
||||
"gemini_files_binary_upload",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
Err(GeminiFilesRequestBodyError::BodyRulesApplyFailed) => {
|
||||
mark_skipped_local_gemini_files_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_body_rules_apply_failed",
|
||||
CandidateFailureDiagnostic::body_rules_apply_failed(
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
GEMINI_FILES_CANDIDATE_API_FORMAT,
|
||||
"gemini_files_body_rules",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
let Some(provider_request_headers) = build_gemini_files_headers(GeminiFilesHeadersInput {
|
||||
headers: &parts.headers,
|
||||
auth_header: &auth_header,
|
||||
auth_value: &auth_value,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
provider_request_body: body_parts.provider_request_body.as_ref(),
|
||||
provider_request_body_base64: body_parts.provider_request_body_base64.as_deref(),
|
||||
original_request_body_json: body_json,
|
||||
original_body_is_empty: body_is_empty,
|
||||
}) else {
|
||||
mark_skipped_local_gemini_files_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
CandidateFailureDiagnostic::header_rules_apply_failed(
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
GEMINI_FILES_CANDIDATE_API_FORMAT,
|
||||
"gemini_files_header_rules",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let file_name = parts
|
||||
.uri
|
||||
.path()
|
||||
.trim_start_matches("/v1beta/")
|
||||
.trim()
|
||||
.to_string();
|
||||
|
||||
Some(LocalGeminiFilesCandidatePayloadParts {
|
||||
transport: Arc::clone(transport),
|
||||
auth_header,
|
||||
auth_value,
|
||||
provider_request_headers,
|
||||
provider_request_body: body_parts.provider_request_body,
|
||||
provider_request_body_base64: body_parts.provider_request_body_base64,
|
||||
upstream_url,
|
||||
file_name,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::json;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
mark_skipped_local_execution_candidate,
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic,
|
||||
materialize_local_execution_candidates_with_serving, LocalCandidateResolutionMode,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_metadata::{
|
||||
build_local_execution_candidate_metadata,
|
||||
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
|
||||
};
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
build_local_authenticated_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
resolve_local_decision_execution_runtime_auth_context, CandidateFailureDiagnostic,
|
||||
ExecutionRuntimeAuthContext, GatewayControlDecision, PlannerAppState,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
pub(super) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalGeminiFilesCandidateAttempt;
|
||||
pub(super) use crate::ai_serving::planner::decision_input::LocalAuthenticatedDecisionInput as LocalGeminiFilesDecisionInput;
|
||||
|
||||
pub(super) const GEMINI_FILES_CANDIDATE_API_FORMAT: &str = "gemini:files";
|
||||
pub(super) const GEMINI_FILES_CLIENT_API_FORMAT: &str = "gemini:files";
|
||||
pub(super) const GEMINI_FILES_REQUIRED_CAPABILITY: &str = "gemini_files";
|
||||
|
||||
pub(super) async fn resolve_local_gemini_files_decision_input(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<LocalGeminiFilesDecisionInput> {
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let explicit_required_capabilities = json!({ "gemini_files": true });
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
auth_context,
|
||||
None,
|
||||
Some(&explicit_required_capabilities),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => return None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local gemini files decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(build_local_authenticated_decision_input(resolved_input))
|
||||
}
|
||||
|
||||
pub(super) async fn materialize_local_gemini_files_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalGeminiFilesDecisionInput,
|
||||
) -> Result<Vec<LocalGeminiFilesCandidateAttempt>, GatewayError> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::GeminiFilesDecision,
|
||||
);
|
||||
let candidates = planner_state
|
||||
.list_selectable_candidates_for_required_capability_without_requested_model(
|
||||
GEMINI_FILES_CANDIDATE_API_FORMAT,
|
||||
GEMINI_FILES_REQUIRED_CAPABILITY,
|
||||
false,
|
||||
Some(&input.auth_snapshot),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
let outcome = materialize_local_execution_candidates_with_serving(
|
||||
planner_state,
|
||||
trace_id,
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
None,
|
||||
Some(&input.auth_snapshot),
|
||||
input.required_capabilities.as_ref(),
|
||||
None,
|
||||
persistence_policy,
|
||||
candidates,
|
||||
Vec::new(),
|
||||
LocalCandidateResolutionMode::WithoutTransportPairGate,
|
||||
|eligible| {
|
||||
let mut extra_fields = serde_json::Map::new();
|
||||
extra_fields.insert(
|
||||
"candidate_api_format".to_string(),
|
||||
json!(GEMINI_FILES_CANDIDATE_API_FORMAT),
|
||||
);
|
||||
Some(build_local_execution_candidate_metadata(
|
||||
LocalExecutionCandidateMetadataParts {
|
||||
eligible,
|
||||
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
extra_fields,
|
||||
},
|
||||
))
|
||||
},
|
||||
|mut skipped_candidate| {
|
||||
let mut extra_fields = serde_json::Map::new();
|
||||
extra_fields.insert(
|
||||
"candidate_api_format".to_string(),
|
||||
json!(GEMINI_FILES_CANDIDATE_API_FORMAT),
|
||||
);
|
||||
skipped_candidate.extra_data =
|
||||
Some(build_local_execution_candidate_metadata_for_candidate(
|
||||
&skipped_candidate.candidate,
|
||||
skipped_candidate.transport_ref(),
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
extra_fields,
|
||||
));
|
||||
skipped_candidate
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
Ok(outcome.attempts)
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_gemini_files_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalGeminiFilesDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::GeminiFilesDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_gemini_files_candidate_with_failure_diagnostic(
|
||||
state: &AppState,
|
||||
input: &LocalGeminiFilesDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
diagnostic: CandidateFailureDiagnostic,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::GeminiFilesDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
diagnostic,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
334
apps/aether-gateway/src/ai_serving/planner/specialized/image.rs
Normal file
334
apps/aether-gateway/src/ai_serving/planner/specialized/image.rs
Normal file
@@ -0,0 +1,334 @@
|
||||
mod decision;
|
||||
mod request;
|
||||
mod support;
|
||||
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::plan_builders::{
|
||||
build_passthrough_sync_plan_from_decision, build_standard_stream_plan_from_decision,
|
||||
AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::ai_serving::{
|
||||
resolve_local_image_stream_spec as resolve_stream_spec,
|
||||
resolve_local_image_sync_spec as resolve_sync_spec,
|
||||
};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
use self::decision::maybe_build_local_openai_image_decision_payload_for_candidate;
|
||||
use self::support::{
|
||||
list_local_openai_image_candidate_attempts, resolve_local_openai_image_decision_input,
|
||||
};
|
||||
|
||||
pub(super) use crate::ai_serving::LocalOpenAiImageSpec;
|
||||
|
||||
pub(crate) async fn build_local_image_sync_plan_and_reports_for_kind(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
build_local_sync_plan_and_reports(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_image_stream_plan_and_reports_for_kind(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
build_local_stream_plan_and_reports(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_sync_local_image_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let spec_metadata = local_openai_image_spec_metadata(spec);
|
||||
|
||||
let Some(input) = resolve_local_openai_image_decision_input(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let Some(attempts) = list_local_openai_image_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
&input,
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_stream_local_image_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let spec_metadata = local_openai_image_spec_metadata(spec);
|
||||
|
||||
let Some(input) = resolve_local_openai_image_decision_input(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let Some(attempts) = list_local_openai_image_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
&input,
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
async fn build_local_sync_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
spec: LocalOpenAiImageSpec,
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
let spec_metadata = local_openai_image_spec_metadata(spec);
|
||||
let Some(input) = resolve_local_openai_image_decision_input(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let Some(attempts) = list_local_openai_image_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
&input,
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_passthrough_sync_plan_from_decision(parts, payload) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = spec_metadata.decision_kind,
|
||||
error = ?err,
|
||||
"gateway local openai image sync decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
async fn build_local_stream_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
spec: LocalOpenAiImageSpec,
|
||||
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
|
||||
let spec_metadata = local_openai_image_spec_metadata(spec);
|
||||
let Some(input) = resolve_local_openai_image_decision_input(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let Some(attempts) = list_local_openai_image_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
&input,
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_standard_stream_plan_from_decision(parts, body_json, payload, false) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = spec_metadata.decision_kind,
|
||||
error = ?err,
|
||||
"gateway local openai image stream decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, LocalExecutionReportContextParts,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
|
||||
use crate::ai_serving::planner::{
|
||||
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::ai_serving::{ai_local_execution_contract_for_formats, PlannerAppState};
|
||||
use crate::{AiExecutionDecision, AppState};
|
||||
|
||||
use super::request::resolve_local_openai_image_candidate_payload_parts;
|
||||
use super::support::{LocalOpenAiImageCandidateAttempt, LocalOpenAiImageDecisionInput};
|
||||
use super::LocalOpenAiImageSpec;
|
||||
|
||||
pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
input: &LocalOpenAiImageDecisionInput,
|
||||
attempt: LocalOpenAiImageCandidateAttempt,
|
||||
spec: LocalOpenAiImageSpec,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
let spec_metadata = local_openai_image_spec_metadata(spec);
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let attempt_identity = attempt.attempt_identity();
|
||||
let resolved = resolve_local_openai_image_candidate_payload_parts(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
input,
|
||||
&attempt,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
let LocalOpenAiImageCandidateAttempt {
|
||||
eligible,
|
||||
candidate_id,
|
||||
..
|
||||
} = attempt;
|
||||
let candidate = eligible.candidate;
|
||||
let transport = resolved.transport;
|
||||
let (execution_strategy, conversion_mode) =
|
||||
ai_local_execution_contract_for_formats(spec_metadata.api_format, spec_metadata.api_format);
|
||||
let proxy = planner_state
|
||||
.app()
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
|
||||
.await;
|
||||
let tls_profile = resolve_transport_tls_profile(&transport);
|
||||
let mut extra_fields = serde_json::Map::new();
|
||||
if let Some(proxy_value) = build_request_trace_proxy_value(Some(&transport), proxy.as_ref()) {
|
||||
extra_fields.insert("proxy".to_string(), proxy_value);
|
||||
}
|
||||
extra_fields.insert("image_request".to_string(), resolved.input_summary.clone());
|
||||
let upstream_is_stream = resolved
|
||||
.provider_request_body
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(spec_metadata.require_streaming);
|
||||
let report_context = build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &input.auth_context,
|
||||
request_id: trace_id,
|
||||
candidate_id: &candidate_id,
|
||||
attempt_identity,
|
||||
model: &resolved.requested_model,
|
||||
provider_name: &transport.provider.name,
|
||||
provider_id: &candidate.provider_id,
|
||||
endpoint_id: &candidate.endpoint_id,
|
||||
key_id: &candidate.key_id,
|
||||
key_name: None,
|
||||
model_id: Some(&candidate.model_id),
|
||||
global_model_id: Some(&candidate.global_model_id),
|
||||
global_model_name: Some(&candidate.global_model_name),
|
||||
provider_api_format: spec_metadata.api_format,
|
||||
client_api_format: spec_metadata.api_format,
|
||||
mapped_model: Some(&resolved.mapped_model),
|
||||
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
|
||||
ranking: eligible.ranking.as_ref(),
|
||||
upstream_url: Some(&resolved.upstream_url),
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
body_rules: transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: Some(serde_json::Value::String(parts.method.to_string())),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: body_base64,
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
upstream_is_stream,
|
||||
has_envelope: false,
|
||||
needs_conversion: false,
|
||||
extra_fields,
|
||||
});
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url: resolved.upstream_url,
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(resolved.auth_header),
|
||||
auth_value: Some(resolved.auth_value),
|
||||
provider_api_format: spec_metadata.api_format.to_string(),
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: resolved.requested_model,
|
||||
mapped_model: resolved.mapped_model,
|
||||
prompt_cache_key: None,
|
||||
provider_request_headers: resolved.provider_request_headers,
|
||||
provider_request_body: Some(resolved.provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::planner::candidate_preparation::{
|
||||
prepare_header_authenticated_candidate, OauthPreparationContext,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
|
||||
use crate::ai_serving::transport::{
|
||||
build_openai_image_headers, build_openai_image_upstream_url,
|
||||
openai_image_transport_unsupported_reason, resolve_openai_image_auth,
|
||||
ProviderOpenAiImageHeadersInput,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
|
||||
build_openai_image_provider_request_body, default_model_for_openai_image_operation,
|
||||
normalize_openai_image_request, CandidateFailureDiagnostic, GatewayProviderTransportSnapshot,
|
||||
PlannerAppState,
|
||||
};
|
||||
use crate::AppState;
|
||||
|
||||
use super::support::{
|
||||
mark_skipped_local_openai_image_candidate,
|
||||
mark_skipped_local_openai_image_candidate_with_failure_diagnostic,
|
||||
LocalOpenAiImageCandidateAttempt, LocalOpenAiImageDecisionInput,
|
||||
};
|
||||
use super::LocalOpenAiImageSpec;
|
||||
|
||||
pub(super) use crate::ai_serving::resolve_requested_openai_image_model_for_request as resolve_requested_image_model_for_request;
|
||||
|
||||
pub(super) struct LocalOpenAiImageCandidatePayloadParts {
|
||||
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
pub(super) auth_header: String,
|
||||
pub(super) auth_value: String,
|
||||
pub(super) requested_model: String,
|
||||
pub(super) mapped_model: String,
|
||||
pub(super) provider_request_headers: BTreeMap<String, String>,
|
||||
pub(super) provider_request_body: Value,
|
||||
pub(super) upstream_url: String,
|
||||
pub(super) input_summary: Value,
|
||||
}
|
||||
|
||||
pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
input: &LocalOpenAiImageDecisionInput,
|
||||
attempt: &LocalOpenAiImageCandidateAttempt,
|
||||
spec: LocalOpenAiImageSpec,
|
||||
) -> Option<LocalOpenAiImageCandidatePayloadParts> {
|
||||
let spec_metadata = local_openai_image_spec_metadata(spec);
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let transport = &attempt.eligible.transport;
|
||||
|
||||
if let Some(skip_reason) =
|
||||
openai_image_transport_unsupported_reason(transport, spec_metadata.api_format)
|
||||
{
|
||||
mark_skipped_local_openai_image_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.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: spec_metadata.api_format,
|
||||
operation: "openai_image_candidate_request",
|
||||
},
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(prepared) => prepared,
|
||||
Err(skip_reason) => {
|
||||
mark_skipped_local_openai_image_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let auth_header = prepared_candidate.auth_header;
|
||||
let auth_value = prepared_candidate.auth_value;
|
||||
|
||||
let Some(normalized_request) = normalize_openai_image_request(parts, body_json, body_base64)
|
||||
else {
|
||||
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"provider_request_body_missing",
|
||||
CandidateFailureDiagnostic::provider_request_body_missing(
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.api_format,
|
||||
"openai_image_request_normalize",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let upstream_url = build_openai_image_upstream_url(transport, parts.uri.query());
|
||||
let mut provider_request_body = build_openai_image_provider_request_body(&normalized_request);
|
||||
apply_codex_openai_responses_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
transport.provider.provider_type.as_str(),
|
||||
spec_metadata.api_format,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
Some(candidate.key_id.as_str()),
|
||||
);
|
||||
|
||||
let Some(mut provider_request_headers) =
|
||||
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
|
||||
headers: &parts.headers,
|
||||
auth_header: &auth_header,
|
||||
auth_value: &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_image_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
CandidateFailureDiagnostic::header_rules_apply_failed(
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.api_format,
|
||||
"openai_image_header_rules",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
apply_codex_openai_responses_special_headers(
|
||||
&mut provider_request_headers,
|
||||
&provider_request_body,
|
||||
&parts.headers,
|
||||
transport.provider.provider_type.as_str(),
|
||||
spec_metadata.api_format,
|
||||
Some(trace_id),
|
||||
transport.key.decrypted_auth_config.as_deref(),
|
||||
);
|
||||
let requested_model = normalized_request
|
||||
.requested_model
|
||||
.clone()
|
||||
.unwrap_or_else(|| {
|
||||
default_model_for_openai_image_operation(normalized_request.operation).to_string()
|
||||
});
|
||||
let mapped_model = provider_request_body
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
|
||||
Some(LocalOpenAiImageCandidatePayloadParts {
|
||||
transport: Arc::clone(transport),
|
||||
auth_header,
|
||||
auth_value,
|
||||
requested_model,
|
||||
mapped_model,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
upstream_url,
|
||||
input_summary: normalized_request.summary_json,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,241 @@
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
mark_skipped_local_execution_candidate,
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic,
|
||||
materialize_local_execution_candidates_with_serving, LocalCandidateResolutionMode,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_metadata::{
|
||||
build_local_execution_candidate_metadata,
|
||||
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
|
||||
use crate::ai_serving::{
|
||||
extract_pool_sticky_session_token, resolve_local_decision_execution_runtime_auth_context,
|
||||
CandidateFailureDiagnostic, ExecutionRuntimeAuthContext, GatewayControlDecision,
|
||||
PlannerAppState,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::AppState;
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
pub(super) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalOpenAiImageCandidateAttempt;
|
||||
pub(super) use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput as LocalOpenAiImageDecisionInput;
|
||||
|
||||
use super::request::resolve_requested_image_model_for_request;
|
||||
|
||||
pub(super) async fn resolve_local_openai_image_decision_input(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<LocalOpenAiImageDecisionInput> {
|
||||
let Some(auth_context) = resolve_local_openai_image_auth_context(decision) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let requested_model = resolve_requested_image_model_for_request(parts, body_json, body_base64)?;
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
auth_context,
|
||||
Some(requested_model.as_str()),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => return None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai image decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(build_local_requested_model_decision_input(
|
||||
resolved_input,
|
||||
requested_model,
|
||||
))
|
||||
}
|
||||
|
||||
fn resolve_local_openai_image_auth_context(
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<ExecutionRuntimeAuthContext> {
|
||||
resolve_local_decision_execution_runtime_auth_context(decision)
|
||||
}
|
||||
|
||||
pub(super) async fn list_local_openai_image_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalOpenAiImageDecisionInput,
|
||||
body_json: &serde_json::Value,
|
||||
api_format: &str,
|
||||
decision_kind: &str,
|
||||
) -> Option<Vec<LocalOpenAiImageCandidateAttempt>> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let (candidates, preselection_skipped) = match planner_state
|
||||
.list_selectable_candidates_with_skip_reasons(
|
||||
api_format,
|
||||
&input.requested_model,
|
||||
false,
|
||||
input.required_capabilities.as_ref(),
|
||||
Some(&input.auth_snapshot),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(candidates) => candidates,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind,
|
||||
error = ?err,
|
||||
"gateway local openai image decision scheduler selection failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(
|
||||
materialize_local_openai_image_candidate_attempts(
|
||||
planner_state,
|
||||
trace_id,
|
||||
input,
|
||||
body_json,
|
||||
candidates,
|
||||
preselection_skipped
|
||||
.into_iter()
|
||||
.map(|item| SkippedLocalExecutionCandidate {
|
||||
candidate: item.candidate,
|
||||
skip_reason: item.skip_reason,
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
})
|
||||
.collect(),
|
||||
api_format,
|
||||
)
|
||||
.await,
|
||||
)
|
||||
}
|
||||
|
||||
async fn materialize_local_openai_image_candidate_attempts(
|
||||
state: PlannerAppState<'_>,
|
||||
trace_id: &str,
|
||||
input: &LocalOpenAiImageDecisionInput,
|
||||
body_json: &serde_json::Value,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
preselection_skipped: Vec<SkippedLocalExecutionCandidate>,
|
||||
api_format: &str,
|
||||
) -> Vec<LocalOpenAiImageCandidateAttempt> {
|
||||
let sticky_session_token = extract_pool_sticky_session_token(body_json);
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::ImageDecision,
|
||||
);
|
||||
let outcome = materialize_local_execution_candidates_with_serving(
|
||||
state,
|
||||
trace_id,
|
||||
api_format,
|
||||
Some(&input.requested_model),
|
||||
Some(&input.auth_snapshot),
|
||||
input.required_capabilities.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
persistence_policy,
|
||||
candidates,
|
||||
preselection_skipped,
|
||||
LocalCandidateResolutionMode::Standard,
|
||||
|eligible| {
|
||||
Some(build_local_execution_candidate_metadata(
|
||||
LocalExecutionCandidateMetadataParts {
|
||||
eligible,
|
||||
provider_api_format: api_format,
|
||||
client_api_format: api_format,
|
||||
extra_fields: serde_json::Map::new(),
|
||||
},
|
||||
))
|
||||
},
|
||||
|mut skipped_candidate| {
|
||||
skipped_candidate.extra_data =
|
||||
Some(build_local_execution_candidate_metadata_for_candidate(
|
||||
&skipped_candidate.candidate,
|
||||
skipped_candidate.transport_ref(),
|
||||
api_format,
|
||||
api_format,
|
||||
serde_json::Map::new(),
|
||||
));
|
||||
skipped_candidate
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
outcome.attempts
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_openai_image_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiImageDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::ImageDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiImageDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
diagnostic: CandidateFailureDiagnostic,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::ImageDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
diagnostic,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
//! Non-matrix AI surfaces such as files and video.
|
||||
|
||||
mod files;
|
||||
mod image;
|
||||
mod video;
|
||||
|
||||
pub(crate) use self::files::{
|
||||
build_local_gemini_files_stream_plan_and_reports_for_kind,
|
||||
build_local_gemini_files_sync_plan_and_reports_for_kind,
|
||||
maybe_build_stream_local_gemini_files_decision_payload,
|
||||
maybe_build_sync_local_gemini_files_decision_payload,
|
||||
};
|
||||
pub(crate) use self::image::{
|
||||
build_local_image_stream_plan_and_reports_for_kind,
|
||||
build_local_image_sync_plan_and_reports_for_kind,
|
||||
maybe_build_stream_local_image_decision_payload, maybe_build_sync_local_image_decision_payload,
|
||||
};
|
||||
pub(crate) use self::video::{
|
||||
build_local_video_sync_plan_and_reports_for_kind, maybe_build_sync_local_video_decision_payload,
|
||||
};
|
||||
140
apps/aether-gateway/src/ai_serving/planner/specialized/video.rs
Normal file
140
apps/aether-gateway/src/ai_serving/planner/specialized/video.rs
Normal file
@@ -0,0 +1,140 @@
|
||||
mod decision;
|
||||
mod request;
|
||||
mod support;
|
||||
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::plan_builders::{
|
||||
build_passthrough_sync_plan_from_decision, AiSyncAttempt,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_video_create_spec_metadata;
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::ai_serving::{
|
||||
resolve_local_video_sync_spec as resolve_sync_spec, LocalVideoCreateFamily,
|
||||
LocalVideoCreateSpec,
|
||||
};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
use self::decision::maybe_build_local_video_create_decision_payload_for_candidate;
|
||||
use self::support::{
|
||||
list_local_video_create_candidate_attempts, resolve_local_video_create_decision_input,
|
||||
};
|
||||
|
||||
pub(crate) async fn build_local_video_sync_plan_and_reports_for_kind(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
build_local_sync_plan_and_reports(state, parts, body_json, trace_id, decision, spec).await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_sync_local_video_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let spec_metadata = local_video_create_spec_metadata(spec);
|
||||
|
||||
let Some(input) = resolve_local_video_create_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let Some(attempts) = list_local_video_create_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate(
|
||||
state, parts, body_json, trace_id, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
async fn build_local_sync_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
spec: LocalVideoCreateSpec,
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
let spec_metadata = local_video_create_spec_metadata(spec);
|
||||
let Some(input) = resolve_local_video_create_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let Some(attempts) = list_local_video_create_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate(
|
||||
state, parts, body_json, trace_id, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_passthrough_sync_plan_from_decision(parts, payload) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = spec_metadata.decision_kind,
|
||||
error = ?err,
|
||||
"gateway local video sync decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
@@ -0,0 +1,138 @@
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, LocalExecutionReportContextParts,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_video_create_spec_metadata;
|
||||
use crate::ai_serving::planner::{
|
||||
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::ai_serving::{ai_local_execution_contract_for_formats, PlannerAppState};
|
||||
use crate::{AiExecutionDecision, AppState};
|
||||
|
||||
use super::request::resolve_local_video_create_candidate_payload_parts;
|
||||
use super::support::{LocalVideoCreateCandidateAttempt, LocalVideoCreateDecisionInput};
|
||||
use super::LocalVideoCreateSpec;
|
||||
|
||||
pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
trace_id: &str,
|
||||
input: &LocalVideoCreateDecisionInput,
|
||||
attempt: LocalVideoCreateCandidateAttempt,
|
||||
spec: LocalVideoCreateSpec,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
let spec_metadata = local_video_create_spec_metadata(spec);
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let attempt_identity = attempt.attempt_identity();
|
||||
let resolved = resolve_local_video_create_candidate_payload_parts(
|
||||
state, parts, body_json, trace_id, input, &attempt, spec,
|
||||
)
|
||||
.await?;
|
||||
let LocalVideoCreateCandidateAttempt {
|
||||
eligible,
|
||||
candidate_id,
|
||||
..
|
||||
} = attempt;
|
||||
let candidate = eligible.candidate;
|
||||
let transport = resolved.transport;
|
||||
let (execution_strategy, conversion_mode) =
|
||||
ai_local_execution_contract_for_formats(spec_metadata.api_format, spec_metadata.api_format);
|
||||
let proxy = planner_state
|
||||
.app()
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
|
||||
.await;
|
||||
let tls_profile = resolve_transport_tls_profile(&transport);
|
||||
let mut extra_fields = serde_json::Map::new();
|
||||
if let Some(proxy_value) = build_request_trace_proxy_value(Some(&transport), proxy.as_ref()) {
|
||||
extra_fields.insert("proxy".to_string(), proxy_value);
|
||||
}
|
||||
let report_context = build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &input.auth_context,
|
||||
request_id: trace_id,
|
||||
candidate_id: &candidate_id,
|
||||
attempt_identity,
|
||||
model: &input.requested_model,
|
||||
provider_name: &transport.provider.name,
|
||||
provider_id: &candidate.provider_id,
|
||||
endpoint_id: &candidate.endpoint_id,
|
||||
key_id: &candidate.key_id,
|
||||
key_name: None,
|
||||
model_id: Some(&candidate.model_id),
|
||||
global_model_id: Some(&candidate.global_model_id),
|
||||
global_model_name: Some(&candidate.global_model_name),
|
||||
provider_api_format: spec_metadata.api_format,
|
||||
client_api_format: spec_metadata.api_format,
|
||||
mapped_model: Some(&resolved.mapped_model),
|
||||
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
|
||||
ranking: eligible.ranking.as_ref(),
|
||||
upstream_url: None,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
body_rules: transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: None,
|
||||
provider_request_headers: None,
|
||||
original_headers: &parts.headers,
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: false,
|
||||
upstream_is_stream: false,
|
||||
has_envelope: false,
|
||||
needs_conversion: false,
|
||||
extra_fields,
|
||||
});
|
||||
let super::request::LocalVideoCreateCandidatePayloadParts {
|
||||
transport: _,
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
upstream_url,
|
||||
} = resolved;
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: false,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format: spec_metadata.api_format.to_string(),
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key: None,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: parts
|
||||
.headers
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream: false,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::planner::candidate_preparation::resolve_candidate_mapped_model;
|
||||
use crate::ai_serving::planner::spec_metadata::local_video_create_spec_metadata;
|
||||
use crate::ai_serving::transport::{
|
||||
build_video_create_headers, build_video_create_request_body, build_video_create_upstream_url,
|
||||
resolve_video_create_auth, video_create_transport_unsupported_reason,
|
||||
ProviderVideoCreateFamily, ProviderVideoCreateHeadersInput,
|
||||
};
|
||||
use crate::ai_serving::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
|
||||
use crate::AppState;
|
||||
|
||||
use super::support::{
|
||||
mark_skipped_local_video_candidate, mark_skipped_local_video_candidate_with_failure_diagnostic,
|
||||
LocalVideoCreateCandidateAttempt, LocalVideoCreateDecisionInput,
|
||||
};
|
||||
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
|
||||
|
||||
pub(super) struct LocalVideoCreateCandidatePayloadParts {
|
||||
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
pub(super) auth_header: String,
|
||||
pub(super) auth_value: String,
|
||||
pub(super) mapped_model: String,
|
||||
pub(super) provider_request_headers: BTreeMap<String, String>,
|
||||
pub(super) provider_request_body: Value,
|
||||
pub(super) upstream_url: String,
|
||||
}
|
||||
|
||||
pub(super) async fn resolve_local_video_create_candidate_payload_parts(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
trace_id: &str,
|
||||
input: &LocalVideoCreateDecisionInput,
|
||||
attempt: &LocalVideoCreateCandidateAttempt,
|
||||
spec: LocalVideoCreateSpec,
|
||||
) -> Option<LocalVideoCreateCandidatePayloadParts> {
|
||||
let spec_metadata = local_video_create_spec_metadata(spec);
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let transport = &attempt.eligible.transport;
|
||||
|
||||
let provider_family = provider_video_create_family(spec.family);
|
||||
let transport_unsupported_reason = video_create_transport_unsupported_reason(
|
||||
transport,
|
||||
provider_family,
|
||||
spec_metadata.api_format,
|
||||
);
|
||||
if let Some(skip_reason) = transport_unsupported_reason {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let auth = resolve_video_create_auth(transport, provider_family);
|
||||
let Some((auth_header, auth_value)) = auth else {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let mapped_model = match resolve_candidate_mapped_model(candidate) {
|
||||
Ok(mapped_model) => mapped_model,
|
||||
Err(skip_reason) => {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
let Some(upstream_url) = build_video_create_upstream_url(
|
||||
transport,
|
||||
parts.uri.path(),
|
||||
parts.uri.query(),
|
||||
&mapped_model,
|
||||
provider_family,
|
||||
) else {
|
||||
mark_skipped_local_video_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"upstream_url_missing",
|
||||
CandidateFailureDiagnostic::upstream_url_missing(
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.api_format,
|
||||
"video_upstream_url",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(provider_request_body) = build_video_create_request_body(
|
||||
body_json,
|
||||
provider_family,
|
||||
&mapped_model,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
) else {
|
||||
mark_skipped_local_video_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_body_rules_apply_failed",
|
||||
CandidateFailureDiagnostic::body_rules_apply_failed(
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.api_format,
|
||||
"video_body_rules",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(provider_request_headers) =
|
||||
build_video_create_headers(ProviderVideoCreateHeadersInput {
|
||||
headers: &parts.headers,
|
||||
auth_header: &auth_header,
|
||||
auth_value: &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_video_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
CandidateFailureDiagnostic::header_rules_apply_failed(
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.api_format,
|
||||
"video_header_rules",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
Some(LocalVideoCreateCandidatePayloadParts {
|
||||
transport: Arc::clone(transport),
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
upstream_url,
|
||||
})
|
||||
}
|
||||
|
||||
fn provider_video_create_family(family: LocalVideoCreateFamily) -> ProviderVideoCreateFamily {
|
||||
match family {
|
||||
LocalVideoCreateFamily::OpenAi => ProviderVideoCreateFamily::OpenAi,
|
||||
LocalVideoCreateFamily::Gemini => ProviderVideoCreateFamily::Gemini,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,253 @@
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use tracing::warn;
|
||||
|
||||
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
mark_skipped_local_execution_candidate,
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic,
|
||||
materialize_local_execution_candidates_with_serving, LocalCandidateResolutionMode,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_metadata::{
|
||||
build_local_execution_candidate_metadata,
|
||||
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_video_create_spec_metadata;
|
||||
use crate::ai_serving::{
|
||||
extract_pool_sticky_session_token, resolve_local_decision_execution_runtime_auth_context,
|
||||
CandidateFailureDiagnostic, ExecutionRuntimeAuthContext, GatewayControlDecision,
|
||||
PlannerAppState,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::AppState;
|
||||
|
||||
pub(super) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalVideoCreateCandidateAttempt;
|
||||
pub(super) use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput as LocalVideoCreateDecisionInput;
|
||||
|
||||
pub(super) async fn resolve_local_video_create_decision_input(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalVideoCreateSpec,
|
||||
) -> Option<LocalVideoCreateDecisionInput> {
|
||||
let spec_metadata = local_video_create_spec_metadata(spec);
|
||||
let Some(auth_context) = resolve_local_video_create_auth_context(decision, spec.family) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let requested_model = extract_requested_model_from_request(
|
||||
parts,
|
||||
body_json,
|
||||
spec_metadata
|
||||
.requested_model_family
|
||||
.expect("video specs should declare requested-model family"),
|
||||
)?;
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
auth_context,
|
||||
Some(requested_model.as_str()),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => return None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = spec_metadata.decision_kind,
|
||||
error = ?err,
|
||||
"gateway local video decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(build_local_requested_model_decision_input(
|
||||
resolved_input,
|
||||
requested_model,
|
||||
))
|
||||
}
|
||||
|
||||
fn resolve_local_video_create_auth_context(
|
||||
decision: &GatewayControlDecision,
|
||||
family: LocalVideoCreateFamily,
|
||||
) -> Option<ExecutionRuntimeAuthContext> {
|
||||
let auth_context = resolve_local_decision_execution_runtime_auth_context(decision)?;
|
||||
match family {
|
||||
LocalVideoCreateFamily::OpenAi | LocalVideoCreateFamily::Gemini => Some(auth_context),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn list_local_video_create_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalVideoCreateDecisionInput,
|
||||
body_json: &serde_json::Value,
|
||||
api_format: &str,
|
||||
decision_kind: &str,
|
||||
) -> Option<Vec<LocalVideoCreateCandidateAttempt>> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let (candidates, preselection_skipped) = match planner_state
|
||||
.list_selectable_candidates_with_skip_reasons(
|
||||
api_format,
|
||||
&input.requested_model,
|
||||
false,
|
||||
input.required_capabilities.as_ref(),
|
||||
Some(&input.auth_snapshot),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(candidates) => candidates,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = decision_kind,
|
||||
error = ?err,
|
||||
"gateway local video decision scheduler selection failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(
|
||||
materialize_local_video_create_candidate_attempts(
|
||||
planner_state,
|
||||
trace_id,
|
||||
input,
|
||||
body_json,
|
||||
candidates,
|
||||
preselection_skipped
|
||||
.into_iter()
|
||||
.map(|item| SkippedLocalExecutionCandidate {
|
||||
candidate: item.candidate,
|
||||
skip_reason: item.skip_reason,
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
})
|
||||
.collect(),
|
||||
api_format,
|
||||
)
|
||||
.await,
|
||||
)
|
||||
}
|
||||
|
||||
async fn materialize_local_video_create_candidate_attempts(
|
||||
state: PlannerAppState<'_>,
|
||||
trace_id: &str,
|
||||
input: &LocalVideoCreateDecisionInput,
|
||||
body_json: &serde_json::Value,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
preselection_skipped: Vec<SkippedLocalExecutionCandidate>,
|
||||
api_format: &str,
|
||||
) -> Vec<LocalVideoCreateCandidateAttempt> {
|
||||
let sticky_session_token = extract_pool_sticky_session_token(body_json);
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::VideoDecision,
|
||||
);
|
||||
let outcome = materialize_local_execution_candidates_with_serving(
|
||||
state,
|
||||
trace_id,
|
||||
api_format,
|
||||
Some(&input.requested_model),
|
||||
Some(&input.auth_snapshot),
|
||||
input.required_capabilities.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
persistence_policy,
|
||||
candidates,
|
||||
preselection_skipped,
|
||||
LocalCandidateResolutionMode::Standard,
|
||||
|eligible| {
|
||||
Some(build_local_execution_candidate_metadata(
|
||||
LocalExecutionCandidateMetadataParts {
|
||||
eligible,
|
||||
provider_api_format: api_format,
|
||||
client_api_format: api_format,
|
||||
extra_fields: serde_json::Map::new(),
|
||||
},
|
||||
))
|
||||
},
|
||||
|mut skipped_candidate| {
|
||||
skipped_candidate.extra_data =
|
||||
Some(build_local_execution_candidate_metadata_for_candidate(
|
||||
&skipped_candidate.candidate,
|
||||
skipped_candidate.transport_ref(),
|
||||
api_format,
|
||||
api_format,
|
||||
serde_json::Map::new(),
|
||||
));
|
||||
skipped_candidate
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
outcome.attempts
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_video_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalVideoCreateDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::VideoDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_video_candidate_with_failure_diagnostic(
|
||||
state: &AppState,
|
||||
input: &LocalVideoCreateDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
diagnostic: CandidateFailureDiagnostic,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::VideoDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
diagnostic,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::ai_serving::{
|
||||
resolve_claude_stream_spec as resolve_surface_stream_spec,
|
||||
resolve_claude_sync_spec as resolve_surface_sync_spec,
|
||||
};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
use super::family::{
|
||||
maybe_build_stream_via_standard_family_payload, maybe_build_sync_via_standard_family_payload,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::normalize_claude_request_to_openai_chat_request;
|
||||
|
||||
pub(crate) fn resolve_sync_spec(plan_kind: &str) -> Option<super::family::LocalStandardSpec> {
|
||||
resolve_surface_sync_spec(plan_kind)
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_stream_spec(plan_kind: &str) -> Option<super::family::LocalStandardSpec> {
|
||||
resolve_surface_stream_spec(plan_kind)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_sync_local_claude_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
maybe_build_sync_via_standard_family_payload(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
plan_kind,
|
||||
resolve_sync_spec,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_stream_local_claude_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
maybe_build_stream_via_standard_family_payload(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
plan_kind,
|
||||
resolve_stream_spec,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
#[cfg(test)]
|
||||
#[path = "codex/tests.rs"]
|
||||
mod tests;
|
||||
|
||||
pub(crate) use crate::ai_serving::{
|
||||
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
|
||||
};
|
||||
@@ -0,0 +1,245 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use super::{
|
||||
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
|
||||
};
|
||||
use http::{HeaderMap, HeaderValue};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn applies_codex_defaults_when_body_rules_do_not_handle_fields() {
|
||||
let mut body = json!({
|
||||
"model": "gpt-5",
|
||||
"max_output_tokens": 128,
|
||||
"temperature": 0.3,
|
||||
"top_p": 0.9,
|
||||
"metadata": {"client": "desktop"},
|
||||
"store": true
|
||||
});
|
||||
|
||||
apply_codex_openai_responses_special_body_edits(
|
||||
&mut body,
|
||||
"codex",
|
||||
"openai:responses",
|
||||
None,
|
||||
None,
|
||||
);
|
||||
|
||||
assert!(body.get("max_output_tokens").is_none());
|
||||
assert!(body.get("temperature").is_none());
|
||||
assert!(body.get("top_p").is_none());
|
||||
assert!(body.get("metadata").is_none());
|
||||
assert_eq!(body["store"], false);
|
||||
assert_eq!(body["instructions"], "You are ChatGPT.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_store_for_compact_even_when_body_rules_handle_it() {
|
||||
let body_rules = json!([
|
||||
{"action":"set","path":"store","value":true},
|
||||
{"action":"set","path":"instructions","value":"Keep custom"},
|
||||
{"action":"set","path":"metadata","value":{"client":"desktop","mode":"custom"}},
|
||||
{"action":"set","path":"top_p","value":0.5}
|
||||
]);
|
||||
let mut body = json!({
|
||||
"model": "gpt-5",
|
||||
"max_output_tokens": 128,
|
||||
"metadata": {"client": "desktop", "mode": "custom"},
|
||||
"store": true,
|
||||
"instructions": "Keep custom",
|
||||
"top_p": 0.5
|
||||
});
|
||||
|
||||
apply_codex_openai_responses_special_body_edits(
|
||||
&mut body,
|
||||
"codex",
|
||||
"openai:responses:compact",
|
||||
Some(&body_rules),
|
||||
None,
|
||||
);
|
||||
|
||||
assert!(body.get("max_output_tokens").is_none());
|
||||
assert!(body.get("store").is_none());
|
||||
assert_eq!(body["instructions"], "Keep custom");
|
||||
assert_eq!(body["metadata"]["mode"], "custom");
|
||||
assert_eq!(body["top_p"], 0.5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn injects_stable_prompt_cache_key_for_codex_requests() {
|
||||
let mut body = json!({
|
||||
"model": "gpt-5",
|
||||
"input": "hello",
|
||||
});
|
||||
|
||||
apply_codex_openai_responses_special_body_edits(
|
||||
&mut body,
|
||||
"codex",
|
||||
"openai:responses",
|
||||
None,
|
||||
Some("key-123"),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
body["prompt_cache_key"],
|
||||
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_existing_prompt_cache_key_for_codex_requests() {
|
||||
let mut body = json!({
|
||||
"model": "gpt-5",
|
||||
"input": "hello",
|
||||
"prompt_cache_key": "existing-key",
|
||||
});
|
||||
|
||||
apply_codex_openai_responses_special_body_edits(
|
||||
&mut body,
|
||||
"codex",
|
||||
"openai:responses",
|
||||
None,
|
||||
Some("key-123"),
|
||||
);
|
||||
|
||||
assert_eq!(body["prompt_cache_key"], "existing-key");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn injects_chatgpt_account_id_and_session_headers_for_codex_requests() {
|
||||
let mut headers = BTreeMap::new();
|
||||
let body = json!({
|
||||
"model": "gpt-5",
|
||||
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
|
||||
});
|
||||
|
||||
apply_codex_openai_responses_special_headers(
|
||||
&mut headers,
|
||||
&body,
|
||||
&HeaderMap::new(),
|
||||
"codex",
|
||||
"openai:responses",
|
||||
Some("trace-codex-123"),
|
||||
Some(r#"{"account_id":"acc-123"}"#),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
headers.get("chatgpt-account-id"),
|
||||
Some(&"acc-123".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("x-client-request-id"),
|
||||
Some(&"trace-codex-123".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("user-agent"),
|
||||
Some(
|
||||
&"codex-tui/0.122.0 (Mac OS 15.2.0; arm64) vscode/2.6.11 (codex-tui; 0.122.0)"
|
||||
.to_string()
|
||||
)
|
||||
);
|
||||
assert_eq!(headers.get("originator"), Some(&"codex-tui".to_string()));
|
||||
assert_eq!(
|
||||
headers.get("session_id"),
|
||||
Some(&"ab5ecce4f0d110fe".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("conversation_id"),
|
||||
Some(&"ab5ecce4f0d110fe".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn respects_existing_codex_request_and_session_headers() {
|
||||
let mut headers = BTreeMap::new();
|
||||
headers.insert(
|
||||
"x-client-request-id".to_string(),
|
||||
"kept-by-rule-request".to_string(),
|
||||
);
|
||||
headers.insert("session_id".to_string(), "kept-by-rule".to_string());
|
||||
let body = json!({
|
||||
"model": "gpt-5",
|
||||
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
|
||||
});
|
||||
let mut original_headers = HeaderMap::new();
|
||||
original_headers.insert(
|
||||
"x-client-request-id",
|
||||
HeaderValue::from_static("user-specified-request"),
|
||||
);
|
||||
original_headers.insert(
|
||||
"session_id",
|
||||
HeaderValue::from_static("user-specified-session"),
|
||||
);
|
||||
original_headers.insert(
|
||||
"conversation_id",
|
||||
HeaderValue::from_static("user-specified-conversation"),
|
||||
);
|
||||
original_headers.insert(
|
||||
"user-agent",
|
||||
HeaderValue::from_static("user-specified-agent"),
|
||||
);
|
||||
original_headers.insert(
|
||||
"originator",
|
||||
HeaderValue::from_static("user-specified-originator"),
|
||||
);
|
||||
|
||||
apply_codex_openai_responses_special_headers(
|
||||
&mut headers,
|
||||
&body,
|
||||
&original_headers,
|
||||
"codex",
|
||||
"openai:responses",
|
||||
Some("trace-codex-123"),
|
||||
Some(r#"{"account_id":"acc-123"}"#),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
headers.get("x-client-request-id"),
|
||||
Some(&"kept-by-rule-request".to_string())
|
||||
);
|
||||
assert!(!headers.contains_key("user-agent"));
|
||||
assert!(!headers.contains_key("originator"));
|
||||
assert_eq!(headers.get("session_id"), Some(&"kept-by-rule".to_string()));
|
||||
assert!(!headers.contains_key("conversation_id"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skips_conversation_id_for_compact_codex_requests() {
|
||||
let mut headers = BTreeMap::new();
|
||||
let body = json!({
|
||||
"model": "gpt-5",
|
||||
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
|
||||
});
|
||||
|
||||
apply_codex_openai_responses_special_headers(
|
||||
&mut headers,
|
||||
&body,
|
||||
&HeaderMap::new(),
|
||||
"codex",
|
||||
"openai:responses:compact",
|
||||
Some("trace-codex-compact-123"),
|
||||
Some(r#"{"account_id":"acc-123"}"#),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
headers.get("chatgpt-account-id"),
|
||||
Some(&"acc-123".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("x-client-request-id"),
|
||||
Some(&"trace-codex-compact-123".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("user-agent"),
|
||||
Some(
|
||||
&"codex-tui/0.122.0 (Mac OS 15.2.0; arm64) vscode/2.6.11 (codex-tui; 0.122.0)"
|
||||
.to_string()
|
||||
)
|
||||
);
|
||||
assert_eq!(headers.get("originator"), Some(&"codex-tui".to_string()));
|
||||
assert_eq!(
|
||||
headers.get("session_id"),
|
||||
Some(&"ab5ecce4f0d110fe".to_string())
|
||||
);
|
||||
assert!(!headers.contains_key("conversation_id"));
|
||||
}
|
||||
@@ -0,0 +1,266 @@
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::plan_builders::{AiStreamAttempt, AiSyncAttempt};
|
||||
use crate::ai_serving::planner::runtime_miss::{
|
||||
apply_local_runtime_candidate_evaluation_progress,
|
||||
apply_local_runtime_candidate_terminal_reason, set_local_runtime_miss_diagnostic_reason,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::{
|
||||
build_stream_plan_from_requested_model_family, build_sync_plan_from_requested_model_family,
|
||||
local_standard_spec_metadata,
|
||||
};
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
use super::candidates::{
|
||||
materialize_local_standard_candidate_attempts, resolve_local_standard_decision_input,
|
||||
};
|
||||
use super::payload::maybe_build_local_standard_decision_payload_for_candidate;
|
||||
use super::LocalStandardSpec;
|
||||
|
||||
pub(crate) async fn maybe_build_sync_via_standard_family_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
resolve_sync_spec: fn(&str) -> Option<LocalStandardSpec>,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let spec_metadata = local_standard_spec_metadata(spec);
|
||||
|
||||
let Some(input) =
|
||||
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_sync_plans");
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_stream_via_standard_family_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
resolve_stream_spec: fn(&str) -> Option<LocalStandardSpec>,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let spec_metadata = local_standard_spec_metadata(spec);
|
||||
|
||||
let Some(input) =
|
||||
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_stream_plans");
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
let spec_metadata = local_standard_spec_metadata(spec);
|
||||
let requested_model_family = spec_metadata
|
||||
.requested_model_family
|
||||
.expect("standard spec metadata should include requested-model family");
|
||||
let Some(input) =
|
||||
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
|
||||
.await
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
extract_requested_model_from_request(parts, body_json, requested_model_family)
|
||||
.as_deref(),
|
||||
"decision_input_unavailable",
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
if candidate_count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let built = build_sync_plan_from_requested_model_family(
|
||||
requested_model_family,
|
||||
parts,
|
||||
body_json,
|
||||
payload,
|
||||
);
|
||||
match built {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec_metadata.api_format,
|
||||
error = ?err,
|
||||
"gateway local standard sync plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_sync_plans");
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
|
||||
let spec_metadata = local_standard_spec_metadata(spec);
|
||||
let requested_model_family = spec_metadata
|
||||
.requested_model_family
|
||||
.expect("standard spec metadata should include requested-model family");
|
||||
let Some(input) =
|
||||
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
|
||||
.await
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
extract_requested_model_from_request(parts, body_json, requested_model_family)
|
||||
.as_deref(),
|
||||
"decision_input_unavailable",
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
if candidate_count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let built = build_stream_plan_from_requested_model_family(
|
||||
requested_model_family,
|
||||
parts,
|
||||
body_json,
|
||||
payload,
|
||||
);
|
||||
match built {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec_metadata.api_format,
|
||||
error = ?err,
|
||||
"gateway local standard stream plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_stream_plans");
|
||||
Ok(plans)
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
materialize_local_execution_candidates_with_serving, LocalCandidateResolutionMode,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_metadata::{
|
||||
build_local_execution_candidate_contract_metadata,
|
||||
build_local_execution_candidate_contract_metadata_for_candidate,
|
||||
LocalExecutionCandidateMetadataParts,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_source::{
|
||||
preselect_local_execution_candidates_with_serving, LocalCandidatePreselectionKeyMode,
|
||||
};
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
|
||||
use crate::ai_serving::{
|
||||
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
|
||||
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision, PlannerAppState,
|
||||
};
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
use super::{LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSpec};
|
||||
|
||||
pub(super) async fn resolve_local_standard_decision_input(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Option<LocalStandardDecisionInput> {
|
||||
let spec_metadata = local_standard_spec_metadata(spec);
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let requested_model = extract_requested_model_from_request(
|
||||
parts,
|
||||
body_json,
|
||||
spec_metadata
|
||||
.requested_model_family
|
||||
.expect("standard specs should declare requested-model family"),
|
||||
)?;
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
auth_context,
|
||||
Some(requested_model.as_str()),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => return None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec_metadata.api_format,
|
||||
error = ?err,
|
||||
"gateway local standard decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(build_local_requested_model_decision_input(
|
||||
resolved_input,
|
||||
requested_model,
|
||||
))
|
||||
}
|
||||
|
||||
pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalStandardDecisionInput,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Result<(Vec<LocalStandardCandidateAttempt>, usize), GatewayError> {
|
||||
let spec_metadata = local_standard_spec_metadata(spec);
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let sticky_session_token = extract_pool_sticky_session_token(body_json);
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::StandardDecision,
|
||||
);
|
||||
let preselection = preselect_local_execution_candidates_with_serving(
|
||||
planner_state,
|
||||
spec_metadata.api_format,
|
||||
&input.requested_model,
|
||||
spec_metadata.require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
&input.auth_snapshot,
|
||||
false,
|
||||
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
|
||||
)
|
||||
.await?;
|
||||
let outcome = materialize_local_execution_candidates_with_serving(
|
||||
planner_state,
|
||||
trace_id,
|
||||
spec_metadata.api_format,
|
||||
Some(&input.requested_model),
|
||||
Some(&input.auth_snapshot),
|
||||
input.required_capabilities.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
persistence_policy,
|
||||
preselection.candidates,
|
||||
preselection.skipped_candidates,
|
||||
LocalCandidateResolutionMode::Standard,
|
||||
|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(),
|
||||
))
|
||||
},
|
||||
|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;
|
||||
|
||||
Ok((outcome.attempts, outcome.candidate_count))
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
mod build;
|
||||
mod candidates;
|
||||
mod payload;
|
||||
mod request;
|
||||
|
||||
pub(crate) use self::build::{
|
||||
build_local_stream_plan_and_reports, build_local_sync_plan_and_reports,
|
||||
maybe_build_stream_via_standard_family_payload, maybe_build_sync_via_standard_family_payload,
|
||||
};
|
||||
pub(super) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalStandardCandidateAttempt;
|
||||
pub(super) use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput as LocalStandardDecisionInput;
|
||||
pub(crate) use crate::ai_serving::{
|
||||
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
|
||||
};
|
||||
@@ -0,0 +1,554 @@
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
use crate::ai_serving::planner::passthrough::maybe_build_local_same_format_provider_decision_payload_for_candidate;
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, LocalExecutionReportContextParts,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
|
||||
use crate::ai_serving::planner::CandidateFailureDiagnostic;
|
||||
use crate::ai_serving::planner::{
|
||||
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
ai_local_execution_contract_for_formats, api_format_alias_matches,
|
||||
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
|
||||
};
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, append_local_failover_policy_to_value,
|
||||
AiExecutionDecision, AppState,
|
||||
};
|
||||
|
||||
use super::request::resolve_local_standard_candidate_payload_parts;
|
||||
use super::{LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSpec};
|
||||
|
||||
pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
body_json: &serde_json::Value,
|
||||
input: &LocalStandardDecisionInput,
|
||||
attempt: LocalStandardCandidateAttempt,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
let spec_metadata = local_standard_spec_metadata(spec);
|
||||
if api_format_alias_matches(
|
||||
&attempt.eligible.provider_api_format,
|
||||
spec_metadata.api_format,
|
||||
) {
|
||||
let same_format_spec = if spec_metadata.require_streaming {
|
||||
resolve_local_same_format_stream_spec(spec_metadata.decision_kind)
|
||||
} else {
|
||||
resolve_local_same_format_sync_spec(spec_metadata.decision_kind)
|
||||
};
|
||||
if let Some(same_format_spec) = same_format_spec {
|
||||
return maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
input,
|
||||
attempt,
|
||||
same_format_spec,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
let LocalStandardCandidateAttempt {
|
||||
eligible,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
..
|
||||
} = &attempt;
|
||||
let candidate = &eligible.candidate;
|
||||
let resolved = resolve_local_standard_candidate_payload_parts(
|
||||
state, parts, trace_id, body_json, input, &attempt, spec,
|
||||
)
|
||||
.await?;
|
||||
let proxy = state
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&resolved.transport)
|
||||
.await;
|
||||
let mut extra_fields = serde_json::Map::new();
|
||||
if let Some(proxy_value) =
|
||||
build_request_trace_proxy_value(Some(&resolved.transport), proxy.as_ref())
|
||||
{
|
||||
extra_fields.insert("proxy".to_string(), proxy_value);
|
||||
}
|
||||
if let Some(envelope_name) = resolved.envelope_name {
|
||||
extra_fields.insert(
|
||||
"envelope_name".to_string(),
|
||||
serde_json::Value::String(envelope_name.to_string()),
|
||||
);
|
||||
}
|
||||
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
|
||||
spec_metadata.api_format,
|
||||
resolved.provider_api_format.as_str(),
|
||||
);
|
||||
let report_context = append_local_failover_policy_to_value(
|
||||
append_execution_contract_fields_to_value(
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &input.auth_context,
|
||||
request_id: trace_id,
|
||||
candidate_id,
|
||||
attempt_identity: attempt.attempt_identity(),
|
||||
model: &input.requested_model,
|
||||
provider_name: &candidate.provider_name,
|
||||
provider_id: &candidate.provider_id,
|
||||
endpoint_id: &candidate.endpoint_id,
|
||||
key_id: &candidate.key_id,
|
||||
key_name: Some(&candidate.key_name),
|
||||
model_id: Some(&candidate.model_id),
|
||||
global_model_id: Some(&candidate.global_model_id),
|
||||
global_model_name: Some(&candidate.global_model_name),
|
||||
provider_api_format: &resolved.provider_api_format,
|
||||
client_api_format: spec_metadata.api_format,
|
||||
mapped_model: Some(&resolved.mapped_model),
|
||||
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
|
||||
ranking: eligible.ranking.as_ref(),
|
||||
upstream_url: Some(&resolved.upstream_url),
|
||||
header_rules: resolved.transport.endpoint.header_rules.as_ref(),
|
||||
body_rules: resolved.transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: Some(serde_json::Value::Null),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: body_json
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
upstream_is_stream: resolved.upstream_is_stream,
|
||||
has_envelope: resolved.envelope_name.is_some(),
|
||||
needs_conversion: true,
|
||||
extra_fields,
|
||||
}),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
spec_metadata.api_format,
|
||||
resolved.provider_api_format.as_str(),
|
||||
),
|
||||
&resolved.transport,
|
||||
);
|
||||
let tls_profile = resolve_transport_tls_profile(&resolved.transport);
|
||||
let timeouts = resolve_transport_execution_timeouts(&resolved.transport);
|
||||
let super::request::LocalStandardCandidatePayloadParts {
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
provider_api_format,
|
||||
provider_request_body,
|
||||
provider_request_headers,
|
||||
upstream_url,
|
||||
upstream_is_stream,
|
||||
envelope_name: _,
|
||||
transport,
|
||||
} = resolved;
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.to_string(),
|
||||
provider_name: candidate.provider_name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: None,
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format,
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key: None,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts,
|
||||
upstream_is_stream,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_standard_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalStandardDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::StandardDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn mark_skipped_local_standard_candidate_with_extra_data(
|
||||
state: &AppState,
|
||||
input: &LocalStandardDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
extra_data: Option<serde_json::Value>,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::StandardDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate_with_extra_data(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
extra_data,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn mark_skipped_local_standard_candidate_with_failure_diagnostic(
|
||||
state: &AppState,
|
||||
input: &LocalStandardDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
diagnostic: CandidateFailureDiagnostic,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::StandardDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
diagnostic,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_provider_transport::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::json;
|
||||
|
||||
use super::maybe_build_local_standard_decision_payload_for_candidate;
|
||||
use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt;
|
||||
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput;
|
||||
use crate::ai_serving::{
|
||||
ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, LocalStandardSourceFamily,
|
||||
LocalStandardSourceMode, LocalStandardSpec,
|
||||
};
|
||||
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||
|
||||
fn sample_auth_snapshot() -> GatewayAuthApiKeySnapshot {
|
||||
GatewayAuthApiKeySnapshot {
|
||||
user_id: "user-1".to_string(),
|
||||
username: "alice".to_string(),
|
||||
email: None,
|
||||
user_role: "user".to_string(),
|
||||
user_auth_source: "local".to_string(),
|
||||
user_is_active: true,
|
||||
user_is_deleted: false,
|
||||
user_rate_limit: None,
|
||||
user_allowed_providers: None,
|
||||
user_allowed_api_formats: None,
|
||||
user_allowed_models: None,
|
||||
api_key_id: "api-key-1".to_string(),
|
||||
api_key_name: Some("default".to_string()),
|
||||
api_key_is_active: true,
|
||||
api_key_is_locked: false,
|
||||
api_key_is_standalone: false,
|
||||
api_key_rate_limit: None,
|
||||
api_key_concurrent_limit: None,
|
||||
api_key_expires_at_unix_secs: None,
|
||||
api_key_allowed_providers: None,
|
||||
api_key_allowed_api_formats: None,
|
||||
api_key_allowed_models: None,
|
||||
currently_usable: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_input() -> LocalRequestedModelDecisionInput {
|
||||
LocalRequestedModelDecisionInput {
|
||||
auth_context: ExecutionRuntimeAuthContext {
|
||||
user_id: "user-1".to_string(),
|
||||
api_key_id: "api-key-1".to_string(),
|
||||
username: Some("alice".to_string()),
|
||||
api_key_name: Some("default".to_string()),
|
||||
balance_remaining: Some(10.0),
|
||||
access_allowed: true,
|
||||
api_key_is_standalone: false,
|
||||
},
|
||||
requested_model: "claude-sonnet-4-5".to_string(),
|
||||
auth_snapshot: sample_auth_snapshot(),
|
||||
required_capabilities: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_transport(api_format: &str, endpoint_id: &str) -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-1".to_string(),
|
||||
name: "provider".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: true,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: endpoint_id.to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
api_format: api_format.to_string(),
|
||||
api_family: Some(
|
||||
api_format
|
||||
.split_once(':')
|
||||
.map(|(family, _)| family)
|
||||
.unwrap_or(api_format)
|
||||
.to_string(),
|
||||
),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
is_active: true,
|
||||
base_url: "https://api.example.test".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: if api_format == "openai:chat" {
|
||||
Some(json!({
|
||||
"enabled": true,
|
||||
"accept_formats": ["claude:messages"],
|
||||
}))
|
||||
} else {
|
||||
None
|
||||
},
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
name: "key".to_string(),
|
||||
auth_type: "api_key".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec![
|
||||
"claude:messages".to_string(),
|
||||
"openai:chat".to_string(),
|
||||
]),
|
||||
auth_type_by_format: None,
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: Some(json!({
|
||||
"claude:messages": 1,
|
||||
"openai:chat": 1,
|
||||
})),
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
decrypted_api_key: "sk-upstream".to_string(),
|
||||
decrypted_auth_config: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_candidate(
|
||||
api_format: &str,
|
||||
endpoint_id: &str,
|
||||
) -> SchedulerMinimalCandidateSelectionCandidate {
|
||||
SchedulerMinimalCandidateSelectionCandidate {
|
||||
provider_id: "provider-1".to_string(),
|
||||
provider_name: "provider".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 1,
|
||||
endpoint_id: endpoint_id.to_string(),
|
||||
endpoint_api_format: api_format.to_string(),
|
||||
key_id: "key-1".to_string(),
|
||||
key_name: "key".to_string(),
|
||||
key_auth_type: "api_key".to_string(),
|
||||
key_internal_priority: 1,
|
||||
key_global_priority_for_format: Some(1),
|
||||
key_capabilities: None,
|
||||
model_id: format!("model-{endpoint_id}"),
|
||||
global_model_id: "global-model-1".to_string(),
|
||||
global_model_name: "claude-sonnet-4-5".to_string(),
|
||||
selected_provider_model_name: if api_format == "claude:messages" {
|
||||
"claude-sonnet-4-5-upstream".to_string()
|
||||
} else {
|
||||
"gpt-4o-upstream".to_string()
|
||||
},
|
||||
mapping_matched_model: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_attempt(
|
||||
api_format: &str,
|
||||
endpoint_id: &str,
|
||||
candidate_index: u32,
|
||||
) -> LocalExecutionCandidateAttempt {
|
||||
LocalExecutionCandidateAttempt {
|
||||
eligible: EligibleLocalExecutionCandidate {
|
||||
candidate: sample_candidate(api_format, endpoint_id),
|
||||
transport: Arc::new(sample_transport(api_format, endpoint_id)),
|
||||
provider_api_format: api_format.to_string(),
|
||||
orchestration: LocalExecutionCandidateMetadata::default(),
|
||||
ranking: None,
|
||||
},
|
||||
candidate_index,
|
||||
retry_index: 0,
|
||||
candidate_id: format!("candidate-{candidate_index}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn claude_stream_spec() -> LocalStandardSpec {
|
||||
LocalStandardSpec {
|
||||
api_format: "claude:messages",
|
||||
decision_kind: "claude_chat_stream",
|
||||
report_kind: "claude_chat_stream_success",
|
||||
family: LocalStandardSourceFamily::Standard,
|
||||
mode: LocalStandardSourceMode::Chat,
|
||||
require_streaming: true,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn standard_family_builds_same_format_candidate_before_cross_format_candidate() {
|
||||
let state = crate::AppState::new().expect("state should build");
|
||||
let request = http::Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/messages?beta=true")
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.body(())
|
||||
.expect("request should build");
|
||||
let (parts, _) = request.into_parts();
|
||||
let body_json = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 32,
|
||||
"stream": true
|
||||
});
|
||||
let input = sample_input();
|
||||
|
||||
let payload = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
&state,
|
||||
&parts,
|
||||
"trace-standard-same-format-first",
|
||||
&body_json,
|
||||
&input,
|
||||
sample_attempt("claude:messages", "endpoint-claude", 0),
|
||||
claude_stream_spec(),
|
||||
)
|
||||
.await
|
||||
.expect("same-format candidate should build a standard-family payload");
|
||||
|
||||
assert_eq!(payload.endpoint_id.as_deref(), Some("endpoint-claude"));
|
||||
assert_eq!(
|
||||
payload.execution_strategy.as_deref(),
|
||||
Some("local_same_format")
|
||||
);
|
||||
assert_eq!(payload.conversion_mode.as_deref(), Some("none"));
|
||||
assert_eq!(
|
||||
payload.provider_api_format.as_deref(),
|
||||
Some("claude:messages")
|
||||
);
|
||||
assert_eq!(
|
||||
payload.client_api_format.as_deref(),
|
||||
Some("claude:messages")
|
||||
);
|
||||
assert_eq!(
|
||||
payload
|
||||
.provider_request_body
|
||||
.as_ref()
|
||||
.and_then(|body| body.get("model"))
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some("claude-sonnet-4-5-upstream")
|
||||
);
|
||||
|
||||
let cross_format_payload = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
&state,
|
||||
&parts,
|
||||
"trace-standard-same-format-first",
|
||||
&body_json,
|
||||
&input,
|
||||
sample_attempt("openai:chat", "endpoint-openai-chat", 1),
|
||||
claude_stream_spec(),
|
||||
)
|
||||
.await
|
||||
.expect("cross-format candidate should still build after the same-format candidate");
|
||||
|
||||
assert_eq!(
|
||||
cross_format_payload.endpoint_id.as_deref(),
|
||||
Some("endpoint-openai-chat")
|
||||
);
|
||||
assert_eq!(
|
||||
cross_format_payload.execution_strategy.as_deref(),
|
||||
Some("local_cross_format")
|
||||
);
|
||||
assert_eq!(
|
||||
cross_format_payload.conversion_mode.as_deref(),
|
||||
Some("bidirectional")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,427 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::planner::candidate_preparation::{
|
||||
prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth,
|
||||
OauthPreparationContext,
|
||||
};
|
||||
use crate::ai_serving::planner::common::force_upstream_streaming_for_provider;
|
||||
use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
|
||||
use crate::ai_serving::planner::standard::{
|
||||
apply_codex_openai_responses_special_headers, request_body_build_failure_extra_data,
|
||||
};
|
||||
use crate::ai_serving::transport::kiro::{
|
||||
build_kiro_provider_headers, build_kiro_provider_request_body,
|
||||
is_kiro_claude_messages_transport, KiroProviderHeadersInput, KiroRequestAuth,
|
||||
KIRO_ENVELOPE_NAME,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
build_kiro_cross_format_upstream_url, build_standard_provider_request_headers,
|
||||
StandardProviderRequestHeadersInput,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
|
||||
};
|
||||
use crate::AppState;
|
||||
|
||||
use super::payload::{
|
||||
mark_skipped_local_standard_candidate, mark_skipped_local_standard_candidate_with_extra_data,
|
||||
mark_skipped_local_standard_candidate_with_failure_diagnostic,
|
||||
};
|
||||
use super::{LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSpec};
|
||||
|
||||
pub(crate) struct LocalStandardCandidatePayloadParts {
|
||||
pub(super) auth_header: String,
|
||||
pub(super) auth_value: String,
|
||||
pub(super) mapped_model: String,
|
||||
pub(super) provider_api_format: String,
|
||||
pub(super) provider_request_body: Value,
|
||||
pub(super) provider_request_headers: BTreeMap<String, String>,
|
||||
pub(super) upstream_url: String,
|
||||
pub(super) upstream_is_stream: bool,
|
||||
pub(super) envelope_name: Option<&'static str>,
|
||||
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_local_standard_candidate_payload_parts(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
body_json: &serde_json::Value,
|
||||
input: &LocalStandardDecisionInput,
|
||||
attempt: &LocalStandardCandidateAttempt,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Option<LocalStandardCandidatePayloadParts> {
|
||||
let spec_metadata = local_standard_spec_metadata(spec);
|
||||
let planner_state = crate::ai_serving::PlannerAppState::new(state);
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let transport = &attempt.eligible.transport;
|
||||
let provider_api_format = attempt.eligible.provider_api_format.as_str();
|
||||
let is_kiro_claude_cli = is_kiro_claude_messages_transport(transport, provider_api_format);
|
||||
let Some(conversion_kind) =
|
||||
crate::ai_serving::request_conversion_kind(spec_metadata.api_format, provider_api_format)
|
||||
else {
|
||||
return None;
|
||||
};
|
||||
|
||||
if let Some(skip_reason) = crate::ai_serving::request_conversion_transport_unsupported_reason(
|
||||
transport,
|
||||
conversion_kind,
|
||||
) {
|
||||
mark_skipped_local_standard_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let oauth_context = OauthPreparationContext {
|
||||
trace_id,
|
||||
api_format: provider_api_format,
|
||||
operation: "standard_family_cross_format",
|
||||
};
|
||||
let kiro_auth = if is_kiro_claude_cli {
|
||||
match crate::ai_serving::planner::candidate_preparation::resolve_candidate_oauth_auth(
|
||||
planner_state,
|
||||
transport,
|
||||
oauth_context,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Some(LocalResolvedOAuthRequestAuth::Kiro(auth)) => Some(auth),
|
||||
_ => {
|
||||
mark_skipped_local_standard_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let prepared_candidate = if let Some(kiro_auth) = kiro_auth.as_ref() {
|
||||
match prepare_header_authenticated_candidate_from_auth(
|
||||
candidate,
|
||||
kiro_auth.name.to_string(),
|
||||
kiro_auth.value.clone(),
|
||||
) {
|
||||
Ok(prepared) => prepared,
|
||||
Err(skip_reason) => {
|
||||
mark_skipped_local_standard_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
match prepare_header_authenticated_candidate(
|
||||
planner_state,
|
||||
transport,
|
||||
candidate,
|
||||
crate::ai_serving::request_conversion_direct_auth(transport, conversion_kind),
|
||||
oauth_context,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(prepared) => prepared,
|
||||
Err(skip_reason) => {
|
||||
mark_skipped_local_standard_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let upstream_is_stream = spec_metadata.require_streaming
|
||||
|| force_upstream_streaming_for_provider(
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format,
|
||||
);
|
||||
let provider_request_body =
|
||||
match crate::ai_serving::planner::standard::build_standard_request_body(
|
||||
body_json,
|
||||
spec_metadata.api_format,
|
||||
&prepared_candidate.mapped_model,
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format,
|
||||
parts.uri.path(),
|
||||
upstream_is_stream,
|
||||
if is_kiro_claude_cli {
|
||||
None
|
||||
} else {
|
||||
transport.endpoint.body_rules.as_ref()
|
||||
},
|
||||
Some(input.auth_context.api_key_id.as_str()),
|
||||
) {
|
||||
Some(body) => body,
|
||||
None => {
|
||||
mark_skipped_local_standard_candidate_with_extra_data(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.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 let Some(kiro_auth) = kiro_auth.as_ref() {
|
||||
return build_kiro_cross_format_payload_parts(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
input,
|
||||
attempt,
|
||||
transport,
|
||||
provider_api_format,
|
||||
prepared_candidate.mapped_model,
|
||||
prepared_candidate.auth_header,
|
||||
prepared_candidate.auth_value,
|
||||
provider_request_body,
|
||||
upstream_is_stream,
|
||||
kiro_auth,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let upstream_url = match crate::ai_serving::planner::standard::build_standard_upstream_url(
|
||||
parts,
|
||||
transport,
|
||||
&prepared_candidate.mapped_model,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
) {
|
||||
Some(url) => url,
|
||||
None => {
|
||||
mark_skipped_local_standard_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"upstream_url_missing",
|
||||
CandidateFailureDiagnostic::upstream_url_missing(
|
||||
spec_metadata.api_format,
|
||||
provider_api_format,
|
||||
"standard_family_url",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let Some(resolved_headers) =
|
||||
build_standard_provider_request_headers(StandardProviderRequestHeadersInput {
|
||||
transport,
|
||||
provider_api_format,
|
||||
same_format: false,
|
||||
headers: &parts.headers,
|
||||
auth_header: &prepared_candidate.auth_header,
|
||||
auth_value: &prepared_candidate.auth_value,
|
||||
extra_headers: &BTreeMap::new(),
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
provider_request_body: &provider_request_body,
|
||||
original_request_body: body_json,
|
||||
upstream_is_stream,
|
||||
})
|
||||
else {
|
||||
mark_skipped_local_standard_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
CandidateFailureDiagnostic::header_rules_apply_failed(
|
||||
spec_metadata.api_format,
|
||||
provider_api_format,
|
||||
"standard_family_headers",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
let mut provider_request_headers = resolved_headers.headers;
|
||||
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(),
|
||||
);
|
||||
|
||||
Some(LocalStandardCandidatePayloadParts {
|
||||
auth_header: resolved_headers.auth_header,
|
||||
auth_value: resolved_headers.auth_value,
|
||||
mapped_model: prepared_candidate.mapped_model,
|
||||
provider_api_format: provider_api_format.to_string(),
|
||||
provider_request_body,
|
||||
provider_request_headers,
|
||||
upstream_url,
|
||||
upstream_is_stream,
|
||||
envelope_name: None,
|
||||
transport: Arc::clone(transport),
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn build_kiro_cross_format_payload_parts(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
original_body_json: &serde_json::Value,
|
||||
input: &LocalStandardDecisionInput,
|
||||
attempt: &LocalStandardCandidateAttempt,
|
||||
transport: &Arc<GatewayProviderTransportSnapshot>,
|
||||
provider_api_format: &str,
|
||||
mapped_model: String,
|
||||
auth_header: String,
|
||||
auth_value: String,
|
||||
claude_request_body: Value,
|
||||
upstream_is_stream: bool,
|
||||
kiro_auth: &KiroRequestAuth,
|
||||
) -> Option<LocalStandardCandidatePayloadParts> {
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let provider_request_body = match build_kiro_provider_request_body(
|
||||
&claude_request_body,
|
||||
&mapped_model,
|
||||
&kiro_auth.auth_config,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
) {
|
||||
Some(body) => body,
|
||||
None => {
|
||||
mark_skipped_local_standard_candidate_with_extra_data(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"provider_request_body_build_failed",
|
||||
request_body_build_failure_extra_data(
|
||||
&claude_request_body,
|
||||
provider_api_format,
|
||||
provider_api_format,
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let upstream_url = match build_kiro_cross_format_upstream_url(
|
||||
transport,
|
||||
&mapped_model,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
kiro_auth.auth_config.effective_api_region(),
|
||||
) {
|
||||
Some(url) => url,
|
||||
None => {
|
||||
mark_skipped_local_standard_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"upstream_url_missing",
|
||||
CandidateFailureDiagnostic::upstream_url_missing(
|
||||
provider_api_format,
|
||||
provider_api_format,
|
||||
"standard_family_kiro_url",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let provider_request_headers = match build_kiro_provider_headers(KiroProviderHeadersInput {
|
||||
headers: &parts.headers,
|
||||
provider_request_body: &provider_request_body,
|
||||
original_request_body: original_body_json,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
auth_header: &auth_header,
|
||||
auth_value: &auth_value,
|
||||
auth_config: &kiro_auth.auth_config,
|
||||
machine_id: kiro_auth.machine_id.as_str(),
|
||||
}) {
|
||||
Some(headers) => headers,
|
||||
None => {
|
||||
mark_skipped_local_standard_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
CandidateFailureDiagnostic::header_rules_apply_failed(
|
||||
provider_api_format,
|
||||
provider_api_format,
|
||||
"standard_family_kiro_headers",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(LocalStandardCandidatePayloadParts {
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
provider_api_format: provider_api_format.to_string(),
|
||||
provider_request_body,
|
||||
provider_request_headers,
|
||||
upstream_url,
|
||||
upstream_is_stream,
|
||||
envelope_name: Some(KIRO_ENVELOPE_NAME),
|
||||
transport: Arc::clone(transport),
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::ai_serving::{
|
||||
resolve_gemini_stream_spec as resolve_surface_stream_spec,
|
||||
resolve_gemini_sync_spec as resolve_surface_sync_spec,
|
||||
};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
use super::family::{
|
||||
maybe_build_stream_via_standard_family_payload, maybe_build_sync_via_standard_family_payload,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::normalize_gemini_request_to_openai_chat_request;
|
||||
|
||||
pub(crate) fn resolve_sync_spec(plan_kind: &str) -> Option<super::family::LocalStandardSpec> {
|
||||
resolve_surface_sync_spec(plan_kind)
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_stream_spec(plan_kind: &str) -> Option<super::family::LocalStandardSpec> {
|
||||
resolve_surface_stream_spec(plan_kind)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_sync_local_gemini_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
maybe_build_sync_via_standard_family_payload(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
plan_kind,
|
||||
resolve_sync_spec,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_stream_local_gemini_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
maybe_build_stream_via_standard_family_payload(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
plan_kind,
|
||||
resolve_stream_spec,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::RequestBody;
|
||||
|
||||
use super::{
|
||||
augment_sync_report_context, build_ai_execution_plan_from_decision,
|
||||
generic_decision_missing_exact_provider_request, take_ai_decision_plan_core,
|
||||
take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts,
|
||||
AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
build_standard_plan_fallback_headers, StandardPlanFallbackAcceptPolicy,
|
||||
StandardPlanFallbackHeadersInput,
|
||||
};
|
||||
use crate::{AiExecutionDecision, GatewayError};
|
||||
|
||||
pub(crate) fn build_gemini_sync_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
_body_json: &serde_json::Value,
|
||||
payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
if generic_decision_missing_exact_provider_request(&payload) {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(core) = take_ai_decision_plan_core(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(url) = take_non_empty_string(&mut payload.upstream_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(auth_pair) = take_ai_upstream_auth_pair(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(provider_request_body_value) = payload.provider_request_body.take() else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let mut provider_request_headers =
|
||||
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
|
||||
request_headers: &parts.headers,
|
||||
existing_provider_request_headers: std::mem::take(
|
||||
&mut payload.provider_request_headers,
|
||||
),
|
||||
auth_header: auth_pair.as_ref().map(|pair| pair.header.as_str()),
|
||||
auth_value: auth_pair.as_ref().map(|pair| pair.value.as_str()),
|
||||
extra_headers: &BTreeMap::new(),
|
||||
content_type: payload.content_type.as_deref(),
|
||||
provider_api_format: core.provider_api_format.as_str(),
|
||||
client_api_format: core.client_api_format.as_str(),
|
||||
upstream_is_stream: payload.upstream_is_stream,
|
||||
build_from_request_when_empty: false,
|
||||
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming,
|
||||
});
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| Some("application/json".to_string()));
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context.take(),
|
||||
&provider_request_headers,
|
||||
&provider_request_body_value,
|
||||
)?;
|
||||
let stream = payload.upstream_is_stream;
|
||||
let plan = build_ai_execution_plan_from_decision(
|
||||
&mut payload,
|
||||
AiExecutionPlanFromDecisionParts {
|
||||
core,
|
||||
method: "POST".to_string(),
|
||||
url,
|
||||
headers: std::mem::take(&mut provider_request_headers),
|
||||
content_type,
|
||||
body: RequestBody::from_json(provider_request_body_value),
|
||||
stream,
|
||||
},
|
||||
);
|
||||
|
||||
Ok(Some(AiSyncAttempt {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn build_gemini_stream_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
_body_json: &serde_json::Value,
|
||||
payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
if generic_decision_missing_exact_provider_request(&payload) {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(core) = take_ai_decision_plan_core(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(url) = take_non_empty_string(&mut payload.upstream_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(auth_pair) = take_ai_upstream_auth_pair(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(provider_request_body_value) = payload.provider_request_body.take() else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let mut provider_request_headers =
|
||||
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
|
||||
request_headers: &parts.headers,
|
||||
existing_provider_request_headers: std::mem::take(
|
||||
&mut payload.provider_request_headers,
|
||||
),
|
||||
auth_header: auth_pair.as_ref().map(|pair| pair.header.as_str()),
|
||||
auth_value: auth_pair.as_ref().map(|pair| pair.value.as_str()),
|
||||
extra_headers: &BTreeMap::new(),
|
||||
content_type: payload.content_type.as_deref(),
|
||||
provider_api_format: core.provider_api_format.as_str(),
|
||||
client_api_format: core.client_api_format.as_str(),
|
||||
upstream_is_stream: true,
|
||||
build_from_request_when_empty: false,
|
||||
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamRequired,
|
||||
});
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| Some("application/json".to_string()));
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context.take(),
|
||||
&provider_request_headers,
|
||||
&provider_request_body_value,
|
||||
)?;
|
||||
let plan = build_ai_execution_plan_from_decision(
|
||||
&mut payload,
|
||||
AiExecutionPlanFromDecisionParts {
|
||||
core,
|
||||
method: "POST".to_string(),
|
||||
url,
|
||||
headers: std::mem::take(&mut provider_request_headers),
|
||||
content_type,
|
||||
body: RequestBody::from_json(provider_request_body_value),
|
||||
stream: true,
|
||||
},
|
||||
);
|
||||
|
||||
Ok(Some(AiStreamAttempt {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
320
apps/aether-gateway/src/ai_serving/planner/standard/mod.rs
Normal file
320
apps/aether-gateway/src/ai_serving/planner/standard/mod.rs
Normal file
@@ -0,0 +1,320 @@
|
||||
//! Standard contract planning surface.
|
||||
//!
|
||||
//! This groups the standard planning surface in one place:
|
||||
//! request-side conversion, matrix registry, and decision payload builders.
|
||||
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
mod claude;
|
||||
mod codex;
|
||||
mod family;
|
||||
mod gemini;
|
||||
mod normalize;
|
||||
mod openai;
|
||||
|
||||
pub(crate) use self::codex::apply_codex_openai_responses_special_headers;
|
||||
pub(crate) use self::family::{
|
||||
build_local_stream_plan_and_reports, build_local_sync_plan_and_reports,
|
||||
};
|
||||
pub(crate) use self::normalize::{
|
||||
build_cross_format_openai_chat_request_body, build_cross_format_openai_chat_upstream_url,
|
||||
build_cross_format_openai_responses_request_body,
|
||||
build_cross_format_openai_responses_upstream_url, build_local_openai_chat_request_body,
|
||||
build_local_openai_chat_upstream_url, build_local_openai_responses_request_body,
|
||||
build_local_openai_responses_upstream_url,
|
||||
};
|
||||
pub(crate) use self::openai::{
|
||||
build_local_openai_chat_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_sync_plan_and_reports_for_kind,
|
||||
build_local_openai_responses_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_responses_sync_plan_and_reports_for_kind, copy_request_number_field,
|
||||
copy_request_number_field_as, map_openai_reasoning_effort_to_claude_output,
|
||||
map_openai_reasoning_effort_to_gemini_budget, maybe_build_stream_local_decision_payload,
|
||||
maybe_build_stream_local_openai_responses_decision_payload,
|
||||
maybe_build_sync_local_decision_payload,
|
||||
maybe_build_sync_local_openai_responses_decision_payload, parse_openai_stop_sequences,
|
||||
resolve_openai_chat_max_tokens, set_local_openai_chat_execution_exhausted_diagnostic,
|
||||
value_as_u64,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::normalize_standard_request_to_openai_chat_request;
|
||||
pub(crate) use crate::ai_serving::{
|
||||
build_core_error_body_for_client_format, request_conversion_kind,
|
||||
request_conversion_transport_supported, sync_chat_response_conversion_kind,
|
||||
sync_cli_response_conversion_kind, RequestConversionKind, SyncChatResponseConversionKind,
|
||||
SyncCliResponseConversionKind,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::{
|
||||
build_standard_request_body, convert_openai_chat_request_to_claude_request,
|
||||
convert_openai_chat_request_to_gemini_request,
|
||||
convert_openai_chat_request_to_openai_responses_request, extract_openai_text_content,
|
||||
normalize_openai_responses_request_to_openai_chat_request, parse_openai_tool_result_content,
|
||||
};
|
||||
pub(crate) use aether_ai_serving::{
|
||||
request_body_build_failure_extra_data, same_format_provider_request_body_failure_extra_data,
|
||||
};
|
||||
|
||||
pub(crate) fn build_standard_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &crate::ai_serving::GatewayProviderTransportSnapshot,
|
||||
mapped_model: &str,
|
||||
provider_api_format: &str,
|
||||
upstream_is_stream: bool,
|
||||
) -> Option<String> {
|
||||
crate::ai_serving::build_provider_transport_request_url(
|
||||
transport,
|
||||
provider_api_format,
|
||||
Some(mapped_model),
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_sync_local_standard_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
if let Some(payload) = self::claude::maybe_build_sync_local_claude_decision_payload(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
self::gemini::maybe_build_sync_local_gemini_decision_payload(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_stream_local_standard_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
if let Some(payload) = self::claude::maybe_build_stream_local_claude_decision_payload(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
self::gemini::maybe_build_stream_local_gemini_decision_payload(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::build_standard_request_body;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn builds_openai_chat_request_from_claude_chat_source() {
|
||||
let request = json!({
|
||||
"model": "claude-3-7-sonnet",
|
||||
"system": "You are concise.",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Hello from Claude"}]
|
||||
}
|
||||
],
|
||||
"max_tokens": 128
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"claude:messages",
|
||||
"gpt-5",
|
||||
"openai",
|
||||
"openai:chat",
|
||||
"/v1/messages",
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude chat should convert to openai chat");
|
||||
|
||||
assert_eq!(converted["model"], "gpt-5");
|
||||
assert_eq!(converted["messages"][0]["role"], "system");
|
||||
assert_eq!(converted["messages"][0]["content"], "You are concise.");
|
||||
assert_eq!(converted["messages"][1]["role"], "user");
|
||||
assert_eq!(converted["messages"][1]["content"], "Hello from Claude");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_claude_chat_request_from_gemini_chat_source() {
|
||||
let request = json!({
|
||||
"systemInstruction": {
|
||||
"parts": [{"text": "Be brief."}]
|
||||
},
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [{"text": "Hello from Gemini"}]
|
||||
}
|
||||
]
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"gemini:generate_content",
|
||||
"claude-sonnet-4-5",
|
||||
"anthropic",
|
||||
"claude:messages",
|
||||
"/v1beta/models/gemini-2.5-pro:generateContent",
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("gemini chat should convert to claude chat");
|
||||
|
||||
assert_eq!(converted["model"], "claude-sonnet-4-5");
|
||||
assert_eq!(converted["messages"][0]["role"], "user");
|
||||
assert!(
|
||||
converted["messages"]
|
||||
.to_string()
|
||||
.contains("Hello from Gemini"),
|
||||
"converted claude payload should retain the gemini user text: {converted}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_gemini_cli_request_from_claude_cli_source() {
|
||||
let request = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Need CLI output"}]
|
||||
}
|
||||
],
|
||||
"max_tokens": 64
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"claude:messages",
|
||||
"gemini-2.5-pro",
|
||||
"google",
|
||||
"gemini:generate_content",
|
||||
"/v1/messages",
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude cli should convert to gemini cli");
|
||||
|
||||
assert_eq!(converted["contents"][0]["role"], "user");
|
||||
assert_eq!(
|
||||
converted["contents"][0]["parts"][0]["text"],
|
||||
"Need CLI output"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_openai_responses_request_from_claude_cli_source_with_forced_stream() {
|
||||
let request = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Need OpenAI Responses output"}]
|
||||
}
|
||||
],
|
||||
"max_tokens": 64
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"claude:messages",
|
||||
"gpt-5",
|
||||
"openai",
|
||||
"openai:responses",
|
||||
"/v1/messages",
|
||||
true,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude cli should convert to openai responses");
|
||||
|
||||
assert_eq!(converted["model"], "gpt-5");
|
||||
assert_eq!(converted["input"][0]["role"], "user");
|
||||
assert_eq!(converted["input"][0]["content"][0]["type"], "input_text");
|
||||
assert_eq!(
|
||||
converted["input"][0]["content"][0]["text"],
|
||||
"Need OpenAI Responses output"
|
||||
);
|
||||
assert_eq!(converted["stream"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_metadata_for_codex_openai_responses_requests() {
|
||||
let request = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"metadata": {"trace_id": "abc"},
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Need OpenAI Responses output"}]
|
||||
}],
|
||||
"max_tokens": 64
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"claude:messages",
|
||||
"gpt-5",
|
||||
"codex",
|
||||
"openai:responses",
|
||||
"/v1/messages",
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude cli should convert to codex request");
|
||||
|
||||
assert!(converted.get("metadata").is_none());
|
||||
assert_eq!(converted["store"], false);
|
||||
assert_eq!(converted["instructions"], "You are ChatGPT.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_store_for_openai_responses_compact_requests() {
|
||||
let request = json!({
|
||||
"model": "gpt-5",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": "Hello from OpenAI Chat"
|
||||
}],
|
||||
"store": true
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"openai:chat",
|
||||
"gpt-5",
|
||||
"openai",
|
||||
"openai:responses:compact",
|
||||
"/v1/chat/completions",
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("openai chat should convert to openai responses compact");
|
||||
|
||||
assert_eq!(converted["model"], "gpt-5");
|
||||
assert!(converted.get("store").is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
#[path = "normalize/chat.rs"]
|
||||
mod chat;
|
||||
#[path = "normalize/responses.rs"]
|
||||
mod responses;
|
||||
#[cfg(test)]
|
||||
#[path = "normalize/tests.rs"]
|
||||
mod tests;
|
||||
|
||||
pub(crate) use self::chat::{
|
||||
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,
|
||||
};
|
||||
pub(crate) use self::responses::{
|
||||
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,
|
||||
};
|
||||
@@ -0,0 +1,75 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::transport::apply_standard_provider_request_body_rules;
|
||||
use crate::ai_serving::{
|
||||
apply_codex_openai_responses_special_body_edits,
|
||||
apply_openai_responses_compact_special_body_edits,
|
||||
build_cross_format_openai_chat_request_body as surface_build_cross_format_openai_chat_request_body,
|
||||
build_local_openai_chat_request_body as surface_build_local_openai_chat_request_body,
|
||||
GatewayProviderTransportSnapshot,
|
||||
};
|
||||
|
||||
pub(crate) fn build_local_openai_chat_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
upstream_is_stream: bool,
|
||||
body_rules: Option<&Value>,
|
||||
) -> Option<Value> {
|
||||
let provider_request_body =
|
||||
surface_build_local_openai_chat_request_body(body_json, mapped_model, upstream_is_stream)?;
|
||||
apply_standard_provider_request_body_rules(provider_request_body, body_rules, body_json)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_openai_chat_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<String> {
|
||||
crate::ai_serving::transport::build_local_openai_chat_upstream_url(transport, parts.uri.query())
|
||||
}
|
||||
|
||||
pub(crate) fn build_cross_format_openai_chat_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
upstream_is_stream: bool,
|
||||
body_rules: Option<&Value>,
|
||||
user_api_key_id: Option<&str>,
|
||||
) -> Option<Value> {
|
||||
let provider_request_body = surface_build_cross_format_openai_chat_request_body(
|
||||
body_json,
|
||||
mapped_model,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
)?;
|
||||
let mut provider_request_body =
|
||||
apply_standard_provider_request_body_rules(provider_request_body, body_rules, body_json)?;
|
||||
apply_codex_openai_responses_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_type,
|
||||
provider_api_format,
|
||||
body_rules,
|
||||
user_api_key_id,
|
||||
);
|
||||
apply_openai_responses_compact_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_api_format,
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
pub(crate) fn build_cross_format_openai_chat_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
mapped_model: &str,
|
||||
provider_api_format: &str,
|
||||
upstream_is_stream: bool,
|
||||
) -> Option<String> {
|
||||
crate::ai_serving::transport::build_cross_format_openai_chat_upstream_url(
|
||||
transport,
|
||||
mapped_model,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::transport::apply_standard_provider_request_body_rules;
|
||||
use crate::ai_serving::{
|
||||
apply_codex_openai_responses_special_body_edits,
|
||||
apply_openai_responses_compact_special_body_edits,
|
||||
build_cross_format_openai_responses_request_body as surface_build_cross_format_openai_responses_request_body,
|
||||
build_local_openai_responses_request_body as surface_build_local_openai_responses_request_body,
|
||||
GatewayProviderTransportSnapshot,
|
||||
};
|
||||
|
||||
pub(crate) fn build_local_openai_responses_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
require_streaming: bool,
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
body_rules: Option<&Value>,
|
||||
user_api_key_id: Option<&str>,
|
||||
) -> Option<Value> {
|
||||
let provider_request_body = surface_build_local_openai_responses_request_body(
|
||||
body_json,
|
||||
mapped_model,
|
||||
require_streaming,
|
||||
)?;
|
||||
let mut provider_request_body =
|
||||
apply_standard_provider_request_body_rules(provider_request_body, body_rules, body_json)?;
|
||||
apply_codex_openai_responses_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_type,
|
||||
provider_api_format,
|
||||
body_rules,
|
||||
user_api_key_id,
|
||||
);
|
||||
apply_openai_responses_compact_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_api_format,
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
pub(crate) fn build_cross_format_openai_responses_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
upstream_is_stream: bool,
|
||||
provider_type: &str,
|
||||
body_rules: Option<&Value>,
|
||||
user_api_key_id: Option<&str>,
|
||||
) -> Option<Value> {
|
||||
let provider_request_body = surface_build_cross_format_openai_responses_request_body(
|
||||
body_json,
|
||||
mapped_model,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
)?;
|
||||
let mut provider_request_body =
|
||||
apply_standard_provider_request_body_rules(provider_request_body, body_rules, body_json)?;
|
||||
apply_codex_openai_responses_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_type,
|
||||
provider_api_format,
|
||||
body_rules,
|
||||
user_api_key_id,
|
||||
);
|
||||
apply_openai_responses_compact_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_api_format,
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_openai_responses_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
compact: bool,
|
||||
) -> Option<String> {
|
||||
crate::ai_serving::transport::build_local_openai_responses_upstream_url(
|
||||
transport,
|
||||
compact,
|
||||
parts.uri.query(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_cross_format_openai_responses_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
mapped_model: &str,
|
||||
client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
upstream_is_stream: bool,
|
||||
) -> Option<String> {
|
||||
crate::ai_serving::transport::build_cross_format_openai_responses_upstream_url(
|
||||
transport,
|
||||
mapped_model,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,301 @@
|
||||
use aether_provider_transport::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
use http::Request;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::{
|
||||
build_cross_format_openai_responses_request_body, build_local_openai_responses_request_body,
|
||||
build_local_openai_responses_upstream_url,
|
||||
};
|
||||
|
||||
fn object_keys(value: &Value) -> Vec<&str> {
|
||||
value
|
||||
.as_object()
|
||||
.expect("json object")
|
||||
.keys()
|
||||
.map(String::as_str)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn sample_transport(base_url: &str, api_format: &str) -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-codex".to_string(),
|
||||
name: "codex".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: false,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-codex".to_string(),
|
||||
provider_id: "provider-codex".to_string(),
|
||||
api_format: api_format.to_string(),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
is_active: true,
|
||||
base_url: base_url.to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-codex".to_string(),
|
||||
provider_id: "provider-codex".to_string(),
|
||||
name: "oauth".to_string(),
|
||||
auth_type: "oauth".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec![api_format.to_string()]),
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
decrypted_api_key: "__placeholder__".to_string(),
|
||||
decrypted_auth_config: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_openai_chat_cross_format_request_body_from_openai_responses_source() {
|
||||
let body_json = json!({
|
||||
"model": "gpt-5",
|
||||
"input": "hello",
|
||||
});
|
||||
|
||||
let provider_request_body = build_cross_format_openai_responses_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"openai:responses",
|
||||
"openai:chat",
|
||||
false,
|
||||
"openai",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("openai responses to openai chat body should build");
|
||||
|
||||
assert_eq!(provider_request_body["model"], "gpt-5-upstream");
|
||||
assert_eq!(provider_request_body["messages"][0]["role"], "user");
|
||||
assert_eq!(provider_request_body["messages"][0]["content"], "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_openai_responses_wrapper_preserves_body_order_after_edits() {
|
||||
let body_json: Value = serde_json::from_str(
|
||||
r#"{
|
||||
"text": {"format": {"type": "text"}},
|
||||
"input": [],
|
||||
"model": "gpt-5.4",
|
||||
"store": false,
|
||||
"tools": [],
|
||||
"stream": true,
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
"reasoning": {"effort": "high"},
|
||||
"tool_choice": "auto"
|
||||
}"#,
|
||||
)
|
||||
.expect("request body should parse");
|
||||
|
||||
let provider_request_body = build_local_openai_responses_request_body(
|
||||
&body_json,
|
||||
"gpt-5.4",
|
||||
true,
|
||||
"codex",
|
||||
"openai:responses",
|
||||
None,
|
||||
Some("key-123"),
|
||||
)
|
||||
.expect("local openai responses body should build");
|
||||
|
||||
assert_eq!(
|
||||
object_keys(&provider_request_body),
|
||||
vec![
|
||||
"text",
|
||||
"input",
|
||||
"model",
|
||||
"store",
|
||||
"tools",
|
||||
"stream",
|
||||
"include",
|
||||
"reasoning",
|
||||
"tool_choice",
|
||||
"instructions",
|
||||
"prompt_cache_key",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_openai_responses_compact_wrapper_strips_store_for_same_format_requests() {
|
||||
let body_json = json!({
|
||||
"model": "gpt-5.4",
|
||||
"input": [],
|
||||
"store": true
|
||||
});
|
||||
|
||||
let provider_request_body = build_local_openai_responses_request_body(
|
||||
&body_json,
|
||||
"gpt-5.4",
|
||||
false,
|
||||
"openai",
|
||||
"openai:responses:compact",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("local openai compact body should build");
|
||||
|
||||
assert!(provider_request_body.get("store").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_openai_responses_upstream_url_preserves_codex_base_path() {
|
||||
let request = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/responses")
|
||||
.body(())
|
||||
.expect("request should build");
|
||||
let (parts, _) = request.into_parts();
|
||||
|
||||
let upstream_url = build_local_openai_responses_upstream_url(
|
||||
&parts,
|
||||
&sample_transport("https://tiger.bookapi.cc/codex", "openai:responses"),
|
||||
false,
|
||||
)
|
||||
.expect("openai responses upstream url should build");
|
||||
|
||||
assert_eq!(upstream_url, "https://tiger.bookapi.cc/codex/responses");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_metadata_for_codex_openai_responses_requests() {
|
||||
let body_json = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"metadata": {"trace_id": "abc"},
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "hello"}]
|
||||
}],
|
||||
});
|
||||
|
||||
let provider_request_body = build_cross_format_openai_responses_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"claude:messages",
|
||||
"openai:responses",
|
||||
true,
|
||||
"codex",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude cli to codex request should build");
|
||||
|
||||
assert!(provider_request_body.get("metadata").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn applies_codex_defaults_unless_body_rules_handle_the_field() {
|
||||
let body_json = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "hello"}]
|
||||
}],
|
||||
"metadata": {"trace_id": "abc"},
|
||||
"store": true
|
||||
});
|
||||
let body_rules = json!([
|
||||
{"action":"set","path":"store","value":true},
|
||||
{"action":"set","path":"instructions","value":"Custom instructions"},
|
||||
{"action":"set","path":"metadata","value":{"trace_id":"keep-me"}}
|
||||
]);
|
||||
|
||||
let provider_request_body = build_cross_format_openai_responses_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"claude:messages",
|
||||
"openai:responses",
|
||||
true,
|
||||
"codex",
|
||||
Some(&body_rules),
|
||||
None,
|
||||
)
|
||||
.expect("claude cli to codex request should build");
|
||||
|
||||
assert_eq!(provider_request_body["store"], true);
|
||||
assert_eq!(provider_request_body["instructions"], "Custom instructions");
|
||||
assert_eq!(provider_request_body["metadata"]["trace_id"], "keep-me");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn injects_codex_prompt_cache_key_for_openai_responses_cross_format_requests() {
|
||||
let body_json = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "hello"}]
|
||||
}],
|
||||
});
|
||||
|
||||
let provider_request_body = build_cross_format_openai_responses_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"claude:messages",
|
||||
"openai:responses",
|
||||
true,
|
||||
"codex",
|
||||
None,
|
||||
Some("key-123"),
|
||||
)
|
||||
.expect("claude cli to codex request should build");
|
||||
|
||||
assert_eq!(
|
||||
provider_request_body["prompt_cache_key"],
|
||||
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn injects_codex_prompt_cache_key_for_openai_chat_cross_format_requests() {
|
||||
let body_json = json!({
|
||||
"model": "gpt-5",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": "hello"
|
||||
}],
|
||||
});
|
||||
|
||||
let provider_request_body = super::build_cross_format_openai_chat_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"codex",
|
||||
"openai:responses",
|
||||
false,
|
||||
None,
|
||||
Some("key-123"),
|
||||
)
|
||||
.expect("openai chat to codex request should build");
|
||||
|
||||
assert_eq!(
|
||||
provider_request_body["prompt_cache_key"],
|
||||
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
#[path = "decision/payload.rs"]
|
||||
mod payload;
|
||||
#[path = "decision/request.rs"]
|
||||
mod request;
|
||||
#[path = "decision/support.rs"]
|
||||
mod support;
|
||||
|
||||
pub(super) use self::payload::maybe_build_local_openai_chat_decision_payload_for_candidate;
|
||||
pub(super) use self::support::{
|
||||
materialize_local_openai_chat_candidate_attempts, LocalOpenAiChatCandidateAttempt,
|
||||
LocalOpenAiChatDecisionInput,
|
||||
};
|
||||
@@ -0,0 +1,182 @@
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, insert_provider_stream_event_api_format,
|
||||
LocalExecutionReportContextParts,
|
||||
};
|
||||
use crate::ai_serving::planner::{
|
||||
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, append_local_failover_policy_to_value,
|
||||
AiExecutionDecision, AppState,
|
||||
};
|
||||
|
||||
use super::request::resolve_local_openai_chat_candidate_payload_parts;
|
||||
use super::support::{LocalOpenAiChatCandidateAttempt, LocalOpenAiChatDecisionInput};
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
body_json: &serde_json::Value,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
attempt: LocalOpenAiChatCandidateAttempt,
|
||||
decision_kind: &str,
|
||||
report_kind: &str,
|
||||
upstream_is_stream: bool,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
let attempt_identity = attempt.attempt_identity();
|
||||
let LocalOpenAiChatCandidateAttempt {
|
||||
eligible,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
..
|
||||
} = attempt;
|
||||
let resolved = resolve_local_openai_chat_candidate_payload_parts(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
input,
|
||||
&eligible,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
decision_kind,
|
||||
report_kind,
|
||||
upstream_is_stream,
|
||||
)
|
||||
.await?;
|
||||
let candidate = &eligible.candidate;
|
||||
|
||||
let prompt_cache_key = resolved
|
||||
.provider_request_body
|
||||
.get("prompt_cache_key")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let proxy = state
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&resolved.transport)
|
||||
.await;
|
||||
let tls_profile = resolve_transport_tls_profile(&resolved.transport);
|
||||
let timeouts = resolve_transport_execution_timeouts(&resolved.transport);
|
||||
let mut extra_fields = serde_json::Map::new();
|
||||
if let Some(proxy_value) =
|
||||
build_request_trace_proxy_value(Some(&resolved.transport), proxy.as_ref())
|
||||
{
|
||||
extra_fields.insert("proxy".to_string(), proxy_value);
|
||||
}
|
||||
if let Some(envelope_name) = resolved.envelope_name {
|
||||
extra_fields.insert(
|
||||
"envelope_name".to_string(),
|
||||
serde_json::Value::String(envelope_name.to_string()),
|
||||
);
|
||||
}
|
||||
insert_provider_stream_event_api_format(
|
||||
&mut extra_fields,
|
||||
resolved.transport.provider.provider_type.as_str(),
|
||||
);
|
||||
let report_context = append_local_failover_policy_to_value(
|
||||
append_execution_contract_fields_to_value(
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &input.auth_context,
|
||||
request_id: trace_id,
|
||||
candidate_id: &candidate_id,
|
||||
attempt_identity,
|
||||
model: &input.requested_model,
|
||||
provider_name: &resolved.transport.provider.name,
|
||||
provider_id: &candidate.provider_id,
|
||||
endpoint_id: &candidate.endpoint_id,
|
||||
key_id: &candidate.key_id,
|
||||
key_name: Some(&candidate.key_name),
|
||||
model_id: Some(&candidate.model_id),
|
||||
global_model_id: Some(&candidate.global_model_id),
|
||||
global_model_name: Some(&candidate.global_model_name),
|
||||
provider_api_format: &resolved.provider_api_format,
|
||||
client_api_format: "openai:chat",
|
||||
mapped_model: Some(&resolved.mapped_model),
|
||||
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
|
||||
ranking: eligible.ranking.as_ref(),
|
||||
upstream_url: Some(&resolved.upstream_url),
|
||||
header_rules: resolved.transport.endpoint.header_rules.as_ref(),
|
||||
body_rules: resolved.transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: Some(serde_json::Value::Null),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: body_json
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
upstream_is_stream,
|
||||
has_envelope: resolved.envelope_name.is_some(),
|
||||
needs_conversion: matches!(
|
||||
resolved.conversion_mode,
|
||||
crate::ai_serving::ConversionMode::Bidirectional
|
||||
),
|
||||
extra_fields,
|
||||
}),
|
||||
resolved.execution_strategy,
|
||||
resolved.conversion_mode,
|
||||
"openai:chat",
|
||||
candidate.endpoint_api_format.as_str(),
|
||||
),
|
||||
&resolved.transport,
|
||||
);
|
||||
let super::request::LocalOpenAiChatCandidatePayloadParts {
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
provider_api_format,
|
||||
provider_request_body,
|
||||
provider_request_headers,
|
||||
upstream_url,
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
report_kind,
|
||||
envelope_name: _,
|
||||
transport,
|
||||
} = resolved;
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: upstream_is_stream,
|
||||
decision_kind: decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: None,
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format,
|
||||
client_api_format: "openai:chat".to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts,
|
||||
upstream_is_stream,
|
||||
report_kind: Some(report_kind),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
}
|
||||
@@ -0,0 +1,609 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::planner::candidate_preparation::{
|
||||
prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth,
|
||||
OauthPreparationContext,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::common::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,
|
||||
};
|
||||
use crate::ai_serving::transport::auth::resolve_local_openai_bearer_auth;
|
||||
use crate::ai_serving::transport::kiro::{
|
||||
build_kiro_provider_headers, build_kiro_provider_request_body,
|
||||
is_kiro_claude_messages_transport, KiroProviderHeadersInput, KiroRequestAuth,
|
||||
KIRO_ENVELOPE_NAME,
|
||||
};
|
||||
use crate::ai_serving::transport::local_openai_chat_transport_unsupported_reason;
|
||||
use crate::ai_serving::transport::{
|
||||
build_kiro_cross_format_upstream_url, build_standard_provider_request_headers,
|
||||
StandardProviderRequestHeadersInput,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
ai_local_execution_contract_for_formats, request_conversion_direct_auth,
|
||||
request_conversion_kind, CandidateFailureDiagnostic, GatewayProviderTransportSnapshot,
|
||||
LocalResolvedOAuthRequestAuth,
|
||||
};
|
||||
use crate::ai_serving::{ConversionMode, ExecutionStrategy};
|
||||
use crate::AppState;
|
||||
|
||||
use super::support::{
|
||||
mark_skipped_local_openai_chat_candidate,
|
||||
mark_skipped_local_openai_chat_candidate_with_extra_data,
|
||||
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic, LocalOpenAiChatDecisionInput,
|
||||
};
|
||||
|
||||
pub(crate) struct LocalOpenAiChatCandidatePayloadParts {
|
||||
pub(super) auth_header: String,
|
||||
pub(super) auth_value: String,
|
||||
pub(super) mapped_model: String,
|
||||
pub(super) provider_api_format: String,
|
||||
pub(super) provider_request_body: Value,
|
||||
pub(super) provider_request_headers: BTreeMap<String, String>,
|
||||
pub(super) upstream_url: String,
|
||||
pub(super) execution_strategy: ExecutionStrategy,
|
||||
pub(super) conversion_mode: ConversionMode,
|
||||
pub(super) report_kind: String,
|
||||
pub(super) envelope_name: Option<&'static str>,
|
||||
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn resolve_local_openai_chat_candidate_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,
|
||||
decision_kind: &str,
|
||||
report_kind: &str,
|
||||
upstream_is_stream: bool,
|
||||
) -> Option<LocalOpenAiChatCandidatePayloadParts> {
|
||||
let planner_state = crate::ai_serving::PlannerAppState::new(state);
|
||||
let candidate = &eligible.candidate;
|
||||
let provider_api_format = eligible.provider_api_format.as_str();
|
||||
let transport = &eligible.transport;
|
||||
|
||||
if provider_api_format == "openai:chat" {
|
||||
if let Some(skip_reason) = local_openai_chat_transport_unsupported_reason(transport) {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let prepared_candidate = match prepare_header_authenticated_candidate(
|
||||
planner_state,
|
||||
transport,
|
||||
candidate,
|
||||
resolve_local_openai_bearer_auth(transport),
|
||||
OauthPreparationContext {
|
||||
trace_id,
|
||||
api_format: "openai:chat",
|
||||
operation: "openai_chat_same_format",
|
||||
},
|
||||
)
|
||||
.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 None;
|
||||
}
|
||||
};
|
||||
|
||||
let Some(provider_request_body) = build_local_openai_chat_request_body(
|
||||
body_json,
|
||||
&prepared_candidate.mapped_model,
|
||||
upstream_is_stream,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
) 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 None;
|
||||
};
|
||||
|
||||
let Some(upstream_url) = build_local_openai_chat_upstream_url(parts, transport) else {
|
||||
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"upstream_url_missing",
|
||||
CandidateFailureDiagnostic::upstream_url_missing(
|
||||
"openai:chat",
|
||||
provider_api_format,
|
||||
"openai_chat_same_format_url",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(resolved_headers) =
|
||||
build_standard_provider_request_headers(StandardProviderRequestHeadersInput {
|
||||
transport,
|
||||
provider_api_format,
|
||||
same_format: true,
|
||||
headers: &parts.headers,
|
||||
auth_header: &prepared_candidate.auth_header,
|
||||
auth_value: &prepared_candidate.auth_value,
|
||||
extra_headers: &BTreeMap::new(),
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
provider_request_body: &provider_request_body,
|
||||
original_request_body: body_json,
|
||||
upstream_is_stream,
|
||||
})
|
||||
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_same_format_headers",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
let mut provider_request_headers = resolved_headers.headers;
|
||||
apply_codex_openai_responses_special_headers(
|
||||
&mut provider_request_headers,
|
||||
&provider_request_body,
|
||||
&parts.headers,
|
||||
transport.provider.provider_type.as_str(),
|
||||
transport.endpoint.api_format.as_str(),
|
||||
Some(trace_id),
|
||||
transport.key.decrypted_auth_config.as_deref(),
|
||||
);
|
||||
|
||||
let (execution_strategy, conversion_mode) =
|
||||
ai_local_execution_contract_for_formats("openai:chat", "openai:chat");
|
||||
return Some(LocalOpenAiChatCandidatePayloadParts {
|
||||
auth_header: resolved_headers.auth_header,
|
||||
auth_value: resolved_headers.auth_value,
|
||||
mapped_model: prepared_candidate.mapped_model,
|
||||
provider_api_format: "openai:chat".to_string(),
|
||||
provider_request_body,
|
||||
provider_request_headers,
|
||||
upstream_url,
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
report_kind: report_kind.to_string(),
|
||||
envelope_name: None,
|
||||
transport: Arc::clone(transport),
|
||||
});
|
||||
}
|
||||
|
||||
let provider_api_format = provider_api_format.trim().to_ascii_lowercase();
|
||||
let Some(conversion_kind) =
|
||||
request_conversion_kind("openai:chat", provider_api_format.as_str())
|
||||
else {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_api_format_unsupported",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
if let Some(skip_reason) = crate::ai_serving::request_conversion_transport_unsupported_reason(
|
||||
transport,
|
||||
conversion_kind,
|
||||
) {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
let is_kiro_claude_cli =
|
||||
is_kiro_claude_messages_transport(transport, provider_api_format.as_str());
|
||||
let oauth_context = OauthPreparationContext {
|
||||
trace_id,
|
||||
api_format: provider_api_format.as_str(),
|
||||
operation: "openai_chat_cross_format",
|
||||
};
|
||||
let kiro_auth = if is_kiro_claude_cli {
|
||||
match crate::ai_serving::planner::candidate_preparation::resolve_candidate_oauth_auth(
|
||||
planner_state,
|
||||
transport,
|
||||
oauth_context,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Some(LocalResolvedOAuthRequestAuth::Kiro(auth)) => Some(auth),
|
||||
_ => {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let prepared_candidate = if let Some(kiro_auth) = kiro_auth.as_ref() {
|
||||
match prepare_header_authenticated_candidate_from_auth(
|
||||
candidate,
|
||||
kiro_auth.name.to_string(),
|
||||
kiro_auth.value.clone(),
|
||||
) {
|
||||
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 None;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
match prepare_header_authenticated_candidate(
|
||||
planner_state,
|
||||
transport,
|
||||
candidate,
|
||||
request_conversion_direct_auth(transport, conversion_kind),
|
||||
oauth_context,
|
||||
)
|
||||
.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 None;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let Some(provider_request_body) = build_cross_format_openai_chat_request_body(
|
||||
body_json,
|
||||
&prepared_candidate.mapped_model,
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format.as_str(),
|
||||
upstream_is_stream,
|
||||
if is_kiro_claude_cli {
|
||||
None
|
||||
} else {
|
||||
transport.endpoint.body_rules.as_ref()
|
||||
},
|
||||
Some(input.auth_context.api_key_id.as_str()),
|
||||
) 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.as_str(),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
if let Some(kiro_auth) = kiro_auth.as_ref() {
|
||||
return build_kiro_openai_chat_cross_format_payload_parts(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
input,
|
||||
eligible,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
decision_kind,
|
||||
transport,
|
||||
provider_api_format.as_str(),
|
||||
prepared_candidate.mapped_model,
|
||||
prepared_candidate.auth_header,
|
||||
prepared_candidate.auth_value,
|
||||
provider_request_body,
|
||||
upstream_is_stream,
|
||||
kiro_auth,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let Some(upstream_url) = build_cross_format_openai_chat_upstream_url(
|
||||
parts,
|
||||
transport,
|
||||
&prepared_candidate.mapped_model,
|
||||
provider_api_format.as_str(),
|
||||
upstream_is_stream,
|
||||
) else {
|
||||
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"upstream_url_missing",
|
||||
CandidateFailureDiagnostic::upstream_url_missing(
|
||||
"openai:chat",
|
||||
provider_api_format.as_str(),
|
||||
"openai_chat_cross_format_url",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
let Some(resolved_headers) =
|
||||
build_standard_provider_request_headers(StandardProviderRequestHeadersInput {
|
||||
transport,
|
||||
provider_api_format: provider_api_format.as_str(),
|
||||
same_format: false,
|
||||
headers: &parts.headers,
|
||||
auth_header: &prepared_candidate.auth_header,
|
||||
auth_value: &prepared_candidate.auth_value,
|
||||
extra_headers: &BTreeMap::new(),
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
provider_request_body: &provider_request_body,
|
||||
original_request_body: body_json,
|
||||
upstream_is_stream,
|
||||
})
|
||||
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.as_str(),
|
||||
"openai_chat_cross_format_headers",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
let mut provider_request_headers = resolved_headers.headers;
|
||||
apply_codex_openai_responses_special_headers(
|
||||
&mut provider_request_headers,
|
||||
&provider_request_body,
|
||||
&parts.headers,
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format.as_str(),
|
||||
Some(trace_id),
|
||||
transport.key.decrypted_auth_config.as_deref(),
|
||||
);
|
||||
|
||||
let resolved_report_kind = if decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND {
|
||||
"openai_chat_stream_success".to_string()
|
||||
} else {
|
||||
"openai_chat_sync_finalize".to_string()
|
||||
};
|
||||
let (execution_strategy, conversion_mode) =
|
||||
ai_local_execution_contract_for_formats("openai:chat", provider_api_format.as_str());
|
||||
|
||||
Some(LocalOpenAiChatCandidatePayloadParts {
|
||||
auth_header: resolved_headers.auth_header,
|
||||
auth_value: resolved_headers.auth_value,
|
||||
mapped_model: prepared_candidate.mapped_model,
|
||||
provider_api_format,
|
||||
provider_request_body,
|
||||
provider_request_headers,
|
||||
upstream_url,
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
report_kind: resolved_report_kind,
|
||||
envelope_name: None,
|
||||
transport: Arc::clone(transport),
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn build_kiro_openai_chat_cross_format_payload_parts(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
original_body_json: &serde_json::Value,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
eligible: &EligibleLocalExecutionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
decision_kind: &str,
|
||||
transport: &Arc<GatewayProviderTransportSnapshot>,
|
||||
provider_api_format: &str,
|
||||
mapped_model: String,
|
||||
auth_header: String,
|
||||
auth_value: String,
|
||||
claude_request_body: Value,
|
||||
upstream_is_stream: bool,
|
||||
kiro_auth: &KiroRequestAuth,
|
||||
) -> Option<LocalOpenAiChatCandidatePayloadParts> {
|
||||
let candidate = &eligible.candidate;
|
||||
let provider_request_body = match build_kiro_provider_request_body(
|
||||
&claude_request_body,
|
||||
&mapped_model,
|
||||
&kiro_auth.auth_config,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
) {
|
||||
Some(body) => body,
|
||||
None => {
|
||||
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"provider_request_body_build_failed",
|
||||
CandidateFailureDiagnostic::envelope_build_failed(
|
||||
"openai:chat",
|
||||
provider_api_format,
|
||||
"openai_chat_kiro_envelope",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let upstream_url = match build_kiro_cross_format_upstream_url(
|
||||
transport,
|
||||
&mapped_model,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
kiro_auth.auth_config.effective_api_region(),
|
||||
) {
|
||||
Some(url) => url,
|
||||
None => {
|
||||
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"upstream_url_missing",
|
||||
CandidateFailureDiagnostic::upstream_url_missing(
|
||||
"openai:chat",
|
||||
provider_api_format,
|
||||
"openai_chat_kiro_url",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let provider_request_headers = match build_kiro_provider_headers(KiroProviderHeadersInput {
|
||||
headers: &parts.headers,
|
||||
provider_request_body: &provider_request_body,
|
||||
original_request_body: original_body_json,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
auth_header: &auth_header,
|
||||
auth_value: &auth_value,
|
||||
auth_config: &kiro_auth.auth_config,
|
||||
machine_id: kiro_auth.machine_id.as_str(),
|
||||
}) {
|
||||
Some(headers) => headers,
|
||||
None => {
|
||||
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_kiro_headers",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let resolved_report_kind = if decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND {
|
||||
"openai_chat_stream_success".to_string()
|
||||
} else {
|
||||
"openai_chat_sync_finalize".to_string()
|
||||
};
|
||||
let (execution_strategy, conversion_mode) =
|
||||
ai_local_execution_contract_for_formats("openai:chat", provider_api_format);
|
||||
|
||||
Some(LocalOpenAiChatCandidatePayloadParts {
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
provider_api_format: provider_api_format.to_string(),
|
||||
provider_request_body,
|
||||
provider_request_headers,
|
||||
upstream_url,
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
report_kind: resolved_report_kind,
|
||||
envelope_name: Some(KIRO_ENVELOPE_NAME),
|
||||
transport: Arc::clone(transport),
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic,
|
||||
materialize_local_execution_candidates_with_serving, LocalCandidateResolutionMode,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_metadata::{
|
||||
build_local_execution_candidate_contract_metadata,
|
||||
build_local_execution_candidate_contract_metadata_for_candidate,
|
||||
LocalExecutionCandidateMetadataParts,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
use crate::ai_serving::planner::CandidateFailureDiagnostic;
|
||||
use crate::ai_serving::{
|
||||
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
|
||||
ExecutionRuntimeAuthContext, PlannerAppState,
|
||||
};
|
||||
use crate::AppState;
|
||||
|
||||
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalOpenAiChatCandidateAttempt;
|
||||
pub(crate) use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput as LocalOpenAiChatDecisionInput;
|
||||
|
||||
pub(crate) async fn mark_skipped_local_openai_chat_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::OpenAiChatDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn mark_skipped_local_openai_chat_candidate_with_extra_data(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
extra_data: Option<serde_json::Value>,
|
||||
) {
|
||||
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::OpenAiChatDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate_with_extra_data(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
extra_data,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn mark_skipped_local_openai_chat_candidate_with_failure_diagnostic(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
diagnostic: CandidateFailureDiagnostic,
|
||||
) {
|
||||
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::OpenAiChatDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
diagnostic,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
body_json: &serde_json::Value,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
preselection_skipped: Vec<SkippedLocalExecutionCandidate>,
|
||||
) -> Vec<LocalOpenAiChatCandidateAttempt> {
|
||||
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 outcome = materialize_local_execution_candidates_with_serving(
|
||||
planner_state,
|
||||
trace_id,
|
||||
"openai:chat",
|
||||
Some(&input.requested_model),
|
||||
Some(&input.auth_snapshot),
|
||||
input.required_capabilities.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
persistence_policy,
|
||||
candidates,
|
||||
preselection_skipped,
|
||||
LocalCandidateResolutionMode::Standard,
|
||||
|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;
|
||||
|
||||
outcome.attempts
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
use serde_json::Value;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::common::{
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_serving::planner::runtime_miss::set_local_runtime_execution_exhausted_diagnostic;
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
mod decision;
|
||||
mod plans;
|
||||
|
||||
use self::decision::{
|
||||
materialize_local_openai_chat_candidate_attempts,
|
||||
maybe_build_local_openai_chat_decision_payload_for_candidate, LocalOpenAiChatDecisionInput,
|
||||
};
|
||||
use self::plans::{
|
||||
build_local_openai_chat_stream_plan_and_reports, build_local_openai_chat_sync_plan_and_reports,
|
||||
list_local_openai_chat_candidates, resolve_local_openai_chat_decision_input,
|
||||
set_local_openai_chat_miss_diagnostic,
|
||||
};
|
||||
|
||||
pub(crate) async fn build_local_openai_chat_sync_plan_and_reports_for_kind(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<crate::ai_serving::planner::plan_builders::AiSyncAttempt>, GatewayError> {
|
||||
build_local_openai_chat_sync_plan_and_reports(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_openai_chat_stream_plan_and_reports_for_kind(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<crate::ai_serving::planner::plan_builders::AiStreamAttempt>, GatewayError> {
|
||||
build_local_openai_chat_stream_plan_and_reports(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) fn set_local_openai_chat_execution_exhausted_diagnostic(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
body_json: &serde_json::Value,
|
||||
plan_count: usize,
|
||||
) {
|
||||
warn!(
|
||||
event_name = "local_openai_chat_candidates_exhausted",
|
||||
log_type = "event",
|
||||
trace_id = %trace_id,
|
||||
plan_kind,
|
||||
route_class = decision.route_class.as_deref().unwrap_or("passthrough"),
|
||||
route_family = decision.route_family.as_deref().unwrap_or("unknown"),
|
||||
candidate_count = plan_count,
|
||||
model = body_json.get("model").and_then(|value| value.as_str()).unwrap_or(""),
|
||||
"gateway local openai chat execution exhausted all candidates"
|
||||
);
|
||||
set_local_runtime_execution_exhausted_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
body_json.get("model").and_then(|value| value.as_str()),
|
||||
plan_count,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_sync_local_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
if plan_kind != OPENAI_CHAT_SYNC_PLAN_KIND {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let Some(input) = resolve_local_openai_chat_decision_input(
|
||||
state, trace_id, decision, body_json, plan_kind, false,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let (candidates, skipped_candidates) =
|
||||
match list_local_openai_chat_candidates(state, &input, false).await {
|
||||
Ok(value) => value,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "local_openai_chat_scheduler_selection_failed",
|
||||
log_type = "event",
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat sync decision scheduler selection failed"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
let attempts = materialize_local_openai_chat_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
candidates,
|
||||
skipped_candidates,
|
||||
)
|
||||
.await;
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
&input,
|
||||
attempt,
|
||||
OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
"openai_chat_sync_success",
|
||||
false,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_stream_local_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
if plan_kind != OPENAI_CHAT_STREAM_PLAN_KIND {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let Some(input) = resolve_local_openai_chat_decision_input(
|
||||
state, trace_id, decision, body_json, plan_kind, false,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let (candidates, skipped_candidates) =
|
||||
match list_local_openai_chat_candidates(state, &input, true).await {
|
||||
Ok(value) => value,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "local_openai_chat_scheduler_selection_failed",
|
||||
log_type = "event",
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat stream decision scheduler selection failed"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
let attempts = materialize_local_openai_chat_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
candidates,
|
||||
skipped_candidates,
|
||||
)
|
||||
.await;
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
&input,
|
||||
attempt,
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND,
|
||||
"openai_chat_stream_success",
|
||||
true,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
#[path = "plans/candidates.rs"]
|
||||
mod candidates;
|
||||
#[path = "plans/diagnostic.rs"]
|
||||
mod diagnostic;
|
||||
#[path = "plans/resolve.rs"]
|
||||
mod resolve;
|
||||
#[path = "plans/stream.rs"]
|
||||
mod stream;
|
||||
#[path = "plans/sync.rs"]
|
||||
mod sync;
|
||||
|
||||
pub(super) use self::candidates::list_local_openai_chat_candidates;
|
||||
pub(super) use self::diagnostic::set_local_openai_chat_miss_diagnostic;
|
||||
pub(super) use self::resolve::resolve_local_openai_chat_decision_input;
|
||||
pub(super) use self::stream::build_local_openai_chat_stream_plan_and_reports;
|
||||
pub(super) use self::sync::build_local_openai_chat_sync_plan_and_reports;
|
||||
@@ -0,0 +1,35 @@
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
use super::super::{GatewayError, LocalOpenAiChatDecisionInput};
|
||||
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::candidate_source::{
|
||||
preselect_local_execution_candidates_with_serving, LocalCandidatePreselectionKeyMode,
|
||||
};
|
||||
use crate::ai_serving::PlannerAppState;
|
||||
use crate::AppState;
|
||||
|
||||
pub(crate) async fn list_local_openai_chat_candidates(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
require_streaming: bool,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
Vec<SkippedLocalExecutionCandidate>,
|
||||
),
|
||||
GatewayError,
|
||||
> {
|
||||
let outcome = preselect_local_execution_candidates_with_serving(
|
||||
PlannerAppState::new(state),
|
||||
"openai:chat",
|
||||
&input.requested_model,
|
||||
require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
&input.auth_snapshot,
|
||||
false,
|
||||
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModel,
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok((outcome.candidates, outcome.skipped_candidates))
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
use super::super::GatewayControlDecision;
|
||||
use crate::ai_serving::planner::runtime_miss::{
|
||||
set_local_runtime_candidate_evaluation_diagnostic, set_local_runtime_miss_diagnostic_reason,
|
||||
};
|
||||
use crate::AppState;
|
||||
|
||||
pub(crate) fn set_local_openai_chat_miss_diagnostic(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
reason: &str,
|
||||
) {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
reason,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn set_local_openai_chat_candidate_evaluation_diagnostic(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
set_local_runtime_candidate_evaluation_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
candidate_count,
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
use tracing::warn;
|
||||
|
||||
use super::super::{GatewayControlDecision, LocalOpenAiChatDecisionInput};
|
||||
use super::diagnostic::set_local_openai_chat_miss_diagnostic;
|
||||
use crate::ai_serving::planner::common::extract_standard_requested_model;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::resolve_local_decision_execution_runtime_auth_context;
|
||||
use crate::AppState;
|
||||
|
||||
pub(crate) async fn resolve_local_openai_chat_decision_input(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
record_miss_diagnostic: bool,
|
||||
) -> Option<LocalOpenAiChatDecisionInput> {
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
route_class = ?decision.route_class,
|
||||
route_family = ?decision.route_family,
|
||||
route_kind = ?decision.route_kind,
|
||||
"gateway local openai chat decision skipped: missing_auth_context"
|
||||
);
|
||||
if record_miss_diagnostic {
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
extract_standard_requested_model(body_json).as_deref(),
|
||||
"missing_auth_context",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(requested_model) = extract_standard_requested_model(body_json) else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
"gateway local openai chat decision skipped: missing_requested_model"
|
||||
);
|
||||
if record_miss_diagnostic {
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
None,
|
||||
"missing_requested_model",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
};
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
auth_context.clone(),
|
||||
Some(requested_model.as_str()),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
user_id = %auth_context.user_id,
|
||||
api_key_id = %auth_context.api_key_id,
|
||||
"gateway local openai chat decision skipped: auth_snapshot_missing"
|
||||
);
|
||||
if record_miss_diagnostic {
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(requested_model.as_str()),
|
||||
"auth_snapshot_missing",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat decision auth snapshot read failed"
|
||||
);
|
||||
if record_miss_diagnostic {
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(requested_model.as_str()),
|
||||
"auth_snapshot_read_failed",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(build_local_requested_model_decision_input(
|
||||
resolved_input,
|
||||
requested_model,
|
||||
))
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
use tracing::warn;
|
||||
|
||||
use super::super::{
|
||||
materialize_local_openai_chat_candidate_attempts,
|
||||
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
|
||||
GatewayError,
|
||||
};
|
||||
use super::candidates::list_local_openai_chat_candidates;
|
||||
use super::diagnostic::{
|
||||
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
|
||||
};
|
||||
use super::resolve::resolve_local_openai_chat_decision_input;
|
||||
use crate::ai_serving::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
|
||||
use crate::ai_serving::planner::plan_builders::{
|
||||
build_openai_chat_stream_plan_from_decision, AiStreamAttempt,
|
||||
};
|
||||
use crate::ai_serving::planner::runtime_miss::apply_local_runtime_candidate_terminal_reason;
|
||||
|
||||
pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
|
||||
if plan_kind != OPENAI_CHAT_STREAM_PLAN_KIND {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let Some(input) = resolve_local_openai_chat_decision_input(
|
||||
state, trace_id, decision, body_json, plan_kind, true,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let (candidates, skipped_candidates) =
|
||||
match list_local_openai_chat_candidates(state, &input, true).await {
|
||||
Ok(value) => value,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat stream decision scheduler selection failed"
|
||||
);
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"scheduler_selection_failed",
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
};
|
||||
if candidates.is_empty() && skipped_candidates.is_empty() {
|
||||
set_local_openai_chat_candidate_evaluation_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
0,
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
set_local_openai_chat_candidate_evaluation_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
candidates.len() + skipped_candidates.len(),
|
||||
);
|
||||
|
||||
let attempts = materialize_local_openai_chat_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
candidates,
|
||||
skipped_candidates,
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
&input,
|
||||
attempt,
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND,
|
||||
"openai_chat_stream_success",
|
||||
true,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_openai_chat_stream_plan_from_decision(parts, body_json, payload) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat stream decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_stream_plans");
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
use tracing::warn;
|
||||
|
||||
use super::super::{
|
||||
materialize_local_openai_chat_candidate_attempts,
|
||||
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
|
||||
GatewayError,
|
||||
};
|
||||
use super::candidates::list_local_openai_chat_candidates;
|
||||
use super::diagnostic::{
|
||||
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
|
||||
};
|
||||
use super::resolve::resolve_local_openai_chat_decision_input;
|
||||
use crate::ai_serving::planner::common::{
|
||||
force_upstream_streaming_for_provider, OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_serving::planner::plan_builders::{
|
||||
build_openai_chat_sync_plan_from_decision, AiSyncAttempt,
|
||||
};
|
||||
use crate::ai_serving::planner::runtime_miss::apply_local_runtime_candidate_terminal_reason;
|
||||
|
||||
fn openai_chat_sync_upstream_is_stream_for_candidate(
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
) -> bool {
|
||||
force_upstream_streaming_for_provider(provider_type, provider_api_format)
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
if plan_kind != OPENAI_CHAT_SYNC_PLAN_KIND {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let Some(input) = resolve_local_openai_chat_decision_input(
|
||||
state, trace_id, decision, body_json, plan_kind, true,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let (candidates, skipped_candidates) =
|
||||
match list_local_openai_chat_candidates(state, &input, false).await {
|
||||
Ok(value) => value,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat sync decision scheduler selection failed"
|
||||
);
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"scheduler_selection_failed",
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
};
|
||||
if candidates.is_empty() && skipped_candidates.is_empty() {
|
||||
set_local_openai_chat_candidate_evaluation_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
0,
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
set_local_openai_chat_candidate_evaluation_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
candidates.len() + skipped_candidates.len(),
|
||||
);
|
||||
|
||||
let attempts = materialize_local_openai_chat_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
candidates,
|
||||
skipped_candidates,
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let upstream_is_stream = openai_chat_sync_upstream_is_stream_for_candidate(
|
||||
attempt.eligible.transport.provider.provider_type.as_str(),
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
);
|
||||
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
&input,
|
||||
attempt,
|
||||
OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
"openai_chat_sync_success",
|
||||
upstream_is_stream,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_openai_chat_sync_plan_from_decision(parts, body_json, payload) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat sync decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_sync_plans");
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::openai_chat_sync_upstream_is_stream_for_candidate;
|
||||
|
||||
#[test]
|
||||
fn openai_chat_sync_forces_streaming_for_codex_openai_responses_candidates() {
|
||||
assert!(openai_chat_sync_upstream_is_stream_for_candidate(
|
||||
"codex",
|
||||
"openai:responses"
|
||||
));
|
||||
assert!(!openai_chat_sync_upstream_is_stream_for_candidate(
|
||||
"codex",
|
||||
"openai:responses:compact"
|
||||
));
|
||||
assert!(!openai_chat_sync_upstream_is_stream_for_candidate(
|
||||
"openai",
|
||||
"openai:responses"
|
||||
));
|
||||
assert!(!openai_chat_sync_upstream_is_stream_for_candidate(
|
||||
"codex",
|
||||
"openai:chat"
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
mod chat;
|
||||
mod responses;
|
||||
|
||||
pub(crate) use crate::ai_serving::{
|
||||
copy_request_number_field, copy_request_number_field_as,
|
||||
map_openai_reasoning_effort_to_claude_output, map_openai_reasoning_effort_to_gemini_budget,
|
||||
parse_openai_stop_sequences, resolve_openai_chat_max_tokens, value_as_u64,
|
||||
};
|
||||
pub(crate) use chat::{
|
||||
build_local_openai_chat_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_sync_plan_and_reports_for_kind,
|
||||
maybe_build_stream_local_decision_payload, maybe_build_sync_local_decision_payload,
|
||||
set_local_openai_chat_execution_exhausted_diagnostic,
|
||||
};
|
||||
pub(crate) use responses::{
|
||||
build_local_openai_responses_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_responses_sync_plan_and_reports_for_kind,
|
||||
maybe_build_stream_local_openai_responses_decision_payload,
|
||||
maybe_build_sync_local_openai_responses_decision_payload,
|
||||
};
|
||||
@@ -0,0 +1,11 @@
|
||||
#[path = "plan_builders/stream.rs"]
|
||||
mod stream;
|
||||
#[path = "plan_builders/sync.rs"]
|
||||
mod sync;
|
||||
|
||||
pub(crate) use self::stream::{
|
||||
build_openai_chat_stream_plan_from_decision, build_openai_responses_stream_plan_from_decision,
|
||||
};
|
||||
pub(crate) use self::sync::{
|
||||
build_openai_chat_sync_plan_from_decision, build_openai_responses_sync_plan_from_decision,
|
||||
};
|
||||
@@ -0,0 +1,504 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::RequestBody;
|
||||
use tracing::debug;
|
||||
|
||||
use super::super::{
|
||||
augment_sync_report_context, build_ai_execution_plan_from_decision,
|
||||
generic_decision_missing_exact_provider_request, take_ai_decision_plan_core,
|
||||
take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts,
|
||||
AiStreamAttempt,
|
||||
};
|
||||
use crate::ai_serving::provider_adaptation_requires_eventstream_accept;
|
||||
use crate::ai_serving::transport::{
|
||||
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
|
||||
build_standard_plan_fallback_openai_responses_url, StandardPlanFallbackAcceptPolicy,
|
||||
StandardPlanFallbackHeadersInput,
|
||||
};
|
||||
use crate::{AiExecutionDecision, GatewayError};
|
||||
|
||||
pub(crate) fn build_openai_chat_stream_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
let Some(core) = take_ai_decision_plan_core(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(auth_pair) = take_ai_upstream_auth_pair(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let url = if let Some(upstream_url) = take_non_empty_string(&mut payload.upstream_url) {
|
||||
upstream_url
|
||||
} else {
|
||||
let Some(upstream_base_url) = take_non_empty_string(&mut payload.upstream_base_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
build_standard_plan_fallback_openai_chat_url(&upstream_base_url, parts.uri.query())
|
||||
};
|
||||
let provider_request_body_value = if let Some(body) = payload.provider_request_body.take() {
|
||||
body
|
||||
} else {
|
||||
let Some(request_body_object) = body_json.as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let mut provider_request_body = serde_json::Map::from_iter(
|
||||
request_body_object
|
||||
.iter()
|
||||
.map(|(key, value)| (key.clone(), value.clone())),
|
||||
);
|
||||
if let Some(mapped_model) = take_non_empty_string(&mut payload.mapped_model) {
|
||||
provider_request_body
|
||||
.insert("model".to_string(), serde_json::Value::String(mapped_model));
|
||||
}
|
||||
provider_request_body.insert("stream".to_string(), serde_json::Value::Bool(true));
|
||||
if let Some(prompt_cache_key) = take_non_empty_string(&mut payload.prompt_cache_key) {
|
||||
let existing = provider_request_body
|
||||
.get("prompt_cache_key")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if existing.is_empty() {
|
||||
provider_request_body.insert(
|
||||
"prompt_cache_key".to_string(),
|
||||
serde_json::Value::String(prompt_cache_key),
|
||||
);
|
||||
}
|
||||
}
|
||||
serde_json::Value::Object(provider_request_body)
|
||||
};
|
||||
let extra_headers = std::mem::take(&mut payload.extra_headers);
|
||||
let mut provider_request_headers =
|
||||
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
|
||||
request_headers: &parts.headers,
|
||||
existing_provider_request_headers: std::mem::take(
|
||||
&mut payload.provider_request_headers,
|
||||
),
|
||||
auth_header: auth_pair.as_ref().map(|pair| pair.header.as_str()),
|
||||
auth_value: auth_pair.as_ref().map(|pair| pair.value.as_str()),
|
||||
extra_headers: &extra_headers,
|
||||
content_type: payload.content_type.as_deref(),
|
||||
provider_api_format: core.provider_api_format.as_str(),
|
||||
client_api_format: core.client_api_format.as_str(),
|
||||
upstream_is_stream: true,
|
||||
build_from_request_when_empty: true,
|
||||
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamRequired,
|
||||
});
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| Some("application/json".to_string()));
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context.take(),
|
||||
&provider_request_headers,
|
||||
&provider_request_body_value,
|
||||
)?;
|
||||
let plan = build_ai_execution_plan_from_decision(
|
||||
&mut payload,
|
||||
AiExecutionPlanFromDecisionParts {
|
||||
core,
|
||||
method: "POST".to_string(),
|
||||
url,
|
||||
headers: std::mem::take(&mut provider_request_headers),
|
||||
content_type,
|
||||
body: RequestBody::from_json(provider_request_body_value),
|
||||
stream: true,
|
||||
},
|
||||
);
|
||||
|
||||
Ok(Some(AiStreamAttempt {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn build_openai_responses_stream_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
_body_json: &serde_json::Value,
|
||||
payload: AiExecutionDecision,
|
||||
compact: bool,
|
||||
) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
if generic_decision_missing_exact_provider_request(&payload) {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(core) = take_ai_decision_plan_core(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(auth_pair) = take_ai_upstream_auth_pair(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let (url, url_source) = if let Some(upstream_url) =
|
||||
take_non_empty_string(&mut payload.upstream_url)
|
||||
{
|
||||
(upstream_url, "upstream_url")
|
||||
} else {
|
||||
let Some(upstream_base_url) = take_non_empty_string(&mut payload.upstream_base_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
(
|
||||
build_standard_plan_fallback_openai_responses_url(
|
||||
&upstream_base_url,
|
||||
parts.uri.query(),
|
||||
compact,
|
||||
),
|
||||
"upstream_base_url",
|
||||
)
|
||||
};
|
||||
let Some(provider_request_body_value) = payload.provider_request_body.take() else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let envelope_name = payload
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|context| context.get("envelope_name"))
|
||||
.and_then(serde_json::Value::as_str);
|
||||
let accept_policy = if provider_adaptation_requires_eventstream_accept(
|
||||
envelope_name,
|
||||
core.provider_api_format.as_str(),
|
||||
) {
|
||||
StandardPlanFallbackAcceptPolicy::ProviderEventStreamIfMissing
|
||||
} else {
|
||||
StandardPlanFallbackAcceptPolicy::TextEventStreamRequired
|
||||
};
|
||||
let mut provider_request_headers =
|
||||
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
|
||||
request_headers: &parts.headers,
|
||||
existing_provider_request_headers: std::mem::take(
|
||||
&mut payload.provider_request_headers,
|
||||
),
|
||||
auth_header: auth_pair.as_ref().map(|pair| pair.header.as_str()),
|
||||
auth_value: auth_pair.as_ref().map(|pair| pair.value.as_str()),
|
||||
extra_headers: &BTreeMap::new(),
|
||||
content_type: payload.content_type.as_deref(),
|
||||
provider_api_format: core.provider_api_format.as_str(),
|
||||
client_api_format: core.client_api_format.as_str(),
|
||||
upstream_is_stream: true,
|
||||
build_from_request_when_empty: false,
|
||||
accept_policy,
|
||||
});
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| Some("application/json".to_string()));
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context.take(),
|
||||
&provider_request_headers,
|
||||
&provider_request_body_value,
|
||||
)?;
|
||||
let plan = build_ai_execution_plan_from_decision(
|
||||
&mut payload,
|
||||
AiExecutionPlanFromDecisionParts {
|
||||
core,
|
||||
method: "POST".to_string(),
|
||||
url,
|
||||
headers: std::mem::take(&mut provider_request_headers),
|
||||
content_type,
|
||||
body: RequestBody::from_json(provider_request_body_value),
|
||||
stream: true,
|
||||
},
|
||||
);
|
||||
|
||||
debug!(
|
||||
event_name = "local_openai_responses_stream_plan_built",
|
||||
log_type = "debug",
|
||||
request_id = %plan.request_id,
|
||||
candidate_id = ?plan.candidate_id,
|
||||
provider_id = %plan.provider_id,
|
||||
endpoint_id = %plan.endpoint_id,
|
||||
key_id = %plan.key_id,
|
||||
downstream_path = %parts.uri.path(),
|
||||
downstream_query = ?parts.uri.query(),
|
||||
url_source,
|
||||
decision_upstream_base_url = ?payload.upstream_base_url,
|
||||
decision_upstream_url = ?payload.upstream_url,
|
||||
plan_url = %plan.url,
|
||||
client_api_format = %plan.client_api_format,
|
||||
provider_api_format = %plan.provider_api_format,
|
||||
upstream_is_stream = payload.upstream_is_stream,
|
||||
compact,
|
||||
"gateway built local openai responses stream execution plan"
|
||||
);
|
||||
|
||||
Ok(Some(AiStreamAttempt {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::{
|
||||
build_openai_chat_stream_plan_from_decision,
|
||||
build_openai_responses_stream_plan_from_decision,
|
||||
};
|
||||
use crate::AiExecutionDecision;
|
||||
|
||||
fn object_keys(value: &Value) -> Vec<&str> {
|
||||
value
|
||||
.as_object()
|
||||
.expect("value should be an object")
|
||||
.keys()
|
||||
.map(String::as_str)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn sample_responses_payload() -> AiExecutionDecision {
|
||||
AiExecutionDecision {
|
||||
action: "stream".to_string(),
|
||||
decision_kind: Some("openai_responses_stream".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: None,
|
||||
request_id: Some("req_123".to_string()),
|
||||
candidate_id: Some("cand_123".to_string()),
|
||||
provider_name: Some("Codex".to_string()),
|
||||
provider_id: Some("prov_123".to_string()),
|
||||
endpoint_id: Some("ep_123".to_string()),
|
||||
key_id: Some("key_123".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/responses".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("authorization".to_string()),
|
||||
auth_value: Some("Bearer test".to_string()),
|
||||
provider_api_format: Some("openai:responses".to_string()),
|
||||
client_api_format: Some("openai:responses".to_string()),
|
||||
provider_contract: Some("openai:responses".to_string()),
|
||||
client_contract: Some("openai:responses".to_string()),
|
||||
model_name: Some("gpt-5.4".to_string()),
|
||||
mapped_model: Some("gpt-5.4".to_string()),
|
||||
prompt_cache_key: Some("cache-key".to_string()),
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::from([(
|
||||
"content-type".to_string(),
|
||||
"application/json".to_string(),
|
||||
)]),
|
||||
provider_request_body: Some(json!({
|
||||
"text": {"verbosity": "low"},
|
||||
"input": [],
|
||||
"model": "gpt-5.4",
|
||||
"store": false,
|
||||
"tools": [],
|
||||
"stream": true,
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
"reasoning": {"effort": "high"},
|
||||
"tool_choice": "auto",
|
||||
"instructions": "You are Codex.",
|
||||
"prompt_cache_key": "cache-key"
|
||||
})),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: true,
|
||||
report_kind: Some("openai_responses_stream_success".to_string()),
|
||||
report_context: Some(json!({})),
|
||||
auth_context: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_openai_responses_stream_plan_preserves_provider_request_body_order_in_plan_and_report()
|
||||
{
|
||||
let parts = http::Request::builder()
|
||||
.uri("http://localhost/v1/responses")
|
||||
.body(())
|
||||
.expect("request should build")
|
||||
.into_parts()
|
||||
.0;
|
||||
let payload = sample_responses_payload();
|
||||
|
||||
let built =
|
||||
build_openai_responses_stream_plan_from_decision(&parts, &json!({}), payload, false)
|
||||
.expect("plan build should succeed")
|
||||
.expect("plan should be produced");
|
||||
let plan_body = built
|
||||
.plan
|
||||
.body
|
||||
.json_body
|
||||
.as_ref()
|
||||
.expect("plan json body should exist");
|
||||
assert_eq!(
|
||||
object_keys(plan_body),
|
||||
vec![
|
||||
"text",
|
||||
"input",
|
||||
"model",
|
||||
"store",
|
||||
"tools",
|
||||
"stream",
|
||||
"include",
|
||||
"reasoning",
|
||||
"tool_choice",
|
||||
"instructions",
|
||||
"prompt_cache_key",
|
||||
]
|
||||
);
|
||||
assert!(
|
||||
built
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("provider_request_body"))
|
||||
.is_none(),
|
||||
"report context should not duplicate provider request body"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_openai_chat_stream_plan_fallback_preserves_complete_same_format_headers() {
|
||||
let parts = http::Request::builder()
|
||||
.uri("http://localhost/v1/chat/completions")
|
||||
.header(http::header::AUTHORIZATION, "Bearer client-token")
|
||||
.header("x-stainless-runtime-version", "v24.0.0")
|
||||
.header("x-app", "codex")
|
||||
.body(())
|
||||
.expect("request should build")
|
||||
.into_parts()
|
||||
.0;
|
||||
let payload = AiExecutionDecision {
|
||||
action: "stream".to_string(),
|
||||
decision_kind: Some("openai_chat_stream".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: None,
|
||||
request_id: Some("req_stream_456".to_string()),
|
||||
candidate_id: Some("cand_stream_456".to_string()),
|
||||
provider_name: Some("OpenAI".to_string()),
|
||||
provider_id: Some("prov_stream_456".to_string()),
|
||||
endpoint_id: Some("ep_stream_456".to_string()),
|
||||
key_id: Some("key_stream_456".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/chat/completions".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("authorization".to_string()),
|
||||
auth_value: Some("Bearer upstream-token".to_string()),
|
||||
provider_api_format: Some("openai:chat".to_string()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: Some("openai:chat".to_string()),
|
||||
client_contract: Some("openai:chat".to_string()),
|
||||
model_name: Some("gpt-5.4".to_string()),
|
||||
mapped_model: Some("gpt-5.4".to_string()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::new(),
|
||||
provider_request_body: Some(json!({"model":"gpt-5.4","messages":[],"stream":true})),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: true,
|
||||
report_kind: Some("openai_chat_stream_success".to_string()),
|
||||
report_context: Some(json!({})),
|
||||
auth_context: None,
|
||||
};
|
||||
|
||||
let built = build_openai_chat_stream_plan_from_decision(&parts, &json!({}), payload)
|
||||
.expect("plan build should succeed")
|
||||
.expect("plan should be produced");
|
||||
|
||||
assert_eq!(
|
||||
built.plan.headers.get("authorization").map(String::as_str),
|
||||
Some("Bearer upstream-token")
|
||||
);
|
||||
assert_eq!(
|
||||
built
|
||||
.plan
|
||||
.headers
|
||||
.get("x-stainless-runtime-version")
|
||||
.map(String::as_str),
|
||||
Some("v24.0.0")
|
||||
);
|
||||
assert_eq!(
|
||||
built.plan.headers.get("x-app").map(String::as_str),
|
||||
Some("codex")
|
||||
);
|
||||
assert_eq!(
|
||||
built.plan.headers.get("accept").map(String::as_str),
|
||||
Some("text/event-stream")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_openai_chat_stream_plan_fallback_restores_claude_headers_for_cross_format() {
|
||||
let parts = http::Request::builder()
|
||||
.uri("http://localhost/v1/chat/completions")
|
||||
.header("anthropic-beta", "prompt-caching-2024-07-31")
|
||||
.header("x-stainless-runtime-version", "v24.0.0")
|
||||
.body(())
|
||||
.expect("request should build")
|
||||
.into_parts()
|
||||
.0;
|
||||
let payload = AiExecutionDecision {
|
||||
action: "stream".to_string(),
|
||||
decision_kind: Some("openai_chat_stream".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: Some("format_conversion".to_string()),
|
||||
request_id: Some("req_stream_789".to_string()),
|
||||
candidate_id: Some("cand_stream_789".to_string()),
|
||||
provider_name: Some("Claude".to_string()),
|
||||
provider_id: Some("prov_stream_789".to_string()),
|
||||
endpoint_id: Some("ep_stream_789".to_string()),
|
||||
key_id: Some("key_stream_789".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/messages".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("x-api-key".to_string()),
|
||||
auth_value: Some("sk-upstream-claude".to_string()),
|
||||
provider_api_format: Some("claude:messages".to_string()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: Some("claude:messages".to_string()),
|
||||
client_contract: Some("openai:chat".to_string()),
|
||||
model_name: Some("claude-sonnet-4-5".to_string()),
|
||||
mapped_model: Some("claude-sonnet-4-5".to_string()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::new(),
|
||||
provider_request_body: Some(
|
||||
json!({"model":"claude-sonnet-4-5","messages":[],"stream":true}),
|
||||
),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: true,
|
||||
report_kind: Some("openai_chat_stream_success".to_string()),
|
||||
report_context: Some(json!({})),
|
||||
auth_context: None,
|
||||
};
|
||||
|
||||
let built = build_openai_chat_stream_plan_from_decision(&parts, &json!({}), payload)
|
||||
.expect("plan build should succeed")
|
||||
.expect("plan should be produced");
|
||||
|
||||
assert_eq!(
|
||||
built.plan.headers.get("x-api-key").map(String::as_str),
|
||||
Some("sk-upstream-claude")
|
||||
);
|
||||
assert_eq!(
|
||||
built.plan.headers.get("anthropic-beta").map(String::as_str),
|
||||
Some("prompt-caching-2024-07-31")
|
||||
);
|
||||
assert_eq!(
|
||||
built
|
||||
.plan
|
||||
.headers
|
||||
.get("anthropic-version")
|
||||
.map(String::as_str),
|
||||
Some("2023-06-01")
|
||||
);
|
||||
assert_eq!(
|
||||
built.plan.headers.get("accept").map(String::as_str),
|
||||
Some("text/event-stream")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,482 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::RequestBody;
|
||||
use tracing::debug;
|
||||
|
||||
use super::super::{
|
||||
augment_sync_report_context, build_ai_execution_plan_from_decision,
|
||||
generic_decision_missing_exact_provider_request, take_ai_decision_plan_core,
|
||||
take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts,
|
||||
AiSyncAttempt,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
|
||||
build_standard_plan_fallback_openai_responses_url, StandardPlanFallbackAcceptPolicy,
|
||||
StandardPlanFallbackHeadersInput,
|
||||
};
|
||||
use crate::{AiExecutionDecision, GatewayError};
|
||||
|
||||
pub(crate) fn build_openai_chat_sync_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
let Some(core) = take_ai_decision_plan_core(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(auth_pair) = take_ai_upstream_auth_pair(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let url = if let Some(upstream_url) = take_non_empty_string(&mut payload.upstream_url) {
|
||||
upstream_url
|
||||
} else {
|
||||
let Some(upstream_base_url) = take_non_empty_string(&mut payload.upstream_base_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
build_standard_plan_fallback_openai_chat_url(&upstream_base_url, parts.uri.query())
|
||||
};
|
||||
let provider_request_body_value = if let Some(body) = payload.provider_request_body.take() {
|
||||
body
|
||||
} else {
|
||||
let Some(request_body_object) = body_json.as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let mut provider_request_body = serde_json::Map::from_iter(
|
||||
request_body_object
|
||||
.iter()
|
||||
.map(|(key, value)| (key.clone(), value.clone())),
|
||||
);
|
||||
if let Some(mapped_model) = take_non_empty_string(&mut payload.mapped_model) {
|
||||
provider_request_body
|
||||
.insert("model".to_string(), serde_json::Value::String(mapped_model));
|
||||
}
|
||||
if payload.upstream_is_stream {
|
||||
provider_request_body.insert("stream".to_string(), serde_json::Value::Bool(true));
|
||||
}
|
||||
if let Some(prompt_cache_key) = take_non_empty_string(&mut payload.prompt_cache_key) {
|
||||
let existing = provider_request_body
|
||||
.get("prompt_cache_key")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if existing.is_empty() {
|
||||
provider_request_body.insert(
|
||||
"prompt_cache_key".to_string(),
|
||||
serde_json::Value::String(prompt_cache_key),
|
||||
);
|
||||
}
|
||||
}
|
||||
serde_json::Value::Object(provider_request_body)
|
||||
};
|
||||
let extra_headers = std::mem::take(&mut payload.extra_headers);
|
||||
let mut provider_request_headers =
|
||||
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
|
||||
request_headers: &parts.headers,
|
||||
existing_provider_request_headers: std::mem::take(
|
||||
&mut payload.provider_request_headers,
|
||||
),
|
||||
auth_header: auth_pair.as_ref().map(|pair| pair.header.as_str()),
|
||||
auth_value: auth_pair.as_ref().map(|pair| pair.value.as_str()),
|
||||
extra_headers: &extra_headers,
|
||||
content_type: payload.content_type.as_deref(),
|
||||
provider_api_format: core.provider_api_format.as_str(),
|
||||
client_api_format: core.client_api_format.as_str(),
|
||||
upstream_is_stream: payload.upstream_is_stream,
|
||||
build_from_request_when_empty: true,
|
||||
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming,
|
||||
});
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| Some("application/json".to_string()));
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context.take(),
|
||||
&provider_request_headers,
|
||||
&provider_request_body_value,
|
||||
)?;
|
||||
let stream = payload.upstream_is_stream;
|
||||
let plan = build_ai_execution_plan_from_decision(
|
||||
&mut payload,
|
||||
AiExecutionPlanFromDecisionParts {
|
||||
core,
|
||||
method: "POST".to_string(),
|
||||
url,
|
||||
headers: std::mem::take(&mut provider_request_headers),
|
||||
content_type,
|
||||
body: RequestBody::from_json(provider_request_body_value),
|
||||
stream,
|
||||
},
|
||||
);
|
||||
|
||||
Ok(Some(AiSyncAttempt {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn build_openai_responses_sync_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
_body_json: &serde_json::Value,
|
||||
payload: AiExecutionDecision,
|
||||
compact: bool,
|
||||
) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
if generic_decision_missing_exact_provider_request(&payload) {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(core) = take_ai_decision_plan_core(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(auth_pair) = take_ai_upstream_auth_pair(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let (url, url_source) = if let Some(upstream_url) =
|
||||
take_non_empty_string(&mut payload.upstream_url)
|
||||
{
|
||||
(upstream_url, "upstream_url")
|
||||
} else {
|
||||
let Some(upstream_base_url) = take_non_empty_string(&mut payload.upstream_base_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
(
|
||||
build_standard_plan_fallback_openai_responses_url(
|
||||
&upstream_base_url,
|
||||
parts.uri.query(),
|
||||
compact,
|
||||
),
|
||||
"upstream_base_url",
|
||||
)
|
||||
};
|
||||
let Some(provider_request_body_value) = payload.provider_request_body.take() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let mut provider_request_headers =
|
||||
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
|
||||
request_headers: &parts.headers,
|
||||
existing_provider_request_headers: std::mem::take(
|
||||
&mut payload.provider_request_headers,
|
||||
),
|
||||
auth_header: auth_pair.as_ref().map(|pair| pair.header.as_str()),
|
||||
auth_value: auth_pair.as_ref().map(|pair| pair.value.as_str()),
|
||||
extra_headers: &BTreeMap::new(),
|
||||
content_type: payload.content_type.as_deref(),
|
||||
provider_api_format: core.provider_api_format.as_str(),
|
||||
client_api_format: core.client_api_format.as_str(),
|
||||
upstream_is_stream: payload.upstream_is_stream,
|
||||
build_from_request_when_empty: false,
|
||||
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming,
|
||||
});
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| Some("application/json".to_string()));
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context.take(),
|
||||
&provider_request_headers,
|
||||
&provider_request_body_value,
|
||||
)?;
|
||||
let stream = payload.upstream_is_stream;
|
||||
let plan = build_ai_execution_plan_from_decision(
|
||||
&mut payload,
|
||||
AiExecutionPlanFromDecisionParts {
|
||||
core,
|
||||
method: "POST".to_string(),
|
||||
url,
|
||||
headers: std::mem::take(&mut provider_request_headers),
|
||||
content_type,
|
||||
body: RequestBody::from_json(provider_request_body_value),
|
||||
stream,
|
||||
},
|
||||
);
|
||||
|
||||
debug!(
|
||||
event_name = "local_openai_responses_sync_plan_built",
|
||||
log_type = "debug",
|
||||
request_id = %plan.request_id,
|
||||
candidate_id = ?plan.candidate_id,
|
||||
provider_id = %plan.provider_id,
|
||||
endpoint_id = %plan.endpoint_id,
|
||||
key_id = %plan.key_id,
|
||||
downstream_path = %parts.uri.path(),
|
||||
downstream_query = ?parts.uri.query(),
|
||||
url_source,
|
||||
decision_upstream_base_url = ?payload.upstream_base_url,
|
||||
decision_upstream_url = ?payload.upstream_url,
|
||||
plan_url = %plan.url,
|
||||
client_api_format = %plan.client_api_format,
|
||||
provider_api_format = %plan.provider_api_format,
|
||||
upstream_is_stream = payload.upstream_is_stream,
|
||||
compact,
|
||||
"gateway built local openai responses sync execution plan"
|
||||
);
|
||||
|
||||
Ok(Some(AiSyncAttempt {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::{
|
||||
build_openai_chat_sync_plan_from_decision, build_openai_responses_sync_plan_from_decision,
|
||||
};
|
||||
use crate::AiExecutionDecision;
|
||||
|
||||
fn object_keys(value: &Value) -> Vec<&str> {
|
||||
value
|
||||
.as_object()
|
||||
.expect("value should be an object")
|
||||
.keys()
|
||||
.map(String::as_str)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn sample_responses_payload() -> AiExecutionDecision {
|
||||
AiExecutionDecision {
|
||||
action: "sync".to_string(),
|
||||
decision_kind: Some("openai_responses_sync".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: None,
|
||||
request_id: Some("req_123".to_string()),
|
||||
candidate_id: Some("cand_123".to_string()),
|
||||
provider_name: Some("Codex".to_string()),
|
||||
provider_id: Some("prov_123".to_string()),
|
||||
endpoint_id: Some("ep_123".to_string()),
|
||||
key_id: Some("key_123".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/responses".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("authorization".to_string()),
|
||||
auth_value: Some("Bearer test".to_string()),
|
||||
provider_api_format: Some("openai:responses".to_string()),
|
||||
client_api_format: Some("openai:responses".to_string()),
|
||||
provider_contract: Some("openai:responses".to_string()),
|
||||
client_contract: Some("openai:responses".to_string()),
|
||||
model_name: Some("gpt-5.4".to_string()),
|
||||
mapped_model: Some("gpt-5.4".to_string()),
|
||||
prompt_cache_key: Some("cache-key".to_string()),
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::from([(
|
||||
"content-type".to_string(),
|
||||
"application/json".to_string(),
|
||||
)]),
|
||||
provider_request_body: Some(json!({
|
||||
"text": {"verbosity": "low"},
|
||||
"input": [],
|
||||
"model": "gpt-5.4",
|
||||
"store": false,
|
||||
"tools": [],
|
||||
"stream": true,
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
"reasoning": {"effort": "high"},
|
||||
"tool_choice": "auto",
|
||||
"instructions": "You are Codex.",
|
||||
"prompt_cache_key": "cache-key"
|
||||
})),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: true,
|
||||
report_kind: Some("openai_responses_sync_success".to_string()),
|
||||
report_context: Some(json!({})),
|
||||
auth_context: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_openai_responses_sync_plan_preserves_provider_request_body_order_in_plan_and_report() {
|
||||
let parts = http::Request::builder()
|
||||
.uri("http://localhost/v1/responses")
|
||||
.body(())
|
||||
.expect("request should build")
|
||||
.into_parts()
|
||||
.0;
|
||||
let payload = sample_responses_payload();
|
||||
|
||||
let built =
|
||||
build_openai_responses_sync_plan_from_decision(&parts, &json!({}), payload, false)
|
||||
.expect("plan build should succeed")
|
||||
.expect("plan should be produced");
|
||||
let plan_body = built
|
||||
.plan
|
||||
.body
|
||||
.json_body
|
||||
.as_ref()
|
||||
.expect("plan json body should exist");
|
||||
assert_eq!(
|
||||
object_keys(plan_body),
|
||||
vec![
|
||||
"text",
|
||||
"input",
|
||||
"model",
|
||||
"store",
|
||||
"tools",
|
||||
"stream",
|
||||
"include",
|
||||
"reasoning",
|
||||
"tool_choice",
|
||||
"instructions",
|
||||
"prompt_cache_key",
|
||||
]
|
||||
);
|
||||
assert!(
|
||||
built
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("provider_request_body"))
|
||||
.is_none(),
|
||||
"report context should not duplicate provider request body"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_openai_chat_sync_plan_fallback_preserves_complete_same_format_headers() {
|
||||
let parts = http::Request::builder()
|
||||
.uri("http://localhost/v1/chat/completions")
|
||||
.header(http::header::AUTHORIZATION, "Bearer client-token")
|
||||
.header("x-stainless-runtime-version", "v24.0.0")
|
||||
.header("x-app", "codex")
|
||||
.body(())
|
||||
.expect("request should build")
|
||||
.into_parts()
|
||||
.0;
|
||||
let payload = AiExecutionDecision {
|
||||
action: "sync".to_string(),
|
||||
decision_kind: Some("openai_chat_sync".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: None,
|
||||
request_id: Some("req_456".to_string()),
|
||||
candidate_id: Some("cand_456".to_string()),
|
||||
provider_name: Some("OpenAI".to_string()),
|
||||
provider_id: Some("prov_456".to_string()),
|
||||
endpoint_id: Some("ep_456".to_string()),
|
||||
key_id: Some("key_456".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/chat/completions".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("authorization".to_string()),
|
||||
auth_value: Some("Bearer upstream-token".to_string()),
|
||||
provider_api_format: Some("openai:chat".to_string()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: Some("openai:chat".to_string()),
|
||||
client_contract: Some("openai:chat".to_string()),
|
||||
model_name: Some("gpt-5.4".to_string()),
|
||||
mapped_model: Some("gpt-5.4".to_string()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::new(),
|
||||
provider_request_body: Some(json!({"model":"gpt-5.4","messages":[],"stream":false})),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: false,
|
||||
report_kind: Some("openai_chat_sync_success".to_string()),
|
||||
report_context: Some(json!({})),
|
||||
auth_context: None,
|
||||
};
|
||||
|
||||
let built = build_openai_chat_sync_plan_from_decision(&parts, &json!({}), payload)
|
||||
.expect("plan build should succeed")
|
||||
.expect("plan should be produced");
|
||||
|
||||
assert_eq!(
|
||||
built.plan.headers.get("authorization").map(String::as_str),
|
||||
Some("Bearer upstream-token")
|
||||
);
|
||||
assert_eq!(
|
||||
built
|
||||
.plan
|
||||
.headers
|
||||
.get("x-stainless-runtime-version")
|
||||
.map(String::as_str),
|
||||
Some("v24.0.0")
|
||||
);
|
||||
assert_eq!(
|
||||
built.plan.headers.get("x-app").map(String::as_str),
|
||||
Some("codex")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_openai_chat_sync_plan_fallback_restores_claude_headers_for_cross_format() {
|
||||
let parts = http::Request::builder()
|
||||
.uri("http://localhost/v1/chat/completions")
|
||||
.header("anthropic-beta", "prompt-caching-2024-07-31")
|
||||
.header("x-stainless-runtime-version", "v24.0.0")
|
||||
.body(())
|
||||
.expect("request should build")
|
||||
.into_parts()
|
||||
.0;
|
||||
let payload = AiExecutionDecision {
|
||||
action: "sync".to_string(),
|
||||
decision_kind: Some("openai_chat_sync".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: Some("format_conversion".to_string()),
|
||||
request_id: Some("req_789".to_string()),
|
||||
candidate_id: Some("cand_789".to_string()),
|
||||
provider_name: Some("Claude".to_string()),
|
||||
provider_id: Some("prov_789".to_string()),
|
||||
endpoint_id: Some("ep_789".to_string()),
|
||||
key_id: Some("key_789".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/messages".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("x-api-key".to_string()),
|
||||
auth_value: Some("sk-upstream-claude".to_string()),
|
||||
provider_api_format: Some("claude:messages".to_string()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: Some("claude:messages".to_string()),
|
||||
client_contract: Some("openai:chat".to_string()),
|
||||
model_name: Some("claude-sonnet-4-5".to_string()),
|
||||
mapped_model: Some("claude-sonnet-4-5".to_string()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::new(),
|
||||
provider_request_body: Some(
|
||||
json!({"model":"claude-sonnet-4-5","messages":[],"stream":false}),
|
||||
),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: false,
|
||||
report_kind: Some("openai_chat_sync_success".to_string()),
|
||||
report_context: Some(json!({})),
|
||||
auth_context: None,
|
||||
};
|
||||
|
||||
let built = build_openai_chat_sync_plan_from_decision(&parts, &json!({}), payload)
|
||||
.expect("plan build should succeed")
|
||||
.expect("plan should be produced");
|
||||
|
||||
assert_eq!(
|
||||
built.plan.headers.get("x-api-key").map(String::as_str),
|
||||
Some("sk-upstream-claude")
|
||||
);
|
||||
assert_eq!(
|
||||
built.plan.headers.get("anthropic-beta").map(String::as_str),
|
||||
Some("prompt-caching-2024-07-31")
|
||||
);
|
||||
assert_eq!(
|
||||
built
|
||||
.plan
|
||||
.headers
|
||||
.get("anthropic-version")
|
||||
.map(String::as_str),
|
||||
Some("2023-06-01")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
#[path = "decision/payload.rs"]
|
||||
mod payload;
|
||||
#[path = "decision/request.rs"]
|
||||
mod request;
|
||||
#[path = "decision/support.rs"]
|
||||
mod support;
|
||||
|
||||
pub(super) use self::payload::maybe_build_local_openai_responses_decision_payload_for_candidate;
|
||||
pub(super) use self::support::{
|
||||
materialize_local_openai_responses_candidate_attempts,
|
||||
resolve_local_openai_responses_decision_input, LocalOpenAiResponsesCandidateAttempt,
|
||||
LocalOpenAiResponsesDecisionInput,
|
||||
};
|
||||
pub(super) use crate::ai_serving::LocalOpenAiResponsesSpec;
|
||||
@@ -0,0 +1,205 @@
|
||||
use serde_json::json;
|
||||
use tracing::debug;
|
||||
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, insert_provider_stream_event_api_format,
|
||||
LocalExecutionReportContextParts,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
|
||||
use crate::ai_serving::planner::{
|
||||
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, append_local_failover_policy_to_value,
|
||||
AiExecutionDecision, AppState,
|
||||
};
|
||||
|
||||
use super::request::resolve_local_openai_responses_candidate_payload_parts;
|
||||
use super::support::{LocalOpenAiResponsesCandidateAttempt, LocalOpenAiResponsesDecisionInput};
|
||||
use super::LocalOpenAiResponsesSpec;
|
||||
|
||||
pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
body_json: &serde_json::Value,
|
||||
input: &LocalOpenAiResponsesDecisionInput,
|
||||
attempt: LocalOpenAiResponsesCandidateAttempt,
|
||||
spec: LocalOpenAiResponsesSpec,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
let spec_metadata = local_openai_responses_spec_metadata(spec);
|
||||
let attempt_identity = attempt.attempt_identity();
|
||||
let LocalOpenAiResponsesCandidateAttempt {
|
||||
eligible,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
..
|
||||
} = attempt;
|
||||
let resolved = resolve_local_openai_responses_candidate_payload_parts(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
input,
|
||||
&eligible,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
let candidate = &eligible.candidate;
|
||||
|
||||
let prompt_cache_key = resolved
|
||||
.provider_request_body
|
||||
.get("prompt_cache_key")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let proxy = state
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&resolved.transport)
|
||||
.await;
|
||||
let tls_profile = resolve_transport_tls_profile(&resolved.transport);
|
||||
let timeouts = resolve_transport_execution_timeouts(&resolved.transport);
|
||||
let mut extra_fields = serde_json::Map::new();
|
||||
if let Some(proxy_value) =
|
||||
build_request_trace_proxy_value(Some(&resolved.transport), proxy.as_ref())
|
||||
{
|
||||
extra_fields.insert("proxy".to_string(), proxy_value);
|
||||
}
|
||||
if let Some(envelope_name) = resolved.envelope_name {
|
||||
extra_fields.insert("envelope_name".to_string(), json!(envelope_name));
|
||||
}
|
||||
insert_provider_stream_event_api_format(
|
||||
&mut extra_fields,
|
||||
resolved.transport.provider.provider_type.as_str(),
|
||||
);
|
||||
let report_context = append_local_failover_policy_to_value(
|
||||
append_execution_contract_fields_to_value(
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &input.auth_context,
|
||||
request_id: trace_id,
|
||||
candidate_id: &candidate_id,
|
||||
attempt_identity,
|
||||
model: &input.requested_model,
|
||||
provider_name: &resolved.transport.provider.name,
|
||||
provider_id: &candidate.provider_id,
|
||||
endpoint_id: &candidate.endpoint_id,
|
||||
key_id: &candidate.key_id,
|
||||
key_name: Some(&candidate.key_name),
|
||||
model_id: Some(&candidate.model_id),
|
||||
global_model_id: Some(&candidate.global_model_id),
|
||||
global_model_name: Some(&candidate.global_model_name),
|
||||
provider_api_format: &resolved.provider_api_format,
|
||||
client_api_format: spec_metadata.api_format,
|
||||
mapped_model: Some(&resolved.mapped_model),
|
||||
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
|
||||
ranking: eligible.ranking.as_ref(),
|
||||
upstream_url: Some(&resolved.upstream_url),
|
||||
header_rules: resolved.transport.endpoint.header_rules.as_ref(),
|
||||
body_rules: resolved.transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: Some(serde_json::Value::Null),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: body_json
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
upstream_is_stream: resolved.upstream_is_stream,
|
||||
has_envelope: resolved.envelope_name.is_some(),
|
||||
needs_conversion: matches!(
|
||||
resolved.conversion_mode,
|
||||
crate::ai_serving::ConversionMode::Bidirectional
|
||||
),
|
||||
extra_fields,
|
||||
}),
|
||||
resolved.execution_strategy,
|
||||
resolved.conversion_mode,
|
||||
spec_metadata.api_format,
|
||||
candidate.endpoint_api_format.as_str(),
|
||||
),
|
||||
&resolved.transport,
|
||||
);
|
||||
|
||||
debug!(
|
||||
event_name = "local_openai_responses_decision_payload_built",
|
||||
log_type = "debug",
|
||||
trace_id = %trace_id,
|
||||
candidate_id = %candidate_id,
|
||||
candidate_index,
|
||||
provider_name = %resolved.transport.provider.name,
|
||||
provider_id = %candidate.provider_id,
|
||||
endpoint_id = %candidate.endpoint_id,
|
||||
key_id = %candidate.key_id,
|
||||
decision_kind = spec_metadata.decision_kind,
|
||||
execution_strategy = resolved.execution_strategy.as_str(),
|
||||
conversion_mode = resolved.conversion_mode.as_str(),
|
||||
client_api_format = spec_metadata.api_format,
|
||||
provider_api_format = %resolved.provider_api_format,
|
||||
request_path = %parts.uri.path(),
|
||||
request_query = ?parts.uri.query(),
|
||||
upstream_base_url = %resolved.transport.endpoint.base_url,
|
||||
upstream_url = %resolved.upstream_url,
|
||||
upstream_is_stream = resolved.upstream_is_stream,
|
||||
has_envelope = resolved.envelope_name.is_some(),
|
||||
"gateway built local openai responses decision payload"
|
||||
);
|
||||
let super::request::LocalOpenAiResponsesCandidatePayloadParts {
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
provider_api_format,
|
||||
provider_request_body,
|
||||
provider_request_headers,
|
||||
upstream_url,
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
is_antigravity: _,
|
||||
envelope_name: _,
|
||||
upstream_is_stream,
|
||||
transport,
|
||||
} = resolved;
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: None,
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format,
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts,
|
||||
upstream_is_stream,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
}
|
||||
@@ -0,0 +1,629 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::Value;
|
||||
use tracing::debug;
|
||||
|
||||
use crate::ai_serving::planner::candidate_preparation::{
|
||||
prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth,
|
||||
OauthPreparationContext,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::common::force_upstream_streaming_for_provider;
|
||||
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,
|
||||
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,
|
||||
};
|
||||
use crate::ai_serving::transport::antigravity::{
|
||||
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
|
||||
classify_local_antigravity_request_support, is_antigravity_provider_transport,
|
||||
AntigravityEnvelopeRequestType, AntigravityRequestEnvelopeSupport,
|
||||
AntigravityRequestSideSupport,
|
||||
};
|
||||
use crate::ai_serving::transport::auth::{
|
||||
resolve_local_gemini_auth, resolve_local_openai_bearer_auth, resolve_local_standard_auth,
|
||||
};
|
||||
use crate::ai_serving::transport::kiro::{
|
||||
build_kiro_provider_headers, build_kiro_provider_request_body,
|
||||
is_kiro_claude_messages_transport,
|
||||
local_kiro_request_transport_unsupported_reason_with_network, KiroProviderHeadersInput,
|
||||
KiroRequestAuth, KIRO_ENVELOPE_NAME,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
build_kiro_cross_format_upstream_url, build_standard_provider_request_headers,
|
||||
local_standard_transport_unsupported_reason_with_network, StandardProviderRequestHeadersInput,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
ai_local_execution_contract_for_formats, request_conversion_direct_auth,
|
||||
request_conversion_kind, CandidateFailureDiagnostic, GatewayProviderTransportSnapshot,
|
||||
LocalResolvedOAuthRequestAuth, PlannerAppState,
|
||||
};
|
||||
use crate::ai_serving::{ConversionMode, ExecutionStrategy};
|
||||
use crate::AppState;
|
||||
|
||||
use super::support::{
|
||||
mark_skipped_local_openai_responses_candidate,
|
||||
mark_skipped_local_openai_responses_candidate_with_extra_data,
|
||||
mark_skipped_local_openai_responses_candidate_with_failure_diagnostic,
|
||||
LocalOpenAiResponsesDecisionInput,
|
||||
};
|
||||
use super::LocalOpenAiResponsesSpec;
|
||||
|
||||
const ANTIGRAVITY_ENVELOPE_NAME: &str = "antigravity:v1internal";
|
||||
|
||||
pub(crate) struct LocalOpenAiResponsesCandidatePayloadParts {
|
||||
pub(super) auth_header: String,
|
||||
pub(super) auth_value: String,
|
||||
pub(super) mapped_model: String,
|
||||
pub(super) provider_api_format: String,
|
||||
pub(super) provider_request_body: Value,
|
||||
pub(super) provider_request_headers: BTreeMap<String, String>,
|
||||
pub(super) upstream_url: String,
|
||||
pub(super) execution_strategy: ExecutionStrategy,
|
||||
pub(super) conversion_mode: ConversionMode,
|
||||
pub(super) is_antigravity: bool,
|
||||
pub(super) envelope_name: Option<&'static str>,
|
||||
pub(super) upstream_is_stream: bool,
|
||||
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn resolve_local_openai_responses_candidate_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 client_api_format = spec_metadata.api_format.trim().to_ascii_lowercase();
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let candidate = &eligible.candidate;
|
||||
let provider_api_format = eligible.provider_api_format.as_str();
|
||||
let transport = &eligible.transport;
|
||||
let is_antigravity = is_antigravity_provider_transport(transport);
|
||||
let is_kiro_claude_cli = is_kiro_claude_messages_transport(transport, provider_api_format);
|
||||
|
||||
let same_format = api_format_alias_matches(provider_api_format, &client_api_format);
|
||||
let conversion_kind = request_conversion_kind(spec_metadata.api_format, provider_api_format);
|
||||
let transport_unsupported_reason = if same_format && is_kiro_claude_cli {
|
||||
local_kiro_request_transport_unsupported_reason_with_network(transport)
|
||||
} else if same_format {
|
||||
local_standard_transport_unsupported_reason_with_network(transport, provider_api_format)
|
||||
} else {
|
||||
match conversion_kind {
|
||||
Some(_) if is_antigravity && provider_api_format == "gemini:generate_content" => None,
|
||||
Some(kind) => {
|
||||
crate::ai_serving::request_conversion_transport_unsupported_reason(transport, kind)
|
||||
}
|
||||
None => Some("transport_api_format_unsupported"),
|
||||
}
|
||||
};
|
||||
if let Some(skip_reason) = transport_unsupported_reason {
|
||||
mark_skipped_local_openai_responses_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let oauth_context = OauthPreparationContext {
|
||||
trace_id,
|
||||
api_format: provider_api_format,
|
||||
operation: "openai_responses_candidate_request",
|
||||
};
|
||||
let kiro_auth = if is_kiro_claude_cli {
|
||||
match crate::ai_serving::planner::candidate_preparation::resolve_candidate_oauth_auth(
|
||||
planner_state,
|
||||
transport,
|
||||
oauth_context,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Some(LocalResolvedOAuthRequestAuth::Kiro(auth)) => Some(auth),
|
||||
_ => {
|
||||
mark_skipped_local_openai_responses_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let direct_auth = if kiro_auth.is_some() {
|
||||
None
|
||||
} else if same_format {
|
||||
match crate::ai_serving::normalize_api_format_alias(provider_api_format).as_str() {
|
||||
"gemini:generate_content" => resolve_local_gemini_auth(transport),
|
||||
"claude:messages" => resolve_local_standard_auth(transport),
|
||||
"openai:responses" | "openai:responses:compact" => {
|
||||
resolve_local_openai_bearer_auth(transport)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
} else {
|
||||
conversion_kind.and_then(|kind| request_conversion_direct_auth(transport, kind))
|
||||
};
|
||||
let prepared_candidate = if let Some(kiro_auth) = kiro_auth.as_ref() {
|
||||
match prepare_header_authenticated_candidate_from_auth(
|
||||
candidate,
|
||||
kiro_auth.name.to_string(),
|
||||
kiro_auth.value.clone(),
|
||||
) {
|
||||
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;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
match prepare_header_authenticated_candidate(
|
||||
planner_state,
|
||||
transport,
|
||||
candidate,
|
||||
direct_auth,
|
||||
oauth_context,
|
||||
)
|
||||
.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 auth_header = prepared_candidate.auth_header;
|
||||
let auth_value = prepared_candidate.auth_value;
|
||||
let mapped_model = prepared_candidate.mapped_model;
|
||||
|
||||
let needs_bidirectional_conversion = !same_format && conversion_kind.is_some();
|
||||
let upstream_is_stream = spec_metadata.require_streaming
|
||||
|| is_antigravity
|
||||
|| force_upstream_streaming_for_provider(
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format,
|
||||
);
|
||||
let Some(base_provider_request_body) = (if needs_bidirectional_conversion {
|
||||
build_cross_format_openai_responses_request_body(
|
||||
body_json,
|
||||
&mapped_model,
|
||||
spec_metadata.api_format,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
transport.provider.provider_type.as_str(),
|
||||
if is_kiro_claude_cli {
|
||||
None
|
||||
} else {
|
||||
transport.endpoint.body_rules.as_ref()
|
||||
},
|
||||
Some(input.auth_context.api_key_id.as_str()),
|
||||
)
|
||||
} else {
|
||||
build_local_openai_responses_request_body(
|
||||
body_json,
|
||||
&mapped_model,
|
||||
upstream_is_stream,
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format,
|
||||
if is_kiro_claude_cli {
|
||||
None
|
||||
} else {
|
||||
transport.endpoint.body_rules.as_ref()
|
||||
},
|
||||
Some(input.auth_context.api_key_id.as_str()),
|
||||
)
|
||||
}) 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;
|
||||
};
|
||||
let antigravity_auth = if is_antigravity {
|
||||
match classify_local_antigravity_request_support(
|
||||
transport,
|
||||
&base_provider_request_body,
|
||||
AntigravityEnvelopeRequestType::Agent,
|
||||
) {
|
||||
AntigravityRequestSideSupport::Supported(spec) => Some(spec.auth),
|
||||
AntigravityRequestSideSupport::Unsupported(_) => {
|
||||
mark_skipped_local_openai_responses_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_unsupported",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let provider_request_body = if let Some(antigravity_auth) = antigravity_auth.as_ref() {
|
||||
match build_antigravity_safe_v1internal_request(
|
||||
antigravity_auth,
|
||||
trace_id,
|
||||
&mapped_model,
|
||||
&base_provider_request_body,
|
||||
AntigravityEnvelopeRequestType::Agent,
|
||||
) {
|
||||
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
|
||||
AntigravityRequestEnvelopeSupport::Unsupported(_) => {
|
||||
mark_skipped_local_openai_responses_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"provider_request_body_build_failed",
|
||||
CandidateFailureDiagnostic::envelope_build_failed(
|
||||
spec_metadata.api_format,
|
||||
provider_api_format,
|
||||
"openai_responses_antigravity_envelope",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
base_provider_request_body
|
||||
};
|
||||
|
||||
if let Some(kiro_auth) = kiro_auth.as_ref() {
|
||||
return build_kiro_openai_responses_payload_parts(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
input,
|
||||
eligible,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
spec_metadata.api_format,
|
||||
transport,
|
||||
provider_api_format,
|
||||
mapped_model,
|
||||
auth_header,
|
||||
auth_value,
|
||||
provider_request_body,
|
||||
upstream_is_stream,
|
||||
needs_bidirectional_conversion,
|
||||
kiro_auth,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let Some(upstream_url) = (if needs_bidirectional_conversion {
|
||||
build_cross_format_openai_responses_upstream_url(
|
||||
parts,
|
||||
transport,
|
||||
&mapped_model,
|
||||
spec_metadata.api_format,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
)
|
||||
} else {
|
||||
build_local_openai_responses_upstream_url(
|
||||
parts,
|
||||
transport,
|
||||
api_format_alias_matches(provider_api_format, "openai:responses:compact"),
|
||||
)
|
||||
}) else {
|
||||
mark_skipped_local_openai_responses_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"upstream_url_missing",
|
||||
CandidateFailureDiagnostic::upstream_url_missing(
|
||||
spec_metadata.api_format,
|
||||
provider_api_format,
|
||||
"openai_responses_url",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
let extra_headers = antigravity_auth
|
||||
.as_ref()
|
||||
.map(build_antigravity_static_identity_headers)
|
||||
.unwrap_or_default();
|
||||
let Some(resolved_headers) =
|
||||
build_standard_provider_request_headers(StandardProviderRequestHeadersInput {
|
||||
transport,
|
||||
provider_api_format,
|
||||
same_format,
|
||||
headers: &parts.headers,
|
||||
auth_header: &auth_header,
|
||||
auth_value: &auth_value,
|
||||
extra_headers: &extra_headers,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
provider_request_body: &provider_request_body,
|
||||
original_request_body: body_json,
|
||||
upstream_is_stream,
|
||||
})
|
||||
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_headers",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
let mut provider_request_headers = resolved_headers.headers;
|
||||
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);
|
||||
|
||||
debug!(
|
||||
event_name = "local_openai_responses_upstream_url_resolved",
|
||||
log_type = "debug",
|
||||
trace_id = %trace_id,
|
||||
candidate_id = %candidate_id,
|
||||
candidate_index,
|
||||
provider_id = %candidate.provider_id,
|
||||
endpoint_id = %candidate.endpoint_id,
|
||||
key_id = %candidate.key_id,
|
||||
provider_type = %transport.provider.provider_type,
|
||||
client_api_format = spec_metadata.api_format,
|
||||
provider_api_format = %provider_api_format,
|
||||
execution_strategy = execution_strategy.as_str(),
|
||||
conversion_mode = conversion_mode.as_str(),
|
||||
base_url = %transport.endpoint.base_url,
|
||||
custom_path = ?transport.endpoint.custom_path,
|
||||
request_path = %parts.uri.path(),
|
||||
request_query = ?parts.uri.query(),
|
||||
mapped_model = %mapped_model,
|
||||
upstream_url = %upstream_url,
|
||||
upstream_is_stream,
|
||||
"gateway resolved local openai responses upstream url"
|
||||
);
|
||||
|
||||
Some(LocalOpenAiResponsesCandidatePayloadParts {
|
||||
auth_header: resolved_headers.auth_header,
|
||||
auth_value: resolved_headers.auth_value,
|
||||
mapped_model,
|
||||
provider_api_format: provider_api_format.to_string(),
|
||||
provider_request_body,
|
||||
provider_request_headers,
|
||||
upstream_url,
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
is_antigravity: is_antigravity
|
||||
|| antigravity_auth.is_some() && ANTIGRAVITY_ENVELOPE_NAME == "antigravity:v1internal",
|
||||
envelope_name: if is_antigravity || antigravity_auth.is_some() {
|
||||
Some(ANTIGRAVITY_ENVELOPE_NAME)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
upstream_is_stream,
|
||||
transport: Arc::clone(transport),
|
||||
})
|
||||
}
|
||||
|
||||
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 build_kiro_openai_responses_payload_parts(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
original_body_json: &serde_json::Value,
|
||||
input: &LocalOpenAiResponsesDecisionInput,
|
||||
eligible: &EligibleLocalExecutionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
client_api_format: &str,
|
||||
transport: &Arc<GatewayProviderTransportSnapshot>,
|
||||
provider_api_format: &str,
|
||||
mapped_model: String,
|
||||
auth_header: String,
|
||||
auth_value: String,
|
||||
claude_request_body: Value,
|
||||
upstream_is_stream: bool,
|
||||
needs_bidirectional_conversion: bool,
|
||||
kiro_auth: &KiroRequestAuth,
|
||||
) -> Option<LocalOpenAiResponsesCandidatePayloadParts> {
|
||||
let candidate = &eligible.candidate;
|
||||
let provider_request_body = match build_kiro_provider_request_body(
|
||||
&claude_request_body,
|
||||
&mapped_model,
|
||||
&kiro_auth.auth_config,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
) {
|
||||
Some(body) => body,
|
||||
None => {
|
||||
mark_skipped_local_openai_responses_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"provider_request_body_build_failed",
|
||||
CandidateFailureDiagnostic::envelope_build_failed(
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
"openai_responses_kiro_envelope",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let upstream_url = match build_kiro_cross_format_upstream_url(
|
||||
transport,
|
||||
&mapped_model,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
kiro_auth.auth_config.effective_api_region(),
|
||||
) {
|
||||
Some(url) => url,
|
||||
None => {
|
||||
mark_skipped_local_openai_responses_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"upstream_url_missing",
|
||||
CandidateFailureDiagnostic::upstream_url_missing(
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
"openai_responses_kiro_url",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let provider_request_headers = match build_kiro_provider_headers(KiroProviderHeadersInput {
|
||||
headers: &parts.headers,
|
||||
provider_request_body: &provider_request_body,
|
||||
original_request_body: original_body_json,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
auth_header: &auth_header,
|
||||
auth_value: &auth_value,
|
||||
auth_config: &kiro_auth.auth_config,
|
||||
machine_id: kiro_auth.machine_id.as_str(),
|
||||
}) {
|
||||
Some(headers) => headers,
|
||||
None => {
|
||||
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(
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
"openai_responses_kiro_headers",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let (execution_strategy, conversion_mode) =
|
||||
ai_local_execution_contract_for_formats(client_api_format, provider_api_format);
|
||||
|
||||
debug!(
|
||||
event_name = "local_openai_responses_kiro_upstream_url_resolved",
|
||||
log_type = "debug",
|
||||
trace_id = %trace_id,
|
||||
candidate_id = %candidate_id,
|
||||
candidate_index,
|
||||
provider_id = %candidate.provider_id,
|
||||
endpoint_id = %candidate.endpoint_id,
|
||||
key_id = %candidate.key_id,
|
||||
provider_type = %transport.provider.provider_type,
|
||||
client_api_format = client_api_format,
|
||||
provider_api_format = %provider_api_format,
|
||||
execution_strategy = execution_strategy.as_str(),
|
||||
conversion_mode = conversion_mode.as_str(),
|
||||
upstream_url = %upstream_url,
|
||||
upstream_is_stream,
|
||||
"gateway resolved local openai responses kiro upstream url"
|
||||
);
|
||||
|
||||
Some(LocalOpenAiResponsesCandidatePayloadParts {
|
||||
auth_header,
|
||||
auth_value,
|
||||
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: Some(KIRO_ENVELOPE_NAME),
|
||||
upstream_is_stream,
|
||||
transport: Arc::clone(transport),
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,309 @@
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic,
|
||||
materialize_local_execution_candidates_with_serving, LocalCandidateResolutionMode,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_metadata::{
|
||||
build_local_execution_candidate_contract_metadata,
|
||||
build_local_execution_candidate_contract_metadata_for_candidate,
|
||||
LocalExecutionCandidateMetadataParts,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_source::{
|
||||
preselect_local_execution_candidates_with_serving, LocalCandidatePreselectionKeyMode,
|
||||
};
|
||||
use crate::ai_serving::planner::common::extract_standard_requested_model;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
use crate::ai_serving::planner::runtime_miss::set_local_runtime_miss_diagnostic_reason;
|
||||
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
|
||||
use crate::ai_serving::planner::CandidateFailureDiagnostic;
|
||||
use crate::ai_serving::{
|
||||
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
|
||||
resolve_local_decision_execution_runtime_auth_context, ExecutionRuntimeAuthContext,
|
||||
GatewayControlDecision, PlannerAppState,
|
||||
};
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
use super::LocalOpenAiResponsesSpec;
|
||||
|
||||
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalOpenAiResponsesCandidateAttempt;
|
||||
pub(crate) use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput as LocalOpenAiResponsesDecisionInput;
|
||||
|
||||
pub(crate) async fn resolve_local_openai_responses_decision_input(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Option<LocalOpenAiResponsesDecisionInput> {
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
route_class = ?decision.route_class,
|
||||
route_family = ?decision.route_family,
|
||||
route_kind = ?decision.route_kind,
|
||||
"gateway local openai responses decision skipped: missing_auth_context"
|
||||
);
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
extract_standard_requested_model(body_json).as_deref(),
|
||||
"missing_auth_context",
|
||||
);
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(requested_model) = extract_standard_requested_model(body_json) else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
"gateway local openai responses decision skipped: missing_requested_model"
|
||||
);
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
None,
|
||||
"missing_requested_model",
|
||||
);
|
||||
return None;
|
||||
};
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
auth_context.clone(),
|
||||
Some(requested_model.as_str()),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
user_id = %auth_context.user_id,
|
||||
api_key_id = %auth_context.api_key_id,
|
||||
"gateway local openai responses decision skipped: auth_snapshot_missing"
|
||||
);
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(requested_model.as_str()),
|
||||
"auth_snapshot_missing",
|
||||
);
|
||||
return None;
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai responses decision auth snapshot read failed"
|
||||
);
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(requested_model.as_str()),
|
||||
"auth_snapshot_read_failed",
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(build_local_requested_model_decision_input(
|
||||
resolved_input,
|
||||
requested_model,
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) async fn materialize_local_openai_responses_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalOpenAiResponsesDecisionInput,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalOpenAiResponsesSpec,
|
||||
) -> Result<(Vec<LocalOpenAiResponsesCandidateAttempt>, 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_with_serving(
|
||||
planner_state,
|
||||
spec_metadata.api_format,
|
||||
&input.requested_model,
|
||||
spec_metadata.require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
&input.auth_snapshot,
|
||||
true,
|
||||
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
|
||||
)
|
||||
.await?;
|
||||
let outcome = materialize_local_execution_candidates_with_serving(
|
||||
planner_state,
|
||||
trace_id,
|
||||
spec_metadata.api_format,
|
||||
Some(&input.requested_model),
|
||||
Some(&input.auth_snapshot),
|
||||
input.required_capabilities.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
persistence_policy,
|
||||
preselection.candidates,
|
||||
preselection.skipped_candidates,
|
||||
LocalCandidateResolutionMode::Standard,
|
||||
|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(),
|
||||
))
|
||||
},
|
||||
|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;
|
||||
|
||||
Ok((outcome.attempts, outcome.candidate_count))
|
||||
}
|
||||
|
||||
pub(crate) async fn mark_skipped_local_openai_responses_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiResponsesDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::OpenAiResponsesDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn mark_skipped_local_openai_responses_candidate_with_extra_data(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiResponsesDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
extra_data: Option<serde_json::Value>,
|
||||
) {
|
||||
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::OpenAiResponsesDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate_with_extra_data(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
extra_data,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn mark_skipped_local_openai_responses_candidate_with_failure_diagnostic(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiResponsesDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
diagnostic: CandidateFailureDiagnostic,
|
||||
) {
|
||||
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::OpenAiResponsesDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
diagnostic,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
use crate::ai_serving::planner::plan_builders::{AiStreamAttempt, AiSyncAttempt};
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
mod decision;
|
||||
mod plans;
|
||||
|
||||
use self::decision::{
|
||||
materialize_local_openai_responses_candidate_attempts,
|
||||
maybe_build_local_openai_responses_decision_payload_for_candidate,
|
||||
resolve_local_openai_responses_decision_input,
|
||||
};
|
||||
use self::plans::{
|
||||
build_local_stream_plan_and_reports, build_local_sync_plan_and_reports, resolve_stream_spec,
|
||||
resolve_sync_spec,
|
||||
};
|
||||
|
||||
pub(crate) async fn build_local_openai_responses_sync_plan_and_reports_for_kind(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
build_local_sync_plan_and_reports(state, parts, trace_id, decision, body_json, spec).await
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_openai_responses_stream_plan_and_reports_for_kind(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
build_local_stream_plan_and_reports(state, parts, trace_id, decision, body_json, spec).await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_sync_local_openai_responses_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let Some(input) = resolve_local_openai_responses_decision_input(
|
||||
state, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let (attempts, _) = materialize_local_openai_responses_candidate_attempts(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_stream_local_openai_responses_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let Some(input) = resolve_local_openai_responses_decision_input(
|
||||
state, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let (attempts, _) = materialize_local_openai_responses_candidate_attempts(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
use tracing::warn;
|
||||
|
||||
use super::decision::{
|
||||
materialize_local_openai_responses_candidate_attempts,
|
||||
maybe_build_local_openai_responses_decision_payload_for_candidate,
|
||||
resolve_local_openai_responses_decision_input, LocalOpenAiResponsesSpec,
|
||||
};
|
||||
use crate::ai_serving::planner::plan_builders::{
|
||||
build_openai_responses_stream_plan_from_decision,
|
||||
build_openai_responses_sync_plan_from_decision, AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
use crate::ai_serving::planner::runtime_miss::{
|
||||
apply_local_runtime_candidate_evaluation_progress,
|
||||
apply_local_runtime_candidate_terminal_reason, set_local_runtime_miss_diagnostic_reason,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
pub(crate) use crate::ai_serving::{
|
||||
resolve_openai_responses_stream_spec as resolve_stream_spec,
|
||||
resolve_openai_responses_sync_spec as resolve_sync_spec,
|
||||
};
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
pub(super) async fn build_local_sync_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalOpenAiResponsesSpec,
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
let spec_metadata = local_openai_responses_spec_metadata(spec);
|
||||
let Some(input) = resolve_local_openai_responses_decision_input(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
|
||||
let (attempts, candidate_count) = materialize_local_openai_responses_candidate_attempts(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
if candidate_count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_openai_responses_sync_plan_from_decision(
|
||||
parts,
|
||||
body_json,
|
||||
payload,
|
||||
spec.compact,
|
||||
) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec_metadata.api_format,
|
||||
error = ?err,
|
||||
"gateway local openai responses sync decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_sync_plans");
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
pub(super) async fn build_local_stream_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalOpenAiResponsesSpec,
|
||||
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
|
||||
let spec_metadata = local_openai_responses_spec_metadata(spec);
|
||||
let Some(input) = resolve_local_openai_responses_decision_input(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
|
||||
let (attempts, candidate_count) = materialize_local_openai_responses_candidate_attempts(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
if candidate_count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_openai_responses_stream_plan_from_decision(
|
||||
parts,
|
||||
body_json,
|
||||
payload,
|
||||
spec.compact,
|
||||
) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec_metadata.api_format,
|
||||
error = ?err,
|
||||
"gateway local openai responses stream decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_stream_plans");
|
||||
Ok(plans)
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user