refactor ai serving modules and crates

This commit is contained in:
fawney19
2026-05-02 13:23:54 +08:00
parent 4fc7cecf30
commit c130d0e2c9
309 changed files with 21549 additions and 14265 deletions

View File

@@ -1,13 +1,16 @@
use aether_ai_serving::{
run_ai_attempt_loop, AiAttemptLoopOutcome, AiAttemptLoopPort, AiExecutionAttempt,
};
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::{
parse_request_candidate_report_context, SchedulerRequestCandidateStatusUpdate,
};
use async_trait::async_trait;
use axum::body::Body;
use axum::http::Response;
use tokio::time::{timeout, Duration};
use tracing::{debug, warn, Instrument};
use crate::ai_pipeline_api::{LocalStreamPlanAndReport, LocalSyncPlanAndReport};
use crate::clock::current_unix_ms;
use crate::control::GatewayControlDecision;
use crate::execution_runtime::{execute_execution_runtime_stream, execute_execution_runtime_sync};
@@ -21,42 +24,6 @@ use crate::{AppState, GatewayError};
const DEFAULT_STREAM_CANDIDATE_WATCHDOG_TIMEOUT_MS: u64 = 300_000;
pub(crate) trait LocalPlanAndReport {
fn plan(&self) -> &aether_contracts::ExecutionPlan;
fn report_kind(&self) -> Option<String>;
fn report_context(&self) -> Option<serde_json::Value>;
}
impl LocalPlanAndReport for LocalSyncPlanAndReport {
fn plan(&self) -> &aether_contracts::ExecutionPlan {
&self.plan
}
fn report_kind(&self) -> Option<String> {
self.report_kind.clone()
}
fn report_context(&self) -> Option<serde_json::Value> {
self.report_context.clone()
}
}
impl LocalPlanAndReport for LocalStreamPlanAndReport {
fn plan(&self) -> &aether_contracts::ExecutionPlan {
&self.plan
}
fn report_kind(&self) -> Option<String> {
self.report_kind.clone()
}
fn report_context(&self) -> Option<serde_json::Value> {
self.report_context.clone()
}
}
pub(crate) async fn execute_sync_plan_and_reports<T>(
state: &AppState,
parts: &http::request::Parts,
@@ -66,12 +33,12 @@ pub(crate) async fn execute_sync_plan_and_reports<T>(
plan_and_reports: Vec<T>,
) -> Result<LocalExecutionRequestOutcome, GatewayError>
where
T: LocalPlanAndReport,
T: AiExecutionAttempt + Send + Sync + 'static,
{
let candidate_count = plan_and_reports.len();
let first_provider = plan_and_reports
.first()
.and_then(|item| item.plan().provider_name.as_deref())
.and_then(|item| item.execution_plan().provider_name.as_deref())
.unwrap_or("-")
.to_string();
let span = tracing::debug_span!(
@@ -92,41 +59,75 @@ where
"candidate loop started"
);
let mut remaining = plan_and_reports.into_iter();
let mut last_attempted = None;
while let Some(plan_and_report) = remaining.next() {
last_attempted = Some((
plan_and_report.plan().clone(),
plan_and_report.report_context(),
));
if let Some(response) = execute_execution_runtime_sync(
state,
parts.uri.path(),
plan_and_report.plan().clone(),
trace_id,
decision,
plan_kind,
plan_and_report.report_kind(),
plan_and_report.report_context(),
)
.await?
{
mark_unused_local_candidates(state, remaining.collect()).await;
return Ok(LocalExecutionRequestOutcome::responded(response));
}
}
let Some((plan, report_context)) = last_attempted else {
return Ok(LocalExecutionRequestOutcome::NoPath);
let port = SyncAttemptLoopPort {
state,
parts,
trace_id,
decision,
plan_kind,
};
Ok(LocalExecutionRequestOutcome::Exhausted(
build_local_execution_exhaustion(state, &plan, report_context.as_ref()).await,
))
match run_ai_attempt_loop(&port, plan_and_reports).await? {
AiAttemptLoopOutcome::Responded(response) => {
Ok(LocalExecutionRequestOutcome::responded(response))
}
AiAttemptLoopOutcome::Exhausted(exhaustion) => {
Ok(LocalExecutionRequestOutcome::Exhausted(exhaustion))
}
AiAttemptLoopOutcome::NoPath => Ok(LocalExecutionRequestOutcome::NoPath),
}
}
.instrument(span)
.await
}
struct SyncAttemptLoopPort<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
decision: &'a GatewayControlDecision,
plan_kind: &'a str,
}
#[async_trait]
impl<T> AiAttemptLoopPort<T> for SyncAttemptLoopPort<'_>
where
T: AiExecutionAttempt + Send + Sync + 'static,
{
type Response = Response<Body>;
type Exhaustion = crate::executor::LocalExecutionExhaustion;
type Error = GatewayError;
async fn execute_attempt(&self, attempt: &T) -> Result<Option<Self::Response>, Self::Error> {
execute_execution_runtime_sync(
self.state,
self.parts.uri.path(),
attempt.execution_plan().clone(),
self.trace_id,
self.decision,
self.plan_kind,
attempt.report_kind(),
attempt.report_context(),
)
.await
}
async fn mark_unused_attempts(&self, attempts: Vec<T>) -> Result<(), Self::Error> {
mark_unused_local_candidates(self.state, attempts).await;
Ok(())
}
async fn build_exhaustion(
&self,
last_plan: aether_contracts::ExecutionPlan,
last_report_context: Option<serde_json::Value>,
) -> Result<Self::Exhaustion, Self::Error> {
Ok(
build_local_execution_exhaustion(self.state, &last_plan, last_report_context.as_ref())
.await,
)
}
}
pub(crate) async fn execute_stream_plan_and_reports<T>(
state: &AppState,
trace_id: &str,
@@ -135,12 +136,12 @@ pub(crate) async fn execute_stream_plan_and_reports<T>(
plan_and_reports: Vec<T>,
) -> Result<LocalExecutionRequestOutcome, GatewayError>
where
T: LocalPlanAndReport,
T: AiExecutionAttempt + Send + Sync + 'static,
{
let candidate_count = plan_and_reports.len();
let first_provider = plan_and_reports
.first()
.and_then(|item| item.plan().provider_name.as_deref())
.and_then(|item| item.execution_plan().provider_name.as_deref())
.unwrap_or("-")
.to_string();
let span = tracing::debug_span!(
@@ -161,90 +162,125 @@ where
"candidate loop started"
);
let mut remaining = plan_and_reports.into_iter();
let mut last_attempted = None;
while let Some(plan_and_report) = remaining.next() {
let plan = plan_and_report.plan().clone();
let report_context = plan_and_report.report_context();
let candidate_index = parse_request_candidate_report_context(report_context.as_ref())
.and_then(|context| context.candidate_index)
.map(|value| value.to_string())
.unwrap_or_else(|| "-".to_string());
debug!(
event_name = "candidate_loop_attempt_started",
log_type = "debug",
trace_id = %trace_id,
plan_kind,
request_id = %short_request_id(plan.request_id.as_str()),
candidate_id = ?plan.candidate_id,
provider_name = plan.provider_name.as_deref().unwrap_or("-"),
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
model_name = plan.model_name.as_deref().unwrap_or("-"),
candidate_index = candidate_index.as_str(),
"candidate loop attempting stream execution candidate"
);
last_attempted = Some((plan.clone(), report_context.clone()));
let watchdog_plan = plan.clone();
let watchdog_report_context = report_context.clone();
let execution_state = state.clone();
let execution_trace_id = trace_id.to_string();
let execution_plan_kind = plan_kind.to_string();
let execution_decision = decision.clone();
let execution_report_kind = plan_and_report.report_kind();
if let Some(response) = execute_stream_candidate_with_watchdog(
state,
trace_id,
plan_kind,
&watchdog_plan,
watchdog_report_context.as_ref(),
move || async move {
execute_execution_runtime_stream(
&execution_state,
plan,
execution_trace_id.as_str(),
&execution_decision,
execution_plan_kind.as_str(),
execution_report_kind,
report_context,
)
.await
},
)
.await?
{
mark_unused_local_candidates(state, remaining.collect()).await;
return Ok(LocalExecutionRequestOutcome::responded(response));
}
}
let Some((plan, report_context)) = last_attempted else {
return Ok(LocalExecutionRequestOutcome::NoPath);
};
warn!(
event_name = "candidate_loop_exhausted",
log_type = "ops",
trace_id = %trace_id,
let port = StreamAttemptLoopPort {
state,
trace_id,
decision,
plan_kind,
};
match run_ai_attempt_loop(&port, plan_and_reports).await? {
AiAttemptLoopOutcome::Responded(response) => {
Ok(LocalExecutionRequestOutcome::responded(response))
}
AiAttemptLoopOutcome::Exhausted(exhaustion) => {
Ok(LocalExecutionRequestOutcome::Exhausted(exhaustion))
}
AiAttemptLoopOutcome::NoPath => Ok(LocalExecutionRequestOutcome::NoPath),
}
}
.instrument(span)
.await
}
struct StreamAttemptLoopPort<'a> {
state: &'a AppState,
trace_id: &'a str,
decision: &'a GatewayControlDecision,
plan_kind: &'a str,
}
#[async_trait]
impl<T> AiAttemptLoopPort<T> for StreamAttemptLoopPort<'_>
where
T: AiExecutionAttempt + Send + Sync + 'static,
{
type Response = Response<Body>;
type Exhaustion = crate::executor::LocalExecutionExhaustion;
type Error = GatewayError;
async fn execute_attempt(&self, attempt: &T) -> Result<Option<Self::Response>, Self::Error> {
let plan = attempt.execution_plan().clone();
let report_context = attempt.report_context();
let candidate_index = parse_request_candidate_report_context(report_context.as_ref())
.and_then(|context| context.candidate_index)
.map(|value| value.to_string())
.unwrap_or_else(|| "-".to_string());
debug!(
event_name = "candidate_loop_attempt_started",
log_type = "debug",
trace_id = %self.trace_id,
plan_kind = self.plan_kind,
request_id = %short_request_id(plan.request_id.as_str()),
candidate_id = ?plan.candidate_id,
provider_name = plan.provider_name.as_deref().unwrap_or("-"),
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
model_name = plan.model_name.as_deref().unwrap_or("-"),
candidate_index = candidate_index.as_str(),
"candidate loop attempting stream execution candidate"
);
let watchdog_plan = plan.clone();
let watchdog_report_context = report_context.clone();
let execution_state = self.state.clone();
let execution_trace_id = self.trace_id.to_string();
let execution_plan_kind = self.plan_kind.to_string();
let execution_decision = self.decision.clone();
let execution_report_kind = attempt.report_kind();
execute_stream_candidate_with_watchdog(
self.state,
self.trace_id,
self.plan_kind,
&watchdog_plan,
watchdog_report_context.as_ref(),
move || async move {
execute_execution_runtime_stream(
&execution_state,
plan,
execution_trace_id.as_str(),
&execution_decision,
execution_plan_kind.as_str(),
execution_report_kind,
report_context,
)
.await
},
)
.await
}
async fn mark_unused_attempts(&self, attempts: Vec<T>) -> Result<(), Self::Error> {
mark_unused_local_candidates(self.state, attempts).await;
Ok(())
}
async fn build_exhaustion(
&self,
last_plan: aether_contracts::ExecutionPlan,
last_report_context: Option<serde_json::Value>,
) -> Result<Self::Exhaustion, Self::Error> {
warn!(
event_name = "candidate_loop_exhausted",
log_type = "ops",
trace_id = %self.trace_id,
plan_kind = self.plan_kind,
request_id = %short_request_id(last_plan.request_id.as_str()),
candidate_id = ?last_plan.candidate_id,
provider_name = last_plan.provider_name.as_deref().unwrap_or("-"),
endpoint_id = %last_plan.endpoint_id,
key_id = %last_plan.key_id,
model_name = last_plan.model_name.as_deref().unwrap_or("-"),
"candidate loop exhausted local stream candidates"
);
Ok(LocalExecutionRequestOutcome::Exhausted(
build_local_execution_exhaustion(state, &plan, report_context.as_ref()).await,
))
Ok(
build_local_execution_exhaustion(self.state, &last_plan, last_report_context.as_ref())
.await,
)
}
.instrument(span)
.await
}
pub(crate) async fn mark_unused_local_candidates<T>(state: &AppState, remaining: Vec<T>)
where
T: LocalPlanAndReport,
T: AiExecutionAttempt,
{
for plan_and_report in remaining {
let report_context = plan_and_report.report_context();
@@ -253,7 +289,7 @@ where
}
record_local_request_candidate_status(
state,
plan_and_report.plan(),
plan_and_report.execution_plan(),
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Unused,

View File

@@ -1,4 +1,4 @@
use crate::ai_pipeline_api::{
use crate::ai_serving::api::{
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,
@@ -13,15 +13,15 @@ use crate::ai_pipeline_api::{
parse_direct_request_body, resolve_claude_stream_spec, resolve_claude_sync_spec,
resolve_gemini_stream_spec, resolve_gemini_sync_spec, resolve_local_same_format_stream_spec,
resolve_local_same_format_sync_spec, set_local_openai_chat_execution_exhausted_diagnostic,
LocalStandardSpec, LocalStreamPlanAndReport, LocalSyncPlanAndReport,
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
AiStreamAttempt, AiSyncAttempt, LocalStandardSpec, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
};
use crate::control::GatewayControlDecision;
use crate::executor::candidate_loop::{
execute_stream_plan_and_reports, execute_sync_plan_and_reports,
};
use crate::executor::LocalExecutionRequestOutcome;
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
use crate::{AiExecutionDecision, AppState, GatewayError};
pub(crate) async fn maybe_execute_sync_local_path(
state: &AppState,
@@ -118,7 +118,7 @@ pub(crate) async fn maybe_execute_sync_via_local_openai_responses_decision(
body_json: &serde_json::Value,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports: Vec<LocalSyncPlanAndReport> =
let plan_and_reports: Vec<AiSyncAttempt> =
build_local_openai_responses_sync_plan_and_reports_for_kind(
state, parts, trace_id, decision, body_json, plan_kind,
)
@@ -146,7 +146,7 @@ pub(crate) async fn maybe_execute_stream_via_local_openai_responses_decision(
body_json: &serde_json::Value,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports: Vec<LocalStreamPlanAndReport> =
let plan_and_reports: Vec<AiStreamAttempt> =
build_local_openai_responses_stream_plan_and_reports_for_kind(
state, parts, trace_id, decision, body_json, plan_kind,
)
@@ -171,11 +171,10 @@ pub(crate) async fn maybe_execute_sync_via_standard_family_decision(
return Ok(LocalExecutionRequestOutcome::NoPath);
};
let plan_and_reports: Vec<LocalSyncPlanAndReport> =
build_standard_family_sync_plan_and_reports(
state, parts, trace_id, decision, body_json, spec,
)
.await?;
let plan_and_reports: Vec<AiSyncAttempt> = build_standard_family_sync_plan_and_reports(
state, parts, trace_id, decision, body_json, spec,
)
.await?;
if plan_and_reports.is_empty() {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
@@ -204,11 +203,10 @@ pub(crate) async fn maybe_execute_stream_via_standard_family_decision(
return Ok(LocalExecutionRequestOutcome::NoPath);
};
let plan_and_reports: Vec<LocalStreamPlanAndReport> =
build_standard_family_stream_plan_and_reports(
state, parts, trace_id, decision, body_json, spec,
)
.await?;
let plan_and_reports: Vec<AiStreamAttempt> = build_standard_family_stream_plan_and_reports(
state, parts, trace_id, decision, body_json, spec,
)
.await?;
if plan_and_reports.is_empty() {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
@@ -330,11 +328,10 @@ pub(crate) async fn maybe_execute_sync_via_local_same_format_provider_decision(
return Ok(LocalExecutionRequestOutcome::NoPath);
};
let plan_and_reports: Vec<LocalSyncPlanAndReport> =
build_local_same_format_sync_plan_and_reports(
state, parts, trace_id, decision, body_json, spec,
)
.await?;
let plan_and_reports: Vec<AiSyncAttempt> = build_local_same_format_sync_plan_and_reports(
state, parts, trace_id, decision, body_json, spec,
)
.await?;
if plan_and_reports.is_empty() {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
@@ -362,11 +359,10 @@ pub(crate) async fn maybe_execute_stream_via_local_same_format_provider_decision
return Ok(LocalExecutionRequestOutcome::NoPath);
};
let plan_and_reports: Vec<LocalStreamPlanAndReport> =
build_local_same_format_stream_plan_and_reports(
state, parts, trace_id, decision, body_json, spec,
)
.await?;
let plan_and_reports: Vec<AiStreamAttempt> = build_local_same_format_stream_plan_and_reports(
state, parts, trace_id, decision, body_json, spec,
)
.await?;
if plan_and_reports.is_empty() {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
@@ -384,7 +380,7 @@ pub(crate) async fn maybe_execute_sync_via_local_gemini_files_decision(
decision: &GatewayControlDecision,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports: Vec<LocalSyncPlanAndReport> =
let plan_and_reports: Vec<AiSyncAttempt> =
build_local_gemini_files_sync_plan_and_reports_for_kind(
state,
parts,
@@ -420,17 +416,16 @@ pub(crate) async fn maybe_execute_sync_via_local_image_decision(
decision: &GatewayControlDecision,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports: Vec<LocalSyncPlanAndReport> =
build_local_image_sync_plan_and_reports_for_kind(
state,
parts,
body_json,
body_base64,
trace_id,
decision,
plan_kind,
)
.await?;
let plan_and_reports: Vec<AiSyncAttempt> = build_local_image_sync_plan_and_reports_for_kind(
state,
parts,
body_json,
body_base64,
trace_id,
decision,
plan_kind,
)
.await?;
if plan_and_reports.is_empty() {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
@@ -453,7 +448,7 @@ pub(crate) async fn maybe_execute_stream_via_local_gemini_files_decision(
decision: &GatewayControlDecision,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports: Vec<LocalStreamPlanAndReport> =
let plan_and_reports: Vec<AiStreamAttempt> =
build_local_gemini_files_stream_plan_and_reports_for_kind(
state, parts, trace_id, decision, plan_kind,
)
@@ -474,7 +469,7 @@ pub(crate) async fn maybe_execute_stream_via_local_image_decision(
decision: &GatewayControlDecision,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports: Vec<LocalStreamPlanAndReport> =
let plan_and_reports: Vec<AiStreamAttempt> =
build_local_image_stream_plan_and_reports_for_kind(
state,
parts,
@@ -500,11 +495,10 @@ pub(crate) async fn maybe_execute_sync_via_local_video_decision(
decision: &GatewayControlDecision,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports: Vec<LocalSyncPlanAndReport> =
build_local_video_sync_plan_and_reports_for_kind(
state, parts, body_json, trace_id, decision, plan_kind,
)
.await?;
let plan_and_reports: Vec<AiSyncAttempt> = build_local_video_sync_plan_and_reports_for_kind(
state, parts, body_json, trace_id, decision, plan_kind,
)
.await?;
if plan_and_reports.is_empty() {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
@@ -596,8 +590,6 @@ pub(crate) fn parse_local_request_body(
parse_direct_request_body(parts, body_bytes)
}
pub(crate) fn decision_payload_is_direct_execution(
payload: &GatewayControlSyncDecisionResponse,
) -> bool {
pub(crate) fn decision_payload_is_direct_execution(payload: &AiExecutionDecision) -> bool {
planner_decision_action(payload.action.as_str())
}

View File

@@ -1,12 +1,11 @@
use crate::ai_pipeline_api::{
maybe_build_stream_plan_payload, maybe_build_sync_plan_payload, LocalStreamPlanAndReport,
LocalSyncPlanAndReport,
use crate::ai_serving::api::{
maybe_build_stream_plan_payload, maybe_build_sync_plan_payload, AiStreamAttempt, AiSyncAttempt,
};
use crate::control::GatewayControlDecision;
use crate::executor::{
execute_stream_plan_and_reports, execute_sync_plan_and_reports, LocalExecutionRequestOutcome,
};
use crate::{AppState, GatewayControlPlanResponse, GatewayError, GatewayFallbackReason};
use crate::{AiExecutionPlanPayload, AppState, GatewayError, GatewayFallbackReason};
pub(crate) async fn maybe_execute_sync_via_plan_fallback(
state: &AppState,
@@ -35,7 +34,7 @@ pub(crate) async fn maybe_execute_sync_via_plan_fallback(
return Ok(LocalExecutionRequestOutcome::NoPath);
};
let GatewayControlPlanResponse {
let AiExecutionPlanPayload {
action: _,
plan_kind,
plan,
@@ -54,7 +53,7 @@ pub(crate) async fn maybe_execute_sync_via_plan_fallback(
trace_id,
decision,
plan_kind.as_str(),
vec![LocalSyncPlanAndReport {
vec![AiSyncAttempt {
plan,
report_kind,
report_context,
@@ -87,7 +86,7 @@ pub(crate) async fn maybe_execute_stream_via_plan_fallback(
return Ok(LocalExecutionRequestOutcome::NoPath);
};
let GatewayControlPlanResponse {
let AiExecutionPlanPayload {
action: _,
plan_kind,
plan,
@@ -105,7 +104,7 @@ pub(crate) async fn maybe_execute_stream_via_plan_fallback(
trace_id,
decision,
plan_kind.as_str(),
vec![LocalStreamPlanAndReport {
vec![AiStreamAttempt {
plan,
report_kind,
report_context,

View File

@@ -8,21 +8,19 @@ use axum::body::Bytes;
use crate::control::GatewayControlDecision;
use crate::headers::header_value_str;
use crate::provider_transport::provider_types::is_codex_cli_backend_url;
use crate::{AppState, GatewayControlSyncDecisionResponse};
use crate::{AiExecutionDecision, AppState};
pub(crate) const DIRECT_PLAN_BYPASS_TTL: Duration = Duration::from_secs(30);
pub(crate) const DIRECT_PLAN_BYPASS_MAX_ENTRIES: usize = 512;
pub(crate) fn should_bypass_execution_runtime_decision(
payload: &GatewayControlSyncDecisionResponse,
) -> bool {
pub(crate) fn should_bypass_execution_runtime_decision(payload: &AiExecutionDecision) -> bool {
let provider_api_format = payload
.provider_api_format
.as_deref()
.map(str::trim)
.unwrap_or_default()
.to_string();
if !crate::ai_pipeline::is_openai_responses_family_format(&provider_api_format) {
if !crate::ai_serving::is_openai_responses_family_format(&provider_api_format) {
return false;
}
@@ -36,7 +34,7 @@ pub(crate) fn should_bypass_execution_runtime_decision(
}
pub(crate) fn should_bypass_execution_runtime_plan(plan: &ExecutionPlan) -> bool {
if !crate::ai_pipeline::is_openai_responses_family_format(&plan.provider_api_format) {
if !crate::ai_serving::is_openai_responses_family_format(&plan.provider_api_format) {
return false;
}

View File

@@ -1,11 +1,15 @@
use aether_ai_serving::{
run_ai_stream_execution_path, AiPlanFallbackReason, AiServingExecutionOutcome,
AiStreamExecutionPathPort, AiStreamExecutionStep,
};
use async_trait::async_trait;
use axum::body::{Body, Bytes};
use axum::http::Response;
use std::collections::BTreeMap;
use crate::ai_pipeline_api::{
use crate::ai_serving::api::{
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
supports_stream_scheduler_decision_kind, LocalStreamPlanAndReport,
OPENAI_VIDEO_CONTENT_PLAN_KIND,
supports_stream_scheduler_decision_kind, AiStreamAttempt, OPENAI_VIDEO_CONTENT_PLAN_KIND,
};
use crate::api::response::build_client_response_from_parts;
use crate::control::GatewayControlDecision;
@@ -47,132 +51,200 @@ pub(crate) async fn maybe_execute_via_stream_decision_path(
return Ok(LocalExecutionRequestOutcome::NoPath);
}
let mut exhausted = None;
match maybe_execute_local_video_task_content_stream(state, parts, trace_id, decision, plan_kind)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
if supports_stream_scheduler_decision_kind(plan_kind) {
match maybe_execute_stream_via_local_image_decision(
state,
parts,
&body_json,
body_base64.as_deref(),
trace_id,
decision,
plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
match maybe_execute_stream_via_local_decision(
state, parts, trace_id, decision, &body_json, plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
match maybe_execute_stream_via_local_openai_responses_decision(
state, parts, trace_id, decision, &body_json, plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
match maybe_execute_stream_via_local_standard_decision(
state, parts, trace_id, decision, &body_json, plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
match maybe_execute_stream_via_local_same_format_provider_decision(
state, parts, trace_id, decision, &body_json, plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
match maybe_execute_stream_via_local_gemini_files_decision(
state, parts, trace_id, decision, plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
if let Some(response) = maybe_execute_stream_via_remote_decision(
state, parts, trace_id, decision, &body_json, plan_kind,
)
.await?
{
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
}
match maybe_execute_stream_via_plan_fallback(
let port = GatewayStreamExecutionPathPort {
state,
parts,
trace_id,
decision,
&body_json,
body_json: &body_json,
body_base64,
plan_kind,
bypass_cache_key,
if supports_stream_scheduler_decision_kind(plan_kind) {
GatewayFallbackReason::RemoteDecisionMiss
} else {
GatewayFallbackReason::SchedulerDecisionUnsupported
},
)
.await?
{
scheduler_supported: supports_stream_scheduler_decision_kind(plan_kind),
};
Ok(from_ai_serving_outcome(
run_ai_stream_execution_path(&port).await?,
))
}
struct GatewayStreamExecutionPathPort<'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<String>,
plan_kind: &'a str,
bypass_cache_key: String,
scheduler_supported: bool,
}
#[async_trait]
impl AiStreamExecutionPathPort for GatewayStreamExecutionPathPort<'_> {
type Response = Response<Body>;
type Exhaustion = super::LocalExecutionExhaustion;
type Error = GatewayError;
fn scheduler_decision_supported(&self) -> bool {
self.scheduler_supported
}
async fn execute_stream_step(
&self,
step: AiStreamExecutionStep,
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error> {
let outcome = match step {
AiStreamExecutionStep::LocalVideoContent => {
maybe_execute_local_video_task_content_stream(
self.state,
self.parts,
self.trace_id,
self.decision,
self.plan_kind,
)
.await?
}
AiStreamExecutionStep::LocalImage => {
maybe_execute_stream_via_local_image_decision(
self.state,
self.parts,
self.body_json,
self.body_base64.as_deref(),
self.trace_id,
self.decision,
self.plan_kind,
)
.await?
}
AiStreamExecutionStep::LocalOpenAiChat => {
maybe_execute_stream_via_local_decision(
self.state,
self.parts,
self.trace_id,
self.decision,
self.body_json,
self.plan_kind,
)
.await?
}
AiStreamExecutionStep::LocalOpenAiResponses => {
maybe_execute_stream_via_local_openai_responses_decision(
self.state,
self.parts,
self.trace_id,
self.decision,
self.body_json,
self.plan_kind,
)
.await?
}
AiStreamExecutionStep::LocalStandardFamily => {
maybe_execute_stream_via_local_standard_decision(
self.state,
self.parts,
self.trace_id,
self.decision,
self.body_json,
self.plan_kind,
)
.await?
}
AiStreamExecutionStep::LocalSameFormatProvider => {
maybe_execute_stream_via_local_same_format_provider_decision(
self.state,
self.parts,
self.trace_id,
self.decision,
self.body_json,
self.plan_kind,
)
.await?
}
AiStreamExecutionStep::LocalGeminiFiles => {
maybe_execute_stream_via_local_gemini_files_decision(
self.state,
self.parts,
self.trace_id,
self.decision,
self.plan_kind,
)
.await?
}
AiStreamExecutionStep::RemoteDecision => {
if let Some(response) = maybe_execute_stream_via_remote_decision(
self.state,
self.parts,
self.trace_id,
self.decision,
self.body_json,
self.plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response)
} else {
LocalExecutionRequestOutcome::NoPath
}
}
};
Ok(to_ai_serving_outcome(outcome))
}
async fn execute_stream_plan_fallback(
&self,
reason: AiPlanFallbackReason,
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error> {
let outcome = maybe_execute_stream_via_plan_fallback(
self.state,
self.parts,
self.trace_id,
self.decision,
self.body_json,
self.body_base64.clone(),
self.plan_kind,
self.bypass_cache_key.clone(),
gateway_fallback_reason(reason),
)
.await?;
Ok(to_ai_serving_outcome(outcome))
}
}
fn to_ai_serving_outcome(
outcome: LocalExecutionRequestOutcome,
) -> AiServingExecutionOutcome<Response<Body>, super::LocalExecutionExhaustion> {
match outcome {
LocalExecutionRequestOutcome::Responded(response) => {
Ok(LocalExecutionRequestOutcome::Responded(response))
AiServingExecutionOutcome::Responded(response)
}
LocalExecutionRequestOutcome::Exhausted(outcome) => {
Ok(LocalExecutionRequestOutcome::Exhausted(outcome))
AiServingExecutionOutcome::Exhausted(outcome)
}
LocalExecutionRequestOutcome::NoPath => AiServingExecutionOutcome::NoPath,
}
}
fn from_ai_serving_outcome(
outcome: AiServingExecutionOutcome<Response<Body>, super::LocalExecutionExhaustion>,
) -> LocalExecutionRequestOutcome {
match outcome {
AiServingExecutionOutcome::Responded(response) => {
LocalExecutionRequestOutcome::Responded(response)
}
AiServingExecutionOutcome::Exhausted(outcome) => {
LocalExecutionRequestOutcome::Exhausted(outcome)
}
AiServingExecutionOutcome::NoPath => LocalExecutionRequestOutcome::NoPath,
}
}
fn gateway_fallback_reason(reason: AiPlanFallbackReason) -> GatewayFallbackReason {
match reason {
AiPlanFallbackReason::RemoteDecisionMiss => GatewayFallbackReason::RemoteDecisionMiss,
AiPlanFallbackReason::SchedulerDecisionUnsupported => {
GatewayFallbackReason::SchedulerDecisionUnsupported
}
LocalExecutionRequestOutcome::NoPath => Ok(exhausted
.map(LocalExecutionRequestOutcome::Exhausted)
.unwrap_or(LocalExecutionRequestOutcome::NoPath)),
}
}
@@ -228,7 +300,7 @@ async fn maybe_execute_local_video_task_content_stream(
trace_id,
decision,
plan_kind,
vec![LocalStreamPlanAndReport {
vec![AiStreamAttempt {
plan,
report_kind: None,
report_context: None,

View File

@@ -1,11 +1,16 @@
use aether_ai_serving::{
run_ai_sync_execution_path, AiPlanFallbackReason, AiServingExecutionOutcome,
AiSyncExecutionPathPort, AiSyncExecutionStep,
};
use async_trait::async_trait;
use axum::body::{Body, Bytes};
use axum::http::Response;
use std::collections::BTreeMap;
use crate::ai_pipeline_api::{
use crate::ai_serving::api::{
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
resolve_execution_runtime_sync_plan_kind, supports_sync_scheduler_decision_kind,
LocalSyncPlanAndReport, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
resolve_execution_runtime_sync_plan_kind, supports_sync_scheduler_decision_kind, AiSyncAttempt,
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::api::response::build_client_response_from_parts;
@@ -57,153 +62,217 @@ pub(crate) async fn maybe_execute_via_sync_decision_path(
return Ok(LocalExecutionRequestOutcome::NoPath);
}
let mut exhausted = None;
match maybe_execute_local_video_task_follow_up_sync(
state, parts, &body_json, trace_id, decision, plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
if supports_sync_scheduler_decision_kind(plan_kind) {
match maybe_execute_sync_via_local_video_decision(
state, parts, &body_json, trace_id, decision, plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
match maybe_execute_sync_via_local_image_decision(
state,
parts,
&body_json,
body_base64.as_deref(),
trace_id,
decision,
plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
match maybe_execute_sync_via_local_decision(
state, parts, trace_id, decision, &body_json, plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
match maybe_execute_sync_via_local_openai_responses_decision(
state, parts, trace_id, decision, &body_json, plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
match maybe_execute_sync_via_local_standard_decision(
state, parts, trace_id, decision, &body_json, plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
match maybe_execute_sync_via_local_same_format_provider_decision(
state, parts, trace_id, decision, &body_json, plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
match maybe_execute_sync_via_local_gemini_files_decision(
state,
parts,
&body_json,
body_base64.as_deref(),
body_bytes.is_empty(),
trace_id,
decision,
plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
LocalExecutionRequestOutcome::NoPath => {}
}
if let Some(response) = maybe_execute_sync_via_remote_decision(
state, parts, trace_id, decision, &body_json, plan_kind,
)
.await?
{
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
}
match maybe_execute_sync_via_plan_fallback(
let port = GatewaySyncExecutionPathPort {
state,
parts,
trace_id,
decision,
&body_json,
body_json: &body_json,
body_base64,
body_is_empty: body_bytes.is_empty(),
plan_kind,
bypass_cache_key,
if supports_sync_scheduler_decision_kind(plan_kind) {
GatewayFallbackReason::RemoteDecisionMiss
} else {
GatewayFallbackReason::SchedulerDecisionUnsupported
},
)
.await?
{
scheduler_supported: supports_sync_scheduler_decision_kind(plan_kind),
};
Ok(from_ai_serving_outcome(
run_ai_sync_execution_path(&port).await?,
))
}
struct GatewaySyncExecutionPathPort<'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<String>,
body_is_empty: bool,
plan_kind: &'a str,
bypass_cache_key: String,
scheduler_supported: bool,
}
#[async_trait]
impl AiSyncExecutionPathPort for GatewaySyncExecutionPathPort<'_> {
type Response = Response<Body>;
type Exhaustion = super::LocalExecutionExhaustion;
type Error = GatewayError;
fn scheduler_decision_supported(&self) -> bool {
self.scheduler_supported
}
async fn execute_sync_step(
&self,
step: AiSyncExecutionStep,
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error> {
let outcome = match step {
AiSyncExecutionStep::VideoTaskFollowUp => {
maybe_execute_local_video_task_follow_up_sync(
self.state,
self.parts,
self.body_json,
self.trace_id,
self.decision,
self.plan_kind,
)
.await?
}
AiSyncExecutionStep::LocalVideo => {
maybe_execute_sync_via_local_video_decision(
self.state,
self.parts,
self.body_json,
self.trace_id,
self.decision,
self.plan_kind,
)
.await?
}
AiSyncExecutionStep::LocalImage => {
maybe_execute_sync_via_local_image_decision(
self.state,
self.parts,
self.body_json,
self.body_base64.as_deref(),
self.trace_id,
self.decision,
self.plan_kind,
)
.await?
}
AiSyncExecutionStep::LocalOpenAiChat => {
maybe_execute_sync_via_local_decision(
self.state,
self.parts,
self.trace_id,
self.decision,
self.body_json,
self.plan_kind,
)
.await?
}
AiSyncExecutionStep::LocalOpenAiResponses => {
maybe_execute_sync_via_local_openai_responses_decision(
self.state,
self.parts,
self.trace_id,
self.decision,
self.body_json,
self.plan_kind,
)
.await?
}
AiSyncExecutionStep::LocalStandardFamily => {
maybe_execute_sync_via_local_standard_decision(
self.state,
self.parts,
self.trace_id,
self.decision,
self.body_json,
self.plan_kind,
)
.await?
}
AiSyncExecutionStep::LocalSameFormatProvider => {
maybe_execute_sync_via_local_same_format_provider_decision(
self.state,
self.parts,
self.trace_id,
self.decision,
self.body_json,
self.plan_kind,
)
.await?
}
AiSyncExecutionStep::LocalGeminiFiles => {
maybe_execute_sync_via_local_gemini_files_decision(
self.state,
self.parts,
self.body_json,
self.body_base64.as_deref(),
self.body_is_empty,
self.trace_id,
self.decision,
self.plan_kind,
)
.await?
}
AiSyncExecutionStep::RemoteDecision => {
if let Some(response) = maybe_execute_sync_via_remote_decision(
self.state,
self.parts,
self.trace_id,
self.decision,
self.body_json,
self.plan_kind,
)
.await?
{
LocalExecutionRequestOutcome::Responded(response)
} else {
LocalExecutionRequestOutcome::NoPath
}
}
};
Ok(to_ai_serving_outcome(outcome))
}
async fn execute_sync_plan_fallback(
&self,
reason: AiPlanFallbackReason,
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error> {
let outcome = maybe_execute_sync_via_plan_fallback(
self.state,
self.parts,
self.trace_id,
self.decision,
self.body_json,
self.body_base64.clone(),
self.plan_kind,
self.bypass_cache_key.clone(),
gateway_fallback_reason(reason),
)
.await?;
Ok(to_ai_serving_outcome(outcome))
}
}
fn to_ai_serving_outcome(
outcome: LocalExecutionRequestOutcome,
) -> AiServingExecutionOutcome<Response<Body>, super::LocalExecutionExhaustion> {
match outcome {
LocalExecutionRequestOutcome::Responded(response) => {
Ok(LocalExecutionRequestOutcome::Responded(response))
AiServingExecutionOutcome::Responded(response)
}
LocalExecutionRequestOutcome::Exhausted(outcome) => {
Ok(LocalExecutionRequestOutcome::Exhausted(outcome))
AiServingExecutionOutcome::Exhausted(outcome)
}
LocalExecutionRequestOutcome::NoPath => AiServingExecutionOutcome::NoPath,
}
}
fn from_ai_serving_outcome(
outcome: AiServingExecutionOutcome<Response<Body>, super::LocalExecutionExhaustion>,
) -> LocalExecutionRequestOutcome {
match outcome {
AiServingExecutionOutcome::Responded(response) => {
LocalExecutionRequestOutcome::Responded(response)
}
AiServingExecutionOutcome::Exhausted(outcome) => {
LocalExecutionRequestOutcome::Exhausted(outcome)
}
AiServingExecutionOutcome::NoPath => LocalExecutionRequestOutcome::NoPath,
}
}
fn gateway_fallback_reason(reason: AiPlanFallbackReason) -> GatewayFallbackReason {
match reason {
AiPlanFallbackReason::RemoteDecisionMiss => GatewayFallbackReason::RemoteDecisionMiss,
AiPlanFallbackReason::SchedulerDecisionUnsupported => {
GatewayFallbackReason::SchedulerDecisionUnsupported
}
LocalExecutionRequestOutcome::NoPath => Ok(exhausted
.map(LocalExecutionRequestOutcome::Exhausted)
.unwrap_or(LocalExecutionRequestOutcome::NoPath)),
}
}
@@ -312,7 +381,7 @@ async fn maybe_execute_local_video_task_follow_up_sync(
trace_id,
decision,
plan_kind,
vec![LocalSyncPlanAndReport {
vec![AiSyncAttempt {
plan: follow_up.plan,
report_kind: follow_up.report_kind,
report_context: follow_up.report_context,