1use serde::{Deserialize, Serialize};
2
3use crate::error::EmbedError;
4use crate::traits::TruncationStrategy;
5
6#[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 pub model_path: Option<String>,
18 pub tokenizer_path: Option<String>,
19 pub base_url: Option<String>,
20}
21
22#[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#[derive(Debug, Clone, Serialize, Deserialize)]
33pub enum IndexBuildMode {
34 Incremental,
35 Manual,
36 Deferred { interval_ms: u64 },
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
41pub enum RebuildTrigger {
42 Count(usize),
43 IntervalSecs(u64),
44 ManualOnly,
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct IndexBuildConfig {
50 pub mode: IndexBuildMode,
51 pub rebuild_trigger: RebuildTrigger,
52}
53
54#[derive(Debug, Clone, Serialize, Deserialize)]
56pub struct EmbedConfig {
57 pub default_model: String,
59 pub default_dimensions: usize,
61 pub truncation: TruncationStrategy,
63 pub storage: StorageConfig,
65 pub index: IndexBuildConfig,
67 pub models: Vec<ModelConfig>,
69}
70
71impl EmbedConfig {
72 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 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#[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}