lellm_provider/
router.rs

1//! ModelRouter — 任务分级路由。
2
3use std::collections::HashMap;
4use std::sync::Arc;
5
6use lellm_core::LlmError;
7
8use crate::LlmProvider;
9
10/// 任务级别。
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12pub enum TaskLevel {
13    /// 快速/便宜，如简单问答、格式转换
14    Flash,
15    /// 默认，如一般对话
16    Standard,
17    /// 复杂推理，如代码生成、深度分析
18    Pro,
19}
20
21/// 路由条目 — Provider + Model 的组合
22#[derive(Debug, Clone)]
23pub struct RouteEntry {
24    pub provider_id: String,
25    pub model: String,
26}
27
28/// 解析后的模型 — 从 Registry 中解析 RouteEntry 得到
29#[derive(Clone)]
30pub struct ResolvedModel {
31    pub provider: Arc<dyn LlmProvider>,
32    pub model: String,
33    /// 模型上下文窗口大小（Token），用于 v0.2 自动推导 ContextBudget
34    /// 若未知则设为 `None`，使用 ContextBudget 的固定默认值
35    pub context_window: Option<usize>,
36}
37
38/// 模型路由器 — 根据任务级别选择路由。
39pub struct ModelRouter {
40    routes: HashMap<TaskLevel, RouteEntry>,
41}
42
43impl ModelRouter {
44    pub fn new() -> Self {
45        Self {
46            routes: HashMap::new(),
47        }
48    }
49
50    pub fn add_route(&mut self, level: TaskLevel, entry: RouteEntry) {
51        self.routes.insert(level, entry);
52    }
53
54    /// 根据任务级别解析路由
55    pub fn resolve(&self, level: TaskLevel) -> Option<&RouteEntry> {
56        self.routes.get(&level)
57    }
58}
59
60impl Default for ModelRouter {
61    fn default() -> Self {
62        Self::new()
63    }
64}
65
66/// Provider 注册表 — 持有所有 Provider 实例。
67pub struct ProviderRegistry {
68    providers: HashMap<String, Arc<dyn LlmProvider>>,
69}
70
71impl ProviderRegistry {
72    pub fn new() -> Self {
73        Self {
74            providers: HashMap::new(),
75        }
76    }
77
78    pub fn register(&mut self, id: &str, provider: Arc<dyn LlmProvider>) {
79        self.providers.insert(id.to_string(), provider);
80    }
81
82    pub fn get(&self, id: &str) -> Option<Arc<dyn LlmProvider>> {
83        self.providers.get(id).cloned()
84    }
85
86    /// 从 RouteEntry 解析为 ResolvedModel
87    pub fn resolve(&self, route: &RouteEntry) -> Result<ResolvedModel, LlmError> {
88        let provider = self
89            .get(&route.provider_id)
90            .ok_or_else(|| LlmError::Provider {
91                provider: route.provider_id.clone(),
92                status: None,
93                code: None,
94                message: format!("provider not registered: {}", route.provider_id),
95            })?;
96        Ok(ResolvedModel {
97            provider,
98            model: route.model.clone(),
99            context_window: None,
100        })
101    }
102}
103
104impl Default for ProviderRegistry {
105    fn default() -> Self {
106        Self::new()
107    }
108}
lellm_provider/router.rs

lellm_provider/
router.rs
