use crate::engine::{BatchResult, MemexEngine, MetaFilter, StoreItem};
use crate::rag::SearchResult;
use anyhow::Result;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolResult {
pub success: bool,
pub message: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub data: Option<Value>,
}
impl ToolResult {
pub fn ok(message: impl Into<String>) -> Self {
Self {
success: true,
message: message.into(),
data: None,
}
}
pub fn ok_with_data(message: impl Into<String>, data: Value) -> Self {
Self {
success: true,
message: message.into(),
data: Some(data),
}
}
pub fn error(message: impl Into<String>) -> Self {
Self {
success: false,
message: message.into(),
data: None,
}
}
}
pub async fn store_document(
engine: &MemexEngine,
id: String,
text: String,
metadata: Value,
) -> Result<ToolResult> {
match engine.store(&id, &text, metadata).await {
Ok(()) => Ok(ToolResult::ok(format!(
"Document '{}' stored successfully",
id
))),
Err(e) => Ok(ToolResult::error(format!(
"Failed to store document '{}': {}",
id, e
))),
}
}
pub async fn search_documents(
engine: &MemexEngine,
query: String,
limit: usize,
filter: Option<MetaFilter>,
) -> Result<Vec<SearchResult>> {
match filter {
Some(f) => engine.search_filtered(&query, f, limit).await,
None => engine.search(&query, limit).await,
}
}
pub async fn get_document(engine: &MemexEngine, id: String) -> Result<Option<SearchResult>> {
engine.get(&id).await
}
pub async fn delete_document(engine: &MemexEngine, id: String) -> Result<ToolResult> {
match engine.delete(&id).await {
Ok(true) => Ok(ToolResult::ok(format!(
"Document '{}' deleted successfully",
id
))),
Ok(false) => Ok(ToolResult::ok_with_data(
format!("Document '{}' not found", id),
json!({"deleted": false}),
)),
Err(e) => Ok(ToolResult::error(format!(
"Failed to delete document '{}': {}",
id, e
))),
}
}
pub async fn store_documents_batch(
engine: &MemexEngine,
items: Vec<StoreItem>,
) -> Result<BatchResult> {
engine.store_batch(items).await
}
pub async fn delete_documents_by_filter(
engine: &MemexEngine,
filter: MetaFilter,
) -> Result<ToolResult> {
match engine.delete_by_filter(filter.clone()).await {
Ok(count) => Ok(ToolResult::ok_with_data(
format!("Deleted {} documents matching filter", count),
json!({
"deleted_count": count,
"filter": filter,
}),
)),
Err(e) => Ok(ToolResult::error(format!(
"Failed to delete by filter: {}",
e
))),
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDefinition {
pub name: String,
pub description: String,
#[serde(rename = "inputSchema", alias = "input_schema")]
pub input_schema: Value,
}
impl ToolDefinition {
pub fn new(
name: impl Into<String>,
description: impl Into<String>,
input_schema: Value,
) -> Self {
Self {
name: name.into(),
description: description.into(),
input_schema,
}
}
}
pub fn tool_definitions() -> Vec<ToolDefinition> {
crate::mcp_protocol::shared_tools_list_result()["tools"]
.as_array()
.expect("shared_tools_list_result().tools must be an array")
.iter()
.map(|tool| {
ToolDefinition::new(
tool["name"]
.as_str()
.expect("shared MCP tool definition missing name"),
tool["description"]
.as_str()
.expect("shared MCP tool definition missing description"),
tool.get("inputSchema")
.cloned()
.expect("shared MCP tool definition missing inputSchema"),
)
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tool_result_ok() {
let result = ToolResult::ok("Success");
assert!(result.success);
assert_eq!(result.message, "Success");
assert!(result.data.is_none());
}
#[test]
fn test_tool_result_ok_with_data() {
let result = ToolResult::ok_with_data("Success", json!({"count": 42}));
assert!(result.success);
assert_eq!(result.message, "Success");
assert_eq!(result.data.unwrap()["count"], 42);
}
#[test]
fn test_tool_result_error() {
let result = ToolResult::error("Something went wrong");
assert!(!result.success);
assert_eq!(result.message, "Something went wrong");
assert!(result.data.is_none());
}
#[test]
fn test_tool_definitions_count() {
let tools = tool_definitions();
assert_eq!(tools.len(), 12);
}
#[test]
fn test_tool_definitions_names() {
let tools = tool_definitions();
let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
assert!(names.contains(&"health"));
assert!(names.contains(&"rag_index"));
assert!(names.contains(&"memory_upsert"));
assert!(names.contains(&"memory_search"));
assert!(names.contains(&"memory_get"));
assert!(names.contains(&"memory_delete"));
assert!(names.contains(&"memory_purge_namespace"));
assert!(names.contains(&"namespace_create_token"));
assert!(names.contains(&"namespace_revoke_token"));
assert!(names.contains(&"namespace_list_protected"));
assert!(names.contains(&"namespace_security_status"));
assert!(names.contains(&"dive"));
}
#[test]
fn test_tool_definitions_have_required_fields() {
let tools = tool_definitions();
for tool in tools {
assert!(!tool.name.is_empty(), "Tool name should not be empty");
assert!(
!tool.description.is_empty(),
"Tool description should not be empty"
);
assert!(
tool.input_schema.is_object(),
"Input schema should be an object"
);
assert!(
tool.input_schema.get("type").is_some(),
"Input schema should have a type field"
);
assert!(
tool.input_schema.get("properties").is_some(),
"Input schema should have properties"
);
}
}
#[test]
fn test_tool_definitions_match_shared_mcp_contract() {
let serialized = serde_json::to_value(tool_definitions()).unwrap();
assert_eq!(
serialized,
crate::mcp_protocol::shared_tools_list_result()["tools"]
);
}
#[test]
fn test_tool_result_serialization() {
let result = ToolResult::ok_with_data("Success", json!({"id": "doc-1"}));
let json_str = serde_json::to_string(&result).unwrap();
assert!(json_str.contains("\"success\":true"));
assert!(json_str.contains("\"message\":\"Success\""));
assert!(json_str.contains("\"data\""));
}
#[test]
fn test_tool_definition_creation() {
let tool = ToolDefinition::new(
"test_tool",
"A test tool",
json!({
"type": "object",
"properties": {
"input": { "type": "string" }
}
}),
);
assert_eq!(tool.name, "test_tool");
assert_eq!(tool.description, "A test tool");
assert!(tool.input_schema["properties"]["input"].is_object());
}
#[test]
fn test_tool_definition_serializes_with_mcp_field_name() {
let tool = ToolDefinition::new(
"test_tool",
"A test tool",
json!({
"type": "object",
"properties": {}
}),
);
let value = serde_json::to_value(tool).unwrap();
assert!(value.get("inputSchema").is_some());
assert!(value.get("input_schema").is_none());
}
}