i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-10 11:19:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor pool candidate scheduling

Browse Source
This commit is contained in:
fawney19
2026-05-03 20:14:29 +08:00
parent 8ebee9922c
commit a24e4a793d
55 changed files with 4825 additions and 311 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision.rs
+2 -1
View File
@@ -7,6 +7,7 @@ mod support;
pub(super) use self::payload::maybe_build_local_openai_chat_decision_payload_for_candidate;
pub(super) use self::support::{
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts, LocalOpenAiChatCandidateAttempt,
LocalOpenAiChatDecisionInput,
LocalOpenAiChatCandidateAttemptSource, LocalOpenAiChatDecisionInput,
};
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/support.rs
+80
View File
@@ -1,9 +1,11 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::ai_serving::planner::candidate_materialization::{
build_local_execution_candidate_attempt_source_with_serving,
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,
LocalExecutionCandidateAttemptSource,
};
use crate::ai_serving::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata,
@@ -22,6 +24,7 @@ use crate::ai_serving::{
use crate::AppState;
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalOpenAiChatCandidateAttempt;
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttemptSource as LocalOpenAiChatCandidateAttemptSource;
pub(crate) use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput as LocalOpenAiChatDecisionInput;
pub(crate) async fn mark_skipped_local_openai_chat_candidate(
@@ -189,3 +192,80 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
outcome.attempts
}
pub(crate) async fn build_local_openai_chat_candidate_attempt_source<'a>(
state: &'a AppState,
trace_id: &str,
input: &LocalOpenAiChatDecisionInput,
body_json: &serde_json::Value,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
preselection_skipped: Vec<SkippedLocalExecutionCandidate>,
) -> (LocalOpenAiChatCandidateAttemptSource<'a>, usize) {
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,
);
build_local_execution_candidate_attempt_source_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(),
input.request_auth_channel.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
}
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/mod.rs
+49 -2
View File
@@ -12,11 +12,14 @@ mod decision;
mod plans;
use self::decision::{
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts,
maybe_build_local_openai_chat_decision_payload_for_candidate, LocalOpenAiChatDecisionInput,
maybe_build_local_openai_chat_decision_payload_for_candidate, LocalOpenAiChatCandidateAttempt,
LocalOpenAiChatCandidateAttemptSource, LocalOpenAiChatDecisionInput,
};
use self::plans::{
build_local_openai_chat_stream_plan_and_reports, build_local_openai_chat_sync_plan_and_reports,
build_local_openai_chat_stream_attempt_source, build_local_openai_chat_stream_plan_and_reports,
build_local_openai_chat_sync_attempt_source, 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,
};
@@ -49,6 +52,50 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports_for_kind(
.await
}
pub(crate) async fn build_local_openai_chat_sync_attempt_source_for_kind<'a>(
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
decision: &'a GatewayControlDecision,
body_json: &'a serde_json::Value,
plan_kind: &str,
) -> Result<
Option<(
impl crate::ai_serving::planner::LocalExecutionAttemptSource<
crate::ai_serving::planner::plan_builders::AiSyncAttempt,
> + 'a,
usize,
)>,
GatewayError,
> {
build_local_openai_chat_sync_attempt_source(
state, parts, trace_id, decision, body_json, plan_kind,
)
.await
}
pub(crate) async fn build_local_openai_chat_stream_attempt_source_for_kind<'a>(
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
decision: &'a GatewayControlDecision,
body_json: &'a serde_json::Value,
plan_kind: &str,
) -> Result<
Option<(
impl crate::ai_serving::planner::LocalExecutionAttemptSource<
crate::ai_serving::planner::plan_builders::AiStreamAttempt,
> + 'a,
usize,
)>,
GatewayError,
> {
build_local_openai_chat_stream_attempt_source(
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,
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans.rs
+6 -2
View File
@@ -12,5 +12,9 @@ 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;
pub(super) use self::stream::{
build_local_openai_chat_stream_attempt_source, build_local_openai_chat_stream_plan_and_reports,
};
pub(super) use self::sync::{
build_local_openai_chat_sync_attempt_source, build_local_openai_chat_sync_plan_and_reports,
};
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs
+161 -1
View File
@@ -1,21 +1,181 @@
use async_trait::async_trait;
use tracing::warn;
use super::super::{
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts,
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
GatewayError,
GatewayError, LocalOpenAiChatCandidateAttempt, LocalOpenAiChatCandidateAttemptSource,
LocalOpenAiChatDecisionInput,
};
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::candidate_materialization::LocalExecutionAttemptSource;
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) struct LocalOpenAiChatStreamAttemptSource<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
body_json: &'a serde_json::Value,
input: LocalOpenAiChatDecisionInput,
candidates: LocalOpenAiChatCandidateAttemptSource<'a>,
}
pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
decision: &'a GatewayControlDecision,
body_json: &'a serde_json::Value,
plan_kind: &str,
) -> Result<Option<(LocalOpenAiChatStreamAttemptSource<'a>, usize)>, 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, true,
)
.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!(
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(None);
}
};
let candidate_count = candidates.len() + skipped_candidates.len();
if candidate_count == 0 {
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
decision,
plan_kind,
Some(input.requested_model.as_str()),
0,
);
return Ok(None);
}
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
decision,
plan_kind,
Some(input.requested_model.as_str()),
candidate_count,
);
let (candidates, candidate_count) = build_local_openai_chat_candidate_attempt_source(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
)
.await;
Ok(Some((
LocalOpenAiChatStreamAttemptSource {
state,
parts,
trace_id,
body_json,
input,
candidates,
},
candidate_count,
)))
}
#[async_trait]
impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiChatStreamAttemptSource<'_> {
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await {
match self.build_stream_attempt(attempt).await? {
Some(attempt) => return Ok(Some(attempt)),
None => continue,
}
}
apply_local_runtime_candidate_terminal_reason(
self.state,
self.trace_id,
"no_local_stream_plans",
);
Ok(None)
}
async fn drain_execution_attempts(&mut self) -> Result<Vec<AiStreamAttempt>, GatewayError> {
let mut drained = Vec::new();
for attempt in self.candidates.drain_static_attempts() {
if let Some(attempt) = self.build_stream_attempt(attempt).await? {
drained.push(attempt);
}
}
Ok(drained)
}
}
impl LocalOpenAiChatStreamAttemptSource<'_> {
async fn build_stream_attempt(
&self,
attempt: LocalOpenAiChatCandidateAttempt,
) -> Result<Option<AiStreamAttempt>, GatewayError> {
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
self.state,
self.parts,
self.trace_id,
self.body_json,
&self.input,
attempt,
OPENAI_CHAT_STREAM_PLAN_KIND,
"openai_chat_stream_success",
true,
)
.await
else {
return Ok(None);
};
match build_openai_chat_stream_plan_from_decision(self.parts, self.body_json, payload) {
Ok(value) => Ok(value),
Err(err) => {
warn!(
trace_id = %self.trace_id,
error = ?err,
"gateway local openai chat stream decision plan build failed"
);
Ok(None)
}
}
}
}
pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
state: &AppState,
parts: &http::request::Parts,
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs
+165 -1
View File
@@ -1,15 +1,19 @@
use async_trait::async_trait;
use tracing::warn;
use super::super::{
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts,
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
GatewayError,
GatewayError, LocalOpenAiChatCandidateAttempt, LocalOpenAiChatCandidateAttemptSource,
LocalOpenAiChatDecisionInput,
};
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::candidate_materialization::LocalExecutionAttemptSource;
use crate::ai_serving::planner::common::{
force_upstream_streaming_for_provider, OPENAI_CHAT_SYNC_PLAN_KIND,
};
@@ -18,6 +22,15 @@ use crate::ai_serving::planner::plan_builders::{
};
use crate::ai_serving::planner::runtime_miss::apply_local_runtime_candidate_terminal_reason;
pub(crate) struct LocalOpenAiChatSyncAttemptSource<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
body_json: &'a serde_json::Value,
input: LocalOpenAiChatDecisionInput,
candidates: LocalOpenAiChatCandidateAttemptSource<'a>,
}
fn openai_chat_sync_upstream_is_stream_for_candidate(
provider_type: &str,
provider_api_format: &str,
@@ -25,6 +38,157 @@ fn openai_chat_sync_upstream_is_stream_for_candidate(
force_upstream_streaming_for_provider(provider_type, provider_api_format)
}
pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>(
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
decision: &'a GatewayControlDecision,
body_json: &'a serde_json::Value,
plan_kind: &str,
) -> Result<Option<(LocalOpenAiChatSyncAttemptSource<'a>, usize)>, 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, true,
)
.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!(
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(None);
}
};
let candidate_count = candidates.len() + skipped_candidates.len();
if candidate_count == 0 {
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
decision,
plan_kind,
Some(input.requested_model.as_str()),
0,
);
return Ok(None);
}
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
decision,
plan_kind,
Some(input.requested_model.as_str()),
candidate_count,
);
let (candidates, candidate_count) = build_local_openai_chat_candidate_attempt_source(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
)
.await;
Ok(Some((
LocalOpenAiChatSyncAttemptSource {
state,
parts,
trace_id,
body_json,
input,
candidates,
},
candidate_count,
)))
}
#[async_trait]
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiChatSyncAttemptSource<'_> {
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await {
match self.build_sync_attempt(attempt).await? {
Some(attempt) => return Ok(Some(attempt)),
None => continue,
}
}
apply_local_runtime_candidate_terminal_reason(
self.state,
self.trace_id,
"no_local_sync_plans",
);
Ok(None)
}
async fn drain_execution_attempts(&mut self) -> Result<Vec<AiSyncAttempt>, GatewayError> {
let mut drained = Vec::new();
for attempt in self.candidates.drain_static_attempts() {
if let Some(attempt) = self.build_sync_attempt(attempt).await? {
drained.push(attempt);
}
}
Ok(drained)
}
}
impl LocalOpenAiChatSyncAttemptSource<'_> {
async fn build_sync_attempt(
&self,
attempt: LocalOpenAiChatCandidateAttempt,
) -> Result<Option<AiSyncAttempt>, GatewayError> {
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(
self.state,
self.parts,
self.trace_id,
self.body_json,
&self.input,
attempt,
OPENAI_CHAT_SYNC_PLAN_KIND,
"openai_chat_sync_success",
upstream_is_stream,
)
.await
else {
return Ok(None);
};
match build_openai_chat_sync_plan_from_decision(self.parts, self.body_json, payload) {
Ok(value) => Ok(value),
Err(err) => {
warn!(
trace_id = %self.trace_id,
error = ?err,
"gateway local openai chat sync decision plan build failed"
);
Ok(None)
}
}
}
}
pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
state: &AppState,
parts: &http::request::Parts,
apps/aether-gateway/src/ai_serving/planner/standard/openai/mod.rs
+4
View File
@@ -7,13 +7,17 @@ pub(crate) use crate::ai_serving::{
parse_openai_stop_sequences, resolve_openai_chat_max_tokens, value_as_u64,
};
pub(crate) use chat::{
build_local_openai_chat_stream_attempt_source_for_kind,
build_local_openai_chat_stream_plan_and_reports_for_kind,
build_local_openai_chat_sync_attempt_source_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_attempt_source_for_kind,
build_local_openai_responses_stream_plan_and_reports_for_kind,
build_local_openai_responses_sync_attempt_source_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,
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision.rs
+2 -1
View File
@@ -7,8 +7,9 @@ mod support;
pub(super) use self::payload::maybe_build_local_openai_responses_decision_payload_for_candidate;
pub(super) use self::support::{
build_local_openai_responses_candidate_attempt_source,
materialize_local_openai_responses_candidate_attempts,
resolve_local_openai_responses_decision_input, LocalOpenAiResponsesCandidateAttempt,
LocalOpenAiResponsesDecisionInput,
LocalOpenAiResponsesCandidateAttemptSource, LocalOpenAiResponsesDecisionInput,
};
pub(super) use crate::ai_serving::LocalOpenAiResponsesSpec;
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/support.rs
+95
View File
@@ -2,9 +2,11 @@ use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use tracing::warn;
use crate::ai_serving::planner::candidate_materialization::{
build_local_execution_candidate_attempt_source_with_serving,
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,
LocalExecutionCandidateAttemptSource,
};
use crate::ai_serving::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata,
@@ -34,6 +36,7 @@ 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::candidate_materialization::LocalExecutionCandidateAttemptSource as LocalOpenAiResponsesCandidateAttemptSource;
pub(crate) use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput as LocalOpenAiResponsesDecisionInput;
pub(crate) async fn resolve_local_openai_responses_decision_input(
@@ -221,6 +224,98 @@ pub(crate) async fn materialize_local_openai_responses_candidate_attempts(
Ok((outcome.attempts, outcome.candidate_count))
}
pub(crate) async fn build_local_openai_responses_candidate_attempt_source<'a>(
state: &'a AppState,
trace_id: &str,
input: &LocalOpenAiResponsesDecisionInput,
body_json: &serde_json::Value,
spec: LocalOpenAiResponsesSpec,
) -> Result<(LocalOpenAiResponsesCandidateAttemptSource<'a>, usize), GatewayError> {
let spec_metadata = local_openai_responses_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
let persistence_policy = build_local_candidate_persistence_policy(
auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiResponsesDecision,
);
let preselection = preselect_local_execution_candidates_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?;
Ok(build_local_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
spec_metadata.api_format,
Some(&input.requested_model),
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.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)
}
pub(crate) async fn mark_skipped_local_openai_responses_candidate(
state: &AppState,
input: &LocalOpenAiResponsesDecisionInput,
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs
+44 -1
View File
@@ -11,7 +11,8 @@ use self::decision::{
resolve_local_openai_responses_decision_input,
};
use self::plans::{
build_local_stream_plan_and_reports, build_local_sync_plan_and_reports, resolve_stream_spec,
build_local_stream_attempt_source, build_local_stream_plan_and_reports,
build_local_sync_attempt_source, build_local_sync_plan_and_reports, resolve_stream_spec,
resolve_sync_spec,
};
@@ -45,6 +46,48 @@ pub(crate) async fn build_local_openai_responses_stream_plan_and_reports_for_kin
build_local_stream_plan_and_reports(state, parts, trace_id, decision, body_json, spec).await
}
pub(crate) async fn build_local_openai_responses_sync_attempt_source_for_kind<'a>(
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
decision: &'a GatewayControlDecision,
body_json: &'a serde_json::Value,
plan_kind: &str,
) -> Result<
Option<(
impl crate::ai_serving::planner::LocalExecutionAttemptSource<AiSyncAttempt> + 'a,
usize,
)>,
GatewayError,
> {
let Some(spec) = resolve_sync_spec(plan_kind) else {
return Ok(None);
};
build_local_sync_attempt_source(state, parts, trace_id, decision, body_json, spec).await
}
pub(crate) async fn build_local_openai_responses_stream_attempt_source_for_kind<'a>(
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
decision: &'a GatewayControlDecision,
body_json: &'a serde_json::Value,
plan_kind: &str,
) -> Result<
Option<(
impl crate::ai_serving::planner::LocalExecutionAttemptSource<AiStreamAttempt> + 'a,
usize,
)>,
GatewayError,
> {
let Some(spec) = resolve_stream_spec(plan_kind) else {
return Ok(None);
};
build_local_stream_attempt_source(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,
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/plans.rs
+260 -1
View File
@@ -1,10 +1,15 @@
use async_trait::async_trait;
use tracing::warn;
use super::decision::{
build_local_openai_responses_candidate_attempt_source,
materialize_local_openai_responses_candidate_attempts,
maybe_build_local_openai_responses_decision_payload_for_candidate,
resolve_local_openai_responses_decision_input, LocalOpenAiResponsesSpec,
resolve_local_openai_responses_decision_input, LocalOpenAiResponsesCandidateAttempt,
LocalOpenAiResponsesCandidateAttemptSource, LocalOpenAiResponsesDecisionInput,
LocalOpenAiResponsesSpec,
};
use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
use crate::ai_serving::planner::plan_builders::{
build_openai_responses_stream_plan_from_decision,
build_openai_responses_sync_plan_from_decision, AiStreamAttempt, AiSyncAttempt,
@@ -21,6 +26,260 @@ pub(crate) use crate::ai_serving::{
};
use crate::{AppState, GatewayError};
pub(crate) struct LocalOpenAiResponsesSyncAttemptSource<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
body_json: &'a serde_json::Value,
input: LocalOpenAiResponsesDecisionInput,
spec: LocalOpenAiResponsesSpec,
candidates: LocalOpenAiResponsesCandidateAttemptSource<'a>,
}
pub(crate) struct LocalOpenAiResponsesStreamAttemptSource<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
body_json: &'a serde_json::Value,
input: LocalOpenAiResponsesDecisionInput,
spec: LocalOpenAiResponsesSpec,
candidates: LocalOpenAiResponsesCandidateAttemptSource<'a>,
}
pub(super) async fn build_local_sync_attempt_source<'a>(
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
decision: &'a GatewayControlDecision,
body_json: &'a serde_json::Value,
spec: LocalOpenAiResponsesSpec,
) -> Result<Option<(LocalOpenAiResponsesSyncAttemptSource<'a>, usize)>, 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(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 (candidates, candidate_count) = build_local_openai_responses_candidate_attempt_source(
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(None);
}
Ok(Some((
LocalOpenAiResponsesSyncAttemptSource {
state,
parts,
trace_id,
body_json,
input,
spec,
candidates,
},
candidate_count,
)))
}
pub(super) async fn build_local_stream_attempt_source<'a>(
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
decision: &'a GatewayControlDecision,
body_json: &'a serde_json::Value,
spec: LocalOpenAiResponsesSpec,
) -> Result<Option<(LocalOpenAiResponsesStreamAttemptSource<'a>, usize)>, 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(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 (candidates, candidate_count) = build_local_openai_responses_candidate_attempt_source(
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(None);
}
Ok(Some((
LocalOpenAiResponsesStreamAttemptSource {
state,
parts,
trace_id,
body_json,
input,
spec,
candidates,
},
candidate_count,
)))
}
#[async_trait]
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiResponsesSyncAttemptSource<'_> {
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await {
match self.build_sync_attempt(attempt).await? {
Some(attempt) => return Ok(Some(attempt)),
None => continue,
}
}
apply_local_runtime_candidate_terminal_reason(
self.state,
self.trace_id,
"no_local_sync_plans",
);
Ok(None)
}
async fn drain_execution_attempts(&mut self) -> Result<Vec<AiSyncAttempt>, GatewayError> {
let mut drained = Vec::new();
for attempt in self.candidates.drain_static_attempts() {
if let Some(attempt) = self.build_sync_attempt(attempt).await? {
drained.push(attempt);
}
}
Ok(drained)
}
}
#[async_trait]
impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiResponsesStreamAttemptSource<'_> {
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await {
match self.build_stream_attempt(attempt).await? {
Some(attempt) => return Ok(Some(attempt)),
None => continue,
}
}
apply_local_runtime_candidate_terminal_reason(
self.state,
self.trace_id,
"no_local_stream_plans",
);
Ok(None)
}
async fn drain_execution_attempts(&mut self) -> Result<Vec<AiStreamAttempt>, GatewayError> {
let mut drained = Vec::new();
for attempt in self.candidates.drain_static_attempts() {
if let Some(attempt) = self.build_stream_attempt(attempt).await? {
drained.push(attempt);
}
}
Ok(drained)
}
}
impl LocalOpenAiResponsesSyncAttemptSource<'_> {
async fn build_sync_attempt(
&self,
attempt: LocalOpenAiResponsesCandidateAttempt,
) -> Result<Option<AiSyncAttempt>, GatewayError> {
let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
self.state,
self.parts,
self.trace_id,
self.body_json,
&self.input,
attempt,
self.spec,
)
.await
else {
return Ok(None);
};
match build_openai_responses_sync_plan_from_decision(
self.parts,
self.body_json,
payload,
self.spec.compact,
) {
Ok(value) => Ok(value),
Err(err) => {
warn!(
trace_id = %self.trace_id,
error = ?err,
"gateway local openai responses sync decision plan build failed"
);
Ok(None)
}
}
}
}
impl LocalOpenAiResponsesStreamAttemptSource<'_> {
async fn build_stream_attempt(
&self,
attempt: LocalOpenAiResponsesCandidateAttempt,
) -> Result<Option<AiStreamAttempt>, GatewayError> {
let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
self.state,
self.parts,
self.trace_id,
self.body_json,
&self.input,
attempt,
self.spec,
)
.await
else {
return Ok(None);
};
match build_openai_responses_stream_plan_from_decision(
self.parts,
self.body_json,
payload,
self.spec.compact,
) {
Ok(value) => Ok(value),
Err(err) => {
warn!(
trace_id = %self.trace_id,
error = ?err,
"gateway local openai responses stream decision plan build failed"
);
Ok(None)
}
}
}
}
pub(super) async fn build_local_sync_plan_and_reports(
state: &AppState,
parts: &http::request::Parts,