lxy
3f2cf5fa3a
- run_agentic_loop: 抽 finish_round_exit + emit_ai_completed_once(6 退出点收敛) - compress/title/knowledge_inject/project: router estimated_context+tier 调用点对齐 - provider_pool: 亲和加 model.enabled 过滤
306 lines
13 KiB
Rust
306 lines
13 KiB
Rust
//! 多 Provider 负载均衡池 — F-260614-04
|
|
//!
|
|
//! 与 `df_ai::router`(F-01 阶段4,纯函数选模型)对齐的纯逻辑核心:给定 enabled provider
|
|
//! 列表 + 已选模型 ID + 选择策略,返回**有序候选 provider 列表**(主→备用)。
|
|
//! 零 IO / 零状态,所有状态(DB 读 / 健康度 / 计数器)由调用方持有。
|
|
//!
|
|
//! ## 架构定位(为何放 src-tauri 而非 df-ai)
|
|
//! `AiProviderRecord` 定义在 df-storage。df-ai 当前存储无关(router.rs 操作 df-ai-core 的
|
|
//! `ModelConfig`,非 storage 类型)。把 ProviderPool 放 df-ai 会强制 df-ai→df-storage 依赖,
|
|
//! 破坏 df-ai 的存储无关边界。ProviderPool 与 `prompt.rs::get_active_provider` /
|
|
//! `secret.rs::build_provider_for` 同属"消费 AiProviderRecord 的 app 层逻辑",故放 src-tauri。
|
|
//!
|
|
//! ## 与 router 协同(职责切分,不重叠)
|
|
//! - **router.select_model_id**:在**单个 provider 的 model_configs 池**中选最优模型(F-01)。
|
|
//! - **ProviderPool::select**:在**多个 provider 间**选主 + 列出备用(本模块)。
|
|
//! 调用顺序:先 router 选模型 → 再 ProviderPool 选 provider 实例(模型在多 provider 间共享时)。
|
|
//!
|
|
//! ## 选择策略(论证,见 select 文档)
|
|
//! 1. **模型亲和优先**:router 已选 model_id,优先选**其池中含该模型的 provider**(正确性:不擅自换模型)。
|
|
//! 2. **加权**:同亲和级内按 weight 降序(weight=0 provider 不入选)。
|
|
//! 3. **is_default 兜底**:模型亲和全 miss 时,默认 provider 优先(零变化启动行为)。
|
|
//!
|
|
//! ## fallback 顺序(select 返回的 Vec 即为 fallback 序)
|
|
//! 调用方(agentic loop)按顺序尝试:主 provider 失败(可重试错误)→ 列表下一个 provider 重试。
|
|
//! 不可重试错误(401/400)立即放弃,不浪费备用 provider(对齐 retry.rs Fatal 分类)。
|
|
//!
|
|
//! ## 向后兼容
|
|
//! - enabled provider 仅 1 个 → 返回单元素 Vec,调用方无 fallback 路径(行为同 F-01 前)。
|
|
//! - enabled provider 0 个 → 返回空 Vec,调用方兜底 get_active_provider(行为同 F-01 前)。
|
|
//! - 模型在多 provider 间共享 → 主=含模型且 weight 最高的 provider,备用=其余含模型的 provider。
|
|
|
|
use df_storage::models::AiProviderRecord;
|
|
|
|
/// 多 Provider 负载均衡池(单元结构,无状态)。
|
|
///
|
|
/// `select` 为关联函数:给定 enabled provider 列表 + 已选模型 ID,返回有序候选列表。
|
|
/// 对齐 `ModelRouter`(df_ai::router)的单元结构 + 关联函数风格。
|
|
pub struct ProviderPool;
|
|
|
|
impl ProviderPool {
|
|
/// 在 enabled provider 池中选出**有序候选列表**(主 → 备用),供调用方 fallback。
|
|
///
|
|
/// ## 选择算法(3 步,稳定排序)
|
|
/// 1. **过滤 enabled**:只选 `enabled == true` 的 provider。
|
|
/// 2. **过滤零权重**:weight == 0 的 provider 不参与(weight=0 = 显式禁用主选;
|
|
/// 若需"仅 fallback"配置,用 enabled=false 替代,语义更清晰)。
|
|
/// 3. **排序**(稳定,保输入顺序作 tiebreak):
|
|
/// - **主键:模型亲和**。含 `model_id` 的 provider 排前(模型亲和=true → 排前)。
|
|
/// model_id 为 None(调用方未走 router,如标题/扫描)→ 退化为全亲和,跳过此键。
|
|
/// - **次键:weight 降序**。同亲和级内 weight 高者排前。
|
|
/// - **末键:is_default 兜底**。同亲和同 weight 时 is_default 排前(启动行为不变)。
|
|
///
|
|
/// ## 向后兼容
|
|
/// - enabled provider 1 个 → 单元素 Vec,无 fallback。
|
|
/// - 0 个 → 空 Vec,调用方兜底 get_active_provider。
|
|
/// - model_id None → 全亲和,纯按 weight + is_default 排序(标题/扫描路径用)。
|
|
pub fn select(
|
|
providers: &[AiProviderRecord],
|
|
model_id: Option<&str>,
|
|
) -> Vec<AiProviderRecord> {
|
|
// 步骤 1+2:enabled 且 weight > 0
|
|
let mut candidates: Vec<AiProviderRecord> = providers
|
|
.iter()
|
|
.filter(|p| p.enabled && p.weight > 0)
|
|
.cloned()
|
|
.collect();
|
|
|
|
// 步骤 3:稳定排序(主键模型亲和,次键 weight,末键 is_default)。
|
|
// sort_by 稳定:同 key 保输入顺序(created_at DESC,list_all 返回序)。
|
|
// 反转比较结果:大者排前(亲和 true > false;weight 大 > 小;is_default true > false)。
|
|
candidates.sort_by(|a, b| {
|
|
// 主键:模型亲和。model_id None → 视两方都亲和(退化为全过此键)。
|
|
// 必须过滤 m.enabled:用户禁用所有匹配实例时,该 provider 不应被标为亲和,
|
|
// 否则会选不含可用目标模型的 provider 做主,污染 fallback 链顺序。
|
|
let a_affinity = model_id
|
|
.map_or(true, |mid| a.model_configs.iter().any(|m| m.model_id == mid && m.enabled));
|
|
let b_affinity = model_id
|
|
.map_or(true, |mid| b.model_configs.iter().any(|m| m.model_id == mid && m.enabled));
|
|
let by_affinity = b_affinity.cmp(&a_affinity); // true 排前
|
|
if by_affinity != std::cmp::Ordering::Equal {
|
|
return by_affinity;
|
|
}
|
|
// 次键:weight 降序。
|
|
let by_weight = b.weight.cmp(&a.weight);
|
|
if by_weight != std::cmp::Ordering::Equal {
|
|
return by_weight;
|
|
}
|
|
// 末键:is_default 兜底。
|
|
b.is_default.cmp(&a.is_default)
|
|
});
|
|
|
|
candidates
|
|
}
|
|
}
|
|
|
|
// ============================================================
|
|
// 单测(纯逻辑,零 IO)
|
|
// ============================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use df_ai::df_ai_core::model::ModelConfig;
|
|
|
|
/// 构造可定制 AiProviderRecord(默认 enabled/weight=50/is_default=false,无 model_configs)。
|
|
fn provider(id: &str) -> AiProviderRecord {
|
|
AiProviderRecord {
|
|
id: id.into(),
|
|
name: id.into(),
|
|
provider_type: "openai_compat".into(),
|
|
api_key: String::new(),
|
|
base_url: "https://x".into(),
|
|
default_model: "glm-4-flash".into(),
|
|
models: None,
|
|
model_configs: Vec::new(),
|
|
is_default: false,
|
|
config: None,
|
|
created_at: "0".into(),
|
|
updated_at: "0".into(),
|
|
enabled: true,
|
|
weight: 50,
|
|
}
|
|
}
|
|
|
|
/// ── 向后兼容:单 provider 路径 ──
|
|
|
|
#[test]
|
|
fn empty_pool_returns_empty() {
|
|
// enabled provider 0 个 → 空 Vec,调用方兜底 get_active_provider。
|
|
let pool: Vec<AiProviderRecord> = vec![];
|
|
assert!(ProviderPool::select(&pool, Some("glm-4-flash")).is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn single_provider_returns_singleton() {
|
|
// 单 provider → 单元素 Vec,无 fallback(行为同 F-01 前)。
|
|
let pool = vec![provider("p1")];
|
|
let selected = ProviderPool::select(&pool, Some("glm-4-flash"));
|
|
assert_eq!(selected.len(), 1);
|
|
assert_eq!(selected[0].id, "p1");
|
|
}
|
|
|
|
/// ── 过滤:enabled / weight=0 ──
|
|
|
|
#[test]
|
|
fn disabled_provider_filtered() {
|
|
// enabled=false → 不入选(单 provider 场景兜底由调用方处理)。
|
|
let pool = vec![
|
|
AiProviderRecord { enabled: false, ..provider("disabled") },
|
|
provider("enabled"),
|
|
];
|
|
let selected = ProviderPool::select(&pool, None);
|
|
assert_eq!(selected.len(), 1);
|
|
assert_eq!(selected[0].id, "enabled");
|
|
}
|
|
|
|
#[test]
|
|
fn zero_weight_provider_filtered() {
|
|
// weight=0 → 显式排除(语义:weight=0 = 不参与主选也不作 fallback)。
|
|
let pool = vec![
|
|
AiProviderRecord { weight: 0, ..provider("zero") },
|
|
provider("normal"),
|
|
];
|
|
let selected = ProviderPool::select(&pool, None);
|
|
assert_eq!(selected.len(), 1);
|
|
assert_eq!(selected[0].id, "normal");
|
|
}
|
|
|
|
/// ── 主键:模型亲和 ──
|
|
|
|
#[test]
|
|
fn model_affinity_picks_provider_with_model() {
|
|
// 两 provider,仅 p2 含目标模型 → p2 排前(即使 p1 weight 更高)。
|
|
let pool = vec![
|
|
AiProviderRecord {
|
|
weight: 90,
|
|
..provider("p1")
|
|
},
|
|
AiProviderRecord {
|
|
weight: 30,
|
|
model_configs: vec![ModelConfig::with_defaults("glm-4-flash")],
|
|
..provider("p2")
|
|
},
|
|
];
|
|
let selected = ProviderPool::select(&pool, Some("glm-4-flash"));
|
|
assert_eq!(selected[0].id, "p2", "含目标模型的 provider 应排前");
|
|
assert_eq!(selected[1].id, "p1");
|
|
}
|
|
|
|
#[test]
|
|
fn model_id_none_all_affinity_equal() {
|
|
// model_id None(标题/扫描路径)→ 全亲和,纯按 weight 排序。
|
|
let pool = vec![
|
|
AiProviderRecord { weight: 30, ..provider("low") },
|
|
AiProviderRecord { weight: 90, ..provider("high") },
|
|
];
|
|
let selected = ProviderPool::select(&pool, None);
|
|
assert_eq!(selected[0].id, "high");
|
|
assert_eq!(selected[1].id, "low");
|
|
}
|
|
|
|
/// ── 次键:weight 降序 ──
|
|
|
|
#[test]
|
|
fn higher_weight_wins_same_affinity() {
|
|
// 两 provider 都不含目标模型(同亲和=false),weight 高者排前。
|
|
let pool = vec![
|
|
AiProviderRecord { weight: 50, ..provider("low") },
|
|
AiProviderRecord { weight: 80, ..provider("high") },
|
|
];
|
|
let selected = ProviderPool::select(&pool, Some("glm-4-flash"));
|
|
assert_eq!(selected[0].id, "high");
|
|
}
|
|
|
|
/// ── 末键:is_default 兜底 ──
|
|
|
|
#[test]
|
|
fn is_default_breaks_weight_tie() {
|
|
// 同亲和 + 同 weight → is_default 排前(启动行为不变)。
|
|
let pool = vec![
|
|
provider("normal"), // is_default=false
|
|
AiProviderRecord { is_default: true, ..provider("default") },
|
|
];
|
|
let selected = ProviderPool::select(&pool, None);
|
|
assert_eq!(selected[0].id, "default");
|
|
}
|
|
|
|
/// ── 综合:多 provider 多维度 ──
|
|
|
|
#[test]
|
|
fn full_scenario_orders_correctly() {
|
|
// 候选:
|
|
// a: enabled=true, weight=70, 含目标模型 → 亲和+weight70
|
|
// b: enabled=true, weight=90, 不含目标模型 → 非亲和+weight90(排 a 后)
|
|
// c: enabled=true, weight=70, 含目标模型, is_default=true → 亲和+weight70+default(同 a 但 default 排前)
|
|
// d: enabled=false → 排除
|
|
// e: weight=0 → 排除
|
|
//
|
|
// 预期顺序:c(亲和/70/default) > a(亲和/70) > b(非亲和/90)
|
|
let pool = vec![
|
|
AiProviderRecord {
|
|
weight: 70,
|
|
model_configs: vec![ModelConfig::with_defaults("glm-4-flash")],
|
|
..provider("a")
|
|
},
|
|
AiProviderRecord {
|
|
weight: 90,
|
|
..provider("b")
|
|
},
|
|
AiProviderRecord {
|
|
weight: 70,
|
|
is_default: true,
|
|
model_configs: vec![ModelConfig::with_defaults("glm-4-flash")],
|
|
..provider("c")
|
|
},
|
|
AiProviderRecord {
|
|
enabled: false,
|
|
..provider("d")
|
|
},
|
|
AiProviderRecord {
|
|
weight: 0,
|
|
..provider("e")
|
|
},
|
|
];
|
|
let selected = ProviderPool::select(&pool, Some("glm-4-flash"));
|
|
assert_eq!(
|
|
selected.iter().map(|p| p.id.as_str()).collect::<Vec<_>>(),
|
|
vec!["c", "a", "b"],
|
|
"应按 亲和 > weight > is_default 排序,排除 enabled=false/weight=0"
|
|
);
|
|
}
|
|
|
|
/// ── 模型亲和:enabled 过滤 ──
|
|
|
|
#[test]
|
|
fn model_affinity_excludes_disabled_model() {
|
|
// 回归点:匹配 model 全 disabled 时,provider 不应被标为亲和。
|
|
// 设计:p1 含目标模型但 enabled=false 且 weight=90(高 weight 诱惑),
|
|
// p2 含目标模型且 enabled=true 但 weight=30(低 weight)。
|
|
// 修复后(过滤 m.enabled):p1 非亲和 / p2 亲和 → p2 因亲和主键排前(即使 weight 低)。
|
|
// bug 复现(不过滤 enabled):p1 p2 同亲和 → 按 weight → p1(90)排前。
|
|
// 位置差异精确捕获回归:p2 排前=修复正确,p1 排前=bug 复现。
|
|
let pool = vec![
|
|
AiProviderRecord {
|
|
weight: 90,
|
|
model_configs: vec![ModelConfig {
|
|
enabled: false,
|
|
..ModelConfig::with_defaults("glm-4-flash")
|
|
}],
|
|
..provider("p1_disabled_model")
|
|
},
|
|
AiProviderRecord {
|
|
weight: 30,
|
|
model_configs: vec![ModelConfig::with_defaults("glm-4-flash")],
|
|
..provider("p2_enabled_model")
|
|
},
|
|
];
|
|
let selected = ProviderPool::select(&pool, Some("glm-4-flash"));
|
|
assert_eq!(
|
|
selected.iter().map(|p| p.id.as_str()).collect::<Vec<_>>(),
|
|
vec!["p2_enabled_model", "p1_disabled_model"],
|
|
"匹配 model 全 disabled 的 p1 非亲和,含 enabled model 的 p2 因亲和排前(即使 weight 低)"
|
|
);
|
|
}
|
|
}
|