mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
refactor ai serving modules and crates
This commit is contained in:
96
crates/aether-ai-serving/src/attempt_loop.rs
Normal file
96
crates/aether-ai-serving/src/attempt_loop.rs
Normal file
@@ -0,0 +1,96 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
pub trait AiExecutionAttempt {
|
||||
fn execution_plan(&self) -> &aether_contracts::ExecutionPlan;
|
||||
|
||||
fn report_kind(&self) -> Option<String>;
|
||||
|
||||
fn report_context(&self) -> Option<serde_json::Value>;
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum AiAttemptLoopOutcome<Response, Exhaustion> {
|
||||
Responded(Response),
|
||||
Exhausted(Exhaustion),
|
||||
NoPath,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AiAttemptLoopPort<Attempt>: Send + Sync
|
||||
where
|
||||
Attempt: AiExecutionAttempt + Send + Sync + 'static,
|
||||
{
|
||||
type Response: Send;
|
||||
type Exhaustion: Send;
|
||||
type Error: Send;
|
||||
|
||||
async fn execute_attempt(
|
||||
&self,
|
||||
attempt: &Attempt,
|
||||
) -> Result<Option<Self::Response>, Self::Error>;
|
||||
|
||||
async fn mark_unused_attempts(&self, attempts: Vec<Attempt>) -> Result<(), Self::Error>;
|
||||
|
||||
async fn build_exhaustion(
|
||||
&self,
|
||||
last_plan: aether_contracts::ExecutionPlan,
|
||||
last_report_context: Option<serde_json::Value>,
|
||||
) -> Result<Self::Exhaustion, Self::Error>;
|
||||
}
|
||||
|
||||
pub async fn run_ai_attempt_loop<Port, Attempt>(
|
||||
port: &Port,
|
||||
attempts: Vec<Attempt>,
|
||||
) -> Result<AiAttemptLoopOutcome<Port::Response, Port::Exhaustion>, Port::Error>
|
||||
where
|
||||
Port: AiAttemptLoopPort<Attempt>,
|
||||
Attempt: AiExecutionAttempt + Send + Sync + 'static,
|
||||
{
|
||||
let mut remaining = attempts.into_iter();
|
||||
let mut last_attempted = None;
|
||||
|
||||
while let Some(attempt) = remaining.next() {
|
||||
last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context()));
|
||||
if let Some(response) = port.execute_attempt(&attempt).await? {
|
||||
port.mark_unused_attempts(remaining.collect()).await?;
|
||||
return Ok(AiAttemptLoopOutcome::Responded(response));
|
||||
}
|
||||
}
|
||||
|
||||
let Some((last_plan, last_report_context)) = last_attempted else {
|
||||
return Ok(AiAttemptLoopOutcome::NoPath);
|
||||
};
|
||||
|
||||
Ok(AiAttemptLoopOutcome::Exhausted(
|
||||
port.build_exhaustion(last_plan, last_report_context)
|
||||
.await?,
|
||||
))
|
||||
}
|
||||
|
||||
impl AiExecutionAttempt for crate::dto::AiSyncAttempt {
|
||||
fn execution_plan(&self) -> &aether_contracts::ExecutionPlan {
|
||||
&self.plan
|
||||
}
|
||||
|
||||
fn report_kind(&self) -> Option<String> {
|
||||
self.report_kind.clone()
|
||||
}
|
||||
|
||||
fn report_context(&self) -> Option<serde_json::Value> {
|
||||
self.report_context.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl AiExecutionAttempt for crate::dto::AiStreamAttempt {
|
||||
fn execution_plan(&self) -> &aether_contracts::ExecutionPlan {
|
||||
&self.plan
|
||||
}
|
||||
|
||||
fn report_kind(&self) -> Option<String> {
|
||||
self.report_kind.clone()
|
||||
}
|
||||
|
||||
fn report_context(&self) -> Option<serde_json::Value> {
|
||||
self.report_context.clone()
|
||||
}
|
||||
}
|
||||
523
crates/aether-ai-serving/src/attempt_plan.rs
Normal file
523
crates/aether-ai-serving/src/attempt_plan.rs
Normal file
@@ -0,0 +1,523 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_ai_surfaces::api::ExecutionRuntimeAuthContext;
|
||||
use aether_contracts::{ExecutionPlan, RequestBody};
|
||||
use url::Url;
|
||||
|
||||
use crate::dto::AiExecutionDecision;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AiDecisionPlanCore {
|
||||
pub request_id: String,
|
||||
pub provider_id: String,
|
||||
pub endpoint_id: String,
|
||||
pub key_id: String,
|
||||
pub provider_api_format: String,
|
||||
pub client_api_format: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AiUpstreamAuthPair {
|
||||
pub header: String,
|
||||
pub value: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct AiExecutionPlanFromDecisionParts {
|
||||
pub core: AiDecisionPlanCore,
|
||||
pub method: String,
|
||||
pub url: String,
|
||||
pub headers: BTreeMap<String, String>,
|
||||
pub content_type: Option<String>,
|
||||
pub body: RequestBody,
|
||||
pub stream: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct AiExecutionDecisionFromPlanParts {
|
||||
pub action: String,
|
||||
pub decision_kind: Option<String>,
|
||||
pub request_id: Option<String>,
|
||||
pub upstream_base_url: Option<String>,
|
||||
pub include_auth_pair: bool,
|
||||
pub plan: ExecutionPlan,
|
||||
pub report_kind: Option<String>,
|
||||
pub report_context: Option<serde_json::Value>,
|
||||
pub auth_context: Option<ExecutionRuntimeAuthContext>,
|
||||
}
|
||||
|
||||
pub fn take_ai_non_empty_string(value: &mut Option<String>) -> Option<String> {
|
||||
value.take().filter(|value| !value.trim().is_empty())
|
||||
}
|
||||
|
||||
pub fn trim_ai_owned_non_empty_string(value: String) -> Option<String> {
|
||||
let trimmed = value.trim();
|
||||
if trimmed.is_empty() {
|
||||
return None;
|
||||
}
|
||||
if trimmed.len() == value.len() {
|
||||
return Some(value);
|
||||
}
|
||||
Some(trimmed.to_owned())
|
||||
}
|
||||
|
||||
pub fn take_ai_decision_plan_core(payload: &mut AiExecutionDecision) -> Option<AiDecisionPlanCore> {
|
||||
Some(AiDecisionPlanCore {
|
||||
request_id: take_ai_non_empty_string(&mut payload.request_id)?,
|
||||
provider_id: take_ai_non_empty_string(&mut payload.provider_id)?,
|
||||
endpoint_id: take_ai_non_empty_string(&mut payload.endpoint_id)?,
|
||||
key_id: take_ai_non_empty_string(&mut payload.key_id)?,
|
||||
provider_api_format: take_ai_non_empty_string(&mut payload.provider_api_format)?,
|
||||
client_api_format: take_ai_non_empty_string(&mut payload.client_api_format)?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn take_ai_upstream_auth_pair(
|
||||
payload: &mut AiExecutionDecision,
|
||||
) -> Option<Option<AiUpstreamAuthPair>> {
|
||||
let header = take_ai_non_empty_string(&mut payload.auth_header);
|
||||
let value = take_ai_non_empty_string(&mut payload.auth_value);
|
||||
match (header, value) {
|
||||
(Some(header), Some(value)) => Some(Some(AiUpstreamAuthPair { header, value })),
|
||||
(None, None) => Some(None),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn resolve_ai_passthrough_sync_request_body(
|
||||
provider_request_body: Option<serde_json::Value>,
|
||||
provider_request_body_base64: Option<String>,
|
||||
) -> RequestBody {
|
||||
if let Some(body_bytes_b64) =
|
||||
provider_request_body_base64.and_then(trim_ai_owned_non_empty_string)
|
||||
{
|
||||
return RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: Some(body_bytes_b64),
|
||||
body_ref: None,
|
||||
};
|
||||
}
|
||||
|
||||
match provider_request_body.unwrap_or(serde_json::Value::Null) {
|
||||
serde_json::Value::Null => RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
},
|
||||
other => RequestBody::from_json(other),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_ai_execution_plan_from_decision(
|
||||
payload: &mut AiExecutionDecision,
|
||||
parts: AiExecutionPlanFromDecisionParts,
|
||||
) -> ExecutionPlan {
|
||||
ExecutionPlan {
|
||||
request_id: parts.core.request_id,
|
||||
candidate_id: payload.candidate_id.take(),
|
||||
provider_name: payload.provider_name.take(),
|
||||
provider_id: parts.core.provider_id,
|
||||
endpoint_id: parts.core.endpoint_id,
|
||||
key_id: parts.core.key_id,
|
||||
method: parts.method,
|
||||
url: parts.url,
|
||||
headers: parts.headers,
|
||||
content_type: parts.content_type,
|
||||
content_encoding: None,
|
||||
body: parts.body,
|
||||
stream: parts.stream,
|
||||
client_api_format: parts.core.client_api_format,
|
||||
provider_api_format: parts.core.provider_api_format,
|
||||
model_name: payload.model_name.take(),
|
||||
proxy: payload.proxy.take(),
|
||||
tls_profile: payload.tls_profile.take(),
|
||||
timeouts: payload.timeouts.take(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_ai_execution_decision_from_plan(
|
||||
parts: AiExecutionDecisionFromPlanParts,
|
||||
) -> AiExecutionDecision {
|
||||
let ExecutionPlan {
|
||||
request_id,
|
||||
candidate_id,
|
||||
provider_name,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
method,
|
||||
url,
|
||||
headers,
|
||||
content_type,
|
||||
content_encoding: _content_encoding,
|
||||
body,
|
||||
stream,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
model_name,
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts,
|
||||
} = parts.plan;
|
||||
let auth_pair = parts
|
||||
.include_auth_pair
|
||||
.then(|| extract_ai_auth_header_pair(&headers))
|
||||
.flatten();
|
||||
let provider_contract = provider_api_format.clone();
|
||||
let client_contract = client_api_format.clone();
|
||||
let request_id = parts.request_id.unwrap_or(request_id);
|
||||
let auth_header = auth_pair.map(|(name, _)| name.to_string());
|
||||
let auth_value = auth_pair.map(|(_, value)| value.to_string());
|
||||
let RequestBody {
|
||||
json_body,
|
||||
body_bytes_b64,
|
||||
body_ref: _body_ref,
|
||||
} = body;
|
||||
|
||||
AiExecutionDecision {
|
||||
action: parts.action,
|
||||
decision_kind: parts.decision_kind,
|
||||
execution_strategy: Some(ai_execution_strategy_for_formats(
|
||||
provider_api_format.as_str(),
|
||||
client_api_format.as_str(),
|
||||
)),
|
||||
conversion_mode: Some(ai_conversion_mode_for_formats(
|
||||
provider_api_format.as_str(),
|
||||
client_api_format.as_str(),
|
||||
)),
|
||||
request_id: Some(request_id),
|
||||
candidate_id,
|
||||
provider_name,
|
||||
provider_id: Some(provider_id),
|
||||
endpoint_id: Some(endpoint_id),
|
||||
key_id: Some(key_id),
|
||||
upstream_base_url: parts.upstream_base_url,
|
||||
upstream_url: Some(url),
|
||||
provider_request_method: Some(method),
|
||||
auth_header,
|
||||
auth_value,
|
||||
provider_api_format: Some(provider_api_format),
|
||||
client_api_format: Some(client_api_format),
|
||||
provider_contract: Some(provider_contract),
|
||||
client_contract: Some(client_contract),
|
||||
model_name,
|
||||
mapped_model: None,
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: headers,
|
||||
provider_request_body: json_body,
|
||||
provider_request_body_base64: body_bytes_b64,
|
||||
content_type,
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts,
|
||||
upstream_is_stream: stream,
|
||||
report_kind: parts.report_kind,
|
||||
report_context: parts.report_context,
|
||||
auth_context: parts.auth_context,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn extract_ai_auth_header_pair(headers: &BTreeMap<String, String>) -> Option<(&str, &str)> {
|
||||
[
|
||||
"authorization",
|
||||
"x-api-key",
|
||||
"api-key",
|
||||
"x-goog-api-key",
|
||||
"proxy-authorization",
|
||||
]
|
||||
.into_iter()
|
||||
.find_map(|name| {
|
||||
headers
|
||||
.iter()
|
||||
.find(|(header_name, _)| header_name.eq_ignore_ascii_case(name))
|
||||
.map(|(header_name, value)| (header_name.as_str(), value.as_str()))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn infer_ai_upstream_base_url(upstream_url: &str) -> Option<String> {
|
||||
let parsed = Url::parse(upstream_url).ok()?;
|
||||
let host = parsed.host_str()?;
|
||||
let mut base = format!("{}://{}", parsed.scheme(), host);
|
||||
if let Some(port) = parsed.port() {
|
||||
base.push(':');
|
||||
base.push_str(port.to_string().as_str());
|
||||
}
|
||||
let base_path = infer_ai_upstream_base_path(parsed.path());
|
||||
if !base_path.is_empty() {
|
||||
base.push_str(base_path);
|
||||
}
|
||||
Some(base)
|
||||
}
|
||||
|
||||
fn infer_ai_upstream_base_path(path: &str) -> &str {
|
||||
let trimmed = path.trim_end_matches('/');
|
||||
if trimmed.is_empty() || trimmed == "/" {
|
||||
return "";
|
||||
}
|
||||
|
||||
for suffix in [
|
||||
"/responses/compact",
|
||||
"/responses",
|
||||
"/chat/completions",
|
||||
"/messages",
|
||||
] {
|
||||
if let Some(prefix) = trimmed.strip_suffix(suffix) {
|
||||
return normalize_inferred_ai_base_path(prefix);
|
||||
}
|
||||
}
|
||||
|
||||
for marker in ["/v1/videos", "/v1beta/"] {
|
||||
if let Some((prefix, _)) = trimmed.split_once(marker) {
|
||||
return normalize_inferred_ai_base_path(prefix);
|
||||
}
|
||||
}
|
||||
|
||||
normalize_inferred_ai_base_path(trimmed)
|
||||
}
|
||||
|
||||
fn normalize_inferred_ai_base_path(path: &str) -> &str {
|
||||
let trimmed = path.trim_end_matches('/');
|
||||
if trimmed.is_empty() || trimmed == "/" {
|
||||
""
|
||||
} else {
|
||||
trimmed
|
||||
}
|
||||
}
|
||||
|
||||
fn ai_execution_strategy_for_formats(provider_api_format: &str, client_api_format: &str) -> String {
|
||||
if provider_api_format == client_api_format {
|
||||
"local_same_format"
|
||||
} else {
|
||||
"local_cross_format"
|
||||
}
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn ai_conversion_mode_for_formats(provider_api_format: &str, client_api_format: &str) -> String {
|
||||
if provider_api_format == client_api_format {
|
||||
"none"
|
||||
} else {
|
||||
"bidirectional"
|
||||
}
|
||||
.to_string()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn take_ai_decision_plan_core_consumes_required_non_empty_fields() {
|
||||
let mut payload = test_decision();
|
||||
|
||||
let core =
|
||||
take_ai_decision_plan_core(&mut payload).expect("core fields should be available");
|
||||
|
||||
assert_eq!(core.request_id, "req_1");
|
||||
assert_eq!(core.provider_id, "provider_1");
|
||||
assert_eq!(core.endpoint_id, "endpoint_1");
|
||||
assert_eq!(core.key_id, "key_1");
|
||||
assert_eq!(core.provider_api_format, "openai:chat");
|
||||
assert_eq!(core.client_api_format, "openai:chat");
|
||||
assert!(payload.request_id.is_none());
|
||||
assert!(payload.provider_api_format.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn take_ai_decision_plan_core_rejects_blank_required_fields() {
|
||||
let mut payload = test_decision();
|
||||
payload.endpoint_id = Some(" ".to_string());
|
||||
|
||||
assert!(take_ai_decision_plan_core(&mut payload).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn take_ai_upstream_auth_pair_rejects_incomplete_auth() {
|
||||
let mut payload = test_decision();
|
||||
payload.auth_header = Some("authorization".to_string());
|
||||
payload.auth_value = Some(" ".to_string());
|
||||
|
||||
assert!(take_ai_upstream_auth_pair(&mut payload).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_ai_passthrough_sync_request_body_prefers_trimmed_base64() {
|
||||
let body = resolve_ai_passthrough_sync_request_body(
|
||||
Some(json!({"ignored": true})),
|
||||
Some(" YWJj ".to_string()),
|
||||
);
|
||||
|
||||
assert_eq!(body.body_bytes_b64.as_deref(), Some("YWJj"));
|
||||
assert!(body.json_body.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_ai_passthrough_sync_request_body_uses_json_when_no_base64() {
|
||||
let body = resolve_ai_passthrough_sync_request_body(Some(json!({"ok": true})), None);
|
||||
|
||||
assert_eq!(body.json_body, Some(json!({"ok": true})));
|
||||
assert!(body.body_bytes_b64.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_ai_execution_plan_from_decision_merges_core_and_remaining_payload_fields() {
|
||||
let mut payload = test_decision();
|
||||
let core =
|
||||
take_ai_decision_plan_core(&mut payload).expect("core fields should be available");
|
||||
let plan = build_ai_execution_plan_from_decision(
|
||||
&mut payload,
|
||||
AiExecutionPlanFromDecisionParts {
|
||||
core,
|
||||
method: "POST".to_string(),
|
||||
url: "https://example.com/v1/chat/completions".to_string(),
|
||||
headers: BTreeMap::from([(
|
||||
"content-type".to_string(),
|
||||
"application/json".to_string(),
|
||||
)]),
|
||||
content_type: Some("application/json".to_string()),
|
||||
body: RequestBody::from_json(json!({"model": "gpt-test"})),
|
||||
stream: true,
|
||||
},
|
||||
);
|
||||
|
||||
assert_eq!(plan.request_id, "req_1");
|
||||
assert_eq!(plan.candidate_id.as_deref(), Some("candidate_1"));
|
||||
assert_eq!(plan.provider_id, "provider_1");
|
||||
assert_eq!(plan.endpoint_id, "endpoint_1");
|
||||
assert_eq!(plan.key_id, "key_1");
|
||||
assert!(plan.stream);
|
||||
assert_eq!(plan.provider_api_format, "openai:chat");
|
||||
assert_eq!(plan.client_api_format, "openai:chat");
|
||||
assert_eq!(plan.model_name.as_deref(), Some("gpt-test"));
|
||||
assert!(payload.candidate_id.is_none());
|
||||
assert!(payload.model_name.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn infer_ai_upstream_base_url_preserves_codex_base_path() {
|
||||
assert_eq!(
|
||||
infer_ai_upstream_base_url("https://tiger.bookapi.cc/codex/responses").as_deref(),
|
||||
Some("https://tiger.bookapi.cc/codex")
|
||||
);
|
||||
assert_eq!(
|
||||
infer_ai_upstream_base_url("https://chatgpt.com/backend-api/codex/responses")
|
||||
.as_deref(),
|
||||
Some("https://chatgpt.com/backend-api/codex")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn infer_ai_upstream_base_url_preserves_nested_v1_prefix() {
|
||||
assert_eq!(
|
||||
infer_ai_upstream_base_url(
|
||||
"https://api.openai.example/custom/v1/chat/completions?mode=1"
|
||||
)
|
||||
.as_deref(),
|
||||
Some("https://api.openai.example/custom/v1")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn infer_ai_upstream_base_url_strips_video_operation_path() {
|
||||
assert_eq!(
|
||||
infer_ai_upstream_base_url("https://video.example/nested/v1/videos/task-123/content")
|
||||
.as_deref(),
|
||||
Some("https://video.example/nested")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_ai_execution_decision_from_plan_maps_plan_fields() {
|
||||
let plan = ExecutionPlan {
|
||||
request_id: "plan-request".to_string(),
|
||||
candidate_id: Some("candidate-1".to_string()),
|
||||
provider_name: Some("provider".to_string()),
|
||||
provider_id: "provider-1".to_string(),
|
||||
endpoint_id: "endpoint-1".to_string(),
|
||||
key_id: "key-1".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: "https://api.example.com/v1/chat/completions".to_string(),
|
||||
headers: BTreeMap::from([("Authorization".to_string(), "Bearer secret".to_string())]),
|
||||
content_type: Some("application/json".to_string()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"model": "mapped"})),
|
||||
stream: false,
|
||||
client_api_format: "openai:chat".to_string(),
|
||||
provider_api_format: "claude:messages".to_string(),
|
||||
model_name: Some("mapped".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
};
|
||||
|
||||
let decision = build_ai_execution_decision_from_plan(AiExecutionDecisionFromPlanParts {
|
||||
action: "execution_runtime.sync_decision".to_string(),
|
||||
decision_kind: Some("openai_chat_sync".to_string()),
|
||||
request_id: Some("trace-1".to_string()),
|
||||
upstream_base_url: Some("https://api.example.com".to_string()),
|
||||
include_auth_pair: true,
|
||||
plan,
|
||||
report_kind: Some("report".to_string()),
|
||||
report_context: Some(json!({"candidate_index": 0})),
|
||||
auth_context: None,
|
||||
});
|
||||
|
||||
assert_eq!(decision.request_id.as_deref(), Some("trace-1"));
|
||||
assert_eq!(
|
||||
decision.execution_strategy.as_deref(),
|
||||
Some("local_cross_format")
|
||||
);
|
||||
assert_eq!(decision.conversion_mode.as_deref(), Some("bidirectional"));
|
||||
assert_eq!(decision.auth_header.as_deref(), Some("Authorization"));
|
||||
assert_eq!(decision.auth_value.as_deref(), Some("Bearer secret"));
|
||||
assert_eq!(
|
||||
decision.provider_request_body,
|
||||
Some(json!({"model": "mapped"}))
|
||||
);
|
||||
assert_eq!(decision.report_kind.as_deref(), Some("report"));
|
||||
}
|
||||
|
||||
fn test_decision() -> AiExecutionDecision {
|
||||
AiExecutionDecision {
|
||||
action: "sync".to_string(),
|
||||
decision_kind: Some("test".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: None,
|
||||
request_id: Some("req_1".to_string()),
|
||||
candidate_id: Some("candidate_1".to_string()),
|
||||
provider_name: Some("provider".to_string()),
|
||||
provider_id: Some("provider_1".to_string()),
|
||||
endpoint_id: Some("endpoint_1".to_string()),
|
||||
key_id: Some("key_1".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/chat/completions".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("authorization".to_string()),
|
||||
auth_value: Some("Bearer token".to_string()),
|
||||
provider_api_format: Some("openai:chat".to_string()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: Some("openai:chat".to_string()),
|
||||
client_contract: Some("openai:chat".to_string()),
|
||||
model_name: Some("gpt-test".to_string()),
|
||||
mapped_model: Some("gpt-test".to_string()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::new(),
|
||||
provider_request_body: None,
|
||||
provider_request_body_base64: None,
|
||||
content_type: None,
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: false,
|
||||
report_kind: None,
|
||||
report_context: None,
|
||||
auth_context: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
159
crates/aether-ai-serving/src/candidate_materialization.rs
Normal file
159
crates/aether-ai-serving/src/candidate_materialization.rs
Normal file
@@ -0,0 +1,159 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AiCandidateMaterializationOutcome<Attempt> {
|
||||
pub attempts: Vec<Attempt>,
|
||||
pub candidate_count: usize,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AiCandidateMaterializationPort: Send + Sync {
|
||||
type Candidate: Send;
|
||||
type Eligible: Send + Sync;
|
||||
type Skipped: Send;
|
||||
type Attempt: Send;
|
||||
type Error: Send;
|
||||
|
||||
async fn resolve_and_rank_candidates(
|
||||
&self,
|
||||
candidates: Vec<Self::Candidate>,
|
||||
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error>;
|
||||
|
||||
fn decorate_skipped_candidate(&self, skipped: Self::Skipped) -> Self::Skipped {
|
||||
skipped
|
||||
}
|
||||
|
||||
fn remember_first_candidate_affinity(&self, candidates: &[Self::Eligible]);
|
||||
|
||||
async fn persist_available_candidates(
|
||||
&self,
|
||||
candidates: Vec<Self::Eligible>,
|
||||
) -> Result<Vec<Self::Attempt>, Self::Error>;
|
||||
|
||||
async fn persist_skipped_candidates(
|
||||
&self,
|
||||
starting_candidate_index: u32,
|
||||
skipped_candidates: Vec<Self::Skipped>,
|
||||
) -> Result<(), Self::Error>;
|
||||
}
|
||||
|
||||
pub async fn run_ai_candidate_materialization<Port>(
|
||||
port: &Port,
|
||||
candidates: Vec<Port::Candidate>,
|
||||
preselection_skipped_candidates: Vec<Port::Skipped>,
|
||||
) -> Result<AiCandidateMaterializationOutcome<Port::Attempt>, Port::Error>
|
||||
where
|
||||
Port: AiCandidateMaterializationPort,
|
||||
{
|
||||
let (candidates, skipped_candidates) = port.resolve_and_rank_candidates(candidates).await?;
|
||||
let skipped_candidates = preselection_skipped_candidates
|
||||
.into_iter()
|
||||
.chain(skipped_candidates)
|
||||
.map(|candidate| port.decorate_skipped_candidate(candidate))
|
||||
.collect::<Vec<_>>();
|
||||
let candidate_count = candidates.len() + skipped_candidates.len();
|
||||
|
||||
port.remember_first_candidate_affinity(&candidates);
|
||||
let available_candidate_count = candidates.len() as u32;
|
||||
let attempts = port.persist_available_candidates(candidates).await?;
|
||||
port.persist_skipped_candidates(available_candidate_count, skipped_candidates)
|
||||
.await?;
|
||||
|
||||
Ok(AiCandidateMaterializationOutcome {
|
||||
attempts,
|
||||
candidate_count,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::Mutex;
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestPort {
|
||||
calls: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiCandidateMaterializationPort for TestPort {
|
||||
type Candidate = &'static str;
|
||||
type Eligible = &'static str;
|
||||
type Skipped = &'static str;
|
||||
type Attempt = &'static str;
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
async fn resolve_and_rank_candidates(
|
||||
&self,
|
||||
candidates: Vec<Self::Candidate>,
|
||||
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error> {
|
||||
self.calls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(format!("resolve:{}", candidates.join(",")));
|
||||
Ok((vec!["eligible-a", "eligible-b"], vec!["resolved-skip"]))
|
||||
}
|
||||
|
||||
fn decorate_skipped_candidate(&self, skipped: Self::Skipped) -> Self::Skipped {
|
||||
self.calls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(format!("decorate:{skipped}"));
|
||||
skipped
|
||||
}
|
||||
|
||||
fn remember_first_candidate_affinity(&self, candidates: &[Self::Eligible]) {
|
||||
self.calls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(format!("affinity:{}", candidates.join(",")));
|
||||
}
|
||||
|
||||
async fn persist_available_candidates(
|
||||
&self,
|
||||
candidates: Vec<Self::Eligible>,
|
||||
) -> Result<Vec<Self::Attempt>, Self::Error> {
|
||||
self.calls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(format!("available:{}", candidates.join(",")));
|
||||
Ok(vec!["attempt-a", "attempt-b"])
|
||||
}
|
||||
|
||||
async fn persist_skipped_candidates(
|
||||
&self,
|
||||
starting_candidate_index: u32,
|
||||
skipped_candidates: Vec<Self::Skipped>,
|
||||
) -> Result<(), Self::Error> {
|
||||
self.calls.lock().unwrap().push(format!(
|
||||
"skipped:{starting_candidate_index}:{}",
|
||||
skipped_candidates.join(",")
|
||||
));
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn materialization_runs_in_serving_order_and_counts_all_candidates() {
|
||||
let port = TestPort::default();
|
||||
|
||||
let outcome =
|
||||
run_ai_candidate_materialization(&port, vec!["candidate-a"], vec!["pre-skip"])
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(outcome.attempts, ["attempt-a", "attempt-b"]);
|
||||
assert_eq!(outcome.candidate_count, 4);
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
[
|
||||
"resolve:candidate-a",
|
||||
"decorate:pre-skip",
|
||||
"decorate:resolved-skip",
|
||||
"affinity:eligible-a,eligible-b",
|
||||
"available:eligible-a,eligible-b",
|
||||
"skipped:2:pre-skip,resolved-skip",
|
||||
]
|
||||
);
|
||||
}
|
||||
}
|
||||
210
crates/aether-ai-serving/src/candidate_metadata.rs
Normal file
210
crates/aether-ai-serving/src/candidate_metadata.rs
Normal file
@@ -0,0 +1,210 @@
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{ConversionMode, ExecutionStrategy};
|
||||
|
||||
pub struct AiCandidateMetadataParts<'a> {
|
||||
pub provider_api_format: &'a str,
|
||||
pub client_api_format: &'a str,
|
||||
pub global_model_id: &'a str,
|
||||
pub global_model_name: &'a str,
|
||||
pub model_id: &'a str,
|
||||
pub selected_provider_model_name: &'a str,
|
||||
pub mapping_matched_model: Option<&'a str>,
|
||||
pub provider_name: &'a str,
|
||||
pub key_name: &'a str,
|
||||
pub extra_fields: Map<String, Value>,
|
||||
}
|
||||
|
||||
pub fn build_ai_candidate_metadata(parts: AiCandidateMetadataParts<'_>) -> Value {
|
||||
let mut object = Map::new();
|
||||
object.insert(
|
||||
"provider_api_format".to_string(),
|
||||
Value::String(parts.provider_api_format.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"client_api_format".to_string(),
|
||||
Value::String(parts.client_api_format.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"global_model_id".to_string(),
|
||||
Value::String(parts.global_model_id.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"global_model_name".to_string(),
|
||||
Value::String(parts.global_model_name.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"model_id".to_string(),
|
||||
Value::String(parts.model_id.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"selected_provider_model_name".to_string(),
|
||||
Value::String(parts.selected_provider_model_name.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"mapping_matched_model".to_string(),
|
||||
parts
|
||||
.mapping_matched_model
|
||||
.map(|value| Value::String(value.to_string()))
|
||||
.unwrap_or(Value::Null),
|
||||
);
|
||||
object.insert(
|
||||
"provider_name".to_string(),
|
||||
Value::String(parts.provider_name.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"key_name".to_string(),
|
||||
Value::String(parts.key_name.to_string()),
|
||||
);
|
||||
object.extend(parts.extra_fields);
|
||||
Value::Object(object)
|
||||
}
|
||||
|
||||
pub fn build_ai_candidate_metadata_from_candidate(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
provider_api_format: &str,
|
||||
client_api_format: &str,
|
||||
extra_fields: Map<String, Value>,
|
||||
) -> Value {
|
||||
build_ai_candidate_metadata(AiCandidateMetadataParts {
|
||||
provider_api_format,
|
||||
client_api_format,
|
||||
global_model_id: candidate.global_model_id.as_str(),
|
||||
global_model_name: candidate.global_model_name.as_str(),
|
||||
model_id: candidate.model_id.as_str(),
|
||||
selected_provider_model_name: candidate.selected_provider_model_name.as_str(),
|
||||
mapping_matched_model: candidate.mapping_matched_model.as_deref(),
|
||||
provider_name: candidate.provider_name.as_str(),
|
||||
key_name: candidate.key_name.as_str(),
|
||||
extra_fields,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn append_ai_execution_contract_fields_to_value(
|
||||
value: Value,
|
||||
execution_strategy: &str,
|
||||
conversion_mode: &str,
|
||||
client_contract: &str,
|
||||
provider_contract: &str,
|
||||
) -> Value {
|
||||
match value {
|
||||
Value::Object(mut object) => {
|
||||
object.insert(
|
||||
"execution_strategy".to_string(),
|
||||
Value::String(execution_strategy.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"conversion_mode".to_string(),
|
||||
Value::String(conversion_mode.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"client_contract".to_string(),
|
||||
Value::String(client_contract.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"provider_contract".to_string(),
|
||||
Value::String(provider_contract.to_string()),
|
||||
);
|
||||
Value::Object(object)
|
||||
}
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn ai_local_execution_contract_for_formats(
|
||||
client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
) -> (ExecutionStrategy, ConversionMode) {
|
||||
if aether_ai_formats::api_format_alias_matches(client_api_format, provider_api_format) {
|
||||
return (ExecutionStrategy::LocalSameFormat, ConversionMode::None);
|
||||
}
|
||||
|
||||
let conversion_mode =
|
||||
if aether_ai_formats::request_conversion_kind(client_api_format, provider_api_format)
|
||||
.is_some()
|
||||
{
|
||||
ConversionMode::Bidirectional
|
||||
} else {
|
||||
ConversionMode::None
|
||||
};
|
||||
(ExecutionStrategy::LocalCrossFormat, conversion_mode)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn candidate_metadata_builds_base_candidate_fields_and_extra_data() {
|
||||
let mut extra_fields = Map::new();
|
||||
extra_fields.insert("source".to_string(), json!("test"));
|
||||
|
||||
let metadata = build_ai_candidate_metadata(AiCandidateMetadataParts {
|
||||
provider_api_format: "openai:responses",
|
||||
client_api_format: "claude:messages",
|
||||
global_model_id: "global-1",
|
||||
global_model_name: "gpt-5.4",
|
||||
model_id: "model-1",
|
||||
selected_provider_model_name: "gpt-5.4",
|
||||
mapping_matched_model: Some("gpt-5"),
|
||||
provider_name: "RightCode",
|
||||
key_name: "key-a",
|
||||
extra_fields,
|
||||
});
|
||||
|
||||
assert_eq!(metadata["provider_api_format"], "openai:responses");
|
||||
assert_eq!(metadata["client_api_format"], "claude:messages");
|
||||
assert_eq!(metadata["global_model_id"], "global-1");
|
||||
assert_eq!(metadata["mapping_matched_model"], "gpt-5");
|
||||
assert_eq!(metadata["source"], "test");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn execution_contract_fields_append_to_object_values_only() {
|
||||
let value = append_ai_execution_contract_fields_to_value(
|
||||
json!({"existing": true}),
|
||||
"local_cross_format",
|
||||
"bidirectional",
|
||||
"openai:chat",
|
||||
"claude:messages",
|
||||
);
|
||||
|
||||
assert_eq!(value["existing"], true);
|
||||
assert_eq!(value["execution_strategy"], "local_cross_format");
|
||||
assert_eq!(value["conversion_mode"], "bidirectional");
|
||||
assert_eq!(value["client_contract"], "openai:chat");
|
||||
assert_eq!(value["provider_contract"], "claude:messages");
|
||||
|
||||
assert_eq!(
|
||||
append_ai_execution_contract_fields_to_value(
|
||||
Value::Null,
|
||||
"local_same_format",
|
||||
"none",
|
||||
"openai:chat",
|
||||
"openai:chat",
|
||||
),
|
||||
Value::Null
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_execution_contract_is_derived_from_client_and_provider_formats() {
|
||||
assert_eq!(
|
||||
ai_local_execution_contract_for_formats(" OPENAI:CHAT ", "openai:chat"),
|
||||
(ExecutionStrategy::LocalSameFormat, ConversionMode::None)
|
||||
);
|
||||
assert_eq!(
|
||||
ai_local_execution_contract_for_formats("openai:chat", "claude:messages"),
|
||||
(
|
||||
ExecutionStrategy::LocalCrossFormat,
|
||||
ConversionMode::Bidirectional,
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
ai_local_execution_contract_for_formats("openai:chat", "unknown:format"),
|
||||
(ExecutionStrategy::LocalCrossFormat, ConversionMode::None)
|
||||
);
|
||||
}
|
||||
}
|
||||
433
crates/aether-ai-serving/src/candidate_persistence.rs
Normal file
433
crates/aether-ai-serving/src/candidate_persistence.rs
Normal file
@@ -0,0 +1,433 @@
|
||||
use aether_scheduler_core::SchedulerRankingOutcome;
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
#[async_trait]
|
||||
pub trait AiAvailableCandidatePersistencePort: Send + Sync {
|
||||
type Candidate: Clone + Send + Sync;
|
||||
type Attempt: Send;
|
||||
type ExtraData: Clone + Send + Sync;
|
||||
type Error: Send;
|
||||
|
||||
fn attempt_slot_count(&self, candidate: &Self::Candidate) -> u32;
|
||||
|
||||
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData>;
|
||||
|
||||
fn generate_candidate_id(&self) -> String;
|
||||
|
||||
fn should_persist_available_candidate(&self, candidate: &Self::Candidate) -> bool;
|
||||
|
||||
async fn persist_available_candidate(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
candidate_index: u32,
|
||||
retry_index: u32,
|
||||
generated_candidate_id: &str,
|
||||
extra_data: Option<Self::ExtraData>,
|
||||
) -> Result<String, Self::Error>;
|
||||
|
||||
fn build_attempt(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
candidate_index: u32,
|
||||
retry_index: u32,
|
||||
candidate_id: String,
|
||||
) -> Self::Attempt;
|
||||
}
|
||||
|
||||
pub async fn run_ai_available_candidate_persistence<Port>(
|
||||
port: &Port,
|
||||
candidates: Vec<Port::Candidate>,
|
||||
) -> Result<Vec<Port::Attempt>, Port::Error>
|
||||
where
|
||||
Port: AiAvailableCandidatePersistencePort,
|
||||
{
|
||||
let total_attempts = candidates
|
||||
.iter()
|
||||
.map(|candidate| port.attempt_slot_count(candidate) as usize)
|
||||
.sum();
|
||||
let mut materialized = Vec::with_capacity(total_attempts);
|
||||
|
||||
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
|
||||
let candidate_index = candidate_index as u32;
|
||||
let attempt_slots = port.attempt_slot_count(&candidate).max(1);
|
||||
let extra_data = port.build_extra_data(&candidate);
|
||||
let mut owned_candidate = Some(candidate);
|
||||
|
||||
for retry_index in 0..attempt_slots {
|
||||
let candidate = owned_candidate
|
||||
.as_ref()
|
||||
.expect("candidate should remain available until final retry");
|
||||
let generated_candidate_id = port.generate_candidate_id();
|
||||
let candidate_id = if port.should_persist_available_candidate(candidate) {
|
||||
port.persist_available_candidate(
|
||||
candidate,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
generated_candidate_id.as_str(),
|
||||
extra_data.clone(),
|
||||
)
|
||||
.await?
|
||||
} else {
|
||||
generated_candidate_id
|
||||
};
|
||||
|
||||
let candidate = if retry_index + 1 == attempt_slots {
|
||||
owned_candidate
|
||||
.take()
|
||||
.expect("final retry should consume owned candidate")
|
||||
} else {
|
||||
candidate.clone()
|
||||
};
|
||||
materialized.push(port.build_attempt(
|
||||
candidate,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
candidate_id,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(materialized)
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AiSkippedCandidatePersistencePort: Send + Sync {
|
||||
type Skipped: Send + Sync;
|
||||
type ExtraData: Send + Sync;
|
||||
type Error: Send;
|
||||
|
||||
fn should_persist_skipped_candidate(&self, candidate: &Self::Skipped) -> bool;
|
||||
|
||||
fn build_extra_data(&self, candidate: &Self::Skipped) -> Option<Self::ExtraData>;
|
||||
|
||||
fn generate_candidate_id(&self) -> String;
|
||||
|
||||
async fn persist_skipped_candidate(
|
||||
&self,
|
||||
candidate: &Self::Skipped,
|
||||
candidate_index: u32,
|
||||
generated_candidate_id: &str,
|
||||
extra_data: Option<Self::ExtraData>,
|
||||
) -> Result<(), Self::Error>;
|
||||
}
|
||||
|
||||
pub async fn run_ai_skipped_candidate_persistence<Port>(
|
||||
port: &Port,
|
||||
starting_candidate_index: u32,
|
||||
skipped_candidates: Vec<Port::Skipped>,
|
||||
) -> Result<(), Port::Error>
|
||||
where
|
||||
Port: AiSkippedCandidatePersistencePort,
|
||||
{
|
||||
let mut next_candidate_index = starting_candidate_index;
|
||||
for skipped_candidate in skipped_candidates {
|
||||
if !port.should_persist_skipped_candidate(&skipped_candidate) {
|
||||
continue;
|
||||
}
|
||||
let generated_candidate_id = port.generate_candidate_id();
|
||||
let extra_data = port.build_extra_data(&skipped_candidate);
|
||||
port.persist_skipped_candidate(
|
||||
&skipped_candidate,
|
||||
next_candidate_index,
|
||||
generated_candidate_id.as_str(),
|
||||
extra_data,
|
||||
)
|
||||
.await?;
|
||||
next_candidate_index = next_candidate_index.saturating_add(1);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn ai_should_persist_available_candidate_for_pool_key(pool_key_index: Option<u32>) -> bool {
|
||||
pool_key_index.is_none_or(|index| index == 0)
|
||||
}
|
||||
|
||||
pub fn ai_should_persist_skipped_candidate_for_pool_membership(is_pool_candidate: bool) -> bool {
|
||||
!is_pool_candidate
|
||||
}
|
||||
|
||||
pub fn ai_candidate_extra_data_with_ranking(
|
||||
extra_data: Option<Value>,
|
||||
ranking: Option<&SchedulerRankingOutcome>,
|
||||
) -> Option<Value> {
|
||||
let Some(ranking) = ranking else {
|
||||
return extra_data;
|
||||
};
|
||||
|
||||
let mut object = match extra_data {
|
||||
Some(Value::Object(object)) => object,
|
||||
Some(value) => {
|
||||
let mut object = serde_json::Map::new();
|
||||
object.insert("extra".to_string(), value);
|
||||
object
|
||||
}
|
||||
None => serde_json::Map::new(),
|
||||
};
|
||||
crate::append_ai_ranking_metadata_to_object(&mut object, ranking);
|
||||
Some(Value::Object(object))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::Mutex;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct TestCandidate {
|
||||
id: &'static str,
|
||||
attempt_slots: u32,
|
||||
persist: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct TestAttempt {
|
||||
id: &'static str,
|
||||
candidate_index: u32,
|
||||
retry_index: u32,
|
||||
candidate_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct TestSkipped {
|
||||
id: &'static str,
|
||||
persist: bool,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestPort {
|
||||
next_id: Mutex<u32>,
|
||||
calls: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
impl TestPort {
|
||||
fn next_candidate_id(&self) -> String {
|
||||
let mut next_id = self.next_id.lock().unwrap();
|
||||
*next_id += 1;
|
||||
format!("candidate-{next_id}")
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiAvailableCandidatePersistencePort for TestPort {
|
||||
type Candidate = TestCandidate;
|
||||
type Attempt = TestAttempt;
|
||||
type ExtraData = String;
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
fn attempt_slot_count(&self, candidate: &Self::Candidate) -> u32 {
|
||||
candidate.attempt_slots
|
||||
}
|
||||
|
||||
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData> {
|
||||
Some(format!("extra:{}", candidate.id))
|
||||
}
|
||||
|
||||
fn generate_candidate_id(&self) -> String {
|
||||
self.next_candidate_id()
|
||||
}
|
||||
|
||||
fn should_persist_available_candidate(&self, candidate: &Self::Candidate) -> bool {
|
||||
candidate.persist
|
||||
}
|
||||
|
||||
async fn persist_available_candidate(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
candidate_index: u32,
|
||||
retry_index: u32,
|
||||
generated_candidate_id: &str,
|
||||
extra_data: Option<Self::ExtraData>,
|
||||
) -> Result<String, Self::Error> {
|
||||
self.calls.lock().unwrap().push(format!(
|
||||
"available:{}:{candidate_index}:{retry_index}:{generated_candidate_id}:{}",
|
||||
candidate.id,
|
||||
extra_data.unwrap_or_default()
|
||||
));
|
||||
Ok(format!("stored-{generated_candidate_id}"))
|
||||
}
|
||||
|
||||
fn build_attempt(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
candidate_index: u32,
|
||||
retry_index: u32,
|
||||
candidate_id: String,
|
||||
) -> Self::Attempt {
|
||||
TestAttempt {
|
||||
id: candidate.id,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
candidate_id,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiSkippedCandidatePersistencePort for TestPort {
|
||||
type Skipped = TestSkipped;
|
||||
type ExtraData = String;
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
fn should_persist_skipped_candidate(&self, candidate: &Self::Skipped) -> bool {
|
||||
candidate.persist
|
||||
}
|
||||
|
||||
fn build_extra_data(&self, candidate: &Self::Skipped) -> Option<Self::ExtraData> {
|
||||
Some(format!("extra:{}", candidate.id))
|
||||
}
|
||||
|
||||
fn generate_candidate_id(&self) -> String {
|
||||
self.next_candidate_id()
|
||||
}
|
||||
|
||||
async fn persist_skipped_candidate(
|
||||
&self,
|
||||
candidate: &Self::Skipped,
|
||||
candidate_index: u32,
|
||||
generated_candidate_id: &str,
|
||||
extra_data: Option<Self::ExtraData>,
|
||||
) -> Result<(), Self::Error> {
|
||||
self.calls.lock().unwrap().push(format!(
|
||||
"skipped:{}:{candidate_index}:{generated_candidate_id}:{}",
|
||||
candidate.id,
|
||||
extra_data.unwrap_or_default()
|
||||
));
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn available_persistence_expands_candidates_into_retry_attempts() {
|
||||
let port = TestPort::default();
|
||||
|
||||
let attempts = run_ai_available_candidate_persistence(
|
||||
&port,
|
||||
vec![
|
||||
TestCandidate {
|
||||
id: "a",
|
||||
attempt_slots: 2,
|
||||
persist: true,
|
||||
},
|
||||
TestCandidate {
|
||||
id: "b",
|
||||
attempt_slots: 1,
|
||||
persist: false,
|
||||
},
|
||||
],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
attempts,
|
||||
[
|
||||
TestAttempt {
|
||||
id: "a",
|
||||
candidate_index: 0,
|
||||
retry_index: 0,
|
||||
candidate_id: "stored-candidate-1".to_string(),
|
||||
},
|
||||
TestAttempt {
|
||||
id: "a",
|
||||
candidate_index: 0,
|
||||
retry_index: 1,
|
||||
candidate_id: "stored-candidate-2".to_string(),
|
||||
},
|
||||
TestAttempt {
|
||||
id: "b",
|
||||
candidate_index: 1,
|
||||
retry_index: 0,
|
||||
candidate_id: "candidate-3".to_string(),
|
||||
},
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
[
|
||||
"available:a:0:0:candidate-1:extra:a",
|
||||
"available:a:0:1:candidate-2:extra:a",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn skipped_persistence_keeps_indices_for_persisted_candidates_only() {
|
||||
let port = TestPort::default();
|
||||
|
||||
run_ai_skipped_candidate_persistence(
|
||||
&port,
|
||||
3,
|
||||
vec![
|
||||
TestSkipped {
|
||||
id: "ignored",
|
||||
persist: false,
|
||||
},
|
||||
TestSkipped {
|
||||
id: "a",
|
||||
persist: true,
|
||||
},
|
||||
TestSkipped {
|
||||
id: "b",
|
||||
persist: true,
|
||||
},
|
||||
],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
[
|
||||
"skipped:a:3:candidate-1:extra:a",
|
||||
"skipped:b:4:candidate-2:extra:b",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pool_candidate_persistence_policy_persists_representatives_only() {
|
||||
assert!(ai_should_persist_available_candidate_for_pool_key(None));
|
||||
assert!(ai_should_persist_available_candidate_for_pool_key(Some(0)));
|
||||
assert!(!ai_should_persist_available_candidate_for_pool_key(Some(1)));
|
||||
|
||||
assert!(ai_should_persist_skipped_candidate_for_pool_membership(
|
||||
false
|
||||
));
|
||||
assert!(!ai_should_persist_skipped_candidate_for_pool_membership(
|
||||
true
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_extra_data_with_ranking_preserves_existing_shapes() {
|
||||
let ranking = SchedulerRankingOutcome {
|
||||
original_index: 1,
|
||||
ranking_index: 0,
|
||||
priority_mode: aether_scheduler_core::SchedulerPriorityMode::Provider,
|
||||
ranking_mode: aether_scheduler_core::SchedulerRankingMode::CacheAffinity,
|
||||
priority_slot: 3,
|
||||
promoted_by: Some("cached_affinity"),
|
||||
demoted_by: None,
|
||||
};
|
||||
|
||||
let object = ai_candidate_extra_data_with_ranking(
|
||||
Some(serde_json::json!({"source": "test"})),
|
||||
Some(&ranking),
|
||||
)
|
||||
.expect("extra data should exist");
|
||||
assert_eq!(object.get("source"), Some(&serde_json::json!("test")));
|
||||
assert_eq!(object.get("ranking_index"), Some(&serde_json::json!(0)));
|
||||
assert_eq!(
|
||||
object.get("promoted_by"),
|
||||
Some(&serde_json::json!("cached_affinity"))
|
||||
);
|
||||
|
||||
let scalar =
|
||||
ai_candidate_extra_data_with_ranking(Some(serde_json::json!("raw")), Some(&ranking))
|
||||
.expect("scalar extra data should be wrapped");
|
||||
assert_eq!(scalar.get("extra"), Some(&serde_json::json!("raw")));
|
||||
assert_eq!(scalar.get("priority_slot"), Some(&serde_json::json!(3)));
|
||||
}
|
||||
}
|
||||
117
crates/aether-ai-serving/src/candidate_persistence_policy.rs
Normal file
117
crates/aether-ai-serving/src/candidate_persistence_policy.rs
Normal file
@@ -0,0 +1,117 @@
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[allow(clippy::enum_variant_names)]
|
||||
pub enum AiCandidatePersistencePolicyKind {
|
||||
StandardDecision,
|
||||
SameFormatProviderDecision,
|
||||
OpenAiChatDecision,
|
||||
OpenAiResponsesDecision,
|
||||
ImageDecision,
|
||||
GeminiFilesDecision,
|
||||
VideoDecision,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct AiCandidatePersistencePolicySpec {
|
||||
pub available_error_context: &'static str,
|
||||
pub skipped_error_context: &'static str,
|
||||
pub record_runtime_miss_diagnostic: bool,
|
||||
}
|
||||
|
||||
pub fn ai_candidate_persistence_policy_spec(
|
||||
kind: AiCandidatePersistencePolicyKind,
|
||||
) -> AiCandidatePersistencePolicySpec {
|
||||
match kind {
|
||||
AiCandidatePersistencePolicyKind::StandardDecision => AiCandidatePersistencePolicySpec {
|
||||
available_error_context:
|
||||
"gateway local standard decision request candidate upsert failed",
|
||||
skipped_error_context:
|
||||
"gateway local standard decision failed to persist skipped candidate",
|
||||
record_runtime_miss_diagnostic: true,
|
||||
},
|
||||
AiCandidatePersistencePolicyKind::SameFormatProviderDecision => {
|
||||
AiCandidatePersistencePolicySpec {
|
||||
available_error_context:
|
||||
"gateway local same-format decision request candidate upsert failed",
|
||||
skipped_error_context:
|
||||
"gateway local same-format decision failed to persist skipped candidate",
|
||||
record_runtime_miss_diagnostic: true,
|
||||
}
|
||||
}
|
||||
AiCandidatePersistencePolicyKind::OpenAiChatDecision => AiCandidatePersistencePolicySpec {
|
||||
available_error_context:
|
||||
"gateway local openai chat decision request candidate upsert failed",
|
||||
skipped_error_context:
|
||||
"gateway local openai chat decision failed to persist skipped candidate",
|
||||
record_runtime_miss_diagnostic: true,
|
||||
},
|
||||
AiCandidatePersistencePolicyKind::OpenAiResponsesDecision => {
|
||||
AiCandidatePersistencePolicySpec {
|
||||
available_error_context:
|
||||
"gateway local openai responses decision request candidate upsert failed",
|
||||
skipped_error_context:
|
||||
"gateway local openai responses decision failed to persist skipped candidate",
|
||||
record_runtime_miss_diagnostic: true,
|
||||
}
|
||||
}
|
||||
AiCandidatePersistencePolicyKind::ImageDecision => AiCandidatePersistencePolicySpec {
|
||||
available_error_context:
|
||||
"gateway local openai image decision request candidate upsert failed",
|
||||
skipped_error_context:
|
||||
"gateway local openai image decision failed to persist skipped candidate",
|
||||
record_runtime_miss_diagnostic: false,
|
||||
},
|
||||
AiCandidatePersistencePolicyKind::GeminiFilesDecision => AiCandidatePersistencePolicySpec {
|
||||
available_error_context: "gateway local gemini files request candidate upsert failed",
|
||||
skipped_error_context: "gateway local gemini files failed to persist skipped candidate",
|
||||
record_runtime_miss_diagnostic: false,
|
||||
},
|
||||
AiCandidatePersistencePolicyKind::VideoDecision => AiCandidatePersistencePolicySpec {
|
||||
available_error_context: "gateway local video decision request candidate upsert failed",
|
||||
skipped_error_context:
|
||||
"gateway local video decision failed to persist skipped candidate",
|
||||
record_runtime_miss_diagnostic: false,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn decision_policies_define_candidate_persistence_side_effects() {
|
||||
let standard = ai_candidate_persistence_policy_spec(
|
||||
AiCandidatePersistencePolicyKind::StandardDecision,
|
||||
);
|
||||
assert_eq!(
|
||||
standard.available_error_context,
|
||||
"gateway local standard decision request candidate upsert failed"
|
||||
);
|
||||
assert_eq!(
|
||||
standard.skipped_error_context,
|
||||
"gateway local standard decision failed to persist skipped candidate"
|
||||
);
|
||||
assert!(standard.record_runtime_miss_diagnostic);
|
||||
|
||||
let same_format = ai_candidate_persistence_policy_spec(
|
||||
AiCandidatePersistencePolicyKind::SameFormatProviderDecision,
|
||||
);
|
||||
assert_eq!(
|
||||
same_format.available_error_context,
|
||||
"gateway local same-format decision request candidate upsert failed"
|
||||
);
|
||||
assert_eq!(
|
||||
same_format.skipped_error_context,
|
||||
"gateway local same-format decision failed to persist skipped candidate"
|
||||
);
|
||||
assert!(same_format.record_runtime_miss_diagnostic);
|
||||
|
||||
let image =
|
||||
ai_candidate_persistence_policy_spec(AiCandidatePersistencePolicyKind::ImageDecision);
|
||||
assert_eq!(
|
||||
image.available_error_context,
|
||||
"gateway local openai image decision request candidate upsert failed"
|
||||
);
|
||||
assert!(!image.record_runtime_miss_diagnostic);
|
||||
}
|
||||
}
|
||||
91
crates/aether-ai-serving/src/candidate_preparation.rs
Normal file
91
crates/aether-ai-serving/src/candidate_preparation.rs
Normal file
@@ -0,0 +1,91 @@
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AiPreparedHeaderAuthenticatedCandidate {
|
||||
pub auth_header: String,
|
||||
pub auth_value: String,
|
||||
pub mapped_model: String,
|
||||
}
|
||||
|
||||
pub fn prepare_ai_header_authenticated_candidate(
|
||||
direct_auth: Option<(String, String)>,
|
||||
oauth_header_auth: Option<(String, String)>,
|
||||
selected_provider_model_name: &str,
|
||||
) -> Result<AiPreparedHeaderAuthenticatedCandidate, &'static str> {
|
||||
let Some((auth_header, auth_value)) = direct_auth.or(oauth_header_auth) else {
|
||||
return Err("transport_auth_unavailable");
|
||||
};
|
||||
let mapped_model = resolve_ai_candidate_mapped_model(selected_provider_model_name)?;
|
||||
|
||||
Ok(AiPreparedHeaderAuthenticatedCandidate {
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn resolve_ai_candidate_mapped_model(
|
||||
selected_provider_model_name: &str,
|
||||
) -> Result<String, &'static str> {
|
||||
let mapped_model = selected_provider_model_name.trim().to_string();
|
||||
if mapped_model.is_empty() {
|
||||
return Err("mapped_model_missing");
|
||||
}
|
||||
|
||||
Ok(mapped_model)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{prepare_ai_header_authenticated_candidate, resolve_ai_candidate_mapped_model};
|
||||
|
||||
#[test]
|
||||
fn mapped_model_trims_selected_provider_model_name() {
|
||||
assert_eq!(
|
||||
resolve_ai_candidate_mapped_model(" gpt-test-upstream "),
|
||||
Ok("gpt-test-upstream".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mapped_model_rejects_empty_selected_provider_model_name() {
|
||||
assert_eq!(
|
||||
resolve_ai_candidate_mapped_model(" "),
|
||||
Err("mapped_model_missing")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn header_auth_preparation_prefers_direct_auth_and_allows_empty_value() {
|
||||
let prepared = prepare_ai_header_authenticated_candidate(
|
||||
Some(("authorization".to_string(), String::new())),
|
||||
Some(("x-oauth".to_string(), "oauth".to_string())),
|
||||
"gpt-test-upstream",
|
||||
)
|
||||
.expect("direct auth should prepare candidate");
|
||||
|
||||
assert_eq!(prepared.auth_header, "authorization");
|
||||
assert_eq!(prepared.auth_value, "");
|
||||
assert_eq!(prepared.mapped_model, "gpt-test-upstream");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn header_auth_preparation_falls_back_to_oauth_header_auth() {
|
||||
let prepared = prepare_ai_header_authenticated_candidate(
|
||||
None,
|
||||
Some(("authorization".to_string(), "Bearer token".to_string())),
|
||||
" gpt-test-upstream ",
|
||||
)
|
||||
.expect("oauth header auth should prepare candidate");
|
||||
|
||||
assert_eq!(prepared.auth_header, "authorization");
|
||||
assert_eq!(prepared.auth_value, "Bearer token");
|
||||
assert_eq!(prepared.mapped_model, "gpt-test-upstream");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn header_auth_preparation_requires_some_auth() {
|
||||
assert_eq!(
|
||||
prepare_ai_header_authenticated_candidate(None, None, "gpt-test-upstream"),
|
||||
Err("transport_auth_unavailable")
|
||||
);
|
||||
}
|
||||
}
|
||||
185
crates/aether-ai-serving/src/candidate_preselection.rs
Normal file
185
crates/aether-ai-serving/src/candidate_preselection.rs
Normal file
@@ -0,0 +1,185 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AiCandidatePreselectionOutcome<Candidate, Skipped> {
|
||||
pub candidates: Vec<Candidate>,
|
||||
pub skipped_candidates: Vec<Skipped>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AiCandidatePreselectionPort: Send + Sync {
|
||||
type Candidate: Send;
|
||||
type Skipped: Send;
|
||||
type Error: Send;
|
||||
|
||||
fn candidate_api_formats(&self) -> Vec<String>;
|
||||
|
||||
fn candidate_api_format_matches_client(&self, candidate_api_format: &str) -> bool;
|
||||
|
||||
async fn list_candidates_for_api_format(
|
||||
&self,
|
||||
candidate_api_format: &str,
|
||||
matches_client_format: bool,
|
||||
) -> Result<(Vec<Self::Candidate>, Vec<Self::Skipped>), Self::Error>;
|
||||
|
||||
fn candidate_allowed(
|
||||
&self,
|
||||
_candidate: &Self::Candidate,
|
||||
_candidate_api_format: &str,
|
||||
_matches_client_format: bool,
|
||||
) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn skipped_candidate_allowed(
|
||||
&self,
|
||||
_skipped_candidate: &Self::Skipped,
|
||||
_candidate_api_format: &str,
|
||||
_matches_client_format: bool,
|
||||
) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn candidate_key(&self, candidate: &Self::Candidate) -> String;
|
||||
|
||||
fn skipped_candidate_key(&self, skipped_candidate: &Self::Skipped) -> String;
|
||||
}
|
||||
|
||||
pub async fn run_ai_candidate_preselection<Port>(
|
||||
port: &Port,
|
||||
) -> Result<AiCandidatePreselectionOutcome<Port::Candidate, Port::Skipped>, Port::Error>
|
||||
where
|
||||
Port: AiCandidatePreselectionPort,
|
||||
{
|
||||
let mut candidates = Vec::new();
|
||||
let mut skipped_candidates = Vec::new();
|
||||
let mut seen_candidates = BTreeSet::new();
|
||||
let mut seen_skipped_candidates = BTreeSet::new();
|
||||
|
||||
for candidate_api_format in port.candidate_api_formats() {
|
||||
let matches_client_format =
|
||||
port.candidate_api_format_matches_client(candidate_api_format.as_str());
|
||||
let (selected, skipped) = port
|
||||
.list_candidates_for_api_format(candidate_api_format.as_str(), matches_client_format)
|
||||
.await?;
|
||||
|
||||
for skipped_candidate in skipped {
|
||||
if !port.skipped_candidate_allowed(
|
||||
&skipped_candidate,
|
||||
candidate_api_format.as_str(),
|
||||
matches_client_format,
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
let candidate_key = port.skipped_candidate_key(&skipped_candidate);
|
||||
if seen_skipped_candidates.insert(candidate_key) {
|
||||
skipped_candidates.push(skipped_candidate);
|
||||
}
|
||||
}
|
||||
|
||||
for candidate in selected {
|
||||
if !port.candidate_allowed(
|
||||
&candidate,
|
||||
candidate_api_format.as_str(),
|
||||
matches_client_format,
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
let candidate_key = port.candidate_key(&candidate);
|
||||
if seen_candidates.insert(candidate_key) {
|
||||
candidates.push(candidate);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(AiCandidatePreselectionOutcome {
|
||||
candidates,
|
||||
skipped_candidates,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::Mutex;
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestPort {
|
||||
calls: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiCandidatePreselectionPort for TestPort {
|
||||
type Candidate = &'static str;
|
||||
type Skipped = &'static str;
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
fn candidate_api_formats(&self) -> Vec<String> {
|
||||
vec!["same".to_string(), "cross".to_string()]
|
||||
}
|
||||
|
||||
fn candidate_api_format_matches_client(&self, candidate_api_format: &str) -> bool {
|
||||
candidate_api_format == "same"
|
||||
}
|
||||
|
||||
async fn list_candidates_for_api_format(
|
||||
&self,
|
||||
candidate_api_format: &str,
|
||||
matches_client_format: bool,
|
||||
) -> Result<(Vec<Self::Candidate>, Vec<Self::Skipped>), Self::Error> {
|
||||
self.calls.lock().unwrap().push(format!(
|
||||
"list:{candidate_api_format}:{matches_client_format}"
|
||||
));
|
||||
Ok(match candidate_api_format {
|
||||
"same" => (vec!["candidate-a"], vec!["skip-a"]),
|
||||
"cross" => (
|
||||
vec!["candidate-a", "candidate-b", "blocked-candidate"],
|
||||
vec!["skip-a", "skip-b", "blocked-skip"],
|
||||
),
|
||||
_ => (Vec::new(), Vec::new()),
|
||||
})
|
||||
}
|
||||
|
||||
fn candidate_allowed(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
_candidate_api_format: &str,
|
||||
matches_client_format: bool,
|
||||
) -> bool {
|
||||
matches_client_format || *candidate != "blocked-candidate"
|
||||
}
|
||||
|
||||
fn skipped_candidate_allowed(
|
||||
&self,
|
||||
skipped_candidate: &Self::Skipped,
|
||||
_candidate_api_format: &str,
|
||||
matches_client_format: bool,
|
||||
) -> bool {
|
||||
matches_client_format || *skipped_candidate != "blocked-skip"
|
||||
}
|
||||
|
||||
fn candidate_key(&self, candidate: &Self::Candidate) -> String {
|
||||
(*candidate).to_string()
|
||||
}
|
||||
|
||||
fn skipped_candidate_key(&self, skipped_candidate: &Self::Skipped) -> String {
|
||||
(*skipped_candidate).to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn preselection_runs_formats_in_order_filters_cross_format_and_dedupes() {
|
||||
let port = TestPort::default();
|
||||
|
||||
let outcome = run_ai_candidate_preselection(&port).await.unwrap();
|
||||
|
||||
assert_eq!(outcome.candidates, ["candidate-a", "candidate-b"]);
|
||||
assert_eq!(outcome.skipped_candidates, ["skip-a", "skip-b"]);
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
["list:same:true", "list:cross:false"]
|
||||
);
|
||||
}
|
||||
}
|
||||
301
crates/aether-ai-serving/src/candidate_ranking.rs
Normal file
301
crates/aether-ai-serving/src/candidate_ranking.rs
Normal file
@@ -0,0 +1,301 @@
|
||||
use aether_scheduler_core::{
|
||||
apply_scheduler_candidate_ranking, requested_capability_priority_for_candidate,
|
||||
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode, SchedulerRankableCandidate,
|
||||
SchedulerRankingContext, SchedulerRankingMode, SchedulerRankingOutcome,
|
||||
SchedulerTunnelAffinityBucket,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AiRankingSchedulingMode {
|
||||
FixedOrder,
|
||||
CacheAffinity,
|
||||
LoadBalance,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct AiRankingContextConfig {
|
||||
pub priority_mode: SchedulerPriorityMode,
|
||||
pub scheduling_mode: AiRankingSchedulingMode,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub struct AiRankableCandidateParts<'a> {
|
||||
pub candidate: &'a SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub original_index: usize,
|
||||
pub normalized_client_api_format: &'a str,
|
||||
pub provider_api_format: &'a str,
|
||||
pub required_capabilities: Option<&'a serde_json::Value>,
|
||||
pub cached_affinity_match: bool,
|
||||
pub tunnel_bucket: SchedulerTunnelAffinityBucket,
|
||||
pub keep_priority_on_conversion: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AiCandidateRankingPort: Send + Sync {
|
||||
type Candidate: Send + Sync;
|
||||
type AffinityTarget: Send + Sync;
|
||||
type Error: Send;
|
||||
|
||||
fn affinity_requested_model(&self, candidates: &[Self::Candidate]) -> Option<String>;
|
||||
|
||||
async fn read_cached_affinity_target(
|
||||
&self,
|
||||
normalized_client_api_format: &str,
|
||||
affinity_requested_model: Option<&str>,
|
||||
) -> Result<Option<Self::AffinityTarget>, Self::Error>;
|
||||
|
||||
fn cached_affinity_matches(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
target: &Self::AffinityTarget,
|
||||
) -> bool;
|
||||
|
||||
async fn build_rankable_candidate(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
original_index: usize,
|
||||
normalized_client_api_format: &str,
|
||||
cached_affinity_match: bool,
|
||||
) -> Result<SchedulerRankableCandidate, Self::Error>;
|
||||
|
||||
fn ranking_context(&self) -> SchedulerRankingContext;
|
||||
|
||||
fn apply_ranking_outcome(
|
||||
&self,
|
||||
candidate: &mut Self::Candidate,
|
||||
outcome: SchedulerRankingOutcome,
|
||||
);
|
||||
}
|
||||
|
||||
pub fn build_ai_rankable_candidate(
|
||||
parts: AiRankableCandidateParts<'_>,
|
||||
) -> SchedulerRankableCandidate {
|
||||
let is_same_format = aether_ai_formats::api_format_alias_matches(
|
||||
parts.provider_api_format,
|
||||
parts.normalized_client_api_format,
|
||||
);
|
||||
let format_preference = aether_ai_formats::request_candidate_api_format_preference(
|
||||
parts.normalized_client_api_format,
|
||||
parts.provider_api_format,
|
||||
)
|
||||
.unwrap_or((u8::MAX, u8::MAX));
|
||||
|
||||
let mut rankable =
|
||||
SchedulerRankableCandidate::from_candidate(parts.candidate, parts.original_index);
|
||||
// The scheduler order is the upstream tie-breaker; AI serving only adds transport facts.
|
||||
rankable.provider_id.clear();
|
||||
rankable.endpoint_id.clear();
|
||||
rankable.key_id.clear();
|
||||
rankable.selected_provider_model_name.clear();
|
||||
|
||||
rankable
|
||||
.with_capability_priority(requested_capability_priority_for_candidate(
|
||||
parts.required_capabilities,
|
||||
parts.candidate,
|
||||
))
|
||||
.with_cached_affinity_match(parts.cached_affinity_match)
|
||||
.with_tunnel_bucket(parts.tunnel_bucket)
|
||||
.with_format_state(
|
||||
!is_same_format && !parts.keep_priority_on_conversion,
|
||||
format_preference,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn ai_ranking_context(config: AiRankingContextConfig) -> SchedulerRankingContext {
|
||||
SchedulerRankingContext {
|
||||
priority_mode: config.priority_mode,
|
||||
ranking_mode: ai_ranking_mode(config.scheduling_mode),
|
||||
include_health: false,
|
||||
load_balance_seed: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn ai_ranking_mode(mode: AiRankingSchedulingMode) -> SchedulerRankingMode {
|
||||
match mode {
|
||||
AiRankingSchedulingMode::FixedOrder => SchedulerRankingMode::FixedOrder,
|
||||
AiRankingSchedulingMode::CacheAffinity => SchedulerRankingMode::CacheAffinity,
|
||||
AiRankingSchedulingMode::LoadBalance => SchedulerRankingMode::LoadBalance,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_ai_candidate_ranking<Port>(
|
||||
port: &Port,
|
||||
mut candidates: Vec<Port::Candidate>,
|
||||
normalized_client_api_format: &str,
|
||||
) -> Result<Vec<Port::Candidate>, Port::Error>
|
||||
where
|
||||
Port: AiCandidateRankingPort,
|
||||
{
|
||||
let affinity_requested_model = port.affinity_requested_model(&candidates);
|
||||
let cached_affinity_target = port
|
||||
.read_cached_affinity_target(
|
||||
normalized_client_api_format,
|
||||
affinity_requested_model.as_deref(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let mut rankables = Vec::with_capacity(candidates.len());
|
||||
for (original_index, candidate) in candidates.iter().enumerate() {
|
||||
let cached_affinity_match = cached_affinity_target
|
||||
.as_ref()
|
||||
.is_some_and(|target| port.cached_affinity_matches(candidate, target));
|
||||
rankables.push(
|
||||
port.build_rankable_candidate(
|
||||
candidate,
|
||||
original_index,
|
||||
normalized_client_api_format,
|
||||
cached_affinity_match,
|
||||
)
|
||||
.await?,
|
||||
);
|
||||
}
|
||||
|
||||
let outcomes =
|
||||
apply_scheduler_candidate_ranking(&mut candidates, &rankables, port.ranking_context());
|
||||
for outcome in outcomes {
|
||||
let ranking_index = outcome.ranking_index;
|
||||
if let Some(candidate) = candidates.get_mut(ranking_index) {
|
||||
port.apply_ranking_outcome(candidate, outcome);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(candidates)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use aether_scheduler_core::{SchedulerPriorityMode, SchedulerRankingMode};
|
||||
use std::sync::Mutex;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct TestCandidate {
|
||||
id: &'static str,
|
||||
priority: i32,
|
||||
ranking_index: Option<usize>,
|
||||
cached_affinity: bool,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestPort {
|
||||
calls: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiCandidateRankingPort for TestPort {
|
||||
type Candidate = TestCandidate;
|
||||
type AffinityTarget = &'static str;
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
fn affinity_requested_model(&self, candidates: &[Self::Candidate]) -> Option<String> {
|
||||
candidates.first().map(|_| "model-a".to_string())
|
||||
}
|
||||
|
||||
async fn read_cached_affinity_target(
|
||||
&self,
|
||||
normalized_client_api_format: &str,
|
||||
affinity_requested_model: Option<&str>,
|
||||
) -> Result<Option<Self::AffinityTarget>, Self::Error> {
|
||||
self.calls.lock().unwrap().push(format!(
|
||||
"affinity:{normalized_client_api_format}:{}",
|
||||
affinity_requested_model.unwrap_or_default()
|
||||
));
|
||||
Ok(Some("candidate-b"))
|
||||
}
|
||||
|
||||
fn cached_affinity_matches(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
target: &Self::AffinityTarget,
|
||||
) -> bool {
|
||||
candidate.id == *target
|
||||
}
|
||||
|
||||
async fn build_rankable_candidate(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
original_index: usize,
|
||||
_normalized_client_api_format: &str,
|
||||
cached_affinity_match: bool,
|
||||
) -> Result<SchedulerRankableCandidate, Self::Error> {
|
||||
self.calls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(format!("rankable:{}:{cached_affinity_match}", candidate.id));
|
||||
Ok(SchedulerRankableCandidate {
|
||||
provider_id: candidate.id.to_string(),
|
||||
endpoint_id: String::new(),
|
||||
key_id: String::new(),
|
||||
selected_provider_model_name: String::new(),
|
||||
provider_priority: candidate.priority,
|
||||
key_internal_priority: 0,
|
||||
key_global_priority_for_format: None,
|
||||
capability_priority: (0, 0),
|
||||
cached_affinity_match,
|
||||
affinity_hash: None,
|
||||
tunnel_bucket: Default::default(),
|
||||
demote_cross_format: false,
|
||||
format_preference: (0, 0),
|
||||
health_bucket: None,
|
||||
health_score: 1.0,
|
||||
original_index,
|
||||
})
|
||||
}
|
||||
|
||||
fn ranking_context(&self) -> SchedulerRankingContext {
|
||||
SchedulerRankingContext {
|
||||
priority_mode: SchedulerPriorityMode::Provider,
|
||||
ranking_mode: SchedulerRankingMode::CacheAffinity,
|
||||
include_health: false,
|
||||
load_balance_seed: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_ranking_outcome(
|
||||
&self,
|
||||
candidate: &mut Self::Candidate,
|
||||
outcome: SchedulerRankingOutcome,
|
||||
) {
|
||||
candidate.ranking_index = Some(outcome.ranking_index);
|
||||
candidate.cached_affinity = outcome.promoted_by.is_some();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ranking_builds_rankables_applies_scheduler_order_and_writes_outcomes() {
|
||||
let port = TestPort::default();
|
||||
let candidates = vec![
|
||||
TestCandidate {
|
||||
id: "candidate-a",
|
||||
priority: 10,
|
||||
ranking_index: None,
|
||||
cached_affinity: false,
|
||||
},
|
||||
TestCandidate {
|
||||
id: "candidate-b",
|
||||
priority: 20,
|
||||
ranking_index: None,
|
||||
cached_affinity: false,
|
||||
},
|
||||
];
|
||||
|
||||
let ranked = run_ai_candidate_ranking(&port, candidates, "openai:chat")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(ranked[0].id, "candidate-b");
|
||||
assert_eq!(ranked[0].ranking_index, Some(0));
|
||||
assert!(ranked[0].cached_affinity);
|
||||
assert_eq!(ranked[1].id, "candidate-a");
|
||||
assert_eq!(ranked[1].ranking_index, Some(1));
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
[
|
||||
"affinity:openai:chat:model-a",
|
||||
"rankable:candidate-a:false",
|
||||
"rankable:candidate-b:true",
|
||||
]
|
||||
);
|
||||
}
|
||||
}
|
||||
415
crates/aether-ai-serving/src/candidate_resolution.rs
Normal file
415
crates/aether-ai-serving/src/candidate_resolution.rs
Normal file
@@ -0,0 +1,415 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AiCandidateResolutionMode {
|
||||
Standard,
|
||||
WithoutTransportPairGate,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct AiCandidateResolutionRequest<'a> {
|
||||
pub client_api_format: &'a str,
|
||||
pub requested_model: Option<&'a str>,
|
||||
pub mode: AiCandidateResolutionMode,
|
||||
}
|
||||
|
||||
impl<'a> AiCandidateResolutionRequest<'a> {
|
||||
pub fn standard(client_api_format: &'a str, requested_model: Option<&'a str>) -> Self {
|
||||
Self {
|
||||
client_api_format,
|
||||
requested_model,
|
||||
mode: AiCandidateResolutionMode::Standard,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn without_transport_pair_gate(
|
||||
client_api_format: &'a str,
|
||||
requested_model: Option<&'a str>,
|
||||
) -> Self {
|
||||
Self {
|
||||
client_api_format,
|
||||
requested_model,
|
||||
mode: AiCandidateResolutionMode::WithoutTransportPairGate,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AiCandidateResolutionOutcome<Eligible, Skipped> {
|
||||
pub eligible_candidates: Vec<Eligible>,
|
||||
pub skipped_candidates: Vec<Skipped>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AiCandidateResolutionPort: Send + Sync {
|
||||
type Candidate: Send;
|
||||
type Transport: Send + Sync;
|
||||
type Eligible: Send + Sync;
|
||||
type Skipped: Send;
|
||||
type Error: Send;
|
||||
|
||||
async fn read_candidate_transport(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
) -> Result<Option<Self::Transport>, Self::Error>;
|
||||
|
||||
fn build_missing_transport_skipped_candidate(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
) -> Self::Skipped;
|
||||
|
||||
fn candidate_common_skip_reason(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
transport: &Self::Transport,
|
||||
requested_model: Option<&str>,
|
||||
) -> Option<&'static str>;
|
||||
|
||||
fn candidate_transport_pair_skip_reason(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
transport: &Self::Transport,
|
||||
normalized_client_api_format: &str,
|
||||
requested_model: &str,
|
||||
) -> Option<&'static str>;
|
||||
|
||||
fn build_skipped_candidate(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
transport: Self::Transport,
|
||||
skip_reason: &'static str,
|
||||
) -> Self::Skipped;
|
||||
|
||||
fn build_eligible_candidate(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
transport: Self::Transport,
|
||||
) -> Self::Eligible;
|
||||
|
||||
async fn rank_eligible_candidates(
|
||||
&self,
|
||||
candidates: Vec<Self::Eligible>,
|
||||
normalized_client_api_format: &str,
|
||||
) -> Result<Vec<Self::Eligible>, Self::Error>;
|
||||
|
||||
async fn apply_pool_scheduler(
|
||||
&self,
|
||||
candidates: Vec<Self::Eligible>,
|
||||
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error>;
|
||||
}
|
||||
|
||||
pub async fn run_ai_candidate_resolution<Port>(
|
||||
port: &Port,
|
||||
candidates: Vec<Port::Candidate>,
|
||||
request: AiCandidateResolutionRequest<'_>,
|
||||
) -> Result<AiCandidateResolutionOutcome<Port::Eligible, Port::Skipped>, Port::Error>
|
||||
where
|
||||
Port: AiCandidateResolutionPort,
|
||||
{
|
||||
let normalized_client_api_format = request.client_api_format.trim().to_ascii_lowercase();
|
||||
let requested_model = request
|
||||
.requested_model
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
let mut eligible = Vec::with_capacity(candidates.len());
|
||||
let mut skipped = Vec::with_capacity(candidates.len());
|
||||
|
||||
for candidate in candidates {
|
||||
let Some(transport) = port.read_candidate_transport(&candidate).await? else {
|
||||
skipped.push(port.build_missing_transport_skipped_candidate(candidate));
|
||||
continue;
|
||||
};
|
||||
|
||||
match candidate_skip_reason_for_mode(
|
||||
port,
|
||||
request.mode,
|
||||
&candidate,
|
||||
&transport,
|
||||
normalized_client_api_format.as_str(),
|
||||
requested_model,
|
||||
) {
|
||||
Some(skip_reason) => {
|
||||
skipped.push(port.build_skipped_candidate(candidate, transport, skip_reason));
|
||||
}
|
||||
None => {
|
||||
eligible.push(port.build_eligible_candidate(candidate, transport));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let ranked = port
|
||||
.rank_eligible_candidates(eligible, normalized_client_api_format.as_str())
|
||||
.await?;
|
||||
let (ranked, pool_skipped) = port.apply_pool_scheduler(ranked).await?;
|
||||
skipped.extend(pool_skipped);
|
||||
|
||||
Ok(AiCandidateResolutionOutcome {
|
||||
eligible_candidates: ranked,
|
||||
skipped_candidates: skipped,
|
||||
})
|
||||
}
|
||||
|
||||
fn candidate_skip_reason_for_mode<Port>(
|
||||
port: &Port,
|
||||
mode: AiCandidateResolutionMode,
|
||||
candidate: &Port::Candidate,
|
||||
transport: &Port::Transport,
|
||||
normalized_client_api_format: &str,
|
||||
requested_model: Option<&str>,
|
||||
) -> Option<&'static str>
|
||||
where
|
||||
Port: AiCandidateResolutionPort,
|
||||
{
|
||||
port.candidate_common_skip_reason(candidate, transport, requested_model)
|
||||
.or_else(|| match mode {
|
||||
AiCandidateResolutionMode::Standard => port.candidate_transport_pair_skip_reason(
|
||||
candidate,
|
||||
transport,
|
||||
normalized_client_api_format,
|
||||
requested_model.unwrap_or_default(),
|
||||
),
|
||||
AiCandidateResolutionMode::WithoutTransportPairGate => None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn extract_ai_pool_sticky_session_token(body_json: &serde_json::Value) -> Option<String> {
|
||||
fn non_empty_str(value: Option<&serde_json::Value>) -> Option<&str> {
|
||||
value
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
let object = body_json.as_object()?;
|
||||
|
||||
non_empty_str(object.get("prompt_cache_key"))
|
||||
.or_else(|| non_empty_str(object.get("conversation_id")))
|
||||
.or_else(|| non_empty_str(object.get("conversationId")))
|
||||
.or_else(|| non_empty_str(object.get("session_id")))
|
||||
.or_else(|| non_empty_str(object.get("sessionId")))
|
||||
.or_else(|| {
|
||||
object
|
||||
.get("metadata")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|metadata| {
|
||||
non_empty_str(metadata.get("session_id"))
|
||||
.or_else(|| non_empty_str(metadata.get("conversation_id")))
|
||||
})
|
||||
})
|
||||
.or_else(|| {
|
||||
object
|
||||
.get("conversationState")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|state| {
|
||||
non_empty_str(state.get("conversationId"))
|
||||
.or_else(|| non_empty_str(state.get("sessionId")))
|
||||
})
|
||||
})
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
use std::sync::Mutex;
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestPort {
|
||||
calls: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiCandidateResolutionPort for TestPort {
|
||||
type Candidate = &'static str;
|
||||
type Transport = &'static str;
|
||||
type Eligible = String;
|
||||
type Skipped = String;
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
async fn read_candidate_transport(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
) -> Result<Option<Self::Transport>, Self::Error> {
|
||||
self.calls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(format!("transport:{candidate}"));
|
||||
Ok(match *candidate {
|
||||
"missing" => None,
|
||||
"inactive" => Some("inactive-transport"),
|
||||
_ => Some("active-transport"),
|
||||
})
|
||||
}
|
||||
|
||||
fn build_missing_transport_skipped_candidate(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
) -> Self::Skipped {
|
||||
format!("{candidate}:transport_snapshot_missing")
|
||||
}
|
||||
|
||||
fn candidate_common_skip_reason(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
_transport: &Self::Transport,
|
||||
requested_model: Option<&str>,
|
||||
) -> Option<&'static str> {
|
||||
self.calls.lock().unwrap().push(format!(
|
||||
"common:{candidate}:{}",
|
||||
requested_model.unwrap_or_default()
|
||||
));
|
||||
(*candidate == "inactive").then_some("provider_inactive")
|
||||
}
|
||||
|
||||
fn candidate_transport_pair_skip_reason(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
_transport: &Self::Transport,
|
||||
normalized_client_api_format: &str,
|
||||
requested_model: &str,
|
||||
) -> Option<&'static str> {
|
||||
self.calls.lock().unwrap().push(format!(
|
||||
"pair:{candidate}:{normalized_client_api_format}:{requested_model}"
|
||||
));
|
||||
(*candidate == "unsupported").then_some("transport_unsupported")
|
||||
}
|
||||
|
||||
fn build_skipped_candidate(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
_transport: Self::Transport,
|
||||
skip_reason: &'static str,
|
||||
) -> Self::Skipped {
|
||||
format!("{candidate}:{skip_reason}")
|
||||
}
|
||||
|
||||
fn build_eligible_candidate(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
_transport: Self::Transport,
|
||||
) -> Self::Eligible {
|
||||
format!("eligible:{candidate}")
|
||||
}
|
||||
|
||||
async fn rank_eligible_candidates(
|
||||
&self,
|
||||
mut candidates: Vec<Self::Eligible>,
|
||||
normalized_client_api_format: &str,
|
||||
) -> Result<Vec<Self::Eligible>, Self::Error> {
|
||||
self.calls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(format!("rank:{normalized_client_api_format}"));
|
||||
candidates.reverse();
|
||||
Ok(candidates)
|
||||
}
|
||||
|
||||
async fn apply_pool_scheduler(
|
||||
&self,
|
||||
candidates: Vec<Self::Eligible>,
|
||||
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error> {
|
||||
self.calls.lock().unwrap().push("pool".to_string());
|
||||
Ok((candidates, vec!["pool:cooldown".to_string()]))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolution_reads_transport_gates_candidates_then_ranks_and_applies_pool() {
|
||||
let port = TestPort::default();
|
||||
|
||||
let outcome = run_ai_candidate_resolution(
|
||||
&port,
|
||||
vec!["first", "missing", "inactive", "unsupported", "second"],
|
||||
AiCandidateResolutionRequest::standard(" OpenAI:Chat ", Some(" gpt-4.1 ")),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
outcome.eligible_candidates,
|
||||
["eligible:second", "eligible:first"]
|
||||
);
|
||||
assert_eq!(
|
||||
outcome.skipped_candidates,
|
||||
[
|
||||
"missing:transport_snapshot_missing",
|
||||
"inactive:provider_inactive",
|
||||
"unsupported:transport_unsupported",
|
||||
"pool:cooldown",
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
[
|
||||
"transport:first",
|
||||
"common:first:gpt-4.1",
|
||||
"pair:first:openai:chat:gpt-4.1",
|
||||
"transport:missing",
|
||||
"transport:inactive",
|
||||
"common:inactive:gpt-4.1",
|
||||
"transport:unsupported",
|
||||
"common:unsupported:gpt-4.1",
|
||||
"pair:unsupported:openai:chat:gpt-4.1",
|
||||
"transport:second",
|
||||
"common:second:gpt-4.1",
|
||||
"pair:second:openai:chat:gpt-4.1",
|
||||
"rank:openai:chat",
|
||||
"pool",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolution_mode_can_skip_transport_pair_gate() {
|
||||
let port = TestPort::default();
|
||||
|
||||
let outcome = run_ai_candidate_resolution(
|
||||
&port,
|
||||
vec!["unsupported"],
|
||||
AiCandidateResolutionRequest::without_transport_pair_gate("openai:chat", None),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(outcome.eligible_candidates, ["eligible:unsupported"]);
|
||||
assert_eq!(outcome.skipped_candidates, ["pool:cooldown"]);
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
[
|
||||
"transport:unsupported",
|
||||
"common:unsupported:",
|
||||
"rank:openai:chat",
|
||||
"pool",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sticky_session_token_is_extracted_from_known_request_fields() {
|
||||
assert_eq!(
|
||||
extract_ai_pool_sticky_session_token(&json!({
|
||||
"prompt_cache_key": " cache-a ",
|
||||
"conversation_id": "conversation-b"
|
||||
}))
|
||||
.as_deref(),
|
||||
Some("cache-a")
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
extract_ai_pool_sticky_session_token(&json!({
|
||||
"metadata": {"conversation_id": " conversation-c "}
|
||||
}))
|
||||
.as_deref(),
|
||||
Some("conversation-c")
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
extract_ai_pool_sticky_session_token(&json!({
|
||||
"conversationState": {"sessionId": " session-d "}
|
||||
}))
|
||||
.as_deref(),
|
||||
Some("session-d")
|
||||
);
|
||||
}
|
||||
}
|
||||
196
crates/aether-ai-serving/src/decision_input.rs
Normal file
196
crates/aether-ai-serving/src/decision_input.rs
Normal file
@@ -0,0 +1,196 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[async_trait]
|
||||
pub trait AiAuthenticatedDecisionInputPort: Send + Sync {
|
||||
type AuthContext: Send + Sync;
|
||||
type AuthSnapshot: Send;
|
||||
type RequiredCapabilities: Send + Sync;
|
||||
type ResolvedInput: Send;
|
||||
type Error: Send;
|
||||
|
||||
async fn read_auth_snapshot(
|
||||
&self,
|
||||
auth_context: &Self::AuthContext,
|
||||
) -> Result<Option<Self::AuthSnapshot>, Self::Error>;
|
||||
|
||||
async fn resolve_required_capabilities(
|
||||
&self,
|
||||
auth_context: &Self::AuthContext,
|
||||
requested_model: Option<&str>,
|
||||
explicit_required_capabilities: Option<&Self::RequiredCapabilities>,
|
||||
) -> Result<Option<Self::RequiredCapabilities>, Self::Error>;
|
||||
|
||||
fn build_resolved_input(
|
||||
&self,
|
||||
auth_context: Self::AuthContext,
|
||||
auth_snapshot: Self::AuthSnapshot,
|
||||
required_capabilities: Option<Self::RequiredCapabilities>,
|
||||
) -> Self::ResolvedInput;
|
||||
}
|
||||
|
||||
pub async fn run_ai_authenticated_decision_input<Port>(
|
||||
port: &Port,
|
||||
auth_context: Port::AuthContext,
|
||||
requested_model: Option<&str>,
|
||||
explicit_required_capabilities: Option<&Port::RequiredCapabilities>,
|
||||
) -> Result<Option<Port::ResolvedInput>, Port::Error>
|
||||
where
|
||||
Port: AiAuthenticatedDecisionInputPort,
|
||||
{
|
||||
let auth_snapshot = match port.read_auth_snapshot(&auth_context).await? {
|
||||
Some(snapshot) => snapshot,
|
||||
None => return Ok(None),
|
||||
};
|
||||
|
||||
let required_capabilities = port
|
||||
.resolve_required_capabilities(
|
||||
&auth_context,
|
||||
requested_model,
|
||||
explicit_required_capabilities,
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(Some(port.build_resolved_input(
|
||||
auth_context,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
)))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::Mutex;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct TestAuthContext {
|
||||
user_id: &'static str,
|
||||
api_key_id: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct TestResolvedInput {
|
||||
auth_context: TestAuthContext,
|
||||
auth_snapshot: &'static str,
|
||||
required_capabilities: Option<String>,
|
||||
}
|
||||
|
||||
struct TestPort {
|
||||
auth_snapshot: Option<&'static str>,
|
||||
calls: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiAuthenticatedDecisionInputPort for TestPort {
|
||||
type AuthContext = TestAuthContext;
|
||||
type AuthSnapshot = &'static str;
|
||||
type RequiredCapabilities = String;
|
||||
type ResolvedInput = TestResolvedInput;
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
async fn read_auth_snapshot(
|
||||
&self,
|
||||
auth_context: &Self::AuthContext,
|
||||
) -> Result<Option<Self::AuthSnapshot>, Self::Error> {
|
||||
self.calls.lock().unwrap().push(format!(
|
||||
"snapshot:{}:{}",
|
||||
auth_context.user_id, auth_context.api_key_id
|
||||
));
|
||||
Ok(self.auth_snapshot)
|
||||
}
|
||||
|
||||
async fn resolve_required_capabilities(
|
||||
&self,
|
||||
auth_context: &Self::AuthContext,
|
||||
requested_model: Option<&str>,
|
||||
explicit_required_capabilities: Option<&Self::RequiredCapabilities>,
|
||||
) -> Result<Option<Self::RequiredCapabilities>, Self::Error> {
|
||||
self.calls.lock().unwrap().push(format!(
|
||||
"capabilities:{}:{}:{}",
|
||||
auth_context.user_id,
|
||||
requested_model.unwrap_or_default(),
|
||||
explicit_required_capabilities
|
||||
.map(String::as_str)
|
||||
.unwrap_or_default()
|
||||
));
|
||||
Ok(Some("merged-capabilities".to_string()))
|
||||
}
|
||||
|
||||
fn build_resolved_input(
|
||||
&self,
|
||||
auth_context: Self::AuthContext,
|
||||
auth_snapshot: Self::AuthSnapshot,
|
||||
required_capabilities: Option<Self::RequiredCapabilities>,
|
||||
) -> Self::ResolvedInput {
|
||||
TestResolvedInput {
|
||||
auth_context,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn authenticated_decision_input_resolves_snapshot_and_capabilities() {
|
||||
let port = TestPort {
|
||||
auth_snapshot: Some("snapshot-a"),
|
||||
calls: Mutex::new(Vec::new()),
|
||||
};
|
||||
let auth_context = TestAuthContext {
|
||||
user_id: "user-a",
|
||||
api_key_id: "key-a",
|
||||
};
|
||||
let explicit = "explicit-capability".to_string();
|
||||
|
||||
let resolved = run_ai_authenticated_decision_input(
|
||||
&port,
|
||||
auth_context.clone(),
|
||||
Some("model-a"),
|
||||
Some(&explicit),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
resolved,
|
||||
Some(TestResolvedInput {
|
||||
auth_context,
|
||||
auth_snapshot: "snapshot-a",
|
||||
required_capabilities: Some("merged-capabilities".to_string()),
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
[
|
||||
"snapshot:user-a:key-a",
|
||||
"capabilities:user-a:model-a:explicit-capability",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn authenticated_decision_input_stops_when_snapshot_is_missing() {
|
||||
let port = TestPort {
|
||||
auth_snapshot: None,
|
||||
calls: Mutex::new(Vec::new()),
|
||||
};
|
||||
|
||||
let resolved = run_ai_authenticated_decision_input(
|
||||
&port,
|
||||
TestAuthContext {
|
||||
user_id: "user-a",
|
||||
api_key_id: "key-a",
|
||||
},
|
||||
Some("model-a"),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(resolved, None);
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
["snapshot:user-a:key-a"]
|
||||
);
|
||||
}
|
||||
}
|
||||
215
crates/aether-ai-serving/src/decision_path.rs
Normal file
215
crates/aether-ai-serving/src/decision_path.rs
Normal file
@@ -0,0 +1,215 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||
pub enum AiSyncDecisionStep {
|
||||
VideoTaskFollowUp,
|
||||
LocalVideo,
|
||||
LocalImage,
|
||||
LocalOpenAiChat,
|
||||
LocalOpenAiResponses,
|
||||
LocalStandardFamily,
|
||||
LocalSameFormatProvider,
|
||||
LocalGeminiFiles,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AiStreamDecisionStep {
|
||||
LocalVideoContent,
|
||||
LocalImage,
|
||||
LocalOpenAiChat,
|
||||
LocalOpenAiResponses,
|
||||
LocalStandardFamily,
|
||||
LocalSameFormatProvider,
|
||||
LocalGeminiFiles,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AiSyncDecisionPathPort: Send + Sync {
|
||||
type Decision: Send;
|
||||
type Error: Send;
|
||||
|
||||
fn sync_decision_step_enabled(&self, step: AiSyncDecisionStep) -> bool;
|
||||
|
||||
async fn build_sync_decision_step(
|
||||
&self,
|
||||
step: AiSyncDecisionStep,
|
||||
) -> Result<Option<Self::Decision>, Self::Error>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AiStreamDecisionPathPort: Send + Sync {
|
||||
type Decision: Send;
|
||||
type Error: Send;
|
||||
|
||||
async fn build_stream_decision_step(
|
||||
&self,
|
||||
step: AiStreamDecisionStep,
|
||||
) -> Result<Option<Self::Decision>, Self::Error>;
|
||||
}
|
||||
|
||||
pub async fn run_ai_sync_decision_path<Port>(
|
||||
port: &Port,
|
||||
) -> Result<Option<Port::Decision>, Port::Error>
|
||||
where
|
||||
Port: AiSyncDecisionPathPort,
|
||||
{
|
||||
for step in [
|
||||
AiSyncDecisionStep::VideoTaskFollowUp,
|
||||
AiSyncDecisionStep::LocalVideo,
|
||||
AiSyncDecisionStep::LocalImage,
|
||||
AiSyncDecisionStep::LocalOpenAiChat,
|
||||
AiSyncDecisionStep::LocalOpenAiResponses,
|
||||
AiSyncDecisionStep::LocalStandardFamily,
|
||||
AiSyncDecisionStep::LocalSameFormatProvider,
|
||||
AiSyncDecisionStep::LocalGeminiFiles,
|
||||
] {
|
||||
if !port.sync_decision_step_enabled(step) {
|
||||
continue;
|
||||
}
|
||||
if let Some(decision) = port.build_sync_decision_step(step).await? {
|
||||
return Ok(Some(decision));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub async fn run_ai_stream_decision_path<Port>(
|
||||
port: &Port,
|
||||
) -> Result<Option<Port::Decision>, Port::Error>
|
||||
where
|
||||
Port: AiStreamDecisionPathPort,
|
||||
{
|
||||
for step in [
|
||||
AiStreamDecisionStep::LocalVideoContent,
|
||||
AiStreamDecisionStep::LocalImage,
|
||||
AiStreamDecisionStep::LocalOpenAiChat,
|
||||
AiStreamDecisionStep::LocalOpenAiResponses,
|
||||
AiStreamDecisionStep::LocalStandardFamily,
|
||||
AiStreamDecisionStep::LocalSameFormatProvider,
|
||||
AiStreamDecisionStep::LocalGeminiFiles,
|
||||
] {
|
||||
if let Some(decision) = port.build_stream_decision_step(step).await? {
|
||||
return Ok(Some(decision));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::collections::{BTreeSet, VecDeque};
|
||||
use std::sync::Mutex;
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestSyncDecisionPort {
|
||||
disabled: BTreeSet<AiSyncDecisionStep>,
|
||||
outcomes: Mutex<VecDeque<Option<&'static str>>>,
|
||||
calls: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiSyncDecisionPathPort for TestSyncDecisionPort {
|
||||
type Decision = &'static str;
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
fn sync_decision_step_enabled(&self, step: AiSyncDecisionStep) -> bool {
|
||||
!self.disabled.contains(&step)
|
||||
}
|
||||
|
||||
async fn build_sync_decision_step(
|
||||
&self,
|
||||
step: AiSyncDecisionStep,
|
||||
) -> Result<Option<Self::Decision>, Self::Error> {
|
||||
self.calls.lock().unwrap().push(format!("{step:?}"));
|
||||
Ok(self.outcomes.lock().unwrap().pop_front().unwrap_or(None))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestStreamDecisionPort {
|
||||
outcomes: Mutex<VecDeque<Option<&'static str>>>,
|
||||
calls: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiStreamDecisionPathPort for TestStreamDecisionPort {
|
||||
type Decision = &'static str;
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
async fn build_stream_decision_step(
|
||||
&self,
|
||||
step: AiStreamDecisionStep,
|
||||
) -> Result<Option<Self::Decision>, Self::Error> {
|
||||
self.calls.lock().unwrap().push(format!("{step:?}"));
|
||||
Ok(self.outcomes.lock().unwrap().pop_front().unwrap_or(None))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_decision_path_runs_steps_in_serving_order() {
|
||||
let port = TestSyncDecisionPort::default();
|
||||
|
||||
let decision = run_ai_sync_decision_path(&port).await.unwrap();
|
||||
|
||||
assert_eq!(decision, None);
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
[
|
||||
"VideoTaskFollowUp",
|
||||
"LocalVideo",
|
||||
"LocalImage",
|
||||
"LocalOpenAiChat",
|
||||
"LocalOpenAiResponses",
|
||||
"LocalStandardFamily",
|
||||
"LocalSameFormatProvider",
|
||||
"LocalGeminiFiles",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_decision_path_skips_disabled_steps_and_stops_at_first_decision() {
|
||||
let port = TestSyncDecisionPort {
|
||||
disabled: BTreeSet::from([AiSyncDecisionStep::LocalGeminiFiles]),
|
||||
outcomes: Mutex::new(VecDeque::from([
|
||||
None,
|
||||
None,
|
||||
Some("image_decision"),
|
||||
Some("should_not_run"),
|
||||
])),
|
||||
calls: Mutex::default(),
|
||||
};
|
||||
|
||||
let decision = run_ai_sync_decision_path(&port).await.unwrap();
|
||||
|
||||
assert_eq!(decision, Some("image_decision"));
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
["VideoTaskFollowUp", "LocalVideo", "LocalImage"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stream_decision_path_stops_at_first_decision() {
|
||||
let port = TestStreamDecisionPort {
|
||||
outcomes: Mutex::new(VecDeque::from([
|
||||
None,
|
||||
None,
|
||||
Some("chat_decision"),
|
||||
Some("should_not_run"),
|
||||
])),
|
||||
calls: Mutex::default(),
|
||||
};
|
||||
|
||||
let decision = run_ai_stream_decision_path(&port).await.unwrap();
|
||||
|
||||
assert_eq!(decision, Some("chat_decision"));
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
["LocalVideoContent", "LocalImage", "LocalOpenAiChat"]
|
||||
);
|
||||
}
|
||||
}
|
||||
92
crates/aether-ai-serving/src/decision_payload.rs
Normal file
92
crates/aether-ai-serving/src/decision_payload.rs
Normal file
@@ -0,0 +1,92 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_ai_surfaces::api::{
|
||||
ExecutionRuntimeAuthContext, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
|
||||
};
|
||||
use aether_contracts::{ExecutionTimeouts, ProxySnapshot};
|
||||
|
||||
use crate::{AiExecutionDecision, ConversionMode, ExecutionStrategy};
|
||||
|
||||
pub struct AiExecutionDecisionResponseParts {
|
||||
pub decision_is_stream: bool,
|
||||
pub decision_kind: String,
|
||||
pub execution_strategy: ExecutionStrategy,
|
||||
pub conversion_mode: ConversionMode,
|
||||
pub request_id: String,
|
||||
pub candidate_id: String,
|
||||
pub provider_name: String,
|
||||
pub provider_id: String,
|
||||
pub endpoint_id: String,
|
||||
pub key_id: String,
|
||||
pub upstream_base_url: String,
|
||||
pub upstream_url: String,
|
||||
pub provider_request_method: Option<String>,
|
||||
pub auth_header: Option<String>,
|
||||
pub auth_value: Option<String>,
|
||||
pub provider_api_format: String,
|
||||
pub client_api_format: String,
|
||||
pub model_name: String,
|
||||
pub mapped_model: String,
|
||||
pub prompt_cache_key: Option<String>,
|
||||
pub provider_request_headers: BTreeMap<String, String>,
|
||||
pub provider_request_body: Option<serde_json::Value>,
|
||||
pub provider_request_body_base64: Option<String>,
|
||||
pub content_type: Option<String>,
|
||||
pub proxy: Option<ProxySnapshot>,
|
||||
pub tls_profile: Option<String>,
|
||||
pub timeouts: Option<ExecutionTimeouts>,
|
||||
pub upstream_is_stream: bool,
|
||||
pub report_kind: Option<String>,
|
||||
pub report_context: Option<serde_json::Value>,
|
||||
pub auth_context: ExecutionRuntimeAuthContext,
|
||||
}
|
||||
|
||||
pub fn build_ai_execution_decision_response(
|
||||
parts: AiExecutionDecisionResponseParts,
|
||||
) -> AiExecutionDecision {
|
||||
AiExecutionDecision {
|
||||
action: ai_execution_decision_action(parts.decision_is_stream).to_string(),
|
||||
decision_kind: Some(parts.decision_kind),
|
||||
execution_strategy: Some(parts.execution_strategy.as_str().to_string()),
|
||||
conversion_mode: Some(parts.conversion_mode.as_str().to_string()),
|
||||
request_id: Some(parts.request_id),
|
||||
candidate_id: Some(parts.candidate_id),
|
||||
provider_name: Some(parts.provider_name),
|
||||
provider_id: Some(parts.provider_id),
|
||||
endpoint_id: Some(parts.endpoint_id),
|
||||
key_id: Some(parts.key_id),
|
||||
upstream_base_url: Some(parts.upstream_base_url),
|
||||
upstream_url: Some(parts.upstream_url),
|
||||
provider_request_method: parts.provider_request_method,
|
||||
auth_header: parts.auth_header,
|
||||
auth_value: parts.auth_value,
|
||||
provider_api_format: Some(parts.provider_api_format.clone()),
|
||||
client_api_format: Some(parts.client_api_format.clone()),
|
||||
provider_contract: Some(parts.provider_api_format),
|
||||
client_contract: Some(parts.client_api_format),
|
||||
model_name: Some(parts.model_name),
|
||||
mapped_model: Some(parts.mapped_model),
|
||||
prompt_cache_key: parts.prompt_cache_key,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: parts.provider_request_headers,
|
||||
provider_request_body: parts.provider_request_body,
|
||||
provider_request_body_base64: parts.provider_request_body_base64,
|
||||
content_type: parts.content_type,
|
||||
proxy: parts.proxy,
|
||||
tls_profile: parts.tls_profile,
|
||||
timeouts: parts.timeouts,
|
||||
upstream_is_stream: parts.upstream_is_stream,
|
||||
report_kind: parts.report_kind,
|
||||
report_context: parts.report_context,
|
||||
auth_context: Some(parts.auth_context),
|
||||
}
|
||||
}
|
||||
|
||||
pub const fn ai_execution_decision_action(decision_is_stream: bool) -> &'static str {
|
||||
if decision_is_stream {
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION
|
||||
} else {
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION
|
||||
}
|
||||
}
|
||||
250
crates/aether-ai-serving/src/dto.rs
Normal file
250
crates/aether-ai-serving/src/dto.rs
Normal file
@@ -0,0 +1,250 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_ai_surfaces::api::ExecutionRuntimeAuthContext;
|
||||
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, ProxySnapshot};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ExecutionStrategy {
|
||||
GatewayAffinityForward,
|
||||
RawPublicProxy,
|
||||
LocalSameFormat,
|
||||
LocalCrossFormat,
|
||||
}
|
||||
|
||||
impl ExecutionStrategy {
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::GatewayAffinityForward => "gateway_affinity_forward",
|
||||
Self::RawPublicProxy => "raw_public_proxy",
|
||||
Self::LocalSameFormat => "local_same_format",
|
||||
Self::LocalCrossFormat => "local_cross_format",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ConversionMode {
|
||||
None,
|
||||
RequestOnly,
|
||||
ResponseOnly,
|
||||
Bidirectional,
|
||||
}
|
||||
|
||||
impl ConversionMode {
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::None => "none",
|
||||
Self::RequestOnly => "request_only",
|
||||
Self::ResponseOnly => "response_only",
|
||||
Self::Bidirectional => "bidirectional",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct AiExecutionPlanPayload {
|
||||
pub action: String,
|
||||
#[serde(default)]
|
||||
pub plan_kind: Option<String>,
|
||||
#[serde(default)]
|
||||
pub plan: Option<ExecutionPlan>,
|
||||
#[serde(default)]
|
||||
pub report_kind: Option<String>,
|
||||
#[serde(default)]
|
||||
pub report_context: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub auth_context: Option<ExecutionRuntimeAuthContext>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct AiExecutionDecision {
|
||||
pub action: String,
|
||||
#[serde(default)]
|
||||
pub decision_kind: Option<String>,
|
||||
#[serde(default)]
|
||||
pub execution_strategy: Option<String>,
|
||||
#[serde(default)]
|
||||
pub conversion_mode: Option<String>,
|
||||
#[serde(default)]
|
||||
pub request_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub candidate_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub provider_name: Option<String>,
|
||||
#[serde(default)]
|
||||
pub provider_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub endpoint_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub key_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub upstream_base_url: Option<String>,
|
||||
#[serde(default)]
|
||||
pub upstream_url: Option<String>,
|
||||
#[serde(default)]
|
||||
pub provider_request_method: Option<String>,
|
||||
#[serde(default)]
|
||||
pub auth_header: Option<String>,
|
||||
#[serde(default)]
|
||||
pub auth_value: Option<String>,
|
||||
#[serde(default)]
|
||||
pub provider_api_format: Option<String>,
|
||||
#[serde(default)]
|
||||
pub client_api_format: Option<String>,
|
||||
#[serde(default)]
|
||||
pub provider_contract: Option<String>,
|
||||
#[serde(default)]
|
||||
pub client_contract: Option<String>,
|
||||
#[serde(default)]
|
||||
pub model_name: Option<String>,
|
||||
#[serde(default)]
|
||||
pub mapped_model: Option<String>,
|
||||
#[serde(default)]
|
||||
pub prompt_cache_key: Option<String>,
|
||||
#[serde(default)]
|
||||
pub extra_headers: BTreeMap<String, String>,
|
||||
#[serde(default)]
|
||||
pub provider_request_headers: BTreeMap<String, String>,
|
||||
#[serde(default)]
|
||||
pub provider_request_body: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub provider_request_body_base64: Option<String>,
|
||||
#[serde(default)]
|
||||
pub content_type: Option<String>,
|
||||
#[serde(default)]
|
||||
pub proxy: Option<ProxySnapshot>,
|
||||
#[serde(default)]
|
||||
pub tls_profile: Option<String>,
|
||||
#[serde(default)]
|
||||
pub timeouts: Option<ExecutionTimeouts>,
|
||||
#[serde(default)]
|
||||
pub upstream_is_stream: bool,
|
||||
#[serde(default)]
|
||||
pub report_kind: Option<String>,
|
||||
#[serde(default)]
|
||||
pub report_context: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub auth_context: Option<ExecutionRuntimeAuthContext>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct AiSyncAttempt {
|
||||
pub plan: ExecutionPlan,
|
||||
pub report_kind: Option<String>,
|
||||
pub report_context: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct AiStreamAttempt {
|
||||
pub plan: ExecutionPlan,
|
||||
pub report_kind: Option<String>,
|
||||
pub report_context: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
pub fn augment_sync_report_context(
|
||||
report_context: Option<serde_json::Value>,
|
||||
provider_request_headers: &BTreeMap<String, String>,
|
||||
_provider_request_body: &serde_json::Value,
|
||||
) -> serde_json::Result<Option<serde_json::Value>> {
|
||||
let mut report_context = match report_context {
|
||||
Some(serde_json::Value::Object(map)) => map,
|
||||
Some(_) => serde_json::Map::new(),
|
||||
None => serde_json::Map::new(),
|
||||
};
|
||||
|
||||
report_context.insert(
|
||||
"provider_request_headers".to_string(),
|
||||
serde_json::to_value(provider_request_headers)?,
|
||||
);
|
||||
|
||||
Ok(Some(serde_json::Value::Object(report_context)))
|
||||
}
|
||||
|
||||
fn decision_has_exact_provider_request(payload: &AiExecutionDecision) -> bool {
|
||||
!payload.provider_request_headers.is_empty()
|
||||
&& (payload.provider_request_body.is_some()
|
||||
|| payload
|
||||
.provider_request_body_base64
|
||||
.as_ref()
|
||||
.map(|value| !value.trim().is_empty())
|
||||
.unwrap_or(false))
|
||||
}
|
||||
|
||||
pub fn generic_decision_missing_exact_provider_request(payload: &AiExecutionDecision) -> bool {
|
||||
!decision_has_exact_provider_request(payload)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use super::{
|
||||
augment_sync_report_context, generic_decision_missing_exact_provider_request,
|
||||
AiExecutionDecision,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn generic_decision_detects_missing_exact_provider_request() {
|
||||
let payload = AiExecutionDecision {
|
||||
action: "local".to_string(),
|
||||
decision_kind: Some("sync".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: None,
|
||||
request_id: None,
|
||||
candidate_id: None,
|
||||
provider_name: None,
|
||||
provider_id: None,
|
||||
endpoint_id: None,
|
||||
key_id: None,
|
||||
upstream_base_url: None,
|
||||
upstream_url: None,
|
||||
provider_request_method: None,
|
||||
auth_header: None,
|
||||
auth_value: None,
|
||||
provider_api_format: None,
|
||||
client_api_format: None,
|
||||
provider_contract: None,
|
||||
client_contract: None,
|
||||
model_name: None,
|
||||
mapped_model: None,
|
||||
prompt_cache_key: None,
|
||||
extra_headers: Default::default(),
|
||||
provider_request_headers: Default::default(),
|
||||
provider_request_body: None,
|
||||
provider_request_body_base64: None,
|
||||
content_type: None,
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: false,
|
||||
report_kind: None,
|
||||
report_context: None,
|
||||
auth_context: None,
|
||||
};
|
||||
|
||||
assert!(generic_decision_missing_exact_provider_request(&payload));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn augment_sync_report_context_attaches_provider_request_headers_only() {
|
||||
let report_context = augment_sync_report_context(
|
||||
Some(serde_json::json!({"trace_id": "abc"})),
|
||||
&BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
|
||||
&serde_json::json!({"model": "gpt-5"}),
|
||||
)
|
||||
.expect("context should serialize")
|
||||
.expect("context should exist");
|
||||
|
||||
assert_eq!(
|
||||
report_context["provider_request_headers"]["content-type"],
|
||||
"application/json"
|
||||
);
|
||||
assert!(
|
||||
report_context.get("provider_request_body").is_none(),
|
||||
"provider request body should not be copied into report context"
|
||||
);
|
||||
}
|
||||
}
|
||||
408
crates/aether-ai-serving/src/execution_path.rs
Normal file
408
crates/aether-ai-serving/src/execution_path.rs
Normal file
@@ -0,0 +1,408 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum AiServingExecutionOutcome<Response, Exhaustion> {
|
||||
Responded(Response),
|
||||
Exhausted(Exhaustion),
|
||||
NoPath,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AiPlanFallbackReason {
|
||||
RemoteDecisionMiss,
|
||||
SchedulerDecisionUnsupported,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AiSyncExecutionStep {
|
||||
VideoTaskFollowUp,
|
||||
LocalVideo,
|
||||
LocalImage,
|
||||
LocalOpenAiChat,
|
||||
LocalOpenAiResponses,
|
||||
LocalStandardFamily,
|
||||
LocalSameFormatProvider,
|
||||
LocalGeminiFiles,
|
||||
RemoteDecision,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AiStreamExecutionStep {
|
||||
LocalVideoContent,
|
||||
LocalImage,
|
||||
LocalOpenAiChat,
|
||||
LocalOpenAiResponses,
|
||||
LocalStandardFamily,
|
||||
LocalSameFormatProvider,
|
||||
LocalGeminiFiles,
|
||||
RemoteDecision,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AiSyncExecutionPathPort: Send + Sync {
|
||||
type Response: Send;
|
||||
type Exhaustion: Send;
|
||||
type Error: Send;
|
||||
|
||||
fn scheduler_decision_supported(&self) -> bool;
|
||||
|
||||
async fn execute_sync_step(
|
||||
&self,
|
||||
step: AiSyncExecutionStep,
|
||||
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>;
|
||||
|
||||
async fn execute_sync_plan_fallback(
|
||||
&self,
|
||||
reason: AiPlanFallbackReason,
|
||||
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AiStreamExecutionPathPort: Send + Sync {
|
||||
type Response: Send;
|
||||
type Exhaustion: Send;
|
||||
type Error: Send;
|
||||
|
||||
fn scheduler_decision_supported(&self) -> bool;
|
||||
|
||||
async fn execute_stream_step(
|
||||
&self,
|
||||
step: AiStreamExecutionStep,
|
||||
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>;
|
||||
|
||||
async fn execute_stream_plan_fallback(
|
||||
&self,
|
||||
reason: AiPlanFallbackReason,
|
||||
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>;
|
||||
}
|
||||
|
||||
pub async fn run_ai_sync_execution_path<Port>(
|
||||
port: &Port,
|
||||
) -> Result<AiServingExecutionOutcome<Port::Response, Port::Exhaustion>, Port::Error>
|
||||
where
|
||||
Port: AiSyncExecutionPathPort,
|
||||
{
|
||||
let mut exhausted = None;
|
||||
|
||||
if let Some(response) =
|
||||
absorb_sync_step(port, AiSyncExecutionStep::VideoTaskFollowUp, &mut exhausted).await?
|
||||
{
|
||||
return Ok(response);
|
||||
}
|
||||
|
||||
if port.scheduler_decision_supported() {
|
||||
for step in [
|
||||
AiSyncExecutionStep::LocalVideo,
|
||||
AiSyncExecutionStep::LocalImage,
|
||||
AiSyncExecutionStep::LocalOpenAiChat,
|
||||
AiSyncExecutionStep::LocalOpenAiResponses,
|
||||
AiSyncExecutionStep::LocalStandardFamily,
|
||||
AiSyncExecutionStep::LocalSameFormatProvider,
|
||||
AiSyncExecutionStep::LocalGeminiFiles,
|
||||
AiSyncExecutionStep::RemoteDecision,
|
||||
] {
|
||||
if let Some(response) = absorb_sync_step(port, step, &mut exhausted).await? {
|
||||
return Ok(response);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let fallback_reason = if port.scheduler_decision_supported() {
|
||||
AiPlanFallbackReason::RemoteDecisionMiss
|
||||
} else {
|
||||
AiPlanFallbackReason::SchedulerDecisionUnsupported
|
||||
};
|
||||
match port.execute_sync_plan_fallback(fallback_reason).await? {
|
||||
AiServingExecutionOutcome::Responded(response) => {
|
||||
Ok(AiServingExecutionOutcome::Responded(response))
|
||||
}
|
||||
AiServingExecutionOutcome::Exhausted(outcome) => {
|
||||
Ok(AiServingExecutionOutcome::Exhausted(outcome))
|
||||
}
|
||||
AiServingExecutionOutcome::NoPath => Ok(exhausted
|
||||
.map(AiServingExecutionOutcome::Exhausted)
|
||||
.unwrap_or(AiServingExecutionOutcome::NoPath)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_ai_stream_execution_path<Port>(
|
||||
port: &Port,
|
||||
) -> Result<AiServingExecutionOutcome<Port::Response, Port::Exhaustion>, Port::Error>
|
||||
where
|
||||
Port: AiStreamExecutionPathPort,
|
||||
{
|
||||
let mut exhausted = None;
|
||||
|
||||
if let Some(response) = absorb_stream_step(
|
||||
port,
|
||||
AiStreamExecutionStep::LocalVideoContent,
|
||||
&mut exhausted,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(response);
|
||||
}
|
||||
|
||||
if port.scheduler_decision_supported() {
|
||||
for step in [
|
||||
AiStreamExecutionStep::LocalImage,
|
||||
AiStreamExecutionStep::LocalOpenAiChat,
|
||||
AiStreamExecutionStep::LocalOpenAiResponses,
|
||||
AiStreamExecutionStep::LocalStandardFamily,
|
||||
AiStreamExecutionStep::LocalSameFormatProvider,
|
||||
AiStreamExecutionStep::LocalGeminiFiles,
|
||||
AiStreamExecutionStep::RemoteDecision,
|
||||
] {
|
||||
if let Some(response) = absorb_stream_step(port, step, &mut exhausted).await? {
|
||||
return Ok(response);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let fallback_reason = if port.scheduler_decision_supported() {
|
||||
AiPlanFallbackReason::RemoteDecisionMiss
|
||||
} else {
|
||||
AiPlanFallbackReason::SchedulerDecisionUnsupported
|
||||
};
|
||||
match port.execute_stream_plan_fallback(fallback_reason).await? {
|
||||
AiServingExecutionOutcome::Responded(response) => {
|
||||
Ok(AiServingExecutionOutcome::Responded(response))
|
||||
}
|
||||
AiServingExecutionOutcome::Exhausted(outcome) => {
|
||||
Ok(AiServingExecutionOutcome::Exhausted(outcome))
|
||||
}
|
||||
AiServingExecutionOutcome::NoPath => Ok(exhausted
|
||||
.map(AiServingExecutionOutcome::Exhausted)
|
||||
.unwrap_or(AiServingExecutionOutcome::NoPath)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn absorb_sync_step<Port>(
|
||||
port: &Port,
|
||||
step: AiSyncExecutionStep,
|
||||
exhausted: &mut Option<Port::Exhaustion>,
|
||||
) -> Result<Option<AiServingExecutionOutcome<Port::Response, Port::Exhaustion>>, Port::Error>
|
||||
where
|
||||
Port: AiSyncExecutionPathPort,
|
||||
{
|
||||
match port.execute_sync_step(step).await? {
|
||||
AiServingExecutionOutcome::Responded(response) => {
|
||||
Ok(Some(AiServingExecutionOutcome::Responded(response)))
|
||||
}
|
||||
AiServingExecutionOutcome::Exhausted(outcome) => {
|
||||
*exhausted = Some(outcome);
|
||||
Ok(None)
|
||||
}
|
||||
AiServingExecutionOutcome::NoPath => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn absorb_stream_step<Port>(
|
||||
port: &Port,
|
||||
step: AiStreamExecutionStep,
|
||||
exhausted: &mut Option<Port::Exhaustion>,
|
||||
) -> Result<Option<AiServingExecutionOutcome<Port::Response, Port::Exhaustion>>, Port::Error>
|
||||
where
|
||||
Port: AiStreamExecutionPathPort,
|
||||
{
|
||||
match port.execute_stream_step(step).await? {
|
||||
AiServingExecutionOutcome::Responded(response) => {
|
||||
Ok(Some(AiServingExecutionOutcome::Responded(response)))
|
||||
}
|
||||
AiServingExecutionOutcome::Exhausted(outcome) => {
|
||||
*exhausted = Some(outcome);
|
||||
Ok(None)
|
||||
}
|
||||
AiServingExecutionOutcome::NoPath => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Mutex;
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestSyncPort {
|
||||
scheduler_supported: bool,
|
||||
outcomes: Mutex<VecDeque<AiServingExecutionOutcome<&'static str, &'static str>>>,
|
||||
calls: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiSyncExecutionPathPort for TestSyncPort {
|
||||
type Response = &'static str;
|
||||
type Exhaustion = &'static str;
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
fn scheduler_decision_supported(&self) -> bool {
|
||||
self.scheduler_supported
|
||||
}
|
||||
|
||||
async fn execute_sync_step(
|
||||
&self,
|
||||
step: AiSyncExecutionStep,
|
||||
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>
|
||||
{
|
||||
self.calls.lock().unwrap().push(format!("{step:?}"));
|
||||
Ok(self
|
||||
.outcomes
|
||||
.lock()
|
||||
.unwrap()
|
||||
.pop_front()
|
||||
.unwrap_or(AiServingExecutionOutcome::NoPath))
|
||||
}
|
||||
|
||||
async fn execute_sync_plan_fallback(
|
||||
&self,
|
||||
reason: AiPlanFallbackReason,
|
||||
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>
|
||||
{
|
||||
self.calls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(format!("Fallback:{reason:?}"));
|
||||
Ok(self
|
||||
.outcomes
|
||||
.lock()
|
||||
.unwrap()
|
||||
.pop_front()
|
||||
.unwrap_or(AiServingExecutionOutcome::NoPath))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestStreamPort {
|
||||
scheduler_supported: bool,
|
||||
outcomes: Mutex<VecDeque<AiServingExecutionOutcome<&'static str, &'static str>>>,
|
||||
calls: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiStreamExecutionPathPort for TestStreamPort {
|
||||
type Response = &'static str;
|
||||
type Exhaustion = &'static str;
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
fn scheduler_decision_supported(&self) -> bool {
|
||||
self.scheduler_supported
|
||||
}
|
||||
|
||||
async fn execute_stream_step(
|
||||
&self,
|
||||
step: AiStreamExecutionStep,
|
||||
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>
|
||||
{
|
||||
self.calls.lock().unwrap().push(format!("{step:?}"));
|
||||
Ok(self
|
||||
.outcomes
|
||||
.lock()
|
||||
.unwrap()
|
||||
.pop_front()
|
||||
.unwrap_or(AiServingExecutionOutcome::NoPath))
|
||||
}
|
||||
|
||||
async fn execute_stream_plan_fallback(
|
||||
&self,
|
||||
reason: AiPlanFallbackReason,
|
||||
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>
|
||||
{
|
||||
self.calls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(format!("Fallback:{reason:?}"));
|
||||
Ok(self
|
||||
.outcomes
|
||||
.lock()
|
||||
.unwrap()
|
||||
.pop_front()
|
||||
.unwrap_or(AiServingExecutionOutcome::NoPath))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_path_runs_scheduler_steps_before_remote_and_fallback() {
|
||||
let port = TestSyncPort {
|
||||
scheduler_supported: true,
|
||||
..TestSyncPort::default()
|
||||
};
|
||||
|
||||
let outcome = run_ai_sync_execution_path(&port).await.unwrap();
|
||||
|
||||
assert!(matches!(outcome, AiServingExecutionOutcome::NoPath));
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
[
|
||||
"VideoTaskFollowUp",
|
||||
"LocalVideo",
|
||||
"LocalImage",
|
||||
"LocalOpenAiChat",
|
||||
"LocalOpenAiResponses",
|
||||
"LocalStandardFamily",
|
||||
"LocalSameFormatProvider",
|
||||
"LocalGeminiFiles",
|
||||
"RemoteDecision",
|
||||
"Fallback:RemoteDecisionMiss",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_path_returns_last_exhaustion_when_fallback_has_no_path() {
|
||||
let port = TestSyncPort {
|
||||
scheduler_supported: true,
|
||||
outcomes: Mutex::new(VecDeque::from([
|
||||
AiServingExecutionOutcome::NoPath,
|
||||
AiServingExecutionOutcome::Exhausted("local_video_exhausted"),
|
||||
])),
|
||||
calls: Mutex::default(),
|
||||
};
|
||||
|
||||
let outcome = run_ai_sync_execution_path(&port).await.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
outcome,
|
||||
AiServingExecutionOutcome::Exhausted("local_video_exhausted")
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stream_path_skips_scheduler_steps_when_unsupported() {
|
||||
let port = TestStreamPort {
|
||||
scheduler_supported: false,
|
||||
..TestStreamPort::default()
|
||||
};
|
||||
|
||||
let outcome = run_ai_stream_execution_path(&port).await.unwrap();
|
||||
|
||||
assert!(matches!(outcome, AiServingExecutionOutcome::NoPath));
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
["LocalVideoContent", "Fallback:SchedulerDecisionUnsupported",]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stream_path_stops_at_first_response() {
|
||||
let port = TestStreamPort {
|
||||
scheduler_supported: true,
|
||||
outcomes: Mutex::new(VecDeque::from([
|
||||
AiServingExecutionOutcome::NoPath,
|
||||
AiServingExecutionOutcome::Responded("image_response"),
|
||||
])),
|
||||
calls: Mutex::default(),
|
||||
};
|
||||
|
||||
let outcome = run_ai_stream_execution_path(&port).await.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
outcome,
|
||||
AiServingExecutionOutcome::Responded("image_response")
|
||||
));
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
["LocalVideoContent", "LocalImage"]
|
||||
);
|
||||
}
|
||||
}
|
||||
190
crates/aether-ai-serving/src/failure_diagnostic.rs
Normal file
190
crates/aether-ai-serving/src/failure_diagnostic.rs
Normal file
@@ -0,0 +1,190 @@
|
||||
use serde_json::{json, Value};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum CandidateFailureDiagnosticKind {
|
||||
RequestBodyBuild,
|
||||
RequestConversion,
|
||||
BodyRules,
|
||||
HeaderRules,
|
||||
UrlBuild,
|
||||
TransportAuth,
|
||||
EnvelopeBuild,
|
||||
}
|
||||
|
||||
impl CandidateFailureDiagnosticKind {
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::RequestBodyBuild => "request_body_build",
|
||||
Self::RequestConversion => "request_conversion",
|
||||
Self::BodyRules => "body_rules",
|
||||
Self::HeaderRules => "header_rules",
|
||||
Self::UrlBuild => "url_build",
|
||||
Self::TransportAuth => "transport_auth",
|
||||
Self::EnvelopeBuild => "envelope_build",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CandidateFailureDiagnostic {
|
||||
kind: CandidateFailureDiagnosticKind,
|
||||
path: String,
|
||||
message: String,
|
||||
source: Option<String>,
|
||||
client_api_format: Option<String>,
|
||||
provider_api_format: Option<String>,
|
||||
safe_to_show: bool,
|
||||
}
|
||||
|
||||
impl CandidateFailureDiagnostic {
|
||||
pub fn new(
|
||||
kind: CandidateFailureDiagnosticKind,
|
||||
path: impl Into<String>,
|
||||
message: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
kind,
|
||||
path: path.into(),
|
||||
message: message.into(),
|
||||
source: None,
|
||||
client_api_format: None,
|
||||
provider_api_format: None,
|
||||
safe_to_show: true,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn source(mut self, source: impl Into<String>) -> Self {
|
||||
self.source = Some(source.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn formats(
|
||||
mut self,
|
||||
client_api_format: impl Into<String>,
|
||||
provider_api_format: impl Into<String>,
|
||||
) -> Self {
|
||||
self.client_api_format = Some(client_api_format.into());
|
||||
self.provider_api_format = Some(provider_api_format.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn to_extra_data(&self) -> Value {
|
||||
let diagnostic = self.to_value();
|
||||
let mut extra_data = json!({
|
||||
"failure_diagnostic": diagnostic,
|
||||
});
|
||||
|
||||
// Compatibility for current usage UI and already persisted trace readers.
|
||||
if self.kind == CandidateFailureDiagnosticKind::RequestBodyBuild {
|
||||
if let Some(object) = extra_data.as_object_mut() {
|
||||
object.insert(
|
||||
"request_body_build_error".to_string(),
|
||||
json!({
|
||||
"path": self.path,
|
||||
"message": self.message,
|
||||
"client_api_format": self.client_api_format,
|
||||
"provider_api_format": self.provider_api_format,
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
extra_data
|
||||
}
|
||||
|
||||
pub fn upstream_url_missing(
|
||||
client_api_format: impl Into<String>,
|
||||
provider_api_format: impl Into<String>,
|
||||
source: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
CandidateFailureDiagnosticKind::UrlBuild,
|
||||
"$.endpoint",
|
||||
"无法构建上游请求地址;请检查 base_url、custom_path、API 格式和模型映射",
|
||||
)
|
||||
.formats(client_api_format, provider_api_format)
|
||||
.source(source)
|
||||
}
|
||||
|
||||
pub fn header_rules_apply_failed(
|
||||
client_api_format: impl Into<String>,
|
||||
provider_api_format: impl Into<String>,
|
||||
source: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
CandidateFailureDiagnosticKind::HeaderRules,
|
||||
"$.endpoint.header_rules",
|
||||
"Header 规则应用失败;请检查规则格式、条件配置,或是否试图覆盖受保护认证头",
|
||||
)
|
||||
.formats(client_api_format, provider_api_format)
|
||||
.source(source)
|
||||
}
|
||||
|
||||
pub fn body_rules_apply_failed(
|
||||
client_api_format: impl Into<String>,
|
||||
provider_api_format: impl Into<String>,
|
||||
source: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
CandidateFailureDiagnosticKind::BodyRules,
|
||||
"$.endpoint.body_rules",
|
||||
"Body 规则应用失败;请检查规则格式、条件配置,或规则输出是否仍是当前上游支持的请求体",
|
||||
)
|
||||
.formats(client_api_format, provider_api_format)
|
||||
.source(source)
|
||||
}
|
||||
|
||||
pub fn body_rules_unsupported_for_binary_upload(
|
||||
client_api_format: impl Into<String>,
|
||||
provider_api_format: impl Into<String>,
|
||||
source: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
CandidateFailureDiagnosticKind::BodyRules,
|
||||
"$.endpoint.body_rules",
|
||||
"二进制上传暂不支持本地应用 Body 规则;请移除该 Endpoint 的 Body 规则或改用 JSON 请求体",
|
||||
)
|
||||
.formats(client_api_format, provider_api_format)
|
||||
.source(source)
|
||||
}
|
||||
|
||||
pub fn provider_request_body_missing(
|
||||
client_api_format: impl Into<String>,
|
||||
provider_api_format: impl Into<String>,
|
||||
source: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
CandidateFailureDiagnosticKind::RequestBodyBuild,
|
||||
"$",
|
||||
"无法构建上游请求体;请检查请求体是否为支持的 JSON object,以及该任务类型必需字段是否存在且取值受支持",
|
||||
)
|
||||
.formats(client_api_format, provider_api_format)
|
||||
.source(source)
|
||||
}
|
||||
|
||||
pub fn envelope_build_failed(
|
||||
client_api_format: impl Into<String>,
|
||||
provider_api_format: impl Into<String>,
|
||||
source: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
CandidateFailureDiagnosticKind::EnvelopeBuild,
|
||||
"$",
|
||||
"无法构建上游请求封装;请检查该 Provider 的认证配置、模型映射、Endpoint Body 规则和当前请求体是否兼容",
|
||||
)
|
||||
.formats(client_api_format, provider_api_format)
|
||||
.source(source)
|
||||
}
|
||||
|
||||
fn to_value(&self) -> Value {
|
||||
json!({
|
||||
"kind": self.kind.as_str(),
|
||||
"path": self.path,
|
||||
"message": self.message,
|
||||
"source": self.source,
|
||||
"client_api_format": self.client_api_format,
|
||||
"provider_api_format": self.provider_api_format,
|
||||
"safe_to_show": self.safe_to_show,
|
||||
})
|
||||
}
|
||||
}
|
||||
131
crates/aether-ai-serving/src/lib.rs
Normal file
131
crates/aether-ai-serving/src/lib.rs
Normal file
@@ -0,0 +1,131 @@
|
||||
pub mod attempt_loop;
|
||||
pub mod attempt_plan;
|
||||
pub mod candidate_materialization;
|
||||
pub mod candidate_metadata;
|
||||
pub mod candidate_persistence;
|
||||
pub mod candidate_persistence_policy;
|
||||
pub mod candidate_preparation;
|
||||
pub mod candidate_preselection;
|
||||
pub mod candidate_ranking;
|
||||
pub mod candidate_resolution;
|
||||
pub mod decision_input;
|
||||
pub mod decision_path;
|
||||
pub mod decision_payload;
|
||||
pub mod dto;
|
||||
pub mod execution_path;
|
||||
pub mod failure_diagnostic;
|
||||
pub mod plan_payload;
|
||||
pub mod pool_scheduler;
|
||||
pub mod ports;
|
||||
pub mod ranking_metadata;
|
||||
pub mod report_context;
|
||||
pub mod request_body_diagnostics;
|
||||
pub mod runtime_miss;
|
||||
pub mod surface_spec;
|
||||
|
||||
pub use attempt_loop::{
|
||||
run_ai_attempt_loop, AiAttemptLoopOutcome, AiAttemptLoopPort, AiExecutionAttempt,
|
||||
};
|
||||
pub use attempt_plan::{
|
||||
build_ai_execution_decision_from_plan, build_ai_execution_plan_from_decision,
|
||||
extract_ai_auth_header_pair, infer_ai_upstream_base_url,
|
||||
resolve_ai_passthrough_sync_request_body, take_ai_decision_plan_core, take_ai_non_empty_string,
|
||||
take_ai_upstream_auth_pair, trim_ai_owned_non_empty_string, AiDecisionPlanCore,
|
||||
AiExecutionDecisionFromPlanParts, AiExecutionPlanFromDecisionParts, AiUpstreamAuthPair,
|
||||
};
|
||||
pub use candidate_materialization::{
|
||||
run_ai_candidate_materialization, AiCandidateMaterializationOutcome,
|
||||
AiCandidateMaterializationPort,
|
||||
};
|
||||
pub use candidate_metadata::{
|
||||
ai_local_execution_contract_for_formats, append_ai_execution_contract_fields_to_value,
|
||||
build_ai_candidate_metadata, build_ai_candidate_metadata_from_candidate,
|
||||
AiCandidateMetadataParts,
|
||||
};
|
||||
pub use candidate_persistence::{
|
||||
ai_candidate_extra_data_with_ranking, ai_should_persist_available_candidate_for_pool_key,
|
||||
ai_should_persist_skipped_candidate_for_pool_membership,
|
||||
run_ai_available_candidate_persistence, run_ai_skipped_candidate_persistence,
|
||||
AiAvailableCandidatePersistencePort, AiSkippedCandidatePersistencePort,
|
||||
};
|
||||
pub use candidate_persistence_policy::{
|
||||
ai_candidate_persistence_policy_spec, AiCandidatePersistencePolicyKind,
|
||||
AiCandidatePersistencePolicySpec,
|
||||
};
|
||||
pub use candidate_preparation::{
|
||||
prepare_ai_header_authenticated_candidate, resolve_ai_candidate_mapped_model,
|
||||
AiPreparedHeaderAuthenticatedCandidate,
|
||||
};
|
||||
pub use candidate_preselection::{
|
||||
run_ai_candidate_preselection, AiCandidatePreselectionOutcome, AiCandidatePreselectionPort,
|
||||
};
|
||||
pub use candidate_ranking::{
|
||||
ai_ranking_context, build_ai_rankable_candidate, run_ai_candidate_ranking,
|
||||
AiCandidateRankingPort, AiRankableCandidateParts, AiRankingContextConfig,
|
||||
AiRankingSchedulingMode,
|
||||
};
|
||||
pub use candidate_resolution::{
|
||||
extract_ai_pool_sticky_session_token, run_ai_candidate_resolution, AiCandidateResolutionMode,
|
||||
AiCandidateResolutionOutcome, AiCandidateResolutionPort, AiCandidateResolutionRequest,
|
||||
};
|
||||
pub use decision_input::{run_ai_authenticated_decision_input, AiAuthenticatedDecisionInputPort};
|
||||
pub use decision_path::{
|
||||
run_ai_stream_decision_path, run_ai_sync_decision_path, AiStreamDecisionPathPort,
|
||||
AiStreamDecisionStep, AiSyncDecisionPathPort, AiSyncDecisionStep,
|
||||
};
|
||||
pub use decision_payload::{
|
||||
ai_execution_decision_action, build_ai_execution_decision_response,
|
||||
AiExecutionDecisionResponseParts,
|
||||
};
|
||||
pub use dto::{
|
||||
augment_sync_report_context, generic_decision_missing_exact_provider_request,
|
||||
AiExecutionDecision, AiExecutionPlanPayload, AiStreamAttempt, AiSyncAttempt, ConversionMode,
|
||||
ExecutionStrategy,
|
||||
};
|
||||
pub use execution_path::{
|
||||
run_ai_stream_execution_path, run_ai_sync_execution_path, AiPlanFallbackReason,
|
||||
AiServingExecutionOutcome, AiStreamExecutionPathPort, AiStreamExecutionStep,
|
||||
AiSyncExecutionPathPort, AiSyncExecutionStep,
|
||||
};
|
||||
pub use failure_diagnostic::{CandidateFailureDiagnostic, CandidateFailureDiagnosticKind};
|
||||
pub use plan_payload::{
|
||||
build_ai_stream_execution_plan_payload, build_ai_sync_execution_plan_payload,
|
||||
};
|
||||
pub use pool_scheduler::{
|
||||
normalize_enabled_ai_pool_presets, run_ai_pool_scheduler, AiPoolCandidateFacts,
|
||||
AiPoolCandidateInput, AiPoolCandidateOrchestration, AiPoolCatalogKeyContext,
|
||||
AiPoolRuntimeState, AiPoolScheduledCandidate, AiPoolSchedulerOutcome, AiPoolSchedulingConfig,
|
||||
AiPoolSchedulingPreset, AiPoolSkippedCandidate, AI_POOL_ACCOUNT_BLOCKED_SKIP_REASON,
|
||||
AI_POOL_ACCOUNT_EXHAUSTED_SKIP_REASON, AI_POOL_COOLDOWN_SKIP_REASON,
|
||||
AI_POOL_COST_LIMIT_REACHED_SKIP_REASON,
|
||||
};
|
||||
pub use ranking_metadata::append_ai_ranking_metadata_to_object;
|
||||
pub use report_context::{
|
||||
build_ai_execution_report_context, build_ai_report_context_original_request_echo,
|
||||
insert_provider_stream_event_api_format, provider_stream_event_api_format_for_provider_type,
|
||||
AiExecutionReportContextParts, AiRequestOrigin,
|
||||
};
|
||||
pub use request_body_diagnostics::{
|
||||
request_body_build_failure_extra_data, same_format_provider_request_body_failure_extra_data,
|
||||
};
|
||||
pub use runtime_miss::{
|
||||
apply_ai_runtime_candidate_evaluation_progress,
|
||||
apply_ai_runtime_candidate_evaluation_progress_preserving_candidate_signal,
|
||||
apply_ai_runtime_candidate_evaluation_progress_to_diagnostic,
|
||||
apply_ai_runtime_candidate_terminal_plan_reason_to_diagnostic,
|
||||
apply_ai_runtime_candidate_terminal_reason, build_ai_runtime_candidate_evaluation_diagnostic,
|
||||
build_ai_runtime_execution_exhausted_diagnostic, record_ai_runtime_candidate_skip_reason,
|
||||
record_ai_runtime_candidate_skip_reason_on_diagnostic,
|
||||
set_ai_runtime_candidate_evaluation_diagnostic, set_ai_runtime_execution_exhausted_diagnostic,
|
||||
set_ai_runtime_miss_diagnostic_reason, AiRuntimeMissDiagnosticFields,
|
||||
AiRuntimeMissDiagnosticPort,
|
||||
};
|
||||
pub use surface_spec::{
|
||||
ai_gemini_files_spec_metadata, ai_openai_image_spec_metadata,
|
||||
ai_openai_responses_spec_metadata, ai_requested_model_family_for_same_format_provider,
|
||||
ai_requested_model_family_for_standard_source, ai_requested_model_family_for_video_create,
|
||||
ai_same_format_provider_spec_metadata, ai_standard_spec_metadata,
|
||||
ai_video_create_spec_metadata, extract_ai_gemini_model_from_path,
|
||||
extract_ai_requested_model_from_request_path, extract_ai_standard_requested_model,
|
||||
AiExecutionSurfaceSpecMetadata, AiRequestedModelFamily,
|
||||
};
|
||||
108
crates/aether-ai-serving/src/plan_payload.rs
Normal file
108
crates/aether-ai-serving/src/plan_payload.rs
Normal file
@@ -0,0 +1,108 @@
|
||||
use aether_ai_surfaces::api::{
|
||||
ExecutionRuntimeAuthContext, EXECUTION_RUNTIME_STREAM_ACTION, EXECUTION_RUNTIME_SYNC_ACTION,
|
||||
};
|
||||
|
||||
use crate::dto::{AiExecutionPlanPayload, AiStreamAttempt, AiSyncAttempt};
|
||||
|
||||
pub fn build_ai_sync_execution_plan_payload(
|
||||
plan_kind: &str,
|
||||
attempt: AiSyncAttempt,
|
||||
auth_context: Option<ExecutionRuntimeAuthContext>,
|
||||
) -> AiExecutionPlanPayload {
|
||||
AiExecutionPlanPayload {
|
||||
action: EXECUTION_RUNTIME_SYNC_ACTION.to_string(),
|
||||
plan_kind: Some(plan_kind.to_string()),
|
||||
plan: Some(attempt.plan),
|
||||
report_kind: attempt.report_kind,
|
||||
report_context: attempt.report_context,
|
||||
auth_context,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_ai_stream_execution_plan_payload(
|
||||
plan_kind: &str,
|
||||
attempt: AiStreamAttempt,
|
||||
auth_context: Option<ExecutionRuntimeAuthContext>,
|
||||
) -> AiExecutionPlanPayload {
|
||||
AiExecutionPlanPayload {
|
||||
action: EXECUTION_RUNTIME_STREAM_ACTION.to_string(),
|
||||
plan_kind: Some(plan_kind.to_string()),
|
||||
plan: Some(attempt.plan),
|
||||
report_kind: attempt.report_kind,
|
||||
report_context: attempt.report_context,
|
||||
auth_context,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use aether_contracts::{ExecutionPlan, RequestBody};
|
||||
|
||||
#[test]
|
||||
fn sync_plan_payload_uses_sync_action_and_attempt_report_fields() {
|
||||
let payload = build_ai_sync_execution_plan_payload(
|
||||
"openai_chat_sync",
|
||||
AiSyncAttempt {
|
||||
plan: test_plan(),
|
||||
report_kind: Some("sync_success".to_string()),
|
||||
report_context: Some(serde_json::json!({"candidate_index": 0})),
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
assert_eq!(payload.action, EXECUTION_RUNTIME_SYNC_ACTION);
|
||||
assert_eq!(payload.plan_kind.as_deref(), Some("openai_chat_sync"));
|
||||
assert_eq!(payload.report_kind.as_deref(), Some("sync_success"));
|
||||
assert_eq!(
|
||||
payload
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("candidate_index"))
|
||||
.and_then(serde_json::Value::as_u64),
|
||||
Some(0)
|
||||
);
|
||||
assert!(payload.plan.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stream_plan_payload_uses_stream_action() {
|
||||
let payload = build_ai_stream_execution_plan_payload(
|
||||
"openai_chat_stream",
|
||||
AiStreamAttempt {
|
||||
plan: test_plan(),
|
||||
report_kind: None,
|
||||
report_context: None,
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
assert_eq!(payload.action, EXECUTION_RUNTIME_STREAM_ACTION);
|
||||
assert_eq!(payload.plan_kind.as_deref(), Some("openai_chat_stream"));
|
||||
assert!(payload.plan.is_some());
|
||||
}
|
||||
|
||||
fn test_plan() -> ExecutionPlan {
|
||||
ExecutionPlan {
|
||||
request_id: "req_1".to_string(),
|
||||
candidate_id: None,
|
||||
provider_name: Some("provider".to_string()),
|
||||
provider_id: "provider_id".to_string(),
|
||||
endpoint_id: "endpoint_id".to_string(),
|
||||
key_id: "key_id".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: "https://example.com/v1/chat/completions".to_string(),
|
||||
headers: Default::default(),
|
||||
content_type: Some("application/json".to_string()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(serde_json::json!({"model": "model"})),
|
||||
stream: false,
|
||||
client_api_format: "openai:chat".to_string(),
|
||||
provider_api_format: "openai:chat".to_string(),
|
||||
model_name: Some("model".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
1048
crates/aether-ai-serving/src/pool_scheduler.rs
Normal file
1048
crates/aether-ai-serving/src/pool_scheduler.rs
Normal file
File diff suppressed because it is too large
Load Diff
12
crates/aether-ai-serving/src/ports.rs
Normal file
12
crates/aether-ai-serving/src/ports.rs
Normal file
@@ -0,0 +1,12 @@
|
||||
pub use crate::attempt_loop::AiAttemptLoopPort;
|
||||
pub use crate::candidate_materialization::AiCandidateMaterializationPort;
|
||||
pub use crate::candidate_persistence::{
|
||||
AiAvailableCandidatePersistencePort, AiSkippedCandidatePersistencePort,
|
||||
};
|
||||
pub use crate::candidate_preselection::AiCandidatePreselectionPort;
|
||||
pub use crate::candidate_ranking::AiCandidateRankingPort;
|
||||
pub use crate::candidate_resolution::AiCandidateResolutionPort;
|
||||
pub use crate::decision_input::AiAuthenticatedDecisionInputPort;
|
||||
pub use crate::decision_path::{AiStreamDecisionPathPort, AiSyncDecisionPathPort};
|
||||
pub use crate::execution_path::{AiStreamExecutionPathPort, AiSyncExecutionPathPort};
|
||||
pub use crate::runtime_miss::AiRuntimeMissDiagnosticPort;
|
||||
66
crates/aether-ai-serving/src/ranking_metadata.rs
Normal file
66
crates/aether-ai-serving/src/ranking_metadata.rs
Normal file
@@ -0,0 +1,66 @@
|
||||
use aether_scheduler_core::SchedulerRankingOutcome;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub fn append_ai_ranking_metadata_to_object(
|
||||
object: &mut Map<String, Value>,
|
||||
ranking: &SchedulerRankingOutcome,
|
||||
) {
|
||||
object.insert(
|
||||
"ranking_mode".to_string(),
|
||||
Value::String(format!("{:?}", ranking.ranking_mode)),
|
||||
);
|
||||
object.insert(
|
||||
"priority_mode".to_string(),
|
||||
Value::String(format!("{:?}", ranking.priority_mode)),
|
||||
);
|
||||
object.insert(
|
||||
"ranking_index".to_string(),
|
||||
Value::Number(serde_json::Number::from(ranking.ranking_index as u64)),
|
||||
);
|
||||
object.insert(
|
||||
"priority_slot".to_string(),
|
||||
Value::Number(serde_json::Number::from(i64::from(ranking.priority_slot))),
|
||||
);
|
||||
if let Some(promoted_by) = ranking.promoted_by {
|
||||
object.insert(
|
||||
"promoted_by".to_string(),
|
||||
Value::String(promoted_by.to_string()),
|
||||
);
|
||||
}
|
||||
if let Some(demoted_by) = ranking.demoted_by {
|
||||
object.insert(
|
||||
"demoted_by".to_string(),
|
||||
Value::String(demoted_by.to_string()),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use aether_scheduler_core::{SchedulerPriorityMode, SchedulerRankingMode};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn ranking_metadata_appends_scheduler_outcome_fields() {
|
||||
let ranking = SchedulerRankingOutcome {
|
||||
original_index: 2,
|
||||
ranking_index: 1,
|
||||
priority_mode: SchedulerPriorityMode::Provider,
|
||||
ranking_mode: SchedulerRankingMode::CacheAffinity,
|
||||
priority_slot: 7,
|
||||
promoted_by: Some("cached_affinity"),
|
||||
demoted_by: Some("cross_format"),
|
||||
};
|
||||
let mut object = Map::new();
|
||||
|
||||
append_ai_ranking_metadata_to_object(&mut object, &ranking);
|
||||
|
||||
assert_eq!(object.get("ranking_mode"), Some(&json!("CacheAffinity")));
|
||||
assert_eq!(object.get("priority_mode"), Some(&json!("Provider")));
|
||||
assert_eq!(object.get("ranking_index"), Some(&json!(1)));
|
||||
assert_eq!(object.get("priority_slot"), Some(&json!(7)));
|
||||
assert_eq!(object.get("promoted_by"), Some(&json!("cached_affinity")));
|
||||
assert_eq!(object.get("demoted_by"), Some(&json!("cross_format")));
|
||||
}
|
||||
}
|
||||
388
crates/aether-ai-serving/src/report_context.rs
Normal file
388
crates/aether-ai-serving/src/report_context.rs
Normal file
@@ -0,0 +1,388 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_ai_surfaces::api::ExecutionRuntimeAuthContext;
|
||||
use aether_scheduler_core::SchedulerRankingOutcome;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::append_ai_ranking_metadata_to_object;
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct AiRequestOrigin {
|
||||
pub client_ip: Option<String>,
|
||||
pub user_agent: Option<String>,
|
||||
}
|
||||
|
||||
pub struct AiExecutionReportContextParts<'a> {
|
||||
pub auth_context: &'a ExecutionRuntimeAuthContext,
|
||||
pub request_id: &'a str,
|
||||
pub candidate_id: &'a str,
|
||||
pub candidate_index: u32,
|
||||
pub retry_index: u32,
|
||||
pub pool_key_index: Option<u32>,
|
||||
pub model: &'a str,
|
||||
pub provider_name: &'a str,
|
||||
pub provider_id: &'a str,
|
||||
pub endpoint_id: &'a str,
|
||||
pub key_id: &'a str,
|
||||
pub key_name: Option<&'a str>,
|
||||
pub model_id: Option<&'a str>,
|
||||
pub global_model_id: Option<&'a str>,
|
||||
pub global_model_name: Option<&'a str>,
|
||||
pub provider_api_format: &'a str,
|
||||
pub client_api_format: &'a str,
|
||||
pub mapped_model: Option<&'a str>,
|
||||
pub candidate_group_id: Option<&'a str>,
|
||||
pub ranking: Option<&'a SchedulerRankingOutcome>,
|
||||
pub upstream_url: Option<&'a str>,
|
||||
pub header_rules: Option<&'a Value>,
|
||||
pub body_rules: Option<&'a Value>,
|
||||
pub provider_request_method: Option<Value>,
|
||||
pub provider_request_headers: Option<&'a BTreeMap<String, String>>,
|
||||
pub original_headers: &'a BTreeMap<String, String>,
|
||||
pub original_request_body: Option<Value>,
|
||||
pub request_origin: AiRequestOrigin,
|
||||
pub client_requested_stream: bool,
|
||||
pub upstream_is_stream: bool,
|
||||
pub has_envelope: bool,
|
||||
pub needs_conversion: bool,
|
||||
pub extra_fields: Map<String, Value>,
|
||||
}
|
||||
|
||||
pub fn build_ai_execution_report_context(parts: AiExecutionReportContextParts<'_>) -> Value {
|
||||
let mut object = Map::new();
|
||||
object.insert(
|
||||
"user_id".to_string(),
|
||||
Value::String(parts.auth_context.user_id.clone()),
|
||||
);
|
||||
object.insert(
|
||||
"api_key_id".to_string(),
|
||||
Value::String(parts.auth_context.api_key_id.clone()),
|
||||
);
|
||||
object.insert(
|
||||
"api_key_is_standalone".to_string(),
|
||||
Value::Bool(parts.auth_context.api_key_is_standalone),
|
||||
);
|
||||
object.insert(
|
||||
"username".to_string(),
|
||||
parts
|
||||
.auth_context
|
||||
.username
|
||||
.clone()
|
||||
.map(Value::String)
|
||||
.unwrap_or(Value::Null),
|
||||
);
|
||||
object.insert(
|
||||
"api_key_name".to_string(),
|
||||
parts
|
||||
.auth_context
|
||||
.api_key_name
|
||||
.clone()
|
||||
.map(Value::String)
|
||||
.unwrap_or(Value::Null),
|
||||
);
|
||||
object.insert(
|
||||
"request_id".to_string(),
|
||||
Value::String(parts.request_id.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"candidate_id".to_string(),
|
||||
Value::String(parts.candidate_id.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"candidate_index".to_string(),
|
||||
Value::Number(parts.candidate_index.into()),
|
||||
);
|
||||
object.insert(
|
||||
"retry_index".to_string(),
|
||||
Value::Number(parts.retry_index.into()),
|
||||
);
|
||||
object.insert("model".to_string(), Value::String(parts.model.to_string()));
|
||||
object.insert(
|
||||
"provider_name".to_string(),
|
||||
Value::String(parts.provider_name.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"provider_id".to_string(),
|
||||
Value::String(parts.provider_id.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"endpoint_id".to_string(),
|
||||
Value::String(parts.endpoint_id.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"key_id".to_string(),
|
||||
Value::String(parts.key_id.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"provider_api_format".to_string(),
|
||||
Value::String(parts.provider_api_format.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"client_api_format".to_string(),
|
||||
Value::String(parts.client_api_format.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"original_headers".to_string(),
|
||||
serde_json::to_value(parts.original_headers).expect("control headers should serialize"),
|
||||
);
|
||||
object.insert(
|
||||
"original_request_body".to_string(),
|
||||
parts.original_request_body.unwrap_or(Value::Null),
|
||||
);
|
||||
if let Some(client_ip) = parts.request_origin.client_ip {
|
||||
object.insert("client_ip".to_string(), Value::String(client_ip));
|
||||
}
|
||||
if let Some(user_agent) = parts.request_origin.user_agent {
|
||||
object.insert("user_agent".to_string(), Value::String(user_agent));
|
||||
}
|
||||
object.insert(
|
||||
"client_requested_stream".to_string(),
|
||||
Value::Bool(parts.client_requested_stream),
|
||||
);
|
||||
object.insert(
|
||||
"upstream_is_stream".to_string(),
|
||||
Value::Bool(parts.upstream_is_stream),
|
||||
);
|
||||
object.insert("has_envelope".to_string(), Value::Bool(parts.has_envelope));
|
||||
object.insert(
|
||||
"needs_conversion".to_string(),
|
||||
Value::Bool(parts.needs_conversion),
|
||||
);
|
||||
|
||||
if let Some(key_name) = parts.key_name {
|
||||
object.insert("key_name".to_string(), Value::String(key_name.to_string()));
|
||||
}
|
||||
if let Some(model_id) = parts.model_id {
|
||||
object.insert("model_id".to_string(), Value::String(model_id.to_string()));
|
||||
}
|
||||
if let Some(global_model_id) = parts.global_model_id {
|
||||
object.insert(
|
||||
"global_model_id".to_string(),
|
||||
Value::String(global_model_id.to_string()),
|
||||
);
|
||||
}
|
||||
if let Some(global_model_name) = parts.global_model_name {
|
||||
object.insert(
|
||||
"global_model_name".to_string(),
|
||||
Value::String(global_model_name.to_string()),
|
||||
);
|
||||
}
|
||||
if let Some(mapped_model) = parts.mapped_model {
|
||||
object.insert(
|
||||
"mapped_model".to_string(),
|
||||
Value::String(mapped_model.to_string()),
|
||||
);
|
||||
}
|
||||
if let Some(candidate_group_id) = parts.candidate_group_id {
|
||||
object.insert(
|
||||
"candidate_group_id".to_string(),
|
||||
Value::String(candidate_group_id.to_string()),
|
||||
);
|
||||
}
|
||||
if let Some(ranking) = parts.ranking {
|
||||
append_ai_ranking_metadata_to_object(&mut object, ranking);
|
||||
}
|
||||
if let Some(upstream_url) = parts.upstream_url {
|
||||
object.insert(
|
||||
"upstream_url".to_string(),
|
||||
Value::String(upstream_url.to_string()),
|
||||
);
|
||||
}
|
||||
if let Some(header_rules) = parts.header_rules {
|
||||
object.insert("header_rules".to_string(), header_rules.clone());
|
||||
}
|
||||
if let Some(body_rules) = parts.body_rules {
|
||||
object.insert("body_rules".to_string(), body_rules.clone());
|
||||
}
|
||||
if let Some(provider_request_method) = parts.provider_request_method {
|
||||
object.insert(
|
||||
"provider_request_method".to_string(),
|
||||
provider_request_method,
|
||||
);
|
||||
}
|
||||
if let Some(provider_request_headers) = parts.provider_request_headers {
|
||||
object.insert(
|
||||
"provider_request_headers".to_string(),
|
||||
serde_json::to_value(provider_request_headers)
|
||||
.expect("provider request headers should serialize"),
|
||||
);
|
||||
}
|
||||
if let Some(pool_key_index) = parts.pool_key_index {
|
||||
object.insert(
|
||||
"pool_key_index".to_string(),
|
||||
Value::Number(pool_key_index.into()),
|
||||
);
|
||||
}
|
||||
|
||||
object.extend(parts.extra_fields);
|
||||
Value::Object(object)
|
||||
}
|
||||
|
||||
pub fn provider_stream_event_api_format_for_provider_type(
|
||||
provider_type: &str,
|
||||
) -> Option<&'static str> {
|
||||
match provider_type.trim().to_ascii_lowercase().as_str() {
|
||||
"codex" => Some("openai:responses"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn insert_provider_stream_event_api_format(
|
||||
extra_fields: &mut Map<String, Value>,
|
||||
provider_type: &str,
|
||||
) {
|
||||
if let Some(api_format) = provider_stream_event_api_format_for_provider_type(provider_type) {
|
||||
extra_fields.insert(
|
||||
"provider_stream_event_api_format".to_string(),
|
||||
Value::String(api_format.to_string()),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_ai_report_context_original_request_echo(
|
||||
body_json: Option<&Value>,
|
||||
body_bytes_b64: Option<&str>,
|
||||
) -> Option<Value> {
|
||||
if let Some(body_bytes_b64) = body_bytes_b64
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
return Some(serde_json::json!({ "body_bytes_b64": body_bytes_b64 }));
|
||||
}
|
||||
|
||||
body_json.filter(|body| !body.is_null()).cloned()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use aether_scheduler_core::{SchedulerPriorityMode, SchedulerRankingMode};
|
||||
use serde_json::json;
|
||||
|
||||
fn sample_auth_context() -> ExecutionRuntimeAuthContext {
|
||||
ExecutionRuntimeAuthContext {
|
||||
user_id: "user-1".to_string(),
|
||||
api_key_id: "key-1".to_string(),
|
||||
username: Some("alice".to_string()),
|
||||
api_key_name: Some("primary".to_string()),
|
||||
balance_remaining: Some(42.0),
|
||||
access_allowed: true,
|
||||
api_key_is_standalone: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn report_context_builds_core_execution_fields() {
|
||||
let auth_context = sample_auth_context();
|
||||
let original_headers = BTreeMap::from([("x-trace-id".to_string(), "trace-a".to_string())]);
|
||||
let provider_headers =
|
||||
BTreeMap::from([("authorization".to_string(), "Bearer token".to_string())]);
|
||||
let mut extra_fields = Map::new();
|
||||
extra_fields.insert("extra".to_string(), json!("value"));
|
||||
let ranking = SchedulerRankingOutcome {
|
||||
original_index: 2,
|
||||
ranking_index: 1,
|
||||
priority_mode: SchedulerPriorityMode::Provider,
|
||||
ranking_mode: SchedulerRankingMode::CacheAffinity,
|
||||
priority_slot: 7,
|
||||
promoted_by: Some("cached_affinity"),
|
||||
demoted_by: None,
|
||||
};
|
||||
|
||||
let report = build_ai_execution_report_context(AiExecutionReportContextParts {
|
||||
auth_context: &auth_context,
|
||||
request_id: "trace-a",
|
||||
candidate_id: "candidate-a",
|
||||
candidate_index: 3,
|
||||
retry_index: 1,
|
||||
pool_key_index: Some(0),
|
||||
model: "gpt-5",
|
||||
provider_name: "RightCode",
|
||||
provider_id: "provider-1",
|
||||
endpoint_id: "endpoint-1",
|
||||
key_id: "key-1",
|
||||
key_name: Some("primary"),
|
||||
model_id: Some("model-1"),
|
||||
global_model_id: Some("global-1"),
|
||||
global_model_name: Some("GPT-5"),
|
||||
provider_api_format: "openai:responses",
|
||||
client_api_format: "openai:chat",
|
||||
mapped_model: Some("gpt-5"),
|
||||
candidate_group_id: Some("group-1"),
|
||||
ranking: Some(&ranking),
|
||||
upstream_url: Some("https://example.com/v1/responses"),
|
||||
header_rules: Some(&json!({"set": []})),
|
||||
body_rules: None,
|
||||
provider_request_method: Some(json!("POST")),
|
||||
provider_request_headers: Some(&provider_headers),
|
||||
original_headers: &original_headers,
|
||||
original_request_body: Some(json!({"model": "gpt-5"})),
|
||||
request_origin: AiRequestOrigin {
|
||||
client_ip: Some("127.0.0.1".to_string()),
|
||||
user_agent: Some("test-agent".to_string()),
|
||||
},
|
||||
client_requested_stream: false,
|
||||
upstream_is_stream: true,
|
||||
has_envelope: false,
|
||||
needs_conversion: true,
|
||||
extra_fields,
|
||||
});
|
||||
|
||||
assert_eq!(report["user_id"], "user-1");
|
||||
assert_eq!(report["candidate_index"], 3);
|
||||
assert_eq!(report["retry_index"], 1);
|
||||
assert_eq!(report["pool_key_index"], 0);
|
||||
assert_eq!(report["original_headers"]["x-trace-id"], "trace-a");
|
||||
assert_eq!(report["original_request_body"]["model"], "gpt-5");
|
||||
assert_eq!(report["ranking_index"], 1);
|
||||
assert_eq!(
|
||||
report["provider_request_headers"]["authorization"],
|
||||
"Bearer token"
|
||||
);
|
||||
assert_eq!(report["extra"], "value");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_stream_event_api_format_is_codex_only() {
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type("codex"),
|
||||
Some("openai:responses")
|
||||
);
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type(" CODEX "),
|
||||
Some("openai:responses")
|
||||
);
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type("openai"),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn original_request_echo_preserves_full_request_body() {
|
||||
let body = json!({
|
||||
"messages": [{"role": "user", "content": "large payload should be omitted"}],
|
||||
"service_tier": "default",
|
||||
"instructions": "Be concise.",
|
||||
"thinking": {"type": "enabled", "budget_tokens": 512},
|
||||
"metadata": {"trace": "keep"},
|
||||
"body_bytes_b64": "aGVsbG8=",
|
||||
});
|
||||
|
||||
let echo = build_ai_report_context_original_request_echo(Some(&body), None)
|
||||
.expect("echo should be produced");
|
||||
|
||||
assert_eq!(echo, body);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn original_request_echo_prefers_binary_body_bytes() {
|
||||
let echo = build_ai_report_context_original_request_echo(
|
||||
Some(&json!({"ignored": true})),
|
||||
Some("aGVsbG8="),
|
||||
)
|
||||
.expect("echo should be produced");
|
||||
|
||||
assert_eq!(echo, json!({"body_bytes_b64": "aGVsbG8="}));
|
||||
}
|
||||
}
|
||||
741
crates/aether-ai-serving/src/request_body_diagnostics.rs
Normal file
741
crates/aether-ai-serving/src/request_body_diagnostics.rs
Normal file
@@ -0,0 +1,741 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use aether_ai_surfaces::api::is_openai_responses_family_format;
|
||||
|
||||
use crate::{CandidateFailureDiagnostic, CandidateFailureDiagnosticKind};
|
||||
|
||||
pub fn request_body_build_failure_extra_data(
|
||||
body_json: &Value,
|
||||
client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
) -> Option<Value> {
|
||||
let diagnostic =
|
||||
diagnose_request_body_build_failure(body_json, client_api_format, provider_api_format)?;
|
||||
Some(
|
||||
diagnostic
|
||||
.formats(client_api_format, provider_api_format)
|
||||
.source(request_body_build_source(
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
))
|
||||
.to_extra_data(),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn same_format_provider_request_body_failure_extra_data(
|
||||
body_json: &Value,
|
||||
provider_api_format: &str,
|
||||
body_rules: Option<&Value>,
|
||||
context: &str,
|
||||
) -> Option<Value> {
|
||||
let diagnostic =
|
||||
diagnose_same_format_provider_request_body_failure(body_json, body_rules, context)?;
|
||||
Some(
|
||||
diagnostic
|
||||
.formats(provider_api_format, provider_api_format)
|
||||
.source(context)
|
||||
.to_extra_data(),
|
||||
)
|
||||
}
|
||||
|
||||
type RequestBodyBuildDiagnostic = CandidateFailureDiagnostic;
|
||||
|
||||
fn diagnose_request_body_build_failure(
|
||||
body_json: &Value,
|
||||
client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
if !body_json.is_object() {
|
||||
return Some(diagnostic("$", "请求体必须是 JSON object"));
|
||||
}
|
||||
|
||||
if is_openai_responses_client_format(client_api_format) {
|
||||
if let Some(diagnostic) = diagnose_openai_responses_request(body_json) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
return Some(diagnostic(
|
||||
"$",
|
||||
"OpenAI Responses 请求体初步结构检查通过;失败可能发生在后续跨格式转换或 Body 规则应用",
|
||||
));
|
||||
}
|
||||
|
||||
if client_api_format == "openai:chat"
|
||||
&& (provider_api_format.starts_with("claude:")
|
||||
|| provider_api_format.starts_with("gemini:"))
|
||||
{
|
||||
return diagnose_openai_chat_cross_format_request(body_json, provider_api_format);
|
||||
}
|
||||
|
||||
Some(diagnostic(
|
||||
"$",
|
||||
"请求体转换失败;当前转换器未返回更细的字段路径",
|
||||
))
|
||||
}
|
||||
|
||||
fn is_openai_responses_client_format(client_api_format: &str) -> bool {
|
||||
is_openai_responses_family_format(client_api_format)
|
||||
}
|
||||
|
||||
fn diagnose_same_format_provider_request_body_failure(
|
||||
body_json: &Value,
|
||||
body_rules: Option<&Value>,
|
||||
context: &str,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
if !body_json.is_object() {
|
||||
return Some(diagnostic("$", "反代请求体必须是 JSON object"));
|
||||
}
|
||||
if body_rules.is_some_and(|rules| !rules.is_array()) {
|
||||
return Some(diagnostic(
|
||||
"$.endpoint.body_rules",
|
||||
"Endpoint Body 规则必须是数组,本地反代无法应用该配置",
|
||||
));
|
||||
}
|
||||
match context {
|
||||
"kiro_envelope" => Some(diagnostic(
|
||||
"$",
|
||||
"Kiro 反代请求体包装失败;请检查 Kiro auth_config 与 Endpoint Body 规则",
|
||||
)),
|
||||
"antigravity_envelope" => Some(diagnostic(
|
||||
"$",
|
||||
"Antigravity 反代请求体包装失败;请检查请求体是否满足该传输封装要求",
|
||||
)),
|
||||
_ => Some(diagnostic(
|
||||
"$",
|
||||
"反代请求体构建失败;当前路径未返回更细的字段信息",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnose_openai_chat_cross_format_request(
|
||||
body_json: &Value,
|
||||
provider_api_format: &str,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let request = body_json.as_object()?;
|
||||
|
||||
if let Some(messages) = request.get("messages") {
|
||||
let Some(messages) = messages.as_array() else {
|
||||
return Some(diagnostic(
|
||||
"$.messages",
|
||||
"OpenAI Chat 的 messages 必须是数组",
|
||||
));
|
||||
};
|
||||
for (message_index, message) in messages.iter().enumerate() {
|
||||
let Some(message_object) = message.as_object() else {
|
||||
return Some(diagnostic(
|
||||
format!("$.messages[{message_index}]"),
|
||||
"message 必须是 object",
|
||||
));
|
||||
};
|
||||
let role = message_object
|
||||
.get("role")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
match role.as_str() {
|
||||
"system" | "developer" => {
|
||||
if let Some(diagnostic) = diagnose_openai_text_content(
|
||||
message_object.get("content"),
|
||||
format!("$.messages[{message_index}].content"),
|
||||
) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
}
|
||||
"user" | "assistant" => {
|
||||
if let Some(diagnostic) = diagnose_openai_content_blocks(
|
||||
message_object.get("content"),
|
||||
format!("$.messages[{message_index}].content"),
|
||||
role.as_str(),
|
||||
) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
if role == "assistant" {
|
||||
if let Some(diagnostic) = diagnose_openai_assistant_tool_calls(
|
||||
message_object.get("tool_calls"),
|
||||
format!("$.messages[{message_index}].tool_calls"),
|
||||
) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
}
|
||||
}
|
||||
"tool" => {
|
||||
let valid_tool_call_id = message_object
|
||||
.get("tool_call_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if !valid_tool_call_id {
|
||||
return Some(diagnostic(
|
||||
format!("$.messages[{message_index}].tool_call_id"),
|
||||
"tool 消息必须包含非空 tool_call_id",
|
||||
));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(diagnostic) = diagnose_openai_tools(request.get("tools"), provider_api_format) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
diagnose_openai_tool_choice(request.get("tool_choice"))
|
||||
}
|
||||
|
||||
fn diagnose_openai_responses_request(body_json: &Value) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let request = body_json.as_object()?;
|
||||
|
||||
if let Some(diagnostic) = diagnose_openai_responses_text_content(
|
||||
request.get("instructions"),
|
||||
"$.instructions".to_string(),
|
||||
) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
|
||||
if let Some(diagnostic) = diagnose_openai_responses_input(request.get("input")) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
if let Some(diagnostic) = diagnose_openai_responses_tools(request.get("tools")) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
diagnose_openai_responses_tool_choice(request.get("tool_choice"))
|
||||
}
|
||||
|
||||
fn diagnose_openai_responses_input(input: Option<&Value>) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let input = input?;
|
||||
match input {
|
||||
Value::Null | Value::String(_) => None,
|
||||
Value::Array(items) => {
|
||||
for (item_index, item) in items.iter().enumerate() {
|
||||
if item.is_string() {
|
||||
continue;
|
||||
}
|
||||
let item_path = format!("$.input[{item_index}]");
|
||||
let Some(item_object) = item.as_object() else {
|
||||
return Some(diagnostic(
|
||||
item_path,
|
||||
"OpenAI Responses input 数组项必须是 string 或 object",
|
||||
));
|
||||
};
|
||||
let item_type = item_object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("message")
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
match item_type.as_str() {
|
||||
"message" => {
|
||||
let role = item_object
|
||||
.get("role")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("user")
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
if role == "system" || role == "developer" {
|
||||
if let Some(diagnostic) = diagnose_openai_responses_text_content(
|
||||
item_object.get("content"),
|
||||
format!("{item_path}.content"),
|
||||
) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
} else if let Some(diagnostic) = diagnose_openai_responses_message_content(
|
||||
item_object.get("content"),
|
||||
format!("{item_path}.content"),
|
||||
) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
}
|
||||
"function_call" => {
|
||||
let valid_name = item_object
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if !valid_name {
|
||||
return Some(diagnostic(
|
||||
format!("{item_path}.name"),
|
||||
"function_call 必须包含非空 name",
|
||||
));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
_ => Some(diagnostic(
|
||||
"$.input",
|
||||
"OpenAI Responses input 必须是 string、array 或 null",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnose_openai_responses_text_content(
|
||||
content: Option<&Value>,
|
||||
path: String,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
match content {
|
||||
None | Some(Value::Null) | Some(Value::String(_)) => None,
|
||||
Some(Value::Array(parts)) => {
|
||||
for (part_index, part) in parts.iter().enumerate() {
|
||||
if !part.is_object() {
|
||||
return Some(diagnostic(
|
||||
format!("{path}[{part_index}]"),
|
||||
"文本 content 数组项必须是 object",
|
||||
));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
Some(_) => Some(diagnostic(
|
||||
path,
|
||||
"文本 content 必须是 string、array 或 null",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnose_openai_responses_message_content(
|
||||
content: Option<&Value>,
|
||||
path: String,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
match content {
|
||||
None | Some(Value::Null) | Some(Value::String(_)) => None,
|
||||
Some(Value::Array(parts)) => {
|
||||
for (part_index, part) in parts.iter().enumerate() {
|
||||
let part_path = format!("{path}[{part_index}]");
|
||||
let Some(part_object) = part.as_object() else {
|
||||
return Some(diagnostic(part_path, "message content 数组项必须是 object"));
|
||||
};
|
||||
let part_type = part_object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
if matches!(
|
||||
part_type.as_str(),
|
||||
"input_image" | "output_image" | "image_url"
|
||||
) && image_part_url(part_object).is_none()
|
||||
{
|
||||
return Some(diagnostic(
|
||||
part_path,
|
||||
"图片 content 缺少 image_url/url,无法规范化为 OpenAI Chat 图片内容",
|
||||
));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
Some(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnose_openai_text_content(
|
||||
content: Option<&Value>,
|
||||
path: String,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
match content {
|
||||
None | Some(Value::Null) | Some(Value::String(_)) => None,
|
||||
Some(Value::Array(parts)) => {
|
||||
for (part_index, part) in parts.iter().enumerate() {
|
||||
if !part.is_object() {
|
||||
return Some(diagnostic(
|
||||
format!("{path}[{part_index}]"),
|
||||
"content 数组项必须是 object",
|
||||
));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
Some(_) => Some(diagnostic(path, "content 必须是 string、array 或 null")),
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnose_openai_content_blocks(
|
||||
content: Option<&Value>,
|
||||
path: String,
|
||||
role: &str,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
match content {
|
||||
None | Some(Value::Null) | Some(Value::String(_)) => None,
|
||||
Some(Value::Array(parts)) => {
|
||||
for (part_index, part) in parts.iter().enumerate() {
|
||||
let part_path = format!("{path}[{part_index}]");
|
||||
let Some(part_object) = part.as_object() else {
|
||||
return Some(diagnostic(part_path, "content 数组项必须是 object"));
|
||||
};
|
||||
let part_type = part_object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
if matches!(part_type, "image_url" | "input_image" | "output_image")
|
||||
&& role == "user"
|
||||
&& image_part_url(part_object).is_none()
|
||||
{
|
||||
return Some(diagnostic(
|
||||
part_path,
|
||||
"图片 content 缺少 image_url/url,无法转换为 Claude image block",
|
||||
));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
Some(_) => Some(diagnostic(path, "content 必须是 string、array 或 null")),
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnose_openai_assistant_tool_calls(
|
||||
tool_calls: Option<&Value>,
|
||||
path: String,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let tool_calls = tool_calls?;
|
||||
let Some(tool_calls) = tool_calls.as_array() else {
|
||||
return Some(diagnostic(path, "assistant.tool_calls 必须是数组"));
|
||||
};
|
||||
for (tool_call_index, tool_call) in tool_calls.iter().enumerate() {
|
||||
let tool_call_path = format!("{path}[{tool_call_index}]");
|
||||
let Some(tool_call_object) = tool_call.as_object() else {
|
||||
return Some(diagnostic(tool_call_path, "tool_call 必须是 object"));
|
||||
};
|
||||
let Some(function) = tool_call_object.get("function").and_then(Value::as_object) else {
|
||||
return Some(diagnostic(
|
||||
format!("{tool_call_path}.function"),
|
||||
"tool_call 必须包含 function object",
|
||||
));
|
||||
};
|
||||
let valid_name = function
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if !valid_name {
|
||||
return Some(diagnostic(
|
||||
format!("{tool_call_path}.function.name"),
|
||||
"tool_call.function.name 必须是非空字符串",
|
||||
));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn diagnose_openai_tools(
|
||||
tools: Option<&Value>,
|
||||
provider_api_format: &str,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let tools = tools?;
|
||||
let Some(tools) = tools.as_array() else {
|
||||
return Some(diagnostic("$.tools", "OpenAI Chat 的 tools 必须是数组"));
|
||||
};
|
||||
for (tool_index, tool) in tools.iter().enumerate() {
|
||||
let tool_path = format!("$.tools[{tool_index}]");
|
||||
let Some(tool_object) = tool.as_object() else {
|
||||
return Some(diagnostic(tool_path, "tool 必须是 object"));
|
||||
};
|
||||
if tool_object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| value != "function")
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let Some(function) = tool_object.get("function").and_then(Value::as_object) else {
|
||||
let native_tool_hint = if provider_api_format.starts_with("claude:") {
|
||||
";如果这是 Claude 原生 tool,请改为 OpenAI function tool 格式"
|
||||
} else if provider_api_format.starts_with("gemini:") {
|
||||
";如果这是 Gemini 原生 tool,请改为 OpenAI function tool 格式"
|
||||
} else {
|
||||
""
|
||||
};
|
||||
return Some(diagnostic(
|
||||
format!("{tool_path}.function"),
|
||||
format!("OpenAI tool 必须包含 function object{native_tool_hint}"),
|
||||
));
|
||||
};
|
||||
let valid_name = function
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if !valid_name {
|
||||
return Some(diagnostic(
|
||||
format!("{tool_path}.function.name"),
|
||||
"OpenAI tool 的 function.name 必须是非空字符串",
|
||||
));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn diagnose_openai_responses_tools(tools: Option<&Value>) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let tools = tools?;
|
||||
let tool_values = tools.as_array()?;
|
||||
for (tool_index, tool) in tool_values.iter().enumerate() {
|
||||
let tool_path = format!("$.tools[{tool_index}]");
|
||||
let Some(tool_object) = tool.as_object() else {
|
||||
return Some(diagnostic(tool_path, "OpenAI Responses tool 必须是 object"));
|
||||
};
|
||||
let tool_type = tool_object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("function")
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
if tool_type.starts_with("web_search")
|
||||
|| tool_object.get("function").is_some()
|
||||
|| tool_type != "function"
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let valid_name = tool_object
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if !valid_name {
|
||||
return Some(diagnostic(
|
||||
format!("{tool_path}.name"),
|
||||
"OpenAI Responses function tool 必须包含非空 name",
|
||||
));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn diagnose_openai_responses_tool_choice(
|
||||
tool_choice: Option<&Value>,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let Some(Value::Object(object)) = tool_choice else {
|
||||
return None;
|
||||
};
|
||||
let is_cli_function_choice = object.get("function").is_none()
|
||||
&& object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| value.eq_ignore_ascii_case("function"));
|
||||
if !is_cli_function_choice {
|
||||
return None;
|
||||
}
|
||||
let valid_name = object
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if valid_name {
|
||||
None
|
||||
} else {
|
||||
Some(diagnostic(
|
||||
"$.tool_choice.name",
|
||||
"OpenAI Responses tool_choice 指定 function 时必须包含非空 name",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnose_openai_tool_choice(tool_choice: Option<&Value>) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let Some(Value::Object(object)) = tool_choice else {
|
||||
return None;
|
||||
};
|
||||
let valid_name = object
|
||||
.get("function")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|function| function.get("name"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if valid_name {
|
||||
None
|
||||
} else {
|
||||
Some(diagnostic(
|
||||
"$.tool_choice.function.name",
|
||||
"tool_choice 指定具体工具时必须包含非空 function.name",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn image_part_url(part_object: &serde_json::Map<String, Value>) -> Option<&str> {
|
||||
part_object
|
||||
.get("image_url")
|
||||
.and_then(|value| {
|
||||
value.as_str().or_else(|| {
|
||||
value
|
||||
.as_object()
|
||||
.and_then(|object| object.get("url"))
|
||||
.and_then(Value::as_str)
|
||||
})
|
||||
})
|
||||
.or_else(|| part_object.get("url").and_then(Value::as_str))
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn diagnostic(path: impl Into<String>, message: impl Into<String>) -> RequestBodyBuildDiagnostic {
|
||||
CandidateFailureDiagnostic::new(
|
||||
CandidateFailureDiagnosticKind::RequestBodyBuild,
|
||||
path,
|
||||
message,
|
||||
)
|
||||
}
|
||||
|
||||
fn request_body_build_source(client_api_format: &str, provider_api_format: &str) -> String {
|
||||
format!("{client_api_format}_to_{provider_api_format}")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::request_body_build_failure_extra_data;
|
||||
|
||||
#[test]
|
||||
fn openai_chat_to_claude_reports_claude_native_tool_shape() {
|
||||
let body = json!({
|
||||
"model": "gpt-5.4",
|
||||
"messages": [{ "role": "user", "content": "hello" }],
|
||||
"tools": [{
|
||||
"name": "read_file",
|
||||
"description": "Read a file",
|
||||
"input_schema": { "type": "object" }
|
||||
}]
|
||||
});
|
||||
|
||||
let diagnostic =
|
||||
request_body_build_failure_extra_data(&body, "openai:chat", "claude:messages")
|
||||
.expect("diagnostic");
|
||||
|
||||
assert_eq!(
|
||||
diagnostic["request_body_build_error"]["path"],
|
||||
"$.tools[0].function"
|
||||
);
|
||||
assert_eq!(
|
||||
diagnostic["failure_diagnostic"]["kind"],
|
||||
"request_body_build"
|
||||
);
|
||||
assert_eq!(
|
||||
diagnostic["failure_diagnostic"]["source"],
|
||||
"openai:chat_to_claude:messages"
|
||||
);
|
||||
assert!(diagnostic["request_body_build_error"]["message"]
|
||||
.as_str()
|
||||
.expect("message")
|
||||
.contains("Claude 原生 tool"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_chat_to_claude_reports_invalid_message_content_part() {
|
||||
let body = json!({
|
||||
"model": "gpt-5.4",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": ["not-an-object"]
|
||||
}]
|
||||
});
|
||||
|
||||
let diagnostic =
|
||||
request_body_build_failure_extra_data(&body, "openai:chat", "claude:messages")
|
||||
.expect("diagnostic");
|
||||
|
||||
assert_eq!(
|
||||
diagnostic["request_body_build_error"]["path"],
|
||||
"$.messages[0].content[0]"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_chat_to_gemini_reports_gemini_native_tool_shape() {
|
||||
let body = json!({
|
||||
"model": "gpt-5.4",
|
||||
"messages": [{ "role": "user", "content": "hello" }],
|
||||
"tools": [{
|
||||
"functionDeclarations": [{
|
||||
"name": "search",
|
||||
"parameters": { "type": "object" }
|
||||
}]
|
||||
}]
|
||||
});
|
||||
|
||||
let diagnostic =
|
||||
request_body_build_failure_extra_data(&body, "openai:chat", "gemini:generate_content")
|
||||
.expect("diagnostic");
|
||||
|
||||
assert_eq!(
|
||||
diagnostic["request_body_build_error"]["path"],
|
||||
"$.tools[0].function"
|
||||
);
|
||||
assert!(diagnostic["request_body_build_error"]["message"]
|
||||
.as_str()
|
||||
.expect("message")
|
||||
.contains("Gemini 原生 tool"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_responses_reports_invalid_function_call_name() {
|
||||
let body = json!({
|
||||
"model": "gpt-5.4",
|
||||
"input": [{
|
||||
"type": "function_call",
|
||||
"arguments": "{}"
|
||||
}]
|
||||
});
|
||||
|
||||
let diagnostic =
|
||||
request_body_build_failure_extra_data(&body, "openai:responses", "claude:messages")
|
||||
.expect("diagnostic");
|
||||
|
||||
assert_eq!(
|
||||
diagnostic["request_body_build_error"]["path"],
|
||||
"$.input[0].name"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_responses_reports_invalid_tool_choice_name() {
|
||||
let body = json!({
|
||||
"model": "gpt-5.4",
|
||||
"input": "hello",
|
||||
"tool_choice": { "type": "function" }
|
||||
});
|
||||
|
||||
let diagnostic = request_body_build_failure_extra_data(
|
||||
&body,
|
||||
"openai:responses",
|
||||
"gemini:generate_content",
|
||||
)
|
||||
.expect("diagnostic");
|
||||
|
||||
assert_eq!(
|
||||
diagnostic["request_body_build_error"]["path"],
|
||||
"$.tool_choice.name"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_format_provider_reports_non_object_body() {
|
||||
let diagnostic = super::same_format_provider_request_body_failure_extra_data(
|
||||
&json!("raw"),
|
||||
"openai:chat",
|
||||
None,
|
||||
"same_format",
|
||||
)
|
||||
.expect("diagnostic");
|
||||
|
||||
assert_eq!(diagnostic["request_body_build_error"]["path"], "$");
|
||||
assert!(diagnostic["request_body_build_error"]["message"]
|
||||
.as_str()
|
||||
.expect("message")
|
||||
.contains("反代请求体必须是 JSON object"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_format_provider_reports_invalid_body_rules_shape() {
|
||||
let diagnostic = super::same_format_provider_request_body_failure_extra_data(
|
||||
&json!({ "model": "gpt-5.4" }),
|
||||
"openai:chat",
|
||||
Some(&json!({ "action": "set" })),
|
||||
"same_format",
|
||||
)
|
||||
.expect("diagnostic");
|
||||
|
||||
assert_eq!(
|
||||
diagnostic["request_body_build_error"]["path"],
|
||||
"$.endpoint.body_rules"
|
||||
);
|
||||
}
|
||||
}
|
||||
468
crates/aether-ai-serving/src/runtime_miss.rs
Normal file
468
crates/aether-ai-serving/src/runtime_miss.rs
Normal file
@@ -0,0 +1,468 @@
|
||||
pub trait AiRuntimeMissDiagnosticPort: Send + Sync {
|
||||
type Decision: Send + Sync;
|
||||
type Diagnostic: Send;
|
||||
|
||||
fn build_runtime_miss_diagnostic(
|
||||
&self,
|
||||
decision: &Self::Decision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
reason: &str,
|
||||
) -> Self::Diagnostic;
|
||||
|
||||
fn set_candidate_count(&self, diagnostic: &mut Self::Diagnostic, candidate_count: usize);
|
||||
|
||||
fn apply_candidate_evaluation_progress(
|
||||
&self,
|
||||
diagnostic: &mut Self::Diagnostic,
|
||||
candidate_count: usize,
|
||||
);
|
||||
|
||||
fn apply_candidate_terminal_plan_reason(
|
||||
&self,
|
||||
diagnostic: &mut Self::Diagnostic,
|
||||
no_plan_reason: &'static str,
|
||||
);
|
||||
|
||||
fn record_candidate_skip_reason(
|
||||
&self,
|
||||
diagnostic: &mut Self::Diagnostic,
|
||||
skip_reason: &'static str,
|
||||
);
|
||||
|
||||
fn set_runtime_miss_diagnostic(&self, trace_id: &str, diagnostic: Self::Diagnostic);
|
||||
|
||||
fn mutate_runtime_miss_diagnostic<F>(&self, trace_id: &str, apply: F)
|
||||
where
|
||||
F: FnOnce(&mut Self::Diagnostic) + Send;
|
||||
|
||||
fn runtime_miss_diagnostic_has_candidate_signal(&self, trace_id: &str) -> bool;
|
||||
}
|
||||
|
||||
pub trait AiRuntimeMissDiagnosticFields {
|
||||
fn set_reason(&mut self, reason: String);
|
||||
fn set_candidate_count(&mut self, candidate_count: usize);
|
||||
fn candidate_count(&self) -> Option<usize>;
|
||||
fn skipped_candidate_count(&self) -> Option<usize>;
|
||||
fn skip_reason_count(&self, skip_reason: &str) -> usize;
|
||||
fn skip_reason_len(&self) -> usize;
|
||||
fn record_skip_reason(&mut self, skip_reason: &'static str);
|
||||
}
|
||||
|
||||
pub fn apply_ai_runtime_candidate_evaluation_progress_to_diagnostic<Diagnostic>(
|
||||
diagnostic: &mut Diagnostic,
|
||||
candidate_count: usize,
|
||||
) where
|
||||
Diagnostic: AiRuntimeMissDiagnosticFields,
|
||||
{
|
||||
diagnostic.set_candidate_count(candidate_count);
|
||||
diagnostic.set_reason(if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else {
|
||||
"candidate_evaluation_incomplete".to_string()
|
||||
});
|
||||
}
|
||||
|
||||
pub fn apply_ai_runtime_candidate_terminal_plan_reason_to_diagnostic<Diagnostic>(
|
||||
diagnostic: &mut Diagnostic,
|
||||
no_plan_reason: &'static str,
|
||||
) where
|
||||
Diagnostic: AiRuntimeMissDiagnosticFields,
|
||||
{
|
||||
let candidate_count = diagnostic.candidate_count().unwrap_or(0);
|
||||
let skipped_candidate_count = diagnostic.skipped_candidate_count().unwrap_or(0);
|
||||
diagnostic.set_reason(if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else if skipped_candidate_count >= candidate_count
|
||||
&& diagnostic.skip_reason_len() == 1
|
||||
&& diagnostic.skip_reason_count("api_key_concurrency_limit_reached") > 0
|
||||
{
|
||||
"api_key_concurrency_limit_reached".to_string()
|
||||
} else if skipped_candidate_count >= candidate_count {
|
||||
"all_candidates_skipped".to_string()
|
||||
} else {
|
||||
no_plan_reason.to_string()
|
||||
});
|
||||
}
|
||||
|
||||
pub fn record_ai_runtime_candidate_skip_reason_on_diagnostic<Diagnostic>(
|
||||
diagnostic: &mut Diagnostic,
|
||||
skip_reason: &'static str,
|
||||
) where
|
||||
Diagnostic: AiRuntimeMissDiagnosticFields,
|
||||
{
|
||||
diagnostic.record_skip_reason(skip_reason);
|
||||
}
|
||||
|
||||
pub fn set_ai_runtime_miss_diagnostic_reason<Port>(
|
||||
port: &Port,
|
||||
trace_id: &str,
|
||||
decision: &Port::Decision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
reason: &str,
|
||||
) where
|
||||
Port: AiRuntimeMissDiagnosticPort,
|
||||
{
|
||||
port.set_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
port.build_runtime_miss_diagnostic(decision, plan_kind, requested_model, reason),
|
||||
);
|
||||
}
|
||||
|
||||
pub fn build_ai_runtime_execution_exhausted_diagnostic<Port>(
|
||||
port: &Port,
|
||||
decision: &Port::Decision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) -> Port::Diagnostic
|
||||
where
|
||||
Port: AiRuntimeMissDiagnosticPort,
|
||||
{
|
||||
let mut diagnostic = port.build_runtime_miss_diagnostic(
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
"execution_runtime_candidates_exhausted",
|
||||
);
|
||||
port.set_candidate_count(&mut diagnostic, candidate_count);
|
||||
diagnostic
|
||||
}
|
||||
|
||||
pub fn set_ai_runtime_execution_exhausted_diagnostic<Port>(
|
||||
port: &Port,
|
||||
trace_id: &str,
|
||||
decision: &Port::Decision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) where
|
||||
Port: AiRuntimeMissDiagnosticPort,
|
||||
{
|
||||
port.set_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_ai_runtime_execution_exhausted_diagnostic(
|
||||
port,
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
candidate_count,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
pub fn build_ai_runtime_candidate_evaluation_diagnostic<Port>(
|
||||
port: &Port,
|
||||
decision: &Port::Decision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) -> Port::Diagnostic
|
||||
where
|
||||
Port: AiRuntimeMissDiagnosticPort,
|
||||
{
|
||||
let mut diagnostic = port.build_runtime_miss_diagnostic(
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
port.apply_candidate_evaluation_progress(&mut diagnostic, candidate_count);
|
||||
diagnostic
|
||||
}
|
||||
|
||||
pub fn set_ai_runtime_candidate_evaluation_diagnostic<Port>(
|
||||
port: &Port,
|
||||
trace_id: &str,
|
||||
decision: &Port::Decision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) where
|
||||
Port: AiRuntimeMissDiagnosticPort,
|
||||
{
|
||||
port.set_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_ai_runtime_candidate_evaluation_diagnostic(
|
||||
port,
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
candidate_count,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
pub fn apply_ai_runtime_candidate_evaluation_progress<Port>(
|
||||
port: &Port,
|
||||
trace_id: &str,
|
||||
candidate_count: usize,
|
||||
) where
|
||||
Port: AiRuntimeMissDiagnosticPort,
|
||||
{
|
||||
port.mutate_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
port.apply_candidate_evaluation_progress(diagnostic, candidate_count);
|
||||
});
|
||||
}
|
||||
|
||||
pub fn apply_ai_runtime_candidate_evaluation_progress_preserving_candidate_signal<Port>(
|
||||
port: &Port,
|
||||
trace_id: &str,
|
||||
candidate_count: usize,
|
||||
) where
|
||||
Port: AiRuntimeMissDiagnosticPort,
|
||||
{
|
||||
let preserve_existing_candidate_signal =
|
||||
candidate_count == 0 && port.runtime_miss_diagnostic_has_candidate_signal(trace_id);
|
||||
if preserve_existing_candidate_signal {
|
||||
return;
|
||||
}
|
||||
apply_ai_runtime_candidate_evaluation_progress(port, trace_id, candidate_count);
|
||||
}
|
||||
|
||||
pub fn apply_ai_runtime_candidate_terminal_reason<Port>(
|
||||
port: &Port,
|
||||
trace_id: &str,
|
||||
no_plan_reason: &'static str,
|
||||
) where
|
||||
Port: AiRuntimeMissDiagnosticPort,
|
||||
{
|
||||
port.mutate_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
port.apply_candidate_terminal_plan_reason(diagnostic, no_plan_reason);
|
||||
});
|
||||
}
|
||||
|
||||
pub fn record_ai_runtime_candidate_skip_reason<Port>(
|
||||
port: &Port,
|
||||
trace_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) where
|
||||
Port: AiRuntimeMissDiagnosticPort,
|
||||
{
|
||||
port.mutate_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
port.record_candidate_skip_reason(diagnostic, skip_reason);
|
||||
});
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Mutex;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct TestDecision {
|
||||
id: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
struct TestDiagnostic {
|
||||
decision_id: String,
|
||||
plan_kind: String,
|
||||
requested_model: Option<String>,
|
||||
reason: String,
|
||||
candidate_count: Option<usize>,
|
||||
terminal_reason: Option<&'static str>,
|
||||
skip_reasons: BTreeMap<&'static str, usize>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestPort {
|
||||
diagnostics: Mutex<BTreeMap<String, TestDiagnostic>>,
|
||||
}
|
||||
|
||||
impl AiRuntimeMissDiagnosticPort for TestPort {
|
||||
type Decision = TestDecision;
|
||||
type Diagnostic = TestDiagnostic;
|
||||
|
||||
fn build_runtime_miss_diagnostic(
|
||||
&self,
|
||||
decision: &Self::Decision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
reason: &str,
|
||||
) -> Self::Diagnostic {
|
||||
TestDiagnostic {
|
||||
decision_id: decision.id.to_string(),
|
||||
plan_kind: plan_kind.to_string(),
|
||||
requested_model: requested_model.map(str::to_string),
|
||||
reason: reason.to_string(),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn set_candidate_count(&self, diagnostic: &mut Self::Diagnostic, candidate_count: usize) {
|
||||
diagnostic.candidate_count = Some(candidate_count);
|
||||
}
|
||||
|
||||
fn apply_candidate_evaluation_progress(
|
||||
&self,
|
||||
diagnostic: &mut Self::Diagnostic,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
diagnostic.candidate_count = Some(candidate_count);
|
||||
}
|
||||
|
||||
fn apply_candidate_terminal_plan_reason(
|
||||
&self,
|
||||
diagnostic: &mut Self::Diagnostic,
|
||||
no_plan_reason: &'static str,
|
||||
) {
|
||||
diagnostic.terminal_reason = Some(no_plan_reason);
|
||||
}
|
||||
|
||||
fn record_candidate_skip_reason(
|
||||
&self,
|
||||
diagnostic: &mut Self::Diagnostic,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
*diagnostic.skip_reasons.entry(skip_reason).or_insert(0) += 1;
|
||||
}
|
||||
|
||||
fn set_runtime_miss_diagnostic(&self, trace_id: &str, diagnostic: Self::Diagnostic) {
|
||||
self.diagnostics
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(trace_id.to_string(), diagnostic);
|
||||
}
|
||||
|
||||
fn mutate_runtime_miss_diagnostic<F>(&self, trace_id: &str, apply: F)
|
||||
where
|
||||
F: FnOnce(&mut Self::Diagnostic) + Send,
|
||||
{
|
||||
let mut diagnostics = self.diagnostics.lock().unwrap();
|
||||
let diagnostic = diagnostics.entry(trace_id.to_string()).or_default();
|
||||
apply(diagnostic);
|
||||
}
|
||||
|
||||
fn runtime_miss_diagnostic_has_candidate_signal(&self, trace_id: &str) -> bool {
|
||||
self.diagnostics
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get(trace_id)
|
||||
.is_some_and(|diagnostic| diagnostic.candidate_count.unwrap_or_default() > 0)
|
||||
}
|
||||
}
|
||||
|
||||
impl AiRuntimeMissDiagnosticFields for TestDiagnostic {
|
||||
fn set_reason(&mut self, reason: String) {
|
||||
self.reason = reason;
|
||||
}
|
||||
|
||||
fn set_candidate_count(&mut self, candidate_count: usize) {
|
||||
self.candidate_count = Some(candidate_count);
|
||||
}
|
||||
|
||||
fn candidate_count(&self) -> Option<usize> {
|
||||
self.candidate_count
|
||||
}
|
||||
|
||||
fn skipped_candidate_count(&self) -> Option<usize> {
|
||||
self.skip_reasons.values().copied().sum::<usize>().into()
|
||||
}
|
||||
|
||||
fn skip_reason_count(&self, skip_reason: &str) -> usize {
|
||||
self.skip_reasons.get(skip_reason).copied().unwrap_or(0)
|
||||
}
|
||||
|
||||
fn skip_reason_len(&self) -> usize {
|
||||
self.skip_reasons.len()
|
||||
}
|
||||
|
||||
fn record_skip_reason(&mut self, skip_reason: &'static str) {
|
||||
*self.skip_reasons.entry(skip_reason).or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_miss_builds_and_sets_execution_exhausted_diagnostic() {
|
||||
let port = TestPort::default();
|
||||
|
||||
set_ai_runtime_execution_exhausted_diagnostic(
|
||||
&port,
|
||||
"trace-a",
|
||||
&TestDecision { id: "decision-a" },
|
||||
"openai_chat",
|
||||
Some("gpt-5"),
|
||||
3,
|
||||
);
|
||||
|
||||
let diagnostic = port.diagnostics.lock().unwrap().get("trace-a").cloned();
|
||||
assert_eq!(
|
||||
diagnostic,
|
||||
Some(TestDiagnostic {
|
||||
decision_id: "decision-a".to_string(),
|
||||
plan_kind: "openai_chat".to_string(),
|
||||
requested_model: Some("gpt-5".to_string()),
|
||||
reason: "execution_runtime_candidates_exhausted".to_string(),
|
||||
candidate_count: Some(3),
|
||||
..Default::default()
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_miss_preserves_candidate_signal_and_records_terminal_updates() {
|
||||
let port = TestPort::default();
|
||||
apply_ai_runtime_candidate_evaluation_progress(&port, "trace-a", 2);
|
||||
|
||||
apply_ai_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
&port, "trace-a", 0,
|
||||
);
|
||||
apply_ai_runtime_candidate_terminal_reason(&port, "trace-a", "no_local_sync_plans");
|
||||
record_ai_runtime_candidate_skip_reason(&port, "trace-a", "transport_missing");
|
||||
|
||||
let diagnostic = port.diagnostics.lock().unwrap().get("trace-a").cloned();
|
||||
assert_eq!(
|
||||
diagnostic,
|
||||
Some(TestDiagnostic {
|
||||
candidate_count: Some(2),
|
||||
terminal_reason: Some("no_local_sync_plans"),
|
||||
skip_reasons: BTreeMap::from([("transport_missing", 1)]),
|
||||
..Default::default()
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_miss_diagnostic_field_helpers_apply_candidate_reason_state_machine() {
|
||||
let mut diagnostic = TestDiagnostic::default();
|
||||
|
||||
apply_ai_runtime_candidate_evaluation_progress_to_diagnostic(&mut diagnostic, 0);
|
||||
assert_eq!(diagnostic.candidate_count, Some(0));
|
||||
assert_eq!(diagnostic.reason, "candidate_list_empty");
|
||||
|
||||
apply_ai_runtime_candidate_evaluation_progress_to_diagnostic(&mut diagnostic, 3);
|
||||
assert_eq!(diagnostic.candidate_count, Some(3));
|
||||
assert_eq!(diagnostic.reason, "candidate_evaluation_incomplete");
|
||||
|
||||
record_ai_runtime_candidate_skip_reason_on_diagnostic(
|
||||
&mut diagnostic,
|
||||
"api_key_concurrency_limit_reached",
|
||||
);
|
||||
record_ai_runtime_candidate_skip_reason_on_diagnostic(
|
||||
&mut diagnostic,
|
||||
"api_key_concurrency_limit_reached",
|
||||
);
|
||||
apply_ai_runtime_candidate_terminal_plan_reason_to_diagnostic(
|
||||
&mut diagnostic,
|
||||
"no_local_sync_plans",
|
||||
);
|
||||
assert_eq!(diagnostic.reason, "no_local_sync_plans");
|
||||
|
||||
diagnostic.candidate_count = Some(2);
|
||||
apply_ai_runtime_candidate_terminal_plan_reason_to_diagnostic(
|
||||
&mut diagnostic,
|
||||
"no_local_sync_plans",
|
||||
);
|
||||
assert_eq!(diagnostic.reason, "api_key_concurrency_limit_reached");
|
||||
|
||||
record_ai_runtime_candidate_skip_reason_on_diagnostic(&mut diagnostic, "transport_missing");
|
||||
apply_ai_runtime_candidate_terminal_plan_reason_to_diagnostic(
|
||||
&mut diagnostic,
|
||||
"no_local_sync_plans",
|
||||
);
|
||||
assert_eq!(diagnostic.reason, "all_candidates_skipped");
|
||||
}
|
||||
}
|
||||
236
crates/aether-ai-serving/src/surface_spec.rs
Normal file
236
crates/aether-ai-serving/src/surface_spec.rs
Normal file
@@ -0,0 +1,236 @@
|
||||
use aether_ai_surfaces::api::{
|
||||
LocalGeminiFilesSpec, LocalOpenAiImageSpec, LocalOpenAiResponsesSpec,
|
||||
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, LocalStandardSourceFamily,
|
||||
LocalStandardSpec, LocalVideoCreateFamily, LocalVideoCreateSpec,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AiRequestedModelFamily {
|
||||
Standard,
|
||||
Gemini,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct AiExecutionSurfaceSpecMetadata {
|
||||
pub api_format: &'static str,
|
||||
pub decision_kind: &'static str,
|
||||
pub report_kind: Option<&'static str>,
|
||||
pub require_streaming: bool,
|
||||
pub requested_model_family: Option<AiRequestedModelFamily>,
|
||||
}
|
||||
|
||||
pub const fn ai_requested_model_family_for_standard_source(
|
||||
family: LocalStandardSourceFamily,
|
||||
) -> AiRequestedModelFamily {
|
||||
match family {
|
||||
LocalStandardSourceFamily::Standard => AiRequestedModelFamily::Standard,
|
||||
LocalStandardSourceFamily::Gemini => AiRequestedModelFamily::Gemini,
|
||||
}
|
||||
}
|
||||
|
||||
pub const fn ai_standard_spec_metadata(spec: LocalStandardSpec) -> AiExecutionSurfaceSpecMetadata {
|
||||
AiExecutionSurfaceSpecMetadata {
|
||||
api_format: spec.api_format,
|
||||
decision_kind: spec.decision_kind,
|
||||
report_kind: Some(spec.report_kind),
|
||||
require_streaming: spec.require_streaming,
|
||||
requested_model_family: Some(ai_requested_model_family_for_standard_source(spec.family)),
|
||||
}
|
||||
}
|
||||
|
||||
pub const fn ai_same_format_provider_spec_metadata(
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> AiExecutionSurfaceSpecMetadata {
|
||||
AiExecutionSurfaceSpecMetadata {
|
||||
api_format: spec.api_format,
|
||||
decision_kind: spec.decision_kind,
|
||||
report_kind: Some(spec.report_kind),
|
||||
require_streaming: spec.require_streaming,
|
||||
requested_model_family: Some(ai_requested_model_family_for_same_format_provider(
|
||||
spec.family,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub const fn ai_openai_responses_spec_metadata(
|
||||
spec: LocalOpenAiResponsesSpec,
|
||||
) -> AiExecutionSurfaceSpecMetadata {
|
||||
AiExecutionSurfaceSpecMetadata {
|
||||
api_format: spec.api_format,
|
||||
decision_kind: spec.decision_kind,
|
||||
report_kind: Some(spec.report_kind),
|
||||
require_streaming: spec.require_streaming,
|
||||
requested_model_family: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub const fn ai_gemini_files_spec_metadata(
|
||||
spec: LocalGeminiFilesSpec,
|
||||
) -> AiExecutionSurfaceSpecMetadata {
|
||||
AiExecutionSurfaceSpecMetadata {
|
||||
api_format: "gemini:files",
|
||||
decision_kind: spec.decision_kind,
|
||||
report_kind: spec.report_kind,
|
||||
require_streaming: spec.require_streaming,
|
||||
requested_model_family: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub const fn ai_openai_image_spec_metadata(
|
||||
spec: LocalOpenAiImageSpec,
|
||||
) -> AiExecutionSurfaceSpecMetadata {
|
||||
AiExecutionSurfaceSpecMetadata {
|
||||
api_format: spec.api_format,
|
||||
decision_kind: spec.decision_kind,
|
||||
report_kind: Some(spec.report_kind),
|
||||
require_streaming: spec.require_streaming,
|
||||
requested_model_family: Some(AiRequestedModelFamily::Standard),
|
||||
}
|
||||
}
|
||||
|
||||
pub const fn ai_video_create_spec_metadata(
|
||||
spec: LocalVideoCreateSpec,
|
||||
) -> AiExecutionSurfaceSpecMetadata {
|
||||
AiExecutionSurfaceSpecMetadata {
|
||||
api_format: spec.api_format,
|
||||
decision_kind: spec.decision_kind,
|
||||
report_kind: Some(spec.report_kind),
|
||||
require_streaming: false,
|
||||
requested_model_family: Some(ai_requested_model_family_for_video_create(spec.family)),
|
||||
}
|
||||
}
|
||||
|
||||
pub const fn ai_requested_model_family_for_same_format_provider(
|
||||
family: LocalSameFormatProviderFamily,
|
||||
) -> AiRequestedModelFamily {
|
||||
match family {
|
||||
LocalSameFormatProviderFamily::Standard => AiRequestedModelFamily::Standard,
|
||||
LocalSameFormatProviderFamily::Gemini => AiRequestedModelFamily::Gemini,
|
||||
}
|
||||
}
|
||||
|
||||
pub const fn ai_requested_model_family_for_video_create(
|
||||
family: LocalVideoCreateFamily,
|
||||
) -> AiRequestedModelFamily {
|
||||
match family {
|
||||
LocalVideoCreateFamily::OpenAi => AiRequestedModelFamily::Standard,
|
||||
LocalVideoCreateFamily::Gemini => AiRequestedModelFamily::Gemini,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn extract_ai_gemini_model_from_path(path: &str) -> Option<String> {
|
||||
let (_, suffix) = path.split_once("/models/")?;
|
||||
let model = suffix
|
||||
.split_once(':')
|
||||
.map(|(value, _)| value)
|
||||
.unwrap_or(suffix);
|
||||
let model = model.trim();
|
||||
if model.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(model.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn extract_ai_standard_requested_model(body_json: &Value) -> Option<String> {
|
||||
body_json
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
pub fn extract_ai_requested_model_from_request_path(
|
||||
request_path: &str,
|
||||
body_json: &Value,
|
||||
family: AiRequestedModelFamily,
|
||||
) -> Option<String> {
|
||||
match family {
|
||||
AiRequestedModelFamily::Standard => extract_ai_standard_requested_model(body_json),
|
||||
AiRequestedModelFamily::Gemini => extract_ai_gemini_model_from_path(request_path),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use aether_ai_surfaces::api::{LocalStandardSourceMode, OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND};
|
||||
|
||||
#[test]
|
||||
fn standard_spec_metadata_maps_model_family_and_report_kind() {
|
||||
let metadata = ai_standard_spec_metadata(LocalStandardSpec {
|
||||
mode: LocalStandardSourceMode::Chat,
|
||||
family: LocalStandardSourceFamily::Standard,
|
||||
api_format: "openai:chat",
|
||||
decision_kind: "openai_chat_sync",
|
||||
report_kind: OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND,
|
||||
require_streaming: false,
|
||||
});
|
||||
|
||||
assert_eq!(metadata.api_format, "openai:chat");
|
||||
assert_eq!(
|
||||
metadata.requested_model_family,
|
||||
Some(AiRequestedModelFamily::Standard)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn video_spec_metadata_maps_gemini_family_without_stream_requirement() {
|
||||
let metadata = ai_video_create_spec_metadata(LocalVideoCreateSpec {
|
||||
family: LocalVideoCreateFamily::Gemini,
|
||||
api_format: "gemini:video",
|
||||
decision_kind: "gemini_video_create",
|
||||
report_kind: "gemini_video_create_success",
|
||||
});
|
||||
|
||||
assert!(!metadata.require_streaming);
|
||||
assert_eq!(
|
||||
metadata.requested_model_family,
|
||||
Some(AiRequestedModelFamily::Gemini)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_model_path_parser_trims_method_suffix() {
|
||||
let model = extract_ai_gemini_model_from_path(
|
||||
"/v1beta/models/gemini-2.5-pro:streamGenerateContent",
|
||||
);
|
||||
|
||||
assert_eq!(model.as_deref(), Some("gemini-2.5-pro"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn standard_requested_model_parser_reads_request_body_model() {
|
||||
let requested_model = extract_ai_standard_requested_model(
|
||||
&serde_json::json!({ "model": " claude-sonnet-4 " }),
|
||||
);
|
||||
|
||||
assert_eq!(requested_model.as_deref(), Some("claude-sonnet-4"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn requested_model_parser_delegates_by_family() {
|
||||
let body = serde_json::json!({ "model": " claude-sonnet-4 " });
|
||||
|
||||
assert_eq!(
|
||||
extract_ai_requested_model_from_request_path(
|
||||
"/v1/chat/completions",
|
||||
&body,
|
||||
AiRequestedModelFamily::Standard,
|
||||
)
|
||||
.as_deref(),
|
||||
Some("claude-sonnet-4")
|
||||
);
|
||||
assert_eq!(
|
||||
extract_ai_requested_model_from_request_path(
|
||||
"/v1beta/models/gemini-2.5-pro:generateContent",
|
||||
&serde_json::json!({}),
|
||||
AiRequestedModelFamily::Gemini,
|
||||
)
|
||||
.as_deref(),
|
||||
Some("gemini-2.5-pro")
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user