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| {
75            m.model.as_deref() == Some(name) || m.provider == name
76        })
77    }
78
79    /// 查找默认模型配置
80    pub fn default_model_config(&self) -> Result<&ModelConfig, EmbedError> {
81        self.find_model(&self.default_model)
82            .ok_or_else(|| EmbedError::Config(format!("默认模型未找到: {}", self.default_model)))
83    }
84}
85
86/// 重试配置
87#[derive(Debug, Clone, Serialize, Deserialize)]
88pub struct RetryConfig {
89    pub max_retries: u32,
90    pub initial_backoff_ms: u64,
91    pub max_backoff_ms: u64,
92    pub backoff_multiplier: f64,
93}
94
95impl Default for RetryConfig {
96    fn default() -> Self {
97        Self {
98            max_retries: 3,
99            initial_backoff_ms: 100,
100            max_backoff_ms: 10_000,
101            backoff_multiplier: 2.0,
102        }
103    }
104}