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
@@ -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;
}