mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 09:50:21 +08:00
refactor ai serving modules and crates
This commit is contained in:
fawney19
@@ -0,0 +1,12 @@
|
||||
//! Requests that can stay in the same public/provider contract family.
|
||||
|
||||
mod provider;
|
||||
|
||||
pub(crate) use self::provider::{
|
||||
build_local_stream_plan_and_reports as build_local_same_format_stream_plan_and_reports,
|
||||
build_local_sync_plan_and_reports as build_local_same_format_sync_plan_and_reports,
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate,
|
||||
maybe_build_stream_local_same_format_provider_decision_payload,
|
||||
maybe_build_sync_local_same_format_provider_decision_payload,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::transport::provider_types::provider_type_supports_local_same_format_transport;
|
||||
@@ -0,0 +1,96 @@
|
||||
use aether_contracts::RequestBody;
|
||||
|
||||
use super::{
|
||||
augment_sync_report_context, build_ai_execution_plan_from_decision,
|
||||
resolve_ai_passthrough_sync_request_body, take_ai_decision_plan_core, take_non_empty_string,
|
||||
AiExecutionPlanFromDecisionParts, AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
use crate::{AiExecutionDecision, GatewayError};
|
||||
|
||||
pub(crate) fn build_passthrough_sync_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
let Some(core) = take_ai_decision_plan_core(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(upstream_url) = take_non_empty_string(&mut payload.upstream_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let provider_request_headers = std::mem::take(&mut payload.provider_request_headers);
|
||||
let ignored_provider_request_body = serde_json::Value::Null;
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context.take(),
|
||||
&provider_request_headers,
|
||||
&ignored_provider_request_body,
|
||||
)?;
|
||||
let request_body = resolve_ai_passthrough_sync_request_body(
|
||||
payload.provider_request_body.take(),
|
||||
payload.provider_request_body_base64.take(),
|
||||
);
|
||||
let provider_request_method = take_non_empty_string(&mut payload.provider_request_method);
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| provider_request_headers.get("content-type").cloned());
|
||||
|
||||
let plan = build_ai_execution_plan_from_decision(
|
||||
&mut payload,
|
||||
AiExecutionPlanFromDecisionParts {
|
||||
core,
|
||||
method: provider_request_method.unwrap_or_else(|| parts.method.to_string()),
|
||||
url: upstream_url,
|
||||
headers: provider_request_headers,
|
||||
content_type,
|
||||
body: request_body,
|
||||
stream: false,
|
||||
},
|
||||
);
|
||||
|
||||
Ok(Some(AiSyncAttempt {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn build_passthrough_stream_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
let Some(core) = take_ai_decision_plan_core(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(upstream_url) = take_non_empty_string(&mut payload.upstream_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let provider_request_headers = std::mem::take(&mut payload.provider_request_headers);
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| provider_request_headers.get("content-type").cloned());
|
||||
let plan = build_ai_execution_plan_from_decision(
|
||||
&mut payload,
|
||||
AiExecutionPlanFromDecisionParts {
|
||||
core,
|
||||
method: parts.method.to_string(),
|
||||
url: upstream_url,
|
||||
headers: provider_request_headers,
|
||||
content_type,
|
||||
body: RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
},
|
||||
stream: true,
|
||||
},
|
||||
);
|
||||
|
||||
Ok(Some(AiStreamAttempt {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context: payload.report_context,
|
||||
}))
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use std::collections::BTreeMap;
|
||||
use url::form_urlencoded;
|
||||
|
||||
use aether_data_contracts::repository::candidates::{
|
||||
RequestCandidateStatus, UpsertRequestCandidateRecord,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::{json, Value};
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_serving::planner::common::{
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
|
||||
};
|
||||
use crate::ai_serving::planner::plan_builders::{AiStreamAttempt, AiSyncAttempt};
|
||||
use crate::ai_serving::transport::antigravity::{
|
||||
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
|
||||
build_antigravity_v1internal_url, classify_local_antigravity_request_support,
|
||||
AntigravityEnvelopeRequestType, AntigravityRequestEnvelopeSupport,
|
||||
AntigravityRequestSideSupport, AntigravityRequestUrlAction,
|
||||
};
|
||||
use crate::ai_serving::transport::auth::{
|
||||
build_openai_passthrough_headers, resolve_local_gemini_auth, resolve_local_standard_auth,
|
||||
};
|
||||
use crate::ai_serving::transport::claude_code::{
|
||||
build_claude_code_messages_url, build_claude_code_passthrough_headers,
|
||||
sanitize_claude_code_request_body, supports_local_claude_code_transport_with_network,
|
||||
};
|
||||
use crate::ai_serving::transport::kiro::{
|
||||
build_kiro_generate_assistant_response_url, build_kiro_provider_headers,
|
||||
build_kiro_provider_request_body, supports_local_kiro_request_transport_with_network,
|
||||
KIRO_ENVELOPE_NAME,
|
||||
};
|
||||
use crate::ai_serving::transport::policy::{
|
||||
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
|
||||
};
|
||||
use crate::ai_serving::transport::url::{
|
||||
build_claude_messages_url, build_gemini_content_url, build_passthrough_path_url,
|
||||
};
|
||||
use crate::ai_serving::transport::vertex::{
|
||||
build_vertex_api_key_gemini_content_url, resolve_local_vertex_api_key_query_auth,
|
||||
supports_local_vertex_api_key_gemini_transport_with_network,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
apply_local_body_rules, apply_local_header_rules, build_passthrough_headers,
|
||||
ensure_upstream_auth_header, resolve_transport_execution_timeouts,
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
|
||||
LocalResolvedOAuthRequestAuth,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
collect_control_headers, ConversionMode, ExecutionStrategy, GatewayControlDecision,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, AiExecutionDecision, AppState, GatewayError,
|
||||
};
|
||||
|
||||
mod family;
|
||||
mod plans;
|
||||
mod request;
|
||||
|
||||
pub(crate) use self::family::{
|
||||
materialize_local_same_format_provider_candidate_attempts,
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate,
|
||||
resolve_local_same_format_provider_decision_input, LocalSameFormatProviderFamily,
|
||||
LocalSameFormatProviderSpec,
|
||||
};
|
||||
pub(crate) use self::family::{
|
||||
maybe_build_stream_local_same_format_provider_decision_payload,
|
||||
maybe_build_sync_local_same_format_provider_decision_payload,
|
||||
};
|
||||
pub(crate) use self::plans::{
|
||||
build_local_stream_plan_and_reports, build_local_sync_plan_and_reports,
|
||||
};
|
||||
|
||||
const ANTIGRAVITY_ENVELOPE_NAME: &str = "antigravity:v1internal";
|
||||
@@ -0,0 +1,149 @@
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::runtime_miss::{
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal,
|
||||
apply_local_runtime_candidate_terminal_reason, set_local_runtime_miss_diagnostic_reason,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
use super::super::plans::{resolve_stream_spec, resolve_sync_spec};
|
||||
use super::candidates::{
|
||||
materialize_local_same_format_provider_candidate_attempts,
|
||||
resolve_local_same_format_provider_decision_input,
|
||||
};
|
||||
use super::payload::maybe_build_local_same_format_provider_decision_payload_for_candidate;
|
||||
|
||||
pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let requested_model_family = spec_metadata
|
||||
.requested_model_family
|
||||
.expect("same-format provider spec metadata should include requested-model family");
|
||||
|
||||
let Some(input) = resolve_local_same_format_provider_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
extract_requested_model_from_request(parts, body_json, requested_model_family)
|
||||
.as_deref(),
|
||||
"decision_input_unavailable",
|
||||
);
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
state,
|
||||
trace_id,
|
||||
candidate_count,
|
||||
);
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) =
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_sync_plans");
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let requested_model_family = spec_metadata
|
||||
.requested_model_family
|
||||
.expect("same-format provider spec metadata should include requested-model family");
|
||||
|
||||
let Some(input) = resolve_local_same_format_provider_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
extract_requested_model_from_request(parts, body_json, requested_model_family)
|
||||
.as_deref(),
|
||||
"decision_input_unavailable",
|
||||
);
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
state,
|
||||
trace_id,
|
||||
candidate_count,
|
||||
);
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) =
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_stream_plans");
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
@@ -0,0 +1,171 @@
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
materialize_local_execution_candidates_with_serving, LocalCandidateResolutionMode,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_metadata::{
|
||||
build_local_execution_candidate_contract_metadata,
|
||||
build_local_execution_candidate_contract_metadata_for_candidate,
|
||||
LocalExecutionCandidateMetadataParts,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_serving::{
|
||||
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
|
||||
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision, PlannerAppState,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
use super::{
|
||||
LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput,
|
||||
LocalSameFormatProviderSpec,
|
||||
};
|
||||
|
||||
pub(crate) async fn resolve_local_same_format_provider_decision_input(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<LocalSameFormatProviderDecisionInput> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let requested_model = extract_requested_model_from_request(
|
||||
parts,
|
||||
body_json,
|
||||
spec_metadata
|
||||
.requested_model_family
|
||||
.expect("same-format provider specs should declare requested-model family"),
|
||||
)?;
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
auth_context,
|
||||
Some(requested_model.as_str()),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => return None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec_metadata.api_format,
|
||||
error = ?err,
|
||||
"gateway local same-format decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(build_local_requested_model_decision_input(
|
||||
resolved_input,
|
||||
requested_model,
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Result<(Vec<LocalSameFormatProviderCandidateAttempt>, usize), GatewayError> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let sticky_session_token = extract_pool_sticky_session_token(body_json);
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision,
|
||||
);
|
||||
let (candidates, preselection_skipped) = planner_state
|
||||
.list_selectable_candidates_with_skip_reasons(
|
||||
spec_metadata.api_format,
|
||||
&input.requested_model,
|
||||
spec_metadata.require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
Some(&input.auth_snapshot),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
let outcome = materialize_local_execution_candidates_with_serving(
|
||||
planner_state,
|
||||
trace_id,
|
||||
spec_metadata.api_format,
|
||||
Some(&input.requested_model),
|
||||
Some(&input.auth_snapshot),
|
||||
input.required_capabilities.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
persistence_policy,
|
||||
candidates,
|
||||
preselection_skipped
|
||||
.into_iter()
|
||||
.map(|item| SkippedLocalExecutionCandidate {
|
||||
candidate: item.candidate,
|
||||
skip_reason: item.skip_reason,
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
})
|
||||
.collect(),
|
||||
LocalCandidateResolutionMode::Standard,
|
||||
|eligible| {
|
||||
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.api_format,
|
||||
);
|
||||
Some(build_local_execution_candidate_contract_metadata(
|
||||
LocalExecutionCandidateMetadataParts {
|
||||
eligible,
|
||||
provider_api_format: spec_metadata.api_format,
|
||||
client_api_format: spec_metadata.api_format,
|
||||
extra_fields: serde_json::Map::new(),
|
||||
},
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
spec_metadata.api_format,
|
||||
))
|
||||
},
|
||||
|mut skipped_candidate| {
|
||||
let provider_api_format = skipped_candidate
|
||||
.transport
|
||||
.as_ref()
|
||||
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
|
||||
.unwrap_or_else(|| spec_metadata.api_format.to_string());
|
||||
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
|
||||
spec_metadata.api_format,
|
||||
provider_api_format.as_str(),
|
||||
);
|
||||
skipped_candidate.extra_data = Some(
|
||||
build_local_execution_candidate_contract_metadata_for_candidate(
|
||||
&skipped_candidate.candidate,
|
||||
skipped_candidate.transport_ref(),
|
||||
provider_api_format.as_str(),
|
||||
spec_metadata.api_format,
|
||||
serde_json::Map::new(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
provider_api_format.as_str(),
|
||||
),
|
||||
);
|
||||
skipped_candidate
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
Ok((outcome.attempts, outcome.candidate_count))
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
mod build;
|
||||
mod candidates;
|
||||
mod payload;
|
||||
mod request;
|
||||
|
||||
pub(crate) use self::build::{
|
||||
maybe_build_stream_local_same_format_provider_decision_payload,
|
||||
maybe_build_sync_local_same_format_provider_decision_payload,
|
||||
};
|
||||
pub(crate) use self::candidates::{
|
||||
materialize_local_same_format_provider_candidate_attempts,
|
||||
resolve_local_same_format_provider_decision_input,
|
||||
};
|
||||
pub(crate) use self::payload::maybe_build_local_same_format_provider_decision_payload_for_candidate;
|
||||
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalSameFormatProviderCandidateAttempt;
|
||||
pub(crate) use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput as LocalSameFormatProviderDecisionInput;
|
||||
pub(crate) use crate::ai_serving::{LocalSameFormatProviderFamily, LocalSameFormatProviderSpec};
|
||||
@@ -0,0 +1,266 @@
|
||||
use serde_json::json;
|
||||
|
||||
use crate::ai_serving::ai_local_execution_contract_for_formats;
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, LocalExecutionReportContextParts,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_serving::planner::CandidateFailureDiagnostic;
|
||||
use crate::ai_serving::planner::{
|
||||
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, append_local_failover_policy_to_value,
|
||||
AiExecutionDecision, AppState,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
use super::request::resolve_local_same_format_provider_candidate_payload_parts;
|
||||
use super::{
|
||||
LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput,
|
||||
LocalSameFormatProviderSpec,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
body_json: &serde_json::Value,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
attempt: LocalSameFormatProviderCandidateAttempt,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let LocalSameFormatProviderCandidateAttempt {
|
||||
eligible,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
..
|
||||
} = &attempt;
|
||||
let candidate = &eligible.candidate;
|
||||
let (execution_strategy, conversion_mode) =
|
||||
ai_local_execution_contract_for_formats(spec_metadata.api_format, spec_metadata.api_format);
|
||||
let resolved = resolve_local_same_format_provider_candidate_payload_parts(
|
||||
state, parts, trace_id, body_json, input, &attempt, spec,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let prompt_cache_key = resolved
|
||||
.provider_request_body
|
||||
.get("prompt_cache_key")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let proxy = state
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&resolved.transport)
|
||||
.await;
|
||||
let tls_profile = resolve_transport_tls_profile(&resolved.transport);
|
||||
let mut extra_fields = serde_json::Map::new();
|
||||
if let Some(proxy_value) =
|
||||
build_request_trace_proxy_value(Some(&resolved.transport), proxy.as_ref())
|
||||
{
|
||||
extra_fields.insert("proxy".to_string(), proxy_value);
|
||||
}
|
||||
if resolved.is_kiro {
|
||||
extra_fields.insert(
|
||||
"envelope_name".to_string(),
|
||||
json!(crate::ai_serving::transport::kiro::KIRO_ENVELOPE_NAME),
|
||||
);
|
||||
} else if resolved.is_antigravity {
|
||||
extra_fields.insert(
|
||||
"envelope_name".to_string(),
|
||||
json!(super::super::ANTIGRAVITY_ENVELOPE_NAME),
|
||||
);
|
||||
}
|
||||
let report_context = append_local_failover_policy_to_value(
|
||||
append_execution_contract_fields_to_value(
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &input.auth_context,
|
||||
request_id: trace_id,
|
||||
candidate_id,
|
||||
attempt_identity: attempt.attempt_identity(),
|
||||
model: &input.requested_model,
|
||||
provider_name: &resolved.transport.provider.name,
|
||||
provider_id: &candidate.provider_id,
|
||||
endpoint_id: &candidate.endpoint_id,
|
||||
key_id: &candidate.key_id,
|
||||
key_name: Some(&candidate.key_name),
|
||||
model_id: Some(&candidate.model_id),
|
||||
global_model_id: Some(&candidate.global_model_id),
|
||||
global_model_name: Some(&candidate.global_model_name),
|
||||
provider_api_format: spec_metadata.api_format,
|
||||
client_api_format: spec_metadata.api_format,
|
||||
mapped_model: Some(&resolved.mapped_model),
|
||||
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
|
||||
ranking: eligible.ranking.as_ref(),
|
||||
upstream_url: Some(&resolved.upstream_url),
|
||||
header_rules: resolved.transport.endpoint.header_rules.as_ref(),
|
||||
body_rules: resolved.transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: Some(serde_json::Value::Null),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: body_json
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
upstream_is_stream: resolved.upstream_is_stream,
|
||||
has_envelope: resolved.is_kiro || resolved.is_antigravity,
|
||||
needs_conversion: false,
|
||||
extra_fields,
|
||||
}),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.api_format,
|
||||
),
|
||||
&resolved.transport,
|
||||
);
|
||||
let super::request::LocalSameFormatProviderCandidatePayloadParts {
|
||||
transport,
|
||||
is_antigravity: _,
|
||||
is_kiro: _,
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
report_kind,
|
||||
upstream_is_stream,
|
||||
upstream_url,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
} = resolved;
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.to_string(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: None,
|
||||
auth_header,
|
||||
auth_value,
|
||||
provider_api_format: spec_metadata.api_format.to_string(),
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream,
|
||||
report_kind: Some(report_kind.to_string()),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_same_format_provider_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn mark_skipped_local_same_format_provider_candidate_with_extra_data(
|
||||
state: &AppState,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
extra_data: Option<serde_json::Value>,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate_with_extra_data(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
extra_data,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn mark_skipped_local_same_format_provider_candidate_with_failure_diagnostic(
|
||||
state: &AppState,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
diagnostic: CandidateFailureDiagnostic,
|
||||
) {
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision,
|
||||
);
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
diagnostic,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -0,0 +1,237 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::transport::antigravity::{
|
||||
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
|
||||
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
|
||||
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
build_same_format_provider_headers, SameFormatProviderHeadersInput,
|
||||
};
|
||||
use crate::ai_serving::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
|
||||
use crate::AppState;
|
||||
|
||||
mod policy;
|
||||
mod prepare;
|
||||
|
||||
use self::prepare::prepare_local_same_format_provider_candidate;
|
||||
use super::payload::{
|
||||
mark_skipped_local_same_format_provider_candidate,
|
||||
mark_skipped_local_same_format_provider_candidate_with_extra_data,
|
||||
mark_skipped_local_same_format_provider_candidate_with_failure_diagnostic,
|
||||
};
|
||||
use super::{
|
||||
LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput,
|
||||
LocalSameFormatProviderSpec,
|
||||
};
|
||||
use crate::ai_serving::planner::standard::same_format_provider_request_body_failure_extra_data;
|
||||
|
||||
pub(crate) struct LocalSameFormatProviderCandidatePayloadParts {
|
||||
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
pub(super) is_antigravity: bool,
|
||||
pub(super) is_kiro: bool,
|
||||
pub(super) auth_header: Option<String>,
|
||||
pub(super) auth_value: Option<String>,
|
||||
pub(super) mapped_model: String,
|
||||
pub(super) report_kind: &'static str,
|
||||
pub(super) upstream_is_stream: bool,
|
||||
pub(super) upstream_url: String,
|
||||
pub(super) provider_request_headers: BTreeMap<String, String>,
|
||||
pub(super) provider_request_body: Value,
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
body_json: &serde_json::Value,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
attempt: &LocalSameFormatProviderCandidateAttempt,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<LocalSameFormatProviderCandidatePayloadParts> {
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let prepared = prepare_local_same_format_provider_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
input,
|
||||
&attempt.eligible,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let Some(base_provider_request_body) =
|
||||
super::super::request::build_same_format_provider_request_body(
|
||||
body_json,
|
||||
&prepared.mapped_model,
|
||||
spec,
|
||||
prepared.transport.endpoint.body_rules.as_ref(),
|
||||
prepared.upstream_is_stream,
|
||||
prepared.kiro_auth.as_ref(),
|
||||
prepared.is_claude_code,
|
||||
)
|
||||
else {
|
||||
mark_skipped_local_same_format_provider_candidate_with_extra_data(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"provider_request_body_missing",
|
||||
same_format_provider_request_body_failure_extra_data(
|
||||
body_json,
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
prepared.transport.endpoint.body_rules.as_ref(),
|
||||
if prepared.kiro_auth.is_some() {
|
||||
"kiro_envelope"
|
||||
} else {
|
||||
"same_format"
|
||||
},
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let antigravity_auth = if prepared.is_antigravity {
|
||||
match classify_local_antigravity_request_support(
|
||||
&prepared.transport,
|
||||
&base_provider_request_body,
|
||||
AntigravityEnvelopeRequestType::Agent,
|
||||
) {
|
||||
AntigravityRequestSideSupport::Supported(spec) => Some(spec.auth),
|
||||
AntigravityRequestSideSupport::Unsupported(_) => {
|
||||
mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_unsupported",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let provider_request_body = if let Some(antigravity_auth) = antigravity_auth.as_ref() {
|
||||
match build_antigravity_safe_v1internal_request(
|
||||
antigravity_auth,
|
||||
trace_id,
|
||||
&prepared.mapped_model,
|
||||
&base_provider_request_body,
|
||||
AntigravityEnvelopeRequestType::Agent,
|
||||
) {
|
||||
AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
|
||||
AntigravityRequestEnvelopeSupport::Unsupported(_) => {
|
||||
mark_skipped_local_same_format_provider_candidate_with_extra_data(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"provider_request_body_missing",
|
||||
same_format_provider_request_body_failure_extra_data(
|
||||
body_json,
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
prepared.transport.endpoint.body_rules.as_ref(),
|
||||
"antigravity_envelope",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
base_provider_request_body
|
||||
};
|
||||
|
||||
let Some(upstream_url) = super::super::request::build_same_format_upstream_url(
|
||||
parts,
|
||||
&prepared.transport,
|
||||
&prepared.mapped_model,
|
||||
spec,
|
||||
prepared.upstream_is_stream,
|
||||
prepared.kiro_auth.as_ref(),
|
||||
) else {
|
||||
mark_skipped_local_same_format_provider_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"upstream_url_missing",
|
||||
CandidateFailureDiagnostic::upstream_url_missing(
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
"same_format_provider_url",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let extra_headers = antigravity_auth
|
||||
.as_ref()
|
||||
.map(build_antigravity_static_identity_headers)
|
||||
.unwrap_or_default();
|
||||
let Some(provider_request_headers) =
|
||||
build_same_format_provider_headers(SameFormatProviderHeadersInput {
|
||||
headers: &parts.headers,
|
||||
provider_request_body: &provider_request_body,
|
||||
original_request_body: body_json,
|
||||
header_rules: prepared.transport.endpoint.header_rules.as_ref(),
|
||||
behavior: prepared.behavior,
|
||||
auth_header: prepared.auth_header.as_deref(),
|
||||
auth_value: prepared.auth_value.as_deref(),
|
||||
extra_headers: &extra_headers,
|
||||
key_fingerprint: prepared.transport.key.fingerprint.as_ref(),
|
||||
kiro_auth_config: prepared.kiro_auth.as_ref().map(|auth| &auth.auth_config),
|
||||
kiro_machine_id: prepared
|
||||
.kiro_auth
|
||||
.as_ref()
|
||||
.map(|auth| auth.machine_id.as_str()),
|
||||
})
|
||||
else {
|
||||
mark_skipped_local_same_format_provider_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
CandidateFailureDiagnostic::header_rules_apply_failed(
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
"same_format_provider_headers",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
Some(LocalSameFormatProviderCandidatePayloadParts {
|
||||
transport: prepared.transport,
|
||||
is_antigravity: prepared.is_antigravity,
|
||||
is_kiro: prepared.is_kiro,
|
||||
auth_header: prepared.auth_header,
|
||||
auth_value: prepared.auth_value,
|
||||
mapped_model: prepared.mapped_model,
|
||||
report_kind: prepared.report_kind,
|
||||
upstream_is_stream: prepared.upstream_is_stream,
|
||||
upstream_url,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
use crate::ai_serving::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata;
|
||||
use crate::ai_serving::transport::{
|
||||
classify_same_format_provider_request_behavior as classify_same_format_provider_request_behavior_impl,
|
||||
resolve_same_format_provider_direct_auth as resolve_same_format_provider_direct_auth_impl,
|
||||
same_format_provider_transport_supported as same_format_provider_transport_supported_impl,
|
||||
same_format_provider_transport_unsupported_reason as same_format_provider_transport_unsupported_reason_impl,
|
||||
should_try_same_format_provider_oauth_auth as should_try_same_format_provider_oauth_auth_impl,
|
||||
GatewayProviderTransportSnapshot, SameFormatProviderFamily, SameFormatProviderRequestBehavior,
|
||||
SameFormatProviderRequestBehaviorParams,
|
||||
};
|
||||
|
||||
use super::super::LocalSameFormatProviderFamily;
|
||||
|
||||
pub(super) fn classify_same_format_provider_request_behavior(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
spec_metadata: LocalExecutionSurfaceSpecMetadata,
|
||||
) -> SameFormatProviderRequestBehavior {
|
||||
classify_same_format_provider_request_behavior_impl(
|
||||
transport,
|
||||
SameFormatProviderRequestBehaviorParams {
|
||||
require_streaming: spec_metadata.require_streaming,
|
||||
report_kind: spec_metadata
|
||||
.report_kind
|
||||
.expect("same-format provider specs should declare report kind"),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn same_format_provider_transport_supported(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
same_format_provider_transport_supported_impl(
|
||||
behavior,
|
||||
transport,
|
||||
same_format_provider_family(family),
|
||||
api_format,
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn same_format_provider_transport_unsupported_reason(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
api_format: &str,
|
||||
) -> Option<&'static str> {
|
||||
same_format_provider_transport_unsupported_reason_impl(
|
||||
behavior,
|
||||
transport,
|
||||
same_format_provider_family(family),
|
||||
api_format,
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn should_try_same_format_provider_oauth_auth(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
) -> bool {
|
||||
should_try_same_format_provider_oauth_auth_impl(
|
||||
behavior,
|
||||
transport,
|
||||
same_format_provider_family(family),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn resolve_same_format_provider_direct_auth(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
) -> Option<(String, String)> {
|
||||
resolve_same_format_provider_direct_auth_impl(
|
||||
behavior,
|
||||
transport,
|
||||
same_format_provider_family(family),
|
||||
)
|
||||
}
|
||||
|
||||
fn same_format_provider_family(family: LocalSameFormatProviderFamily) -> SameFormatProviderFamily {
|
||||
match family {
|
||||
LocalSameFormatProviderFamily::Standard => SameFormatProviderFamily::Standard,
|
||||
LocalSameFormatProviderFamily::Gemini => SameFormatProviderFamily::Gemini,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,175 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::ai_serving::planner::candidate_preparation::{
|
||||
resolve_candidate_mapped_model, resolve_candidate_oauth_auth, OauthPreparationContext,
|
||||
};
|
||||
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_serving::transport::kiro::KiroRequestAuth;
|
||||
use crate::ai_serving::transport::vertex::resolve_local_vertex_api_key_query_auth;
|
||||
use crate::ai_serving::transport::SameFormatProviderRequestBehavior;
|
||||
use crate::ai_serving::{
|
||||
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
|
||||
};
|
||||
use crate::AppState;
|
||||
|
||||
use super::super::LocalSameFormatProviderDecisionInput;
|
||||
use super::super::LocalSameFormatProviderSpec;
|
||||
use super::policy::{
|
||||
classify_same_format_provider_request_behavior, resolve_same_format_provider_direct_auth,
|
||||
same_format_provider_transport_supported, same_format_provider_transport_unsupported_reason,
|
||||
should_try_same_format_provider_oauth_auth,
|
||||
};
|
||||
|
||||
pub(super) struct PreparedSameFormatProviderCandidate {
|
||||
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
pub(super) behavior: SameFormatProviderRequestBehavior,
|
||||
pub(super) is_antigravity: bool,
|
||||
pub(super) is_claude_code: bool,
|
||||
pub(super) is_vertex: bool,
|
||||
pub(super) is_kiro: bool,
|
||||
pub(super) kiro_auth: Option<KiroRequestAuth>,
|
||||
pub(super) auth_header: Option<String>,
|
||||
pub(super) auth_value: Option<String>,
|
||||
pub(super) mapped_model: String,
|
||||
pub(super) report_kind: &'static str,
|
||||
pub(super) upstream_is_stream: bool,
|
||||
}
|
||||
|
||||
pub(super) async fn prepare_local_same_format_provider_candidate(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
eligible: &EligibleLocalExecutionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<PreparedSameFormatProviderCandidate> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let candidate = &eligible.candidate;
|
||||
let transport = Arc::clone(&eligible.transport);
|
||||
let behavior = classify_same_format_provider_request_behavior(&transport, spec_metadata);
|
||||
|
||||
if !same_format_provider_transport_supported(
|
||||
&behavior,
|
||||
&transport,
|
||||
spec.family,
|
||||
spec_metadata.api_format,
|
||||
) {
|
||||
let skip_reason = same_format_provider_transport_unsupported_reason(
|
||||
&behavior,
|
||||
&transport,
|
||||
spec.family,
|
||||
spec_metadata.api_format,
|
||||
)
|
||||
.unwrap_or("transport_unsupported");
|
||||
super::super::payload::mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let vertex_query_auth = if behavior.is_vertex {
|
||||
resolve_local_vertex_api_key_query_auth(&transport)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let should_try_oauth_auth =
|
||||
should_try_same_format_provider_oauth_auth(&behavior, &transport, spec.family);
|
||||
let oauth_auth = if should_try_oauth_auth {
|
||||
resolve_candidate_oauth_auth(
|
||||
planner_state,
|
||||
&transport,
|
||||
OauthPreparationContext {
|
||||
trace_id,
|
||||
api_format: spec_metadata.api_format,
|
||||
operation: "same_format_provider_prepare",
|
||||
},
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let kiro_auth = match oauth_auth.as_ref() {
|
||||
Some(LocalResolvedOAuthRequestAuth::Kiro(auth)) => Some(auth.clone()),
|
||||
_ => None,
|
||||
};
|
||||
let auth = if let Some(kiro_auth) = kiro_auth.as_ref() {
|
||||
Some((kiro_auth.name.to_string(), kiro_auth.value.clone()))
|
||||
} else if let Some(LocalResolvedOAuthRequestAuth::Header { name, value }) = oauth_auth.as_ref()
|
||||
{
|
||||
Some((name.clone(), value.clone()))
|
||||
} else {
|
||||
resolve_same_format_provider_direct_auth(&behavior, &transport, spec.family)
|
||||
};
|
||||
let (auth_header, auth_value) = match auth {
|
||||
Some((name, value)) => (Some(name), Some(value)),
|
||||
None if behavior.is_vertex && vertex_query_auth.is_some() => (None, None),
|
||||
None => {
|
||||
super::super::payload::mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
if behavior.is_vertex && vertex_query_auth.is_none() {
|
||||
super::super::payload::mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let mapped_model = match resolve_candidate_mapped_model(candidate) {
|
||||
Ok(mapped_model) => mapped_model,
|
||||
Err(skip_reason) => {
|
||||
super::super::payload::mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(PreparedSameFormatProviderCandidate {
|
||||
transport,
|
||||
behavior,
|
||||
is_antigravity: behavior.is_antigravity,
|
||||
is_claude_code: behavior.is_claude_code,
|
||||
is_vertex: behavior.is_vertex,
|
||||
is_kiro: behavior.is_kiro,
|
||||
kiro_auth,
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
report_kind: behavior.report_kind,
|
||||
upstream_is_stream: behavior.upstream_is_stream,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::runtime_miss::{
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal,
|
||||
apply_local_runtime_candidate_terminal_reason, set_local_runtime_miss_diagnostic_reason,
|
||||
};
|
||||
use crate::ai_serving::planner::spec_metadata::{
|
||||
build_stream_plan_from_requested_model_family, build_sync_plan_from_requested_model_family,
|
||||
local_same_format_provider_spec_metadata,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::{
|
||||
resolve_local_same_format_stream_spec as resolve_stream_spec,
|
||||
resolve_local_same_format_sync_spec as resolve_sync_spec,
|
||||
};
|
||||
|
||||
use super::{
|
||||
materialize_local_same_format_provider_candidate_attempts,
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate,
|
||||
resolve_local_same_format_provider_decision_input, AiStreamAttempt, AiSyncAttempt, AppState,
|
||||
GatewayControlDecision, GatewayError, LocalSameFormatProviderSpec,
|
||||
};
|
||||
|
||||
pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let requested_model_family = spec_metadata
|
||||
.requested_model_family
|
||||
.expect("same-format provider spec metadata should include requested-model family");
|
||||
let Some(input) = resolve_local_same_format_provider_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
extract_requested_model_from_request(parts, body_json, requested_model_family)
|
||||
.as_deref(),
|
||||
"decision_input_unavailable",
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
state,
|
||||
trace_id,
|
||||
candidate_count,
|
||||
);
|
||||
if candidate_count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let built = build_sync_plan_from_requested_model_family(
|
||||
requested_model_family,
|
||||
parts,
|
||||
body_json,
|
||||
payload,
|
||||
);
|
||||
|
||||
match built {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec_metadata.api_format,
|
||||
error = ?err,
|
||||
"gateway local same-format sync decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_sync_plans");
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let requested_model_family = spec_metadata
|
||||
.requested_model_family
|
||||
.expect("same-format provider spec metadata should include requested-model family");
|
||||
let Some(input) = resolve_local_same_format_provider_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
extract_requested_model_from_request(parts, body_json, requested_model_family)
|
||||
.as_deref(),
|
||||
"decision_input_unavailable",
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
spec_metadata.decision_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
state,
|
||||
trace_id,
|
||||
candidate_count,
|
||||
);
|
||||
if candidate_count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let built = build_stream_plan_from_requested_model_family(
|
||||
requested_model_family,
|
||||
parts,
|
||||
body_json,
|
||||
payload,
|
||||
);
|
||||
|
||||
match built {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec_metadata.api_format,
|
||||
error = ?err,
|
||||
"gateway local same-format stream decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
apply_local_runtime_candidate_terminal_reason(state, trace_id, "no_local_stream_plans");
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
mod body;
|
||||
mod url;
|
||||
|
||||
pub(super) use self::body::build_same_format_provider_request_body;
|
||||
pub(super) use self::url::build_same_format_upstream_url;
|
||||
@@ -0,0 +1,36 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use super::super::LocalSameFormatProviderSpec;
|
||||
use crate::ai_serving::transport::{
|
||||
build_same_format_provider_request_body as build_same_format_provider_request_body_impl,
|
||||
SameFormatProviderFamily, SameFormatProviderRequestBodyInput,
|
||||
};
|
||||
|
||||
pub(crate) fn build_same_format_provider_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
body_rules: Option<&Value>,
|
||||
upstream_is_stream: bool,
|
||||
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
|
||||
is_claude_code: bool,
|
||||
) -> Option<Value> {
|
||||
build_same_format_provider_request_body_impl(SameFormatProviderRequestBodyInput {
|
||||
body_json,
|
||||
mapped_model,
|
||||
family: same_format_provider_family(spec.family),
|
||||
body_rules,
|
||||
upstream_is_stream,
|
||||
kiro_auth_config: kiro_auth.map(|auth| &auth.auth_config),
|
||||
is_claude_code,
|
||||
})
|
||||
}
|
||||
|
||||
fn same_format_provider_family(
|
||||
family: super::super::LocalSameFormatProviderFamily,
|
||||
) -> SameFormatProviderFamily {
|
||||
match family {
|
||||
super::super::LocalSameFormatProviderFamily::Standard => SameFormatProviderFamily::Standard,
|
||||
super::super::LocalSameFormatProviderFamily::Gemini => SameFormatProviderFamily::Gemini,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
use crate::ai_serving::transport::{
|
||||
build_same_format_provider_upstream_url as build_same_format_provider_upstream_url_impl,
|
||||
SameFormatProviderUpstreamUrlParams,
|
||||
};
|
||||
use crate::ai_serving::GatewayProviderTransportSnapshot;
|
||||
|
||||
use super::super::LocalSameFormatProviderSpec;
|
||||
|
||||
pub(crate) fn build_same_format_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
mapped_model: &str,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
upstream_is_stream: bool,
|
||||
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
|
||||
) -> Option<String> {
|
||||
build_same_format_provider_upstream_url_impl(
|
||||
transport,
|
||||
SameFormatProviderUpstreamUrlParams {
|
||||
provider_api_format: spec.api_format,
|
||||
mapped_model,
|
||||
upstream_is_stream,
|
||||
request_query: parts.uri.query(),
|
||||
kiro_api_region: kiro_auth.map(|auth| auth.auth_config.effective_api_region()),
|
||||
},
|
||||
)
|
||||
}
|
||||
Reference in New Issue
Block a user