use rucora::rag::{chunk_text, index_chunks, index_text, retrieve};
use rucora_core::retrieval::VectorStore;
use rucora_embed::openai::OpenAiEmbeddingProvider;
use rucora_retrieval::in_memory::InMemoryVectorStore;
use tracing::{Level, info};
use tracing_subscriber::FmtSubscriber;
#[tokio::main]
async fn main() -> anyhow::Result<()> {
dotenv::dotenv().ok();
let subscriber = FmtSubscriber::builder()
.with_max_level(Level::INFO)
.with_target(false)
.finish();
tracing::subscriber::set_global_default(subscriber)?;
info!("╔════════════════════════════════════════╗");
info!("║ rucora RAG（检索增强生成）示例 ║");
info!("╚════════════════════════════════════════╝\n");
if std::env::var("OPENAI_API_KEY").is_err() && std::env::var("OPENAI_BASE_URL").is_err() {
info!("⚠ 未设置 API 配置");
info!(" 使用 OpenAI: export OPENAI_API_KEY=sk-your-key");
info!(" 使用 Ollama: export OPENAI_BASE_URL=http://localhost:11434");
info!("\n注意：RAG 需要 Embedding Provider 来向量化文本");
return Ok(());
}
info!("═══════════════════════════════════════");
info!("RAG 工作流程:");
info!("═══════════════════════════════════════");
info!("1. 文档分块（Chunking）: 将长文档分割成小块");
info!("2. 向量化（Embedding）: 将文本转换为向量");
info!("3. 索引（Indexing）: 将向量存储到向量数据库");
info!("4. 检索（Retrieval）: 根据查询找到相关文档");
info!("5. 增强生成（Generation）: 结合检索结果生成回答");
info!("═══════════════════════════════════════\n");
info!("═══════════════════════════════════════");
info!("示例 1: 文档分块");
info!("═══════════════════════════════════════\n");
let long_document = "
Rust 是一门系统编程语言，专注于安全、并发和性能。它由 Mozilla 研究团队开发，
于 2010 年首次发布。Rust 的最大特点是所有权（Ownership）系统，它能够在编译时
保证内存安全，而不需要垃圾回收机制。
Rust 的核心概念包括：
1. 所有权（Ownership）：每个值都有一个所有者，当所有者离开作用域时，值会被丢弃。
2. 借用（Borrowing）：通过引用来访问值，而不获取所有权。
3. 生命周期（Lifetimes）：确保引用在有效期间内使用。
4. 模式匹配（Pattern Matching）：强大的模式匹配能力。
5. 零成本抽象（Zero-cost Abstractions）：高级抽象不带来运行时开销。
Rust 的应用场景非常广泛：
- 系统编程：操作系统、驱动程序、嵌入式系统
- Web 后端：高性能 Web 服务器和 API
- 命令行工具：快速、可靠的 CLI 工具
- 网络服务：分布式系统、微服务架构
- 游戏开发：游戏引擎、图形渲染
- 区块链：智能合约、加密货币
Rust 的生态系统正在快速增长，Cargo 包管理器提供了便捷的依赖管理和构建工具。
crates.io 上有数以万计的开源库，涵盖了各种应用场景。
";
info!("1.1 原始文档长度：{} 字符\n", long_document.len());
info!("1.2 分块处理（每块 200 字符，重叠 30 字符）...");
let chunks = chunk_text("rust_doc", long_document, 200, 30);
info!("✓ 分块完成，共 {} 个块\n", chunks.len());
for (i, chunk) in chunks.iter().enumerate() {
info!(" 块 {} [{}]:", i + 1, chunk.id);
info!(
" 内容：{}...\n",
chunk.text.chars().take(50).collect::<String>()
);
}
info!("═══════════════════════════════════════");
info!("示例 2: 向量化和索引");
info!("═══════════════════════════════════════\n");
info!("2.1 创建向量存储...");
let vector_store = InMemoryVectorStore::new();
info!("✓ 向量存储创建成功\n");
info!("2.2 创建 Embedding Provider...");
let embedder = OpenAiEmbeddingProvider::from_env()?;
info!("✓ Embedding Provider 创建成功\n");
info!("2.3 索引分块...");
index_chunks(&embedder, &vector_store, &chunks).await?;
info!("✓ 索引完成\n");
let count = vector_store.count().await?;
info!(" 向量存储中当前有 {} 个向量\n", count);
info!("═══════════════════════════════════════");
info!("示例 3: 完整 RAG 流程");
info!("═══════════════════════════════════════\n");
info!("3.1 索引新文档...");
let python_doc = "
Python 是一种高级编程语言，由 Guido van Rossum 于 1991 年创建。
Python 的设计哲学强调代码可读性，使用缩进来表示代码块。
Python 的特点：
1. 简洁易读：语法简单，接近自然语言
2. 解释型语言：不需要编译，直接运行
3. 动态类型：变量类型在运行时确定
4. 丰富的库：标准库和第三方库非常庞大
5. 跨平台：可在 Windows、Linux、macOS 上运行
Python 的应用领域：
- 数据科学：数据分析、机器学习、人工智能
- Web 开发：Django、Flask 等框架
- 自动化：脚本编写、任务自动化
- 科学计算：NumPy、SciPy、Pandas
- 网络爬虫：Scrapy、BeautifulSoup
";
let python_chunks =
index_text(&embedder, &vector_store, "python_doc", python_doc, 150, 20).await?;
info!("✓ 已索引 Python 文档，共 {} 个块\n", python_chunks.len());
info!("3.2 检索相关文档...");
let queries = vec![
("Rust 的所有权是什么？", "Rust 相关"),
("Python 适合做什么？", "Python 相关"),
("系统编程语言有哪些特点？", "跨文档检索"),
];
for (query, description) in queries {
info!(" 查询：\"{}\"（{}）", query, description);
let citations = retrieve(&embedder, &vector_store, query, 3).await?;
info!(" 找到 {} 条相关引用:", citations.len());
for (i, citation) in citations.iter().enumerate() {
info!(
" {}. [{}] {} (相似度：{:.3})",
i + 1,
citation.render(),
citation
.text
.as_ref()
.unwrap_or(&String::new())
.chars()
.take(40)
.collect::<String>(),
citation.score
);
}
info!("");
}
info!("═══════════════════════════════════════");
info!("示例 4: RAG 最佳实践");
info!("═══════════════════════════════════════\n");
info!("4.1 分块大小建议:");
info!(" - 小文本（< 1000 字符）: 不需要分块");
info!(" - 中等文本（1000-5000 字符）: 500 字符/块，50 字符重叠");
info!(" - 大文本（> 5000 字符）: 1000 字符/块，100 字符重叠\n");
info!("4.2 TopK 选择建议:");
info!(" - 精确查询：top_k = 3-5");
info!(" - 模糊查询：top_k = 5-10");
info!(" - 探索性查询：top_k = 10-20\n");
info!("4.3 性能优化建议:");
info!(" - 使用缓存减少重复嵌入");
info!(" - 批量嵌入优于单次嵌入");
info!(" - 定期清理向量存储\n");
info!("═══════════════════════════════════════");
info!("示例 5: RAG + LLM 生成回答");
info!("═══════════════════════════════════════\n");
use rucora::agent::SimpleAgent;
use rucora::prelude::Agent;
use rucora::provider::OpenAiProvider;
info!("5.1 创建 LLM Provider...");
let model_name = std::env::var("MODEL_NAME").unwrap_or_else(|_| "gpt-4o-mini".to_string());
let llm_provider = OpenAiProvider::from_env()?;
info!("✓ LLM Provider 创建成功\n");
info!("5.2 创建 Agent...");
let agent = SimpleAgent::builder()
.provider(llm_provider)
.model(&model_name)
.system_prompt(
"你是一个基于检索结果的问答助手。请根据提供的上下文信息回答问题。\n\
如果上下文中没有相关信息，请诚实地说明。\n\
回答时请引用相关的来源。",
)
.build();
info!("✓ Agent 创建成功\n");
let query = "Rust 的所有权系统是如何工作的？";
info!("5.3 查询：\"{}\"", query);
let citations = retrieve(&embedder, &vector_store, query, 3).await?;
let context = citations
.iter()
.map(|c| {
format!(
"[{}]: {}",
c.render(),
c.text.as_ref().unwrap_or(&String::new())
)
})
.collect::<Vec<_>>()
.join("\n\n");
info!(" 检索到 {} 条相关文档", citations.len());
info!(" 上下文长度：{} 字符\n", context.len());
let enhanced_prompt = format!(
"请根据以下上下文信息回答问题。\n\n\
=== 上下文 ===\n\
{context}\n\n\
=== 问题 ===\n\
{query}\n\n\
请提供详细、准确的回答，并引用相关来源。"
);
info!("5.4 生成回答...");
let output = agent.run(enhanced_prompt.into()).await?;
if let Some(text) = output.text() {
info!(" 回答：\n{}\n", text);
}
info!("═══════════════════════════════════════");
info!("示例完成！");
info!("═══════════════════════════════════════\n");
info!("📝 RAG 核心概念总结：\n");
info!("1. 文档分块（Chunking）:");
info!(" - 将长文档分割成适合向量化的小块");
info!(" - 重叠分块保持上下文连续性\n");
info!("2. 向量化（Embedding）:");
info!(" - 将文本转换为数值向量");
info!(" - 语义相似的文本向量距离更近\n");
info!("3. 索引（Indexing）:");
info!(" - 将向量存储到向量数据库");
info!(" - 支持高效的相似度搜索\n");
info!("4. 检索（Retrieval）:");
info!(" - 将查询向量化");
info!(" - 找到最相关的文档块\n");
info!("5. 增强生成（Generation）:");
info!(" - 将检索结果作为上下文");
info!(" - LLM 基于上下文生成准确回答\n");
info!("💡 使用建议:");
info!(" - 选择合适的分块大小（通常 200-1000 字符）");
info!(" - 设置合理的重叠（通常为分块大小的 10-20%）");
info!(" - 根据查询类型调整 top_k");
info!(" - 使用缓存优化性能\n");
Ok(())
}