PicoBot/src/tools/memory_search.rs

use std::sync::Arc;

use async_trait::async_trait;
use serde_json::json;

use crate::storage::{MemoryRecord, MemoryRepository};
use crate::tools::extract_u64;
use crate::tools::traits::{Tool, ToolContext, ToolResult};

pub struct MemorySearchTool {
    memories: Arc<dyn MemoryRepository>,
}

impl MemorySearchTool {
    pub fn new(memories: Arc<dyn MemoryRepository>) -> Self {
        Self { memories }
    }
}

#[async_trait]
impl Tool for MemorySearchTool {
    fn name(&self) -> &str {
        "memory_search"
    }

    fn description(&self) -> &str {
        "Search and read long-term user memories from the configured memory repository. This is the default entry point for memory retrieval and should usually be the first memory tool you call at the start of a request, unless the request is clearly a simple greeting, a one-off calculation, or a direct fact question that does not depend on user history. Use it to recall prior preferences, stable facts, historical decisions, and ongoing task context. If the request also needs other independent read-only tools, you may call memory_search in the same round alongside them. This tool is read-only and supports three actions: search for multi-keyword recall, get for exact namespace/key lookup, and list for browsing recent memories. Prefer this tool over memory_manage whenever you only need to retrieve memory."
    }

    fn parameters_schema(&self) -> serde_json::Value {
        let namespaces = crate::storage::allowed_namespace_names();
        json!({
            "type": "object",
            "properties": {
                "action": {
                    "type": "string",
                    "enum": ["search", "get", "list"],
                    "description": "检索操作。search 用于多关键词召回，get 用于精确 namespace/key 读取，list 用于浏览最近记忆。"
                },
                "namespace": {
                    "type": "string",
                    "enum": namespaces,
                    "description": "可选的命名空间过滤。get 操作时必填。"
                },
                "queries": {
                    "type": "array",
                    "items": {
                        "type": "string"
                    },
                    "description": "搜索关键词数组。建议提供多个简洁的双语关键词、英文别名和可能的 snake_case memory_key。search 操作时必填。",
                    "minItems": 1
                },
                "key": {
                    "type": "string",
                    "description": "命名空间内的记忆键名。get 操作时必填。"
                },
                "limit": {
                    "type": "integer",
                    "description": "返回记忆的最大数量",
                    "minimum": 1,
                    "default": 10
                }
            },
            "required": ["action"]
        })
    }

    async fn execute(&self, _args: serde_json::Value) -> anyhow::Result<ToolResult> {
        Ok(error_result("memory_search requires tool context"))
    }

    async fn execute_with_context(
        &self,
        context: &ToolContext,
        args: serde_json::Value,
    ) -> anyhow::Result<ToolResult> {
        let action = match args.get("action").and_then(|value| value.as_str()) {
            Some(action) => action,
            None => return Ok(error_result("Missing required parameter: action")),
        };

        let scope_key = match scope_key_from_context(context) {
            Ok(scope_key) => scope_key,
            Err(result) => return Ok(result),
        };

        let namespace = args.get("namespace").and_then(|value| value.as_str());
        let key = args.get("key").and_then(|value| value.as_str());

        let payload = match action {
            "list" => {
                let limit = extract_u64(&args, "limit").unwrap_or(10) as usize;
                let memories = self
                    .memories
                    .list_memories("user", &scope_key, namespace, limit)?;
                json!({
                    "count": memories.len(),
                    "memories": memories.into_iter().map(memory_to_json).collect::<Vec<_>>()
                })
            }
            "search" => {
                let queries = match args.get("queries") {
                    Some(value) => {
                        // 支持两种格式：实际数组 或 字符串化的数组
                        if let Some(arr) = value.as_array() {
                            arr.iter()
                                .filter_map(|v| v.as_str())
                                .map(str::trim)
                                .filter(|v| !v.is_empty())
                                .map(ToOwned::to_owned)
                                .collect::<Vec<_>>()
                        } else if let Some(s) = value.as_str() {
                            // 尝试解析字符串化的 JSON 数组
                            match serde_json::from_str::<Vec<serde_json::Value>>(s) {
                                Ok(arr) => arr
                                    .iter()
                                    .filter_map(|v| v.as_str())
                                    .map(str::trim)
                                    .filter(|v| !v.is_empty())
                                    .map(ToOwned::to_owned)
                                    .collect::<Vec<_>>(),
                                Err(_) => {
                                    // 如果不是 JSON 数组，尝试按逗号分割
                                    s.split(',')
                                        .map(str::trim)
                                        .filter(|v| !v.is_empty())
                                        .map(ToOwned::to_owned)
                                        .collect::<Vec<_>>()
                                }
                            }
                        } else {
                            vec![]
                        }
                    }
                    None => vec![],
                };
                if queries.is_empty() {
                    return Ok(error_result("Missing required parameter: queries"));
                }
                let limit = extract_u64(&args, "limit").unwrap_or(10) as usize;
                let memories = self
                    .memories
                    .search_memories_any("user", &scope_key, &queries, namespace, limit)?;
                json!({
                    "queries": queries,
                    "count": memories.len(),
                    "memories": memories.into_iter().map(memory_to_json).collect::<Vec<_>>()
                })
            }
            "get" => {
                let namespace = match namespace {
                    Some(namespace) => namespace,
                    None => return Ok(error_result("Missing required parameter: namespace")),
                };
                let key = match key {
                    Some(key) => key,
                    None => return Ok(error_result("Missing required parameter: key")),
                };

                match self
                    .memories
                    .get_memory("user", &scope_key, namespace, key)?
                {
                    Some(memory) => memory_to_json(memory),
                    None => {
                        return Ok(error_result(&format!(
                            "memory '{}.{}' not found",
                            namespace, key
                        )));
                    }
                }
            }
            _ => return Ok(error_result("Unsupported action")),
        };

        Ok(ToolResult {
            success: true,
            output: serde_json::to_string_pretty(&payload)?,
            error: None,
        })
    }

    fn read_only(&self) -> bool {
        true
    }
}

fn scope_key_from_context(_context: &ToolContext) -> Result<String, ToolResult> {
    Ok(crate::storage::GLOBAL_SCOPE_KEY.to_string())
}

fn memory_to_json(memory: MemoryRecord) -> serde_json::Value {
    json!({
        "id": memory.id,
        "scope_kind": memory.scope_kind,
        "scope_key": memory.scope_key,
        "namespace": memory.namespace,
        "key": memory.memory_key,
        "content": memory.content,
        "source_type": memory.source_type,
        "source_session_id": memory.source_session_id,
        "source_message_id": memory.source_message_id,
        "source_message_seq": memory.source_message_seq,
        "source_channel_name": memory.source_channel_name,
        "source_chat_id": memory.source_chat_id,
        "created_at": memory.created_at,
        "updated_at": memory.updated_at,
    })
}

fn error_result(message: &str) -> ToolResult {
    ToolResult {
        success: false,
        output: String::new(),
        error: Some(message.to_string()),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::storage::SessionStore;

    const TEST_CHANNEL: &str = "test-channel";

    #[tokio::test]
    async fn test_memory_search_search_and_get() {
        let store = Arc::new(SessionStore::in_memory().unwrap());
        store
            .put_memory(&crate::storage::MemoryUpsert {
                scope_kind: "user".to_string(),
                scope_key: crate::storage::GLOBAL_SCOPE_KEY.to_string(),
                namespace: "user".to_string(),
                memory_key: "language".to_string(),
                content: "User prefers Chinese responses".to_string(),
                source_type: "message".to_string(),
                source_session_id: Some(format!("{}:chat-1", TEST_CHANNEL)),
                source_message_id: Some("msg-1".to_string()),
                source_message_seq: Some(1),
                source_channel_name: Some(TEST_CHANNEL.to_string()),
                source_chat_id: Some("chat-1".to_string()),
            })
            .unwrap();

        let tool = MemorySearchTool::new(store);
        let context = ToolContext {
            channel_name: Some(TEST_CHANNEL.to_string()),
            ..ToolContext::default()
        };

        let search = tool
            .execute_with_context(
                &context,
                json!({
                    "action": "search",
                    "queries": ["Chinese", "language"],
                    "limit": 5
                }),
            )
            .await
            .unwrap();
        assert!(search.success);
        assert!(search.output.contains("language"));

        let get = tool
            .execute_with_context(
                &context,
                json!({
                    "action": "get",
                    "namespace": "user",
                    "key": "language"
                }),
            )
            .await
            .unwrap();
        assert!(get.success);
        assert!(get.output.contains("Chinese"));
    }

    #[tokio::test]
    async fn test_memory_search_is_read_only_and_works_with_default_context() {
        let store = Arc::new(SessionStore::in_memory().unwrap());
        let tool = MemorySearchTool::new(store);

        assert!(tool.read_only());

        // scope_key 已全局统一为 "default"，不再依赖 channel_name
        let result = tool
            .execute_with_context(&ToolContext::default(), json!({ "action": "list" }))
            .await
            .unwrap();
        assert!(result.success);
    }

    #[tokio::test]
    async fn test_memory_search_search_requires_queries() {
        let store = Arc::new(SessionStore::in_memory().unwrap());
        let tool = MemorySearchTool::new(store);
        let context = ToolContext {
            channel_name: Some(TEST_CHANNEL.to_string()),
            ..ToolContext::default()
        };

        let result = tool
            .execute_with_context(&context, json!({ "action": "search", "queries": [] }))
            .await
            .unwrap();
        assert!(!result.success);
        assert!(result.error.unwrap().contains("queries"));
    }
}