Skip to main content

xz_embed/
config.rs

1use serde::{Deserialize, Serialize};
2
3use crate::error::EmbedError;
4use crate::traits::TruncationStrategy;
5
6/// 模型配置
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct ModelConfig {
9    pub provider: String,
10    pub api_key: Option<String>,
11    pub model: Option<String>,
12    pub dimensions: usize,
13    pub max_batch_size: usize,
14    pub max_input_tokens: usize,
15    pub input_per_million: f64,
16    /// 本地模型路径
17    pub model_path: Option<String>,
18    pub tokenizer_path: Option<String>,
19    pub base_url: Option<String>,
20}
21
22/// 存储配置
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct StorageConfig {
25    pub backend: String,
26    pub path: String,
27    pub max_capacity_bytes: Option<u64>,
28    pub table_name: String,
29}
30
31/// 索引构建模式
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub enum IndexBuildMode {
34    Incremental,
35    Manual,
36    Deferred { interval_ms: u64 },
37}
38
39/// 索引重建触发器
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub enum RebuildTrigger {
42    Count(usize),
43    IntervalSecs(u64),
44    ManualOnly,
45}
46
47/// 索引构建配置
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct IndexBuildConfig {
50    pub mode: IndexBuildMode,
51    pub rebuild_trigger: RebuildTrigger,
52}
53
54/// 完整嵌入配置
55#[derive(Debug, Clone, Serialize, Deserialize)]
56pub struct EmbedConfig {
57    /// 默认模型名称
58    pub default_model: String,
59    /// 默认向量维度
60    pub default_dimensions: usize,
61    /// 截断策略
62    pub truncation: TruncationStrategy,
63    /// 存储配置
64    pub storage: StorageConfig,
65    /// 索引配置
66    pub index: IndexBuildConfig,
67    /// 模型列表
68    pub models: Vec<ModelConfig>,
69}
70
71impl EmbedConfig {
72    /// 查找模型配置
73    pub fn find_model(&self, name: &str) -> Option<&ModelConfig> {
74        self.models.iter().find(|m| m.model.as_deref() == Some(name) || m.provider == name)
75    }
76
77    /// 查找默认模型配置
78    pub fn default_model_config(&self) -> Result<&ModelConfig, EmbedError> {
79        self.find_model(&self.default_model)
80            .ok_or_else(|| EmbedError::Config(format!("默认模型未找到: {}", self.default_model)))
81    }
82}
83
84/// 重试配置
85#[derive(Debug, Clone, Serialize, Deserialize)]
86pub struct RetryConfig {
87    pub max_retries: u32,
88    pub initial_backoff_ms: u64,
89    pub max_backoff_ms: u64,
90    pub backoff_multiplier: f64,
91}
92
93impl Default for RetryConfig {
94    fn default() -> Self {
95        Self {
96            max_retries: 3,
97            initial_backoff_ms: 100,
98            max_backoff_ms: 10_000,
99            backoff_multiplier: 2.0,
100        }
101    }
102}