//! 知识库 —— KnowledgePage 存储与关键词检索。 use std::sync::Arc; use time::OffsetDateTime; use crate::memory::error::MemoryError; use crate::memory::store::MemoryStore; use crate::memory::types::{KnowledgePage, MemoryFilter, MemoryItem}; pub use crate::memory::types::PageIndexEntry; /// `MemoryItem.id` 中知识页面前缀。 pub const KNOWLEDGE_PREFIX: &str = "knowledge_"; /// 知识库 —— KnowledgePage CRUD + 关键词检索 + 内容索引。 /// /// 内部以 `MemoryStore` 为后端存储 KnowledgePage(序列化为 JSON), /// 同时维护一个 `Vec` 索引以加速列表遍历。 pub struct KnowledgeStore { store: Arc, index: tokio::sync::Mutex>, } impl KnowledgeStore { /// 创建一个新的 KnowledgeStore。 pub fn new(store: Arc) -> Self { Self { store, index: tokio::sync::Mutex::new(Vec::new()), } } /// 从 MemoryStore 重建索引(修复 index 与 store 的不同步问题)。 pub async fn rebuild_index(&self) -> Result<(), MemoryError> { let items = self .store .list(&MemoryFilter { prefix: Some(KNOWLEDGE_PREFIX.to_string()), ..Default::default() }) .await?; let mut index = self.index.lock().await; index.clear(); for item in items { let page: KnowledgePage = serde_json::from_str(&item.content) .map_err(|e| MemoryError::Serialization(e.to_string()))?; index.push(PageIndexEntry::from(&page)); } Ok(()) } /// 创建一个新的知识页面。 pub async fn add_page(&self, page: KnowledgePage) -> Result<(), MemoryError> { if page.id.is_empty() { return Err(MemoryError::InvalidInput("page.id is empty".into())); } let now = OffsetDateTime::now_utc(); let id = format!("{KNOWLEDGE_PREFIX}{}", page.id); let content = serde_json::to_string(&page) .map_err(|e| MemoryError::Serialization(e.to_string()))?; let item = MemoryItem { id, content, metadata: serde_json::json!({}), created_at: now, }; self.store.save(item).await?; let mut index = self.index.lock().await; // 替换或追加 if let Some(existing) = index.iter_mut().find(|e| e.id == page.id) { *existing = PageIndexEntry::from(&page); } else { index.push(PageIndexEntry::from(&page)); } Ok(()) } /// 根据 page id 获取一个页面。 pub async fn get_page(&self, id: &str) -> Result, MemoryError> { let full_id = format!("{KNOWLEDGE_PREFIX}{id}"); let item = self.store.get(&full_id).await?; match item { None => Ok(None), Some(item) => { let page: KnowledgePage = serde_json::from_str(&item.content) .map_err(|e| MemoryError::Serialization(e.to_string()))?; Ok(Some(page)) } } } /// 更新一个已存在的知识页面。 pub async fn update_page(&self, page: KnowledgePage) -> Result<(), MemoryError> { if page.id.is_empty() { return Err(MemoryError::InvalidInput("page.id is empty".into())); } // 通过 get_page 检查存在性 if self.get_page(&page.id).await?.is_none() { return Err(MemoryError::NotFound(page.id)); } self.add_page(page).await } /// 删除一个知识页面。 pub async fn delete_page(&self, id: &str) -> Result<(), MemoryError> { let full_id = format!("{KNOWLEDGE_PREFIX}{id}"); self.store.delete(&full_id).await?; let mut index = self.index.lock().await; index.retain(|e| e.id != id); Ok(()) } /// 根据关键词搜索知识页面。 /// /// 匹配规则:在 `title` / `summary` / `tags` 中查找子串(不区分大小写)。 /// 全文 `content` 搜索走 `MemoryStore`。 pub async fn search(&self, query: &str) -> Result, MemoryError> { if query.is_empty() { return Ok(Vec::new()); } let needle = query.to_lowercase(); // 锁内仅 clone 匹配的 entry id,避免异步 get_page() 持有 index 锁。 let ids: Vec = { let index = self.index.lock().await; index .iter() .filter(|entry| { entry.title.to_lowercase().contains(&needle) || entry.summary.to_lowercase().contains(&needle) || entry.tags.iter().any(|t| t.to_lowercase().contains(&needle)) }) .map(|entry| entry.id.clone()) .collect() }; let mut results = Vec::new(); for id in &ids { if let Some(page) = self.get_page(id).await? { results.push(page); } } Ok(results) } /// 获取内容目录(所有页面的轻量级索引条目)。 pub async fn get_index(&self) -> Vec { self.index.lock().await.clone() } } #[cfg(test)] mod tests { use super::*; use crate::memory::InMemoryStore; use time::OffsetDateTime; fn make_page(id: &str, title: &str, tags: &[&str]) -> KnowledgePage { let now = OffsetDateTime::now_utc(); KnowledgePage { id: id.to_string(), title: title.to_string(), summary: format!("summary of {title}"), content: format!("full content of {title}"), tags: tags.iter().map(|s| s.to_string()).collect(), references: Vec::new(), created_at: now, updated_at: now, } } #[tokio::test] async fn add_get_delete_page() { let store = Arc::new(InMemoryStore::new()) as Arc; let ks = KnowledgeStore::new(store); ks.add_page(make_page("p1", "LangGraph", &["langgraph", "framework"])) .await .unwrap(); let got = ks.get_page("p1").await.unwrap(); assert!(got.is_some()); assert_eq!(got.unwrap().title, "LangGraph"); ks.delete_page("p1").await.unwrap(); assert!(ks.get_page("p1").await.unwrap().is_none()); } #[tokio::test] async fn add_page_rejects_empty_id() { let store = Arc::new(InMemoryStore::new()) as Arc; let ks = KnowledgeStore::new(store); let result = ks.add_page(make_page("", "NoId", &[])).await; assert!(matches!(result, Err(MemoryError::InvalidInput(_)))); } #[tokio::test] async fn update_page_requires_existing() { let store = Arc::new(InMemoryStore::new()) as Arc; let ks = KnowledgeStore::new(store); let result = ks.update_page(make_page("nope", "Ghost", &[])).await; assert!(matches!(result, Err(MemoryError::NotFound(_)))); } #[tokio::test] async fn search_finds_by_title_summary_tag() { let store = Arc::new(InMemoryStore::new()) as Arc; let ks = KnowledgeStore::new(store); ks.add_page(make_page("p1", "LangGraph StateGraph", &["llm"])) .await .unwrap(); ks.add_page(make_page("p2", "Other", &["knowledge-graph"])) .await .unwrap(); ks.add_page(make_page("p3", "Third", &["unrelated"])) .await .unwrap(); let results = ks.search("stategraph").await.unwrap(); assert_eq!(results.len(), 1); assert_eq!(results[0].id, "p1"); let results = ks.search("knowledge-graph").await.unwrap(); assert_eq!(results.len(), 1); assert_eq!(results[0].id, "p2"); let results = ks.search("nonexistent").await.unwrap(); assert!(results.is_empty()); } #[tokio::test] async fn get_index_returns_all_pages() { let store = Arc::new(InMemoryStore::new()) as Arc; let ks = KnowledgeStore::new(store); ks.add_page(make_page("p1", "A", &[])).await.unwrap(); ks.add_page(make_page("p2", "B", &[])).await.unwrap(); let index = ks.get_index().await; assert_eq!(index.len(), 2); } #[tokio::test] async fn rebuild_index_recovers_from_drift() { let store = Arc::new(InMemoryStore::new()) as Arc; let ks = KnowledgeStore::new(store); // 添加页面 ks.add_page(make_page("p1", "A", &[])).await.unwrap(); ks.add_page(make_page("p2", "B", &[])).await.unwrap(); assert_eq!(ks.get_index().await.len(), 2); // 模拟 index 漂移:清空后重建 ks.index.lock().await.clear(); assert_eq!(ks.get_index().await.len(), 0); ks.rebuild_index().await.unwrap(); assert_eq!(ks.get_index().await.len(), 2); } }