rucora_core/retrieval/
trait.rs

1//! 向量数据库（Vector Store）抽象。
2//!
3//! 提供向量存储、检索和管理的统一接口，支持语义搜索和 RAG 场景。
4
5use async_trait::async_trait;
6use serde::{Deserialize, Serialize};
7
8use crate::error::ProviderError;
9
10/// 向量记录。
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct VectorRecord {
13    /// 唯一标识符。
14    pub id: String,
15    /// 向量数据。
16    pub vector: Vec<f32>,
17    /// 关联的文本内容（可选）。
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub text: Option<String>,
20    /// 元数据（可选）。
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub metadata: Option<serde_json::Value>,
23}
24
25impl VectorRecord {
26    /// 创建新的向量记录。
27    pub fn new(id: impl Into<String>, vector: Vec<f32>) -> Self {
28        assert!(!vector.is_empty(), "VectorRecord vector must not be empty");
29        Self {
30            id: id.into(),
31            vector,
32            text: None,
33            metadata: None,
34        }
35    }
36
37    /// 设置文本内容。
38    pub fn with_text(mut self, text: impl Into<String>) -> Self {
39        self.text = Some(text.into());
40        self
41    }
42
43    /// 设置元数据。
44    pub fn with_metadata(mut self, metadata: serde_json::Value) -> Self {
45        self.metadata = Some(metadata);
46        self
47    }
48}
49
50/// 搜索结果。
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct SearchResult {
53    /// 记录 ID。
54    pub id: String,
55    /// 相似度分数（通常 0-1，越大越相似）。
56    pub score: f32,
57    /// 向量数据（可选，取决于 store 配置）。
58    #[serde(skip_serializing_if = "Option::is_none")]
59    pub vector: Option<Vec<f32>>,
60    /// 关联的文本内容。
61    #[serde(skip_serializing_if = "Option::is_none")]
62    pub text: Option<String>,
63    /// 元数据。
64    #[serde(skip_serializing_if = "Option::is_none")]
65    pub metadata: Option<serde_json::Value>,
66}
67
68/// 查询条件。
69#[derive(Debug, Clone, Default)]
70pub struct VectorQuery {
71    /// 查询向量。
72    pub vector: Vec<f32>,
73    /// 返回结果数量（默认 10）。
74    pub top_k: usize,
75    /// 最小相似度阈值（可选，过滤低相似度结果）。
76    pub score_threshold: Option<f32>,
77    /// 元数据过滤条件（可选，JSON 格式）。
78    pub filter: Option<serde_json::Value>,
79}
80
81impl VectorQuery {
82    /// 创建新的查询。
83    pub fn new(vector: Vec<f32>) -> Self {
84        Self {
85            vector,
86            top_k: 10,
87            score_threshold: None,
88            filter: None,
89        }
90    }
91
92    /// 设置返回数量。
93    pub fn with_top_k(mut self, top_k: usize) -> Self {
94        assert!(top_k > 0, "top_k must be greater than 0");
95        assert!(top_k <= 1000, "top_k must not exceed 1000");
96        self.top_k = top_k;
97        self
98    }
99
100    /// 设置相似度阈值。
101    pub fn with_score_threshold(mut self, threshold: f32) -> Self {
102        self.score_threshold = Some(threshold);
103        self
104    }
105
106    /// 设置过滤条件。
107    pub fn with_filter(mut self, filter: serde_json::Value) -> Self {
108        self.filter = Some(filter);
109        self
110    }
111}
112
113/// 向量数据库抽象。
114///
115/// 支持基本的增删改查和相似度搜索。
116#[async_trait]
117pub trait VectorStore: Send + Sync {
118    /// 添加或更新记录。
119    async fn upsert(&self, records: Vec<VectorRecord>) -> Result<(), ProviderError>;
120
121    /// 根据 ID 删除记录。
122    async fn delete(&self, ids: Vec<String>) -> Result<(), ProviderError>;
123
124    /// 根据 ID 查询记录。
125    async fn get(&self, ids: Vec<String>) -> Result<Vec<VectorRecord>, ProviderError>;
126
127    /// 向量相似度搜索。
128    async fn search(&self, query: VectorQuery) -> Result<Vec<SearchResult>, ProviderError>;
129
130    /// 清空所有数据。
131    async fn clear(&self) -> Result<(), ProviderError>;
132
133    /// 获取记录数量。
134    async fn count(&self) -> Result<usize, ProviderError>;
135}
rucora_core/retrieval/trait.rs

rucora_core/retrieval/
trait.rs
