agcore/src/memory/knowledge.rs

//! 知识库 —— 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<PageIndexEntry>` 索引以加速列表遍历。
pub struct KnowledgeStore {
    store: Arc<dyn MemoryStore>,
    index: tokio::sync::Mutex<Vec<PageIndexEntry>>,
}

impl KnowledgeStore {
    /// 创建一个新的 KnowledgeStore。
    pub fn new(store: Arc<dyn MemoryStore>) -> 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<Option<KnowledgePage>, 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<Vec<KnowledgePage>, MemoryError> {
        if query.is_empty() {
            return Ok(Vec::new());
        }
        let needle = query.to_lowercase();
        // 锁内仅 clone 匹配的 entry id，避免异步 get_page() 持有 index 锁。
        let ids: Vec<String> = {
            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<PageIndexEntry> {
        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<dyn MemoryStore>;
        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<dyn MemoryStore>;
        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<dyn MemoryStore>;
        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<dyn MemoryStore>;
        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<dyn MemoryStore>;
        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<dyn MemoryStore>;
        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);
    }
}