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