mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-11 03:39:49 +08:00
feat: add routing profile scheduling policies
This commit is contained in:
@@ -1,4 +1,5 @@
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, LocalExecutionReportContextParts,
|
||||
};
|
||||
@@ -10,7 +11,9 @@ use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_profile,
|
||||
};
|
||||
use crate::ai_serving::{ai_local_execution_contract_for_formats, PlannerAppState};
|
||||
use crate::{append_execution_contract_fields_to_value, AiExecutionDecision, AppState};
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, AiExecutionDecision, AppState, GatewayError,
|
||||
};
|
||||
|
||||
use super::request::resolve_local_openai_image_candidate_payload_parts;
|
||||
use super::support::{LocalOpenAiImageCandidateAttempt, LocalOpenAiImageDecisionInput};
|
||||
@@ -25,11 +28,11 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
input: &LocalOpenAiImageDecisionInput,
|
||||
attempt: LocalOpenAiImageCandidateAttempt,
|
||||
spec: LocalOpenAiImageSpec,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let spec_metadata = local_openai_image_spec_metadata(spec);
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let attempt_identity = attempt.attempt_identity();
|
||||
let resolved = resolve_local_openai_image_candidate_payload_parts(
|
||||
let Some(resolved) = resolve_local_openai_image_candidate_payload_parts(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
@@ -39,7 +42,10 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
&attempt,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let LocalOpenAiImageCandidateAttempt {
|
||||
eligible,
|
||||
candidate_id,
|
||||
@@ -88,6 +94,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(spec_metadata.require_streaming);
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
let report_context = append_execution_contract_fields_to_value(
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &input.auth_context,
|
||||
@@ -114,7 +121,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
body_rules: transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: Some(serde_json::Value::String(parts.method.to_string())),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
original_headers: effective_headers,
|
||||
request_path: Some(parts.uri.path()),
|
||||
request_query_string: parts.uri.query(),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
@@ -134,39 +141,39 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
provider_api_format.as_str(),
|
||||
);
|
||||
|
||||
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: resolved.upstream_url,
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(resolved.auth_header),
|
||||
auth_value: Some(resolved.auth_value),
|
||||
provider_api_format,
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: resolved.requested_model,
|
||||
mapped_model: resolved.mapped_model,
|
||||
prompt_cache_key: None,
|
||||
provider_request_headers: resolved.provider_request_headers,
|
||||
provider_request_body: Some(resolved.provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
let mut decision = 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: resolved.upstream_url,
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(resolved.auth_header),
|
||||
auth_value: Some(resolved.auth_value),
|
||||
provider_api_format,
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: resolved.requested_model,
|
||||
mapped_model: resolved.mapped_model,
|
||||
prompt_cache_key: None,
|
||||
provider_request_headers: resolved.provider_request_headers,
|
||||
provider_request_body: Some(resolved.provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
});
|
||||
apply_provider_request_routing_policy_to_decision(input, &mut decision)?;
|
||||
Ok(Some(decision))
|
||||
}
|
||||
|
||||
@@ -59,6 +59,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let transport = &attempt.eligible.transport;
|
||||
let provider_api_format = attempt.eligible.provider_api_format.as_str();
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
|
||||
if provider_api_format == "gemini:generate_content" {
|
||||
return resolve_local_openai_image_to_gemini_candidate_payload_parts(
|
||||
@@ -170,7 +171,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
|
||||
|
||||
let Some(mut provider_request_headers) =
|
||||
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
auth_header: &auth_header,
|
||||
auth_value: &auth_value,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
@@ -201,7 +202,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
|
||||
apply_codex_openai_responses_special_headers(
|
||||
&mut provider_request_headers,
|
||||
&provider_request_body,
|
||||
&parts.headers,
|
||||
effective_headers,
|
||||
transport.provider.provider_type.as_str(),
|
||||
spec_metadata.api_format,
|
||||
Some(trace_id),
|
||||
@@ -256,6 +257,7 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let transport = &attempt.eligible.transport;
|
||||
let provider_api_format = "gemini:generate_content";
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
|
||||
let prepared_candidate = match prepare_header_authenticated_candidate(
|
||||
PlannerAppState::new(state),
|
||||
@@ -332,7 +334,7 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
|
||||
converted.body_json,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
body_json,
|
||||
&parts.headers,
|
||||
effective_headers,
|
||||
) {
|
||||
Some(body) => body,
|
||||
None => {
|
||||
@@ -384,7 +386,7 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
|
||||
transport,
|
||||
provider_api_format,
|
||||
same_format: false,
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
auth_header: &prepared_candidate.auth_header,
|
||||
auth_value: &prepared_candidate.auth_value,
|
||||
extra_headers: &BTreeMap::new(),
|
||||
|
||||
@@ -13,6 +13,7 @@ use crate::ai_serving::planner::candidate_metadata::{
|
||||
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
attach_routing_policy_to_local_requested_model_input,
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
@@ -42,12 +43,16 @@ pub(super) async fn resolve_local_openai_image_decision_input(
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<LocalOpenAiImageDecisionInput> {
|
||||
) -> Result<Option<LocalOpenAiImageDecisionInput>, GatewayError> {
|
||||
let Some(auth_context) = resolve_local_openai_image_auth_context(decision) else {
|
||||
return None;
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let requested_model = resolve_requested_image_model_for_request(parts, body_json, body_base64)?;
|
||||
let Some(requested_model) =
|
||||
resolve_requested_image_model_for_request(parts, body_json, body_base64)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
@@ -58,21 +63,37 @@ pub(super) async fn resolve_local_openai_image_decision_input(
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => return None,
|
||||
Ok(None) => return Ok(None),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai image decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
|
||||
let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
|
||||
input.request_auth_channel = decision.request_auth_channel.clone();
|
||||
input.client_session_affinity = client_session_affinity_from_parts(parts, Some(body_json));
|
||||
Some(input)
|
||||
if let Err(err) = attach_routing_policy_to_local_requested_model_input(
|
||||
state,
|
||||
parts,
|
||||
&mut input,
|
||||
body_json,
|
||||
"openai:image",
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai image decision routing profile resolution failed"
|
||||
);
|
||||
return Err(err);
|
||||
}
|
||||
Ok(Some(input))
|
||||
}
|
||||
|
||||
fn resolve_local_openai_image_auth_context(
|
||||
@@ -229,6 +250,7 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
@@ -305,6 +327,7 @@ async fn materialize_local_openai_image_candidate_attempts(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
|
||||
Reference in New Issue
Block a user