Skip to main content

xz_provider/router/
decision.rs

1use crate::types::{CapabilityRequest, ModelInfo};
2
3/// 路由决策上下文 —— 一次路由需要的所有输入
4#[derive(Debug, Clone, Default)]
5pub struct RouteContext {
6    pub capabilities: Option<CapabilityRequest>,
7    pub named_route: Option<String>,
8    pub model: Option<String>,
9    pub cost_preference: CostPreference,
10}
11
12/// 成本偏好
13#[derive(Debug, Clone, Default)]
14pub enum CostPreference {
15    Cheapest,
16    Fastest,
17    Balanced,
18    #[default]
19    NoPreference,
20}
21
22/// 路由决策结果
23#[derive(Debug, Clone)]
24pub struct RouteDecision {
25    pub model: String,
26    pub provider: String,
27    pub fallback_chain: Vec<FallbackEntry>,
28}
29
30/// 回退条目
31#[derive(Debug, Clone)]
32pub struct FallbackEntry {
33    pub model: String,
34    pub provider: String,
35    pub condition: FallbackCondition,
36}
37
38/// 回退触发条件
39#[derive(Debug, Clone)]
40pub enum FallbackCondition {
41    Always,
42    RateLimitOnly,
43    ErrorStatus(Vec<u16>),
44}
45
46impl RouteDecision {
47    pub fn iter_entries(&self) -> impl Iterator<Item = (&str, &str)> {
48        let primary = std::iter::once((self.provider.as_str(), self.model.as_str()));
49        let fallbacks = self.fallback_chain.iter().map(|e| (e.provider.as_str(), e.model.as_str()));
50        primary.chain(fallbacks)
51    }
52}
53
54/// 延迟追踪器 —— 记录每个 model 的历史延迟,供 Fastest 路由决策使用
55#[derive(Debug, Clone)]
56pub struct LatencyTracker {
57    history: std::collections::HashMap<String, std::collections::VecDeque<u64>>,
58    max_history: usize,
59}
60
61impl LatencyTracker {
62    pub fn new(max_history: usize) -> Self {
63        Self { history: std::collections::HashMap::new(), max_history }
64    }
65
66    pub fn record(&mut self, model: &str, latency_ms: u64) {
67        let entry = self.history.entry(model.to_owned()).or_default();
68        entry.push_back(latency_ms);
69        if entry.len() > self.max_history {
70            entry.pop_front();
71        }
72    }
73
74    pub fn record_error(&mut self, model: &str) {
75        let entry = self.history.entry(model.to_owned()).or_default();
76        entry.push_back(u64::MAX / 2);
77        if entry.len() > self.max_history {
78            entry.pop_front();
79        }
80    }
81
82    pub fn avg_latency(&self, model: &str) -> Option<u64> {
83        let history = self.history.get(model)?;
84        if history.is_empty() {
85            return None;
86        }
87        let sum: u64 = history.iter().sum();
88        Some(sum / history.len() as u64)
89    }
90
91    pub fn fastest(&self, models: &[&ModelInfo]) -> Option<String> {
92        models
93            .iter()
94            .filter(|m| self.history.contains_key(&m.name))
95            .min_by_key(|m| self.avg_latency(&m.name).unwrap_or(u64::MAX))
96            .or_else(|| models.first())
97            .map(|m| m.name.clone())
98    }
99}
100
101impl Default for LatencyTracker {
102    fn default() -> Self {
103        Self::new(10)
104    }
105}