agcore/examples/knowledge_search_demo.rs

//! knowledge_search_demo —— 知识页面存储与关键词检索。
//! Required features: cargo run --example knowledge_search_demo --features "memory"
//!
//! 演示：
//! 1. `KnowledgeStore` 存储多个 `KnowledgePage`
//! 2. `MemoryRetriever` 按关键词检索 + TextOverlap (Dice) 评分
//! 3. 评分 [0.0, 1.0] 范围校验
//! 4. `RetrieverConfig::min_score` 阈值过滤
//! 5. `RetrieverConfig::max_results` 截断
//! 6. 空 query 返回空结果
//!
//! 运行：`cargo run --example knowledge_search_demo`

use std::sync::Arc;

use agcore::memory::{
    InMemoryStore, KnowledgePage, KnowledgeStore, MemoryRetriever, MemoryStore, RetrieverConfig,
};
use time::OffsetDateTime;

fn make_page(id: &str, title: &str, content: &str) -> KnowledgePage {
    let now = OffsetDateTime::now_utc();
    KnowledgePage {
        id: id.to_string(),
        title: title.to_string(),
        summary: content.chars().take(30).collect(),
        content: content.to_string(),
        tags: Vec::new(),
        references: Vec::new(),
        created_at: now,
        updated_at: now,
    }
}

#[tokio::main]
async fn main() {
    // 1. 创建知识库 + 批量存储页面
    let store: Arc<dyn MemoryStore> = Arc::new(InMemoryStore::new());
    let ks = KnowledgeStore::new(store);

    let pages = vec![
        make_page(
            "rust-1",
            "Rust 入门",
            "Rust 是一门系统级编程语言，注重安全性与并发。",
        ),
        make_page(
            "python-1",
            "Python 简介",
            "Python 是一门动态类型的高级编程语言。",
        ),
        make_page(
            "langgraph-1",
            "LangGraph 框架",
            "LangGraph 是 LangChain 的状态图扩展，用于构建多步 Agent。",
        ),
        make_page(
            "rust-async",
            "Rust 异步编程",
            "Rust 异步基于 tokio 与 futures 抽象。",
        ),
    ];
    for p in &pages {
        ks.add_page(p.clone()).await.expect("保存页面失败");
    }
    println!("=== 已存储 {} 个知识页面 ===", pages.len());
    let index = ks.get_index().await;
    for entry in &index {
        println!("  - {} ({})", entry.title, entry.id);
    }

    // 2. 关键词检索 —— 期望命中 Rust 相关页面
    println!("\n=== 关键词检索：'Rust 异步' ===");
    let retriever = MemoryRetriever::new(ks, RetrieverConfig::default());
    let result = retriever.retrieve("Rust 异步").await.unwrap();
    println!("query: {}", result.query);
    for item in &result.items {
        match item {
            agcore::memory::RetrievalItem::KnowledgePage { page, score } => {
                println!("  命中: {} (score={:.3})", page.title, score);
                assert!((0.0..=1.0).contains(score), "score 应在 [0, 1] 区间");
            }
            agcore::memory::RetrievalItem::GraphEntity { entity, score, .. } => {
                println!("  命中实体: {} (score={:.3})", entity.name, score);
            }
        }
    }
    assert!(!result.items.is_empty(), "应至少命中一个页面");

    // 3. min_score 阈值过滤
    println!("\n=== min_score=0.5 阈值过滤（无关 query）===");
    let store2: Arc<dyn MemoryStore> = Arc::new(InMemoryStore::new());
    let ks2 = KnowledgeStore::new(store2);
    ks2.add_page(make_page("rust-1", "Rust 入门", "Rust 入门内容。"))
        .await
        .unwrap();
    let cfg = RetrieverConfig {
        max_results: 20,
        min_score: 0.5,
        graph_depth: 2,
    };
    let retriever2 = MemoryRetriever::new(ks2, cfg);
    let result = retriever2.retrieve("完全不相关的火锅配方").await.unwrap();
    println!("无关 query → items.len = {} (期望 0)", result.items.len());
    assert!(result.items.is_empty());

    // 4. max_results 截断
    println!("\n=== max_results=2 截断 ===");
    let store3: Arc<dyn MemoryStore> = Arc::new(InMemoryStore::new());
    let ks3 = KnowledgeStore::new(store3);
    for i in 0..5 {
        ks3.add_page(make_page(
            &format!("rust-{i}"),
            "Rust 主题",
            &format!("第 {i} 篇关于 Rust 的内容"),
        ))
        .await
        .unwrap();
    }
    let cfg = RetrieverConfig {
        max_results: 2,
        min_score: 0.0,
        graph_depth: 2,
    };
    let retriever3 = MemoryRetriever::new(ks3, cfg);
    let result = retriever3.retrieve("Rust").await.unwrap();
    println!(
        "5 个相关页面 → 返回 items.len = {} (期望 2)",
        result.items.len()
    );
    assert_eq!(result.items.len(), 2);

    // 5. 空 query
    println!("\n=== 空 query ===");
    let empty = retriever3.retrieve("").await.unwrap();
    println!("空 query → items.len = {}", empty.items.len());
    assert!(empty.items.is_empty());

    // 6. 停用词过滤：`extract_keywords` 在检索前过滤单字符词与停用词
    println!("\n=== 停用词过滤 ===");
    let mixed = retriever3.retrieve("the Rust is").await.unwrap();
    println!(
        "query='the Rust is' → 命中 {} 个 (停用词 'the'/'is' 被过滤，仅 'rust' 进入搜索)",
        mixed.items.len()
    );
    assert!(
        !mixed.items.is_empty(),
        "非停用词 'rust' 应命中页面（即使 query 中含停用词）"
    );

    let only_stop = retriever3.retrieve("the is are").await.unwrap();
    println!(
        "纯停用词 query='the is are' → 命中 {} 个 (期望 0，所有 token 均被过滤)",
        only_stop.items.len()
    );
    assert!(only_stop.items.is_empty(), "纯停用词 query 必须返回空结果");

    println!("\n✓ knowledge_search_demo 完成");
}