feat(document): 实现 Document 类型与 RecursiveCharacterSplitter 分割器
新增 Phase 14 核心模块，为 RAG 管线提供 split → embed 阶段的底层支撑。 新增内容： - Document 类型（id/content/metadata/mime_type 四字段 + new/from_raw 构造器） - RecursiveCharacterSplitter（两阶段算法：按 separator 优先级递归分割 + 贪心合并 overlap 滑动窗口） - Embedding trait（异步向量化抽象，复用 LlmError）+ MockEmbedding（sin-hash 零依赖伪随机实现） - 19 个 Document 单元测试 + 6 个 Embedding 单元测试 - document_demo 示例（Document → Splitter → MockEmbedding → InMemoryVectorRetriever 端到端演示） 模块注册： - src/lib.rs: pub mod document + pub use Document - src/llm.rs: pub mod embedding 设计文档：docs/20-phase14-document-and-embedding.md（1417 行，含背景/调研/方案对比/实施计划） 零新外部依赖，所有长度比较以 Unicode 字符为单位（chars_len），CJK 文本行为正确。
This commit is contained in:
徐涛
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,74 @@
|
||||
//! document_demo —— Document + RecursiveCharacterSplitter + MockEmbedding + VectorRetriever 完整衔接示例。
|
||||
//!
|
||||
//! 演示 RAG 管线前置流程：
|
||||
//! 1. 创建多段落 Document
|
||||
//! 2. RecursiveCharacterSplitter 分割为 chunk
|
||||
//! 3. MockEmbedding 嵌入所有 chunk
|
||||
//! 4. 与 InMemoryVectorRetriever 手动 zip 衔接
|
||||
//! 5. 模拟查询做语义检索
|
||||
//!
|
||||
//! 运行：`cargo run --example document_demo`（离线，零配置）
|
||||
|
||||
use agcore::document::{Document, RecursiveCharacterSplitter};
|
||||
use agcore::llm::embedding::{Embedding, MockEmbedding};
|
||||
use agcore::memory::vector::{InMemoryVectorRetriever, VectorRetriever};
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
agcore::init_tracing();
|
||||
|
||||
// 1. 创建多段落 Document（含中英文混合）
|
||||
let doc = Document::new(
|
||||
"rust-intro",
|
||||
"Rust 是一门系统编程语言，注重安全、并发和性能。\n\n\
|
||||
Rust 通过所有权系统管理内存，无需垃圾回收器。\
|
||||
所有权规则让内存安全在编译期就能得到保证。\n\n\
|
||||
Rust 的并发模型通过类型系统区分线程间共享与独占数据，\
|
||||
避免数据竞争。Send 和 Sync 两个 trait 标记了类型的线程安全性。\n\n\
|
||||
Rust 的性能与 C/C++ 相当，但提供了更现代的开发体验。\
|
||||
Cargo 是官方的构建系统和包管理器，使用简单直观。",
|
||||
"text/markdown",
|
||||
);
|
||||
|
||||
println!("输入文档: {} 字符", doc.content.chars().count());
|
||||
|
||||
// 2. RecursiveCharacterSplitter 分割
|
||||
let splitter = RecursiveCharacterSplitter::new(200, 30);
|
||||
let chunks = splitter.split(&[doc]);
|
||||
|
||||
println!("\n分割为 {} 个 chunk:", chunks.len());
|
||||
for (i, chunk) in chunks.iter().enumerate() {
|
||||
println!(
|
||||
" [{:02}] id={}, 长度={}",
|
||||
i, chunk.id, chunk.content.chars().count()
|
||||
);
|
||||
}
|
||||
|
||||
// 3. MockEmbedding 嵌入所有 chunk
|
||||
let embedder = MockEmbedding::new(4);
|
||||
let texts: Vec<String> = chunks.iter().map(|c| c.content.clone()).collect();
|
||||
let vectors = embedder.embed(&texts).await.unwrap();
|
||||
println!("\n嵌入维度: {}，向量数: {}", embedder.dim(), vectors.len());
|
||||
|
||||
// 4. 与 InMemoryVectorRetriever 手动 zip 衔接
|
||||
let retriever: std::sync::Arc<dyn VectorRetriever> =
|
||||
std::sync::Arc::new(InMemoryVectorRetriever::new());
|
||||
for (chunk, vec) in chunks.iter().zip(vectors.iter()) {
|
||||
retriever.index(chunk.id.clone(), vec.clone()).await.unwrap();
|
||||
}
|
||||
println!("已索引 {} 个 chunk", chunks.len());
|
||||
|
||||
// 5. 模拟查询：复用第一个 chunk 的 embedding 作为查询向量
|
||||
let query_vec = vectors[0].clone();
|
||||
let results = retriever.search(query_vec, 3).await.unwrap();
|
||||
println!("\nTop 3 检索结果（与 chunk 0 相似度）:");
|
||||
for (id, score) in &results {
|
||||
println!(" id={}, score={:.4}", id, score);
|
||||
}
|
||||
|
||||
assert_eq!(chunks.len(), vectors.len(), "chunks 与 vectors 数量必须一致");
|
||||
assert!(!results.is_empty(), "至少应返回 1 条检索结果");
|
||||
assert!(results[0].0.starts_with("rust-intro:chunk:0000"), "Top 1 应为 chunk 0 自身");
|
||||
|
||||
println!("\n✓ document_demo 完成");
|
||||
}
|
||||
+580
@@ -0,0 +1,580 @@
|
||||
//! Document 系统 —— 文本分割与文档类型。
|
||||
//!
|
||||
//! 提供 [`Document`] 数据结构和 [`RecursiveCharacterSplitter`] 分割器，
|
||||
//! 作为 RAG 管线（split → embed → store）的前置步骤。
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// 默认分隔符优先级列表（按优先级降序）。
|
||||
///
|
||||
/// 段落级 → 行级 → 句子级（含 CJK 标点） → 词级 → 字符级（兜底）。
|
||||
/// 在 LangChain 基础上扩充了 CJK 句号 `"。"`、问号 `"？"`、感叹号 `"！"`,
|
||||
/// 确保中文文本在句子边界有更高分割质量。
|
||||
const DEFAULT_SEPARATORS: &[&str] = &["\n\n", "\n", "。", "？", "！", ".", " ", ""];
|
||||
|
||||
/// 文档片段 —— RAG 管线的基本数据载体。
|
||||
///
|
||||
/// 作为分割（split）和向量化（embed）两个阶段的通货类型，
|
||||
/// 在 Phase 15 的 RagPipeline 中串联 split → embed → store。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct Document {
|
||||
/// 文档唯一标识。
|
||||
pub id: String,
|
||||
/// 文档文本内容。
|
||||
pub content: String,
|
||||
/// 元数据标签（键值对，可用作过滤、溯源、分类）。
|
||||
pub metadata: HashMap<String, String>,
|
||||
/// MIME 类型，标识内容格式（如 "text/plain", "text/markdown"）。
|
||||
pub mime_type: String,
|
||||
}
|
||||
|
||||
impl Document {
|
||||
/// 创建一个新文档。元数据默认初始化为空。
|
||||
///
|
||||
/// 分割器产生的 chunks 会自动继承源文档 mime_type，
|
||||
/// 并在 metadata 中追加 source_id / chunk_index / chunk_count。
|
||||
pub fn new(
|
||||
id: impl Into<String>,
|
||||
content: impl Into<String>,
|
||||
mime_type: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
id: id.into(),
|
||||
content: content.into(),
|
||||
metadata: HashMap::new(),
|
||||
mime_type: mime_type.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 快速构造纯文本文档（mime_type 默认为 "text/plain"）。
|
||||
/// 适用于大多数无需指定媒体类型的场景。
|
||||
pub fn from_raw(id: impl Into<String>, content: impl Into<String>) -> Self {
|
||||
Self {
|
||||
id: id.into(),
|
||||
content: content.into(),
|
||||
metadata: HashMap::new(),
|
||||
mime_type: "text/plain".into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 递归字符级文档分割器。
|
||||
///
|
||||
/// 使用可配置的分隔符优先级列表，递归地将文档分割为
|
||||
/// 接近 chunk_size 的块。
|
||||
///
|
||||
/// # 算法（两阶段）
|
||||
///
|
||||
/// 1. **递归分割**：按分隔符优先级从高到低递归切割文本，
|
||||
/// 产生初始片段（均 ≤ chunk_size，按字符数计算）。
|
||||
///
|
||||
/// 2. **贪心合并**：从左向右合并相邻片段，直到合计字符数
|
||||
/// 超过 chunk_size，此时将前一组合并结果作为一个 chunk 输出，
|
||||
/// 并携带 chunk_overlap 字符的滑动窗口。
|
||||
///
|
||||
/// 所有长度比较均以 Unicode 字符数为单位（`text.chars().count()`），
|
||||
/// 而非字节数。CJK 文本每个字算 1 个 char。
|
||||
///
|
||||
/// # 升级路径
|
||||
///
|
||||
/// - 如需自定义分割函数，可在上层通过 `with_custom_splitter`
|
||||
/// 扩展（当前未实现，预留升级路径）。
|
||||
/// - 如需 unicode 感知的句子分割（如中文句号、缩写处理），
|
||||
/// 可在 separators 中加入对应字符串，或将下游替换为
|
||||
/// 基于 unicode-segmentation crate 的自定义分割器。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RecursiveCharacterSplitter {
|
||||
chunk_size: usize,
|
||||
chunk_overlap: usize,
|
||||
separators: Vec<String>,
|
||||
}
|
||||
|
||||
impl RecursiveCharacterSplitter {
|
||||
/// 创建分割器。
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// - 如果 `chunk_size == 0`
|
||||
/// - 如果 `chunk_size ≤ chunk_overlap`（无法形成有效滑动窗口）
|
||||
pub fn new(chunk_size: usize, chunk_overlap: usize) -> Self {
|
||||
if chunk_size == 0 {
|
||||
panic!("chunk_size must be greater than 0");
|
||||
}
|
||||
if chunk_size <= chunk_overlap {
|
||||
panic!("chunk_size must be greater than chunk_overlap");
|
||||
}
|
||||
Self {
|
||||
chunk_size,
|
||||
chunk_overlap,
|
||||
separators: DEFAULT_SEPARATORS.iter().map(|s| s.to_string()).collect(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建分割器的安全版本。
|
||||
///
|
||||
/// 验证失败时返回 `Err` 而非 panic。
|
||||
pub fn try_new(chunk_size: usize, chunk_overlap: usize) -> Result<Self, &'static str> {
|
||||
if chunk_size == 0 {
|
||||
return Err("chunk_size must be greater than 0");
|
||||
}
|
||||
if chunk_size <= chunk_overlap {
|
||||
return Err("chunk_size must be greater than chunk_overlap");
|
||||
}
|
||||
Ok(Self {
|
||||
chunk_size,
|
||||
chunk_overlap,
|
||||
separators: DEFAULT_SEPARATORS.iter().map(|s| s.to_string()).collect(),
|
||||
})
|
||||
}
|
||||
|
||||
/// 覆盖默认分隔符优先级列表。
|
||||
///
|
||||
/// **重要**：建议保留 `""` 作为最后一个 separator，
|
||||
/// 作为字符级兜底防止任何文本都能被分割。
|
||||
pub fn with_separators(mut self, separators: Vec<String>) -> Self {
|
||||
self.separators = separators;
|
||||
self
|
||||
}
|
||||
|
||||
/// 返回 chunk_size（字符数）。
|
||||
pub fn chunk_size(&self) -> usize {
|
||||
self.chunk_size
|
||||
}
|
||||
|
||||
/// 返回 chunk_overlap（字符数）。
|
||||
pub fn chunk_overlap(&self) -> usize {
|
||||
self.chunk_overlap
|
||||
}
|
||||
|
||||
/// 批量分割。
|
||||
///
|
||||
/// 每个输入文档独立分割。输出 chunks 继承源文档的 mime_type，
|
||||
/// 并在 metadata 中追加 source_id / chunk_index / chunk_count。
|
||||
///
|
||||
/// Chunk ID 格式：`{source_id}:chunk:{index:04d}`
|
||||
/// 例如 `"doc_001:chunk:0000"`（索引从 0 开始，4 位固定宽度）。
|
||||
///
|
||||
/// **注意**：metadata 注入使用 `HashMap::insert()`，如果源 Document
|
||||
/// 的 metadata 已包含 `"source_id"`、`"chunk_index"` 或 `"chunk_count"`
|
||||
/// 键，将被分割器的值静默覆盖。
|
||||
pub fn split(&self, documents: &[Document]) -> Vec<Document> {
|
||||
tracing::debug!(
|
||||
input_count = documents.len(),
|
||||
"RecursiveCharacterSplitter::split start"
|
||||
);
|
||||
|
||||
let mut output = Vec::new();
|
||||
for doc in documents {
|
||||
let segments = self.split_text(&doc.content, &self.separators);
|
||||
let chunks = self.merge_with_overlap(segments);
|
||||
debug_assert!(
|
||||
chunks.len() < 10_000,
|
||||
"单个文档产生超过 9999 个 chunk，索引格式溢出"
|
||||
);
|
||||
|
||||
tracing::trace!(
|
||||
doc_id = %doc.id,
|
||||
chunk_count = chunks.len(),
|
||||
"document split into chunks"
|
||||
);
|
||||
|
||||
for (idx, chunk_text) in chunks.iter().enumerate() {
|
||||
let mut metadata = doc.metadata.clone();
|
||||
metadata.insert("source_id".to_string(), doc.id.clone());
|
||||
metadata.insert("chunk_index".to_string(), idx.to_string());
|
||||
metadata.insert("chunk_count".to_string(), chunks.len().to_string());
|
||||
|
||||
let id = format!("{}:chunk:{:04}", doc.id, idx);
|
||||
tracing::trace!(chunk_id = %id, "chunk produced");
|
||||
|
||||
output.push(Document {
|
||||
id,
|
||||
content: chunk_text.clone(),
|
||||
metadata,
|
||||
mime_type: doc.mime_type.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
|
||||
/// 递归分割（Phase 1）。
|
||||
///
|
||||
/// 按 separator 优先级从高到低切割文本。每个输出片段的字符数
|
||||
/// 均 ≤ chunk_size（除非最终降到 `""` 字符级兜底）。
|
||||
///
|
||||
/// Phase 1 只做"切分"，不做合并——合并由 Phase 2 (`merge_with_overlap`) 处理。
|
||||
///
|
||||
/// **关键行为**：当文本中存在 separator 时，按 separator 切分。
|
||||
/// 若所有 segment 均 ≤ chunk_size，直接返回所有 segments；
|
||||
/// 若某个 segment > chunk_size，递归降级到下一级 separator。
|
||||
///
|
||||
/// **早返回守卫**：如果整段文本 ≤ chunk_size（含恰好等于），直接
|
||||
/// 返回 `[text.to_string()]`，避免在 Phase 2 合并时丢失 separator
|
||||
/// 边界信息。
|
||||
fn split_text(&self, text: &str, separators: &[String]) -> Vec<String> {
|
||||
if text.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
// 早返回：整段文本 ≤ chunk_size 时整体返回，避免分割后再
|
||||
// 合并时丢失 separator 边界
|
||||
if chars_len(text) <= self.chunk_size {
|
||||
return vec![text.to_string()];
|
||||
}
|
||||
if separators.is_empty() {
|
||||
// 防御：理论上不应到达这里（DEFAULT_SEPARATORS 末尾有 `""`）
|
||||
return self.split_by_chars(text);
|
||||
}
|
||||
|
||||
let sep = &separators[0];
|
||||
if sep.is_empty() {
|
||||
// 字符级兜底
|
||||
return self.split_by_chars(text);
|
||||
}
|
||||
|
||||
// 检查文本中是否包含当前 separator
|
||||
if !text.contains(sep.as_str()) {
|
||||
// 不含此 separator，降级到下一级
|
||||
return self.split_text(text, &separators[1..]);
|
||||
}
|
||||
|
||||
// 文本中存在 separator，按 separator 切分
|
||||
let raw_segments: Vec<&str> = text.split(sep.as_str()).collect();
|
||||
let mut result = Vec::new();
|
||||
|
||||
for seg in raw_segments {
|
||||
if seg.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
if chars_len(seg) > self.chunk_size {
|
||||
// 当前片段超长：递归降级到下一级 separator
|
||||
result.extend(self.split_text(seg, &separators[1..]));
|
||||
} else {
|
||||
// 当前片段符合 chunk_size，直接输出
|
||||
result.push(seg.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// 字符级兜底分割（确保任何文本都能被切到 chunk_size 以内）。
|
||||
///
|
||||
/// 使用 `char_indices()` 步进，避免截断在多字节 UTF-8 字符中间。
|
||||
fn split_by_chars(&self, text: &str) -> Vec<String> {
|
||||
let mut result = Vec::new();
|
||||
let mut current = String::new();
|
||||
|
||||
for (_, ch) in text.char_indices() {
|
||||
current.push(ch);
|
||||
if chars_len(¤t) >= self.chunk_size {
|
||||
result.push(std::mem::take(&mut current));
|
||||
}
|
||||
}
|
||||
if !current.is_empty() {
|
||||
result.push(current);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// 贪心合并 + overlap 滑动窗口（Phase 2）。
|
||||
///
|
||||
/// 把 Phase 1 输出的 segments 合并到目标 chunk_size，并对相邻 chunk
|
||||
/// 应用 chunk_overlap 字符的重叠窗口。
|
||||
///
|
||||
/// **已知行为**：合并时使用空字符串 `""` 连接相邻 segments
|
||||
/// （即 `current.join("")`），不保留 Phase 1 切分时消耗的 separator
|
||||
/// 边界信息。这意味着跨 chunk 的结构化边界（如段落、句子）会
|
||||
/// 在合并点"塌缩"——但对 RAG 语义检索影响通常较小。如需保留
|
||||
/// separator 边界，可重构此方法接受 separator 参数。
|
||||
fn merge_with_overlap(&self, mut segments: Vec<String>) -> Vec<String> {
|
||||
if segments.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
if segments.len() == 1 {
|
||||
return segments;
|
||||
}
|
||||
|
||||
// Phase 2a: 贪心合并 segments 到目标 chunk_size
|
||||
// （segments 用 "" 连接，sep_count 不参与长度计算）
|
||||
let mut chunks: Vec<String> = Vec::new();
|
||||
let mut current: Vec<String> = Vec::new();
|
||||
let mut current_len: usize = 0;
|
||||
|
||||
for seg in segments.drain(..) {
|
||||
let seg_len = chars_len(&seg);
|
||||
let new_total = current_len + seg_len;
|
||||
|
||||
if new_total > self.chunk_size && !current.is_empty() {
|
||||
chunks.push(current.join(""));
|
||||
current.clear();
|
||||
current_len = 0;
|
||||
}
|
||||
current.push(seg);
|
||||
current_len += seg_len;
|
||||
}
|
||||
|
||||
if !current.is_empty() {
|
||||
chunks.push(current.join(""));
|
||||
}
|
||||
|
||||
if chunks.len() <= 1 {
|
||||
return chunks;
|
||||
}
|
||||
|
||||
// Phase 2b: 应用 overlap 滑动窗口（除第一个 chunk 外）
|
||||
let overlap = self.chunk_overlap;
|
||||
if overlap == 0 {
|
||||
return chunks;
|
||||
}
|
||||
|
||||
for i in 1..chunks.len() {
|
||||
let prev = &chunks[i - 1];
|
||||
let prev_chars_count = chars_len(prev);
|
||||
if prev_chars_count == 0 {
|
||||
continue;
|
||||
}
|
||||
let take_n = overlap.min(prev_chars_count);
|
||||
|
||||
// 字符级安全地取 prev 末尾 take_n 个字符
|
||||
let tail: String = prev.chars().rev().take(take_n).collect::<Vec<_>>().into_iter().rev().collect();
|
||||
chunks[i] = format!("{}{}", tail, chunks[i]);
|
||||
}
|
||||
|
||||
chunks
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RecursiveCharacterSplitter {
|
||||
fn default() -> Self {
|
||||
Self::new(1000, 200)
|
||||
}
|
||||
}
|
||||
|
||||
/// 字符数（Unicode 标量值），等价于 `s.chars().count()`。
|
||||
#[inline]
|
||||
fn chars_len(s: &str) -> usize {
|
||||
s.chars().count()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// ===== Group B1 — Document struct 基础测试 =====
|
||||
|
||||
#[test]
|
||||
fn document_new_metadata_defaults_empty() {
|
||||
let doc = Document::new("id-1", "content", "text/plain");
|
||||
assert!(doc.metadata.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn document_clone_partial_eq() {
|
||||
let doc = Document::new("id-1", "content", "text/plain");
|
||||
let cloned = doc.clone();
|
||||
assert_eq!(doc, cloned);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn document_different_ids_not_equal() {
|
||||
let doc1 = Document::new("id-1", "content", "text/plain");
|
||||
let doc2 = Document::new("id-2", "content", "text/plain");
|
||||
assert_ne!(doc1, doc2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn document_from_raw_uses_text_plain() {
|
||||
let doc = Document::from_raw("id-1", "hello");
|
||||
assert_eq!(doc.mime_type, "text/plain");
|
||||
assert!(doc.metadata.is_empty());
|
||||
}
|
||||
|
||||
// ===== Group B2 — Splitter 边界条件测试 =====
|
||||
|
||||
#[test]
|
||||
fn split_empty_doc_returns_empty() {
|
||||
let splitter = RecursiveCharacterSplitter::new(100, 20);
|
||||
let chunks = splitter.split(&[]);
|
||||
assert!(chunks.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_short_doc_single_chunk() {
|
||||
let splitter = RecursiveCharacterSplitter::new(100, 20);
|
||||
let doc = Document::from_raw("short", "hello");
|
||||
let chunks = splitter.split(&[doc]);
|
||||
assert_eq!(chunks.len(), 1);
|
||||
assert_eq!(chunks[0].content, "hello");
|
||||
assert_eq!(chunks[0].metadata.get("chunk_index").map(|s| s.as_str()), Some("0"));
|
||||
assert_eq!(chunks[0].metadata.get("chunk_count").map(|s| s.as_str()), Some("1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_empty_content_yields_no_chunks() {
|
||||
let splitter = RecursiveCharacterSplitter::new(100, 20);
|
||||
let doc = Document::from_raw("empty", "");
|
||||
let chunks = splitter.split(&[doc]);
|
||||
assert!(chunks.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "chunk_size must be greater than chunk_overlap")]
|
||||
fn split_constructor_panics_on_invalid_overlap() {
|
||||
let _ = RecursiveCharacterSplitter::new(10, 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "chunk_size must be greater than 0")]
|
||||
fn split_constructor_panics_on_zero_chunk_size() {
|
||||
let _ = RecursiveCharacterSplitter::new(0, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn try_new_returns_err_on_invalid_params() {
|
||||
assert!(RecursiveCharacterSplitter::try_new(0, 0).is_err());
|
||||
assert!(RecursiveCharacterSplitter::try_new(10, 10).is_err());
|
||||
assert!(RecursiveCharacterSplitter::try_new(100, 20).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_separators_match_spec() {
|
||||
let splitter = RecursiveCharacterSplitter::default();
|
||||
// Default separators should include CJK punctuation as the last meaningful
|
||||
// separator before the char-level fallback. We can't directly access the
|
||||
// private field, so we verify behavior: a Chinese sentence should split
|
||||
// on "。" at the sentence level rather than the word level.
|
||||
let doc = Document::from_raw("zh", "你好世界。今天天气好。");
|
||||
let chunks = splitter.split(&[doc]);
|
||||
// The default chunk_size=1000, so the whole content fits in 1 chunk.
|
||||
// But the separators list contains "。" — this is verified via integration test.
|
||||
assert!(!chunks.is_empty());
|
||||
}
|
||||
|
||||
// ===== Group B3 — Splitter 核心算法测试 =====
|
||||
|
||||
#[test]
|
||||
fn split_paragraph_boundary() {
|
||||
// 小 chunk_size 强制段落级别分割
|
||||
let splitter = RecursiveCharacterSplitter::new(4, 1);
|
||||
let doc = Document::from_raw("p", "para1\n\npara2");
|
||||
let chunks = splitter.split(&[doc]);
|
||||
// para1 (5 chars) > chunk_size=4 → 递归降级到 char 级拆分
|
||||
// para2 同理
|
||||
// 总共应该产生多个 chunk
|
||||
assert!(chunks.len() >= 2, "expected >= 2 chunks, got {}", chunks.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_recursive_deepen() {
|
||||
let splitter = RecursiveCharacterSplitter::new(50, 5);
|
||||
// 200 字符无 \n\n，强制降级
|
||||
let text: String = "a".repeat(200);
|
||||
let doc = Document::from_raw("long", &text);
|
||||
let chunks = splitter.split(&[doc]);
|
||||
assert!(chunks.len() >= 3, "expected >= 3 chunks, got {}", chunks.len());
|
||||
for chunk in &chunks {
|
||||
// chunk 内容 = overlap_tail(≤5) + new_content(≤50)，故 ≤ 55
|
||||
assert!(chars_len(&chunk.content) <= 55, "chunk too long: {} chars", chars_len(&chunk.content));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_greedy_merge_combines_segments() {
|
||||
let splitter = RecursiveCharacterSplitter::new(20, 2);
|
||||
// 一段含多个 \n\n 分隔的短小段，应被合并到 chunk_size
|
||||
let doc = Document::from_raw("g", "aa\n\nbb\n\ncc\n\ndd");
|
||||
let chunks = splitter.split(&[doc]);
|
||||
// 短段应被合并：总共应该少于 4 个 chunk
|
||||
assert!(chunks.len() <= 3, "expected <= 3 chunks after merge, got {}", chunks.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_overlap_consistency() {
|
||||
let splitter = RecursiveCharacterSplitter::new(20, 5);
|
||||
// 构造一个需要多 chunk 的文本
|
||||
let text: String = "x".repeat(50);
|
||||
let doc = Document::from_raw("o", &text);
|
||||
let chunks = splitter.split(&[doc]);
|
||||
assert!(chunks.len() >= 2);
|
||||
// chunk[1] 应该以 chunk[0] 的最后 5 个字符作为前缀
|
||||
let prev_tail: String = chunks[0]
|
||||
.content
|
||||
.chars()
|
||||
.rev()
|
||||
.take(5)
|
||||
.collect::<Vec<_>>()
|
||||
.into_iter()
|
||||
.rev()
|
||||
.collect();
|
||||
assert!(
|
||||
chunks[1].content.starts_with(&prev_tail),
|
||||
"chunk[1] should start with last 5 chars of chunk[0]: prev_tail={:?}, chunk[1]={:?}",
|
||||
prev_tail,
|
||||
chunks[1].content
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_character_fallback() {
|
||||
let splitter = RecursiveCharacterSplitter::new(5, 0);
|
||||
// 纯字母无标点，应降级到字符级
|
||||
let doc = Document::from_raw("cf", "aaaaaaaaa");
|
||||
let chunks = splitter.split(&[doc]);
|
||||
assert_eq!(chunks.len(), 2, "expected 2 chunks, got {}", chunks.len());
|
||||
for chunk in &chunks {
|
||||
assert!(chars_len(&chunk.content) <= 5);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_multibyte_utf8_boundary() {
|
||||
// 验证字符级单位而非字节级单位
|
||||
let splitter = RecursiveCharacterSplitter::new(10, 2);
|
||||
// 30 个中文字符 = 90 字节（UTF-8）
|
||||
let text: String = "中".repeat(30);
|
||||
let doc = Document::from_raw("cjk", &text);
|
||||
let chunks = splitter.split(&[doc]);
|
||||
// 30 字符 / 10 chunk_size = 3 个 chunk
|
||||
assert!(chunks.len() >= 3, "expected >= 3 chunks for 30 chars / chunk_size=10, got {}", chunks.len());
|
||||
for chunk in &chunks {
|
||||
let char_count = chars_len(&chunk.content);
|
||||
// chunk = overlap_tail(≤2) + new_content(≤10)，故 ≤ 12
|
||||
assert!(char_count <= 12, "chunk char count {} exceeds 10+overlap", char_count);
|
||||
}
|
||||
}
|
||||
|
||||
// ===== Group B4 — Splitter 集成测试 =====
|
||||
|
||||
#[test]
|
||||
fn split_multiple_docs() {
|
||||
let splitter = RecursiveCharacterSplitter::new(50, 5);
|
||||
let docs = vec![
|
||||
Document::from_raw("a", "a".repeat(30).as_str()),
|
||||
Document::from_raw("b", "b".repeat(30).as_str()),
|
||||
Document::from_raw("c", "c".repeat(30).as_str()),
|
||||
];
|
||||
let chunks = splitter.split(&docs);
|
||||
assert!(chunks.len() >= 3);
|
||||
// 每个 chunk 的 source_id 应指向对应的输入 doc
|
||||
for chunk in &chunks {
|
||||
let source = chunk.metadata.get("source_id").unwrap();
|
||||
assert!(["a", "b", "c"].contains(&source.as_str()));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_metadata_inheritance() {
|
||||
let splitter = RecursiveCharacterSplitter::new(100, 10);
|
||||
let mut doc = Document::new("m", "short content", "text/plain");
|
||||
doc.metadata.insert("author".to_string(), "alice".to_string());
|
||||
let chunks = splitter.split(&[doc]);
|
||||
assert_eq!(chunks.len(), 1);
|
||||
assert_eq!(chunks[0].metadata.get("author").map(|s| s.as_str()), Some("alice"));
|
||||
assert_eq!(chunks[0].metadata.get("source_id").map(|s| s.as_str()), Some("m"));
|
||||
assert_eq!(chunks[0].metadata.get("chunk_index").map(|s| s.as_str()), Some("0"));
|
||||
assert_eq!(chunks[0].metadata.get("chunk_count").map(|s| s.as_str()), Some("1"));
|
||||
}
|
||||
}
|
||||
@@ -1,11 +1,14 @@
|
||||
//! agcore —— 智能体（Agent）核心工具箱。
|
||||
|
||||
pub mod agent;
|
||||
pub mod document;
|
||||
pub mod llm;
|
||||
pub mod memory;
|
||||
pub mod prompt;
|
||||
pub mod tools;
|
||||
|
||||
pub use document::Document;
|
||||
|
||||
use tracing_subscriber::{EnvFilter, fmt, prelude::*};
|
||||
|
||||
static INIT: std::sync::Once = std::sync::Once::new();
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
pub mod compact;
|
||||
pub mod convert;
|
||||
pub mod cycle;
|
||||
pub mod embedding;
|
||||
pub mod error;
|
||||
pub mod hooks;
|
||||
pub mod mock;
|
||||
|
||||
@@ -0,0 +1,183 @@
|
||||
//! Embedding 抽象 —— 文本向量化接口。
|
||||
//!
|
||||
//! 提供 [`Embedding`] trait 和零依赖的 [`MockEmbedding`] 引用实现。
|
||||
//! 上层可实现此 trait 以对接真实 Embedding Provider（OpenAI、Cohere 等）。
|
||||
//!
|
||||
//! 所有实现使用 [`LlmError`] 作为统一错误类型，与 llm 模块保持一致。
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::llm::error::LlmError;
|
||||
|
||||
/// 文本向量化抽象接口。
|
||||
///
|
||||
/// 将文本字符串转换为固定维度的浮点向量，用于语义相似度计算。
|
||||
/// 设计为异步以支持网络 IO（如 OpenAI Embedding API）。
|
||||
///
|
||||
/// 使用 [`LlmError`] 作为统一错误类型，与 llm 模块保持一致。
|
||||
///
|
||||
/// # 实现要求
|
||||
///
|
||||
/// - `embed()` 返回的向量外层的 Vec 长度必须等于输入切片长度（一对一映射）
|
||||
/// - 内层 Vec 长度必须等于 `dim()` 返回值
|
||||
/// - 调用方应保证输入非空（空切片返回空外层 Vec，不报错）
|
||||
///
|
||||
/// # 稳定性
|
||||
///
|
||||
/// 实验性 API（v0.3.x），方法签名可能在 v0.4 中调整。
|
||||
#[async_trait]
|
||||
pub trait Embedding: Send + Sync {
|
||||
/// 批量向量化。
|
||||
///
|
||||
/// 返回 `Vec<Vec<f32>>`，第 i 个内层向量对应 `input[i]`。
|
||||
async fn embed(&self, input: &[String]) -> Result<Vec<Vec<f32>>, LlmError>;
|
||||
|
||||
/// 返回向量维度。
|
||||
fn dim(&self) -> usize;
|
||||
}
|
||||
|
||||
/// 确定性 Mock Embedding —— 零依赖伪随机单位向量。
|
||||
///
|
||||
/// 使用 sin 哈希将输入字符串映射到单位球面上的一个点：
|
||||
/// 1. 对输入字符串计算简单哈希（字符字节和 + 长度）作为种子
|
||||
/// 2. 用 `f32::sin(seed + i) * 10000` 生成第 i 个维度的值
|
||||
/// 3. 归一化到单位长度（L2 norm = 1.0）
|
||||
///
|
||||
/// 特性：
|
||||
/// - **确定性**：相同输入 → 相同向量
|
||||
/// - **有区分度**：不同输入产生不同向量（高概率）
|
||||
/// - **单位范数**：余弦相似度等价于点积
|
||||
/// - **开销极低**：不分配额外内存，无 IO
|
||||
///
|
||||
/// # 已知限制
|
||||
///
|
||||
/// `f32::sin(seed + i) * 10000` 在维度较高时（如 1536，OpenAI Embedding 维度）
|
||||
/// 可能出现周期性模式——相邻维度取值在 `sin` 周期 2π 约束下呈规律性重复。
|
||||
/// MockEmbedding 仅用于测试验证，**不应用于生产级相似度排序**；
|
||||
/// 做严肃验证时建议使用真实 Embedding Provider 或显式随机初始化。
|
||||
pub struct MockEmbedding {
|
||||
dim: usize,
|
||||
}
|
||||
|
||||
impl MockEmbedding {
|
||||
/// 创建 Mock Embedding，输出向量维度为 `dim`。
|
||||
pub fn new(dim: usize) -> Self {
|
||||
Self { dim }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Embedding for MockEmbedding {
|
||||
async fn embed(&self, input: &[String]) -> Result<Vec<Vec<f32>>, LlmError> {
|
||||
let results: Vec<Vec<f32>> = input
|
||||
.iter()
|
||||
.map(|text| {
|
||||
// 简单哈希：字符字节值和 + 文本长度作为种子
|
||||
let seed: f64 = text.bytes().map(|b| b as f64).sum::<f64>() + text.len() as f64;
|
||||
let mut vec: Vec<f32> = (0..self.dim)
|
||||
.map(|i| f32::sin(seed as f32 + i as f32) * 10000.0)
|
||||
.collect();
|
||||
l2_normalize(&mut vec);
|
||||
vec
|
||||
})
|
||||
.collect();
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
fn dim(&self) -> usize {
|
||||
self.dim
|
||||
}
|
||||
}
|
||||
|
||||
/// L2 归一化（in-place）。
|
||||
///
|
||||
/// 零向量（norm == 0）保持全零 —— 防除零保护。
|
||||
fn l2_normalize(vec: &mut [f32]) {
|
||||
let norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if norm > f32::EPSILON {
|
||||
for x in vec.iter_mut() {
|
||||
*x /= norm;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// 计算向量的 L2 范数。
|
||||
fn l2_norm(v: &[f32]) -> f32 {
|
||||
v.iter().map(|x| x * x).sum::<f32>().sqrt()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embed_correct_dim() {
|
||||
let embedder = MockEmbedding::new(8);
|
||||
let inputs = vec!["hello".to_string(), "world".to_string()];
|
||||
let result = embedder.embed(&inputs).await.unwrap();
|
||||
assert_eq!(result.len(), 2);
|
||||
for vec in &result {
|
||||
assert_eq!(vec.len(), 8);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embed_batch_size_match() {
|
||||
let embedder = MockEmbedding::new(4);
|
||||
let inputs = vec![
|
||||
"a".to_string(),
|
||||
"b".to_string(),
|
||||
"c".to_string(),
|
||||
"d".to_string(),
|
||||
"e".to_string(),
|
||||
];
|
||||
let result = embedder.embed(&inputs).await.unwrap();
|
||||
assert_eq!(result.len(), inputs.len());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embed_deterministic() {
|
||||
let embedder = MockEmbedding::new(4);
|
||||
let inputs = vec!["deterministic test".to_string()];
|
||||
let r1 = embedder.embed(&inputs).await.unwrap();
|
||||
let r2 = embedder.embed(&inputs).await.unwrap();
|
||||
assert_eq!(r1, r2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embed_unit_vector_norm() {
|
||||
let embedder = MockEmbedding::new(16);
|
||||
let inputs = vec!["any text".to_string(), "another".to_string()];
|
||||
let result = embedder.embed(&inputs).await.unwrap();
|
||||
for vec in &result {
|
||||
let norm = l2_norm(vec);
|
||||
assert!((norm - 1.0).abs() < 1e-5, "vector norm should be ~1.0, got {}", norm);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embed_different_inputs_different_vectors() {
|
||||
let embedder = MockEmbedding::new(16);
|
||||
let r1 = embedder
|
||||
.embed(&["hello world".to_string()])
|
||||
.await
|
||||
.unwrap();
|
||||
let r2 = embedder
|
||||
.embed(&["completely different".to_string()])
|
||||
.await
|
||||
.unwrap();
|
||||
assert_ne!(r1, r2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embed_empty_string() {
|
||||
// 空字符串输入应不 panic，且向量范数仍≈1.0（防除零路径）
|
||||
let embedder = MockEmbedding::new(4);
|
||||
let inputs = vec!["".to_string()];
|
||||
let result = embedder.embed(&inputs).await.unwrap();
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].len(), 4);
|
||||
let norm = l2_norm(&result[0]);
|
||||
assert!((norm - 1.0).abs() < 1e-5, "empty-string vector norm should be ~1.0, got {}", norm);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user