xz_embed/
traits.rs

1use async_trait::async_trait;
2use serde::{Deserialize, Serialize};
3use std::fmt::Debug;
4
5use crate::error::{EmbedError, StoreError};
6use crate::types::{MetadataFilter, SearchResult, StoreStats, VectorEntry};
7
8/// 文本截断策略
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub enum TruncationStrategy {
11    /// 直接报错（严格要求）
12    Error,
13    /// 从开头截断（保留前 N 个 token）
14    TruncateStart,
15    /// 从末尾截断（保留后 N 个 token）
16    TruncateEnd,
17    /// 保留首尾各 N/2 个 token（适合 LLM 输出）
18    TruncateBoth,
19}
20
21/// 嵌入模型定价信息
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct EmbedPricing {
24    /// 每百万 token 价格（输入）
25    pub input_per_million: f64,
26}
27
28/// 嵌入模型元信息
29#[derive(Debug, Clone, Serialize, Deserialize)]
30pub struct EmbedModelInfo {
31    pub name: String,
32    pub display_name: String,
33    /// 支持的维度列表（None 表示不支持维度选择）
34    pub supported_dimensions: Option<Vec<usize>>,
35    /// 当前维度
36    pub current_dimension: usize,
37    /// 最大输入 token 限制
38    pub max_input_tokens: usize,
39    /// 最大批次大小
40    pub max_batch_size: usize,
41    /// 定价信息
42    pub pricing: EmbedPricing,
43}
44
45/// 统一的文本向量嵌入接口
46#[async_trait]
47pub trait EmbeddingModel: Send + Sync + Debug {
48    /// 核心嵌入方法。对一批文本生成向量。
49    async fn embed(&self, input: &[&str]) -> Result<Vec<Vec<f32>>, EmbedError>;
50
51    /// 便捷方法：对单个文本生成向量
52    async fn embed_single(&self, text: &str) -> Result<Vec<f32>, EmbedError> {
53        let mut results = self.embed(&[text]).await?;
54        if results.is_empty() {
55            return Err(EmbedError::Model("embed_single 返回空结果".into()));
56        }
57        Ok(results.remove(0))
58    }
59
60    /// 模型信息
61    fn model_info(&self) -> &EmbedModelInfo;
62
63    /// 最大批次大小
64    fn max_batch_size(&self) -> usize;
65
66    /// 向量维度
67    fn dimensions(&self) -> usize;
68}
69
70/// 维度约简策略
71#[derive(Debug, Clone, Serialize, Deserialize)]
72pub enum DimensionReducer {
73    /// API 原生支持（OpenAI dimensions 参数）
74    Native(usize),
75    /// 截断（取前 N 维）
76    Truncate(usize),
77}
78
79impl DimensionReducer {
80    /// 应用降维
81    pub fn reduce(&self, vectors: &mut [Vec<f32>]) {
82        match self {
83            DimensionReducer::Native(_) => {
84                // 原生降维由 API 端处理，此处无需操作
85            }
86            DimensionReducer::Truncate(n) => {
87                for v in vectors.iter_mut() {
88                    v.truncate(*n);
89                }
90            }
91        }
92    }
93}
94
95/// 向量存储抽象 — 可插拔后端
96#[async_trait]
97pub trait VectorStore: Send + Sync + Debug {
98    /// 插入单条向量
99    async fn insert(&self, entry: VectorEntry) -> Result<(), StoreError>;
100
101    /// 批量插入向量（推荐的高吞吐写入路径）
102    async fn insert_batch(&self, entries: Vec<VectorEntry>) -> Result<(), StoreError>;
103
104    /// 相似度搜索（cosine 相似度，返回 Top-K）
105    async fn search(&self, query: &[f32], limit: usize) -> Result<Vec<SearchResult>, StoreError>;
106
107    /// 带元数据过滤的相似度搜索
108    async fn search_with_filter(
109        &self,
110        query: &[f32],
111        filter: &MetadataFilter,
112        limit: usize,
113    ) -> Result<Vec<SearchResult>, StoreError>;
114
115    /// 按 ID 批量删除
116    async fn delete(&self, ids: &[String]) -> Result<usize, StoreError>;
117
118    /// 按元数据过滤条件删除
119    async fn delete_by_filter(&self, filter: &MetadataFilter) -> Result<usize, StoreError>;
120
121    /// 清空存储
122    async fn clear(&self) -> Result<(), StoreError>;
123
124    /// 存储中的总条目数
125    async fn count(&self) -> Result<usize, StoreError>;
126
127    /// 创建/重建索引（后台线程，不阻塞写入）
128    async fn rebuild_index(&self) -> Result<(), StoreError>;
129
130    /// 存储统计信息
131    async fn stats(&self) -> Result<StoreStats, StoreError>;
132}
133
134/// VectorStore 生命周期 trait
135#[async_trait]
136pub trait StoreLifecycle: VectorStore {
137    /// 初始化存储（创建表结构等）
138    async fn initialize(&self) -> Result<(), StoreError>;
139
140    /// 优雅关闭
141    async fn close(&self) -> Result<(), StoreError>;
142
143    /// 强制持久化检查点
144    async fn checkpoint(&self) -> Result<(), StoreError>;
145
146    /// 是否健康
147    async fn health_check(&self) -> Result<bool, StoreError>;
148}
149
150/// BM25 关键词检索 trait
151pub trait KeywordSearch: Send + Sync {
152    /// 全文检索，返回 (document_id, bm25_score)
153    fn search(&self, query: &str, limit: usize) -> Vec<(String, f32)>;
154}
xz_embed/traits.rs

xz_embed/
traits.rs
