Initial public release

src/tools/mod.rs

@@ -0,0 +1,106 @@
+//! MCP Tools for OpenBrain
+//!
+//! Provides the core tools for memory storage and retrieval:
+//! - `store`: Store a memory with automatic embedding generation
+//! - `query`: Query memories by semantic similarity
+//! - `purge`: Delete memories by agent_id or time range
+
+pub mod query;
+pub mod store;
+pub mod purge;
+
+use anyhow::Result;
+use serde_json::{json, Value};
+use std::sync::Arc;
+
+use crate::AppState;
+
+/// Get all tool definitions for MCP tools/list
+pub fn get_tool_definitions() -> Vec<Value> {
+    vec![
+        json!({
+            "name": "store",
+            "description": "Store a memory with automatic embedding generation and keyword extraction. The memory will be associated with the agent_id for isolated retrieval.",
+            "inputSchema": {
+                "type": "object",
+                "properties": {
+                    "content": {
+                        "type": "string",
+                        "description": "The text content to store as a memory"
+                    },
+                    "agent_id": {
+                        "type": "string",
+                        "description": "Unique identifier for the agent storing the memory (default: 'default')"
+                    },
+                    "metadata": {
+                        "type": "object",
+                        "description": "Optional metadata to attach to the memory"
+                    }
+                },
+                "required": ["content"]
+            }
+        }),
+        json!({
+            "name": "query",
+            "description": "Query stored memories using semantic similarity search. Returns the most relevant memories based on the query text.",
+            "inputSchema": {
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "string",
+                        "description": "The search query text"
+                    },
+                    "agent_id": {
+                        "type": "string",
+                        "description": "Agent ID to search within (default: 'default')"
+                    },
+                    "limit": {
+                        "type": "integer",
+                        "description": "Maximum number of results to return (default: 10)"
+                    },
+                    "threshold": {
+                        "type": "number",
+                        "description": "Minimum similarity threshold 0.0-1.0 (default: 0.5)"
+                    }
+                },
+                "required": ["query"]
+            }
+        }),
+        json!({
+            "name": "purge",
+            "description": "Delete memories for an agent. Can delete all memories or those before a specific timestamp.",
+            "inputSchema": {
+                "type": "object",
+                "properties": {
+                    "agent_id": {
+                        "type": "string",
+                        "description": "Agent ID whose memories to delete (required)"
+                    },
+                    "before": {
+                        "type": "string",
+                        "description": "Optional ISO8601 timestamp - delete memories created before this time"
+                    },
+                    "confirm": {
+                        "type": "boolean",
+                        "description": "Must be true to confirm deletion"
+                    }
+                },
+                "required": ["agent_id", "confirm"]
+            }
+        })
+    ]
+}
+
+/// Execute a tool by name with given arguments
+pub async fn execute_tool(
+    state: &Arc<AppState>,
+    tool_name: &str,
+    arguments: Value,
+) -> Result<String> {
+    match tool_name {
+        "store" => store::execute(state, arguments).await,
+        "query" => query::execute(state, arguments).await,
+        "purge" => purge::execute(state, arguments).await,
+        _ => anyhow::bail!("Unknown tool: {}", tool_name),
+    }
+}

src/tools/purge.rs

@@ -0,0 +1,79 @@
+//! Purge Tool - Delete memories by agent_id or time range
+
+use anyhow::{bail, Context, Result};
+use chrono::DateTime;
+use serde_json::Value;
+use std::sync::Arc;
+use tracing::{info, warn};
+
+use crate::AppState;
+
+/// Execute the purge tool
+pub async fn execute(state: &Arc<AppState>, arguments: Value) -> Result<String> {
+    // Extract parameters
+    let agent_id = arguments
+        .get("agent_id")
+        .and_then(|v| v.as_str())
+        .context("Missing required parameter: agent_id")?;
+
+    let confirm = arguments
+        .get("confirm")
+        .and_then(|v| v.as_bool())
+        .unwrap_or(false);
+
+    if !confirm {
+        bail!("Purge operation requires 'confirm: true' to proceed");
+    }
+
+    let before = arguments
+        .get("before")
+        .and_then(|v| v.as_str())
+        .map(|s| DateTime::parse_from_rfc3339(s))
+        .transpose()
+        .context("Invalid 'before' timestamp format - use ISO8601/RFC3339")?
+        .map(|dt| dt.with_timezone(&chrono::Utc));
+
+    // Get current count before purge
+    let count_before = state
+        .db
+        .count_memories(agent_id)
+        .await
+        .context("Failed to count memories")?;
+
+    if count_before == 0 {
+        info!("No memories found for agent '{}'", agent_id);
+        return Ok(serde_json::json!({
+            "success": true,
+            "agent_id": agent_id,
+            "deleted": 0,
+            "message": "No memories found to purge"
+        })
+        .to_string());
+    }
+
+    warn!(
+        "Purging memories for agent '{}' (before={:?})",
+        agent_id, before
+    );
+
+    // Execute purge
+    let deleted = state
+        .db
+        .purge_memories(agent_id, before)
+        .await
+        .context("Failed to purge memories")?;
+
+    info!(
+        "Purged {} memories for agent '{}'",
+        deleted, agent_id
+    );
+
+    Ok(serde_json::json!({
+        "success": true,
+        "agent_id": agent_id,
+        "deleted": deleted,
+        "had_before_filter": before.is_some(),
+        "message": format!("Successfully purged {} memories", deleted)
+    })
+    .to_string())
+}

src/tools/query.rs

@@ -0,0 +1,81 @@
+//! Query Tool - Search memories by semantic similarity
+
+use anyhow::{Context, Result, anyhow};
+use serde_json::Value;
+use std::sync::Arc;
+use tracing::info;
+
+use crate::AppState;
+
+/// Execute the query tool
+pub async fn execute(state: &Arc<AppState>, arguments: Value) -> Result<String> {
+    // Get embedding engine, return error if not ready
+    let embedding_engine = state
+        .get_embedding()
+        .await
+        .ok_or_else(|| anyhow!("Embedding engine not ready - service is still initializing"))?;
+
+    // Extract parameters
+    let query_text = arguments
+        .get("query")
+        .and_then(|v| v.as_str())
+        .context("Missing required parameter: query")?;
+
+    let agent_id = arguments
+        .get("agent_id")
+        .and_then(|v| v.as_str())
+        .unwrap_or("default");
+
+    let limit = arguments
+        .get("limit")
+        .and_then(|v| v.as_i64())
+        .unwrap_or(10);
+
+    let threshold = arguments
+        .get("threshold")
+        .and_then(|v| v.as_f64())
+        .unwrap_or(0.5) as f32;
+
+    info!(
+        "Querying memories for agent '{}': '{}' (limit={}, threshold={})",
+        agent_id, query_text, limit, threshold
+    );
+
+    // Generate embedding for query using Arc<EmbeddingEngine>
+    let query_embedding = embedding_engine
+        .embed(query_text)
+        .context("Failed to generate query embedding")?;
+
+    // Search database
+    let matches = state
+        .db
+        .query_memories(agent_id, &query_embedding, limit, threshold)
+        .await
+        .context("Failed to query memories")?;
+
+    info!("Found {} matching memories", matches.len());
+
+    // Format results
+    let results: Vec<Value> = matches
+        .iter()
+        .map(|m| {
+            serde_json::json!({
+                "id": m.record.id.to_string(),
+                "content": m.record.content,
+                "similarity": m.similarity,
+                "keywords": m.record.keywords,
+                "metadata": m.record.metadata,
+                "created_at": m.record.created_at.to_rfc3339()
+            })
+        })
+        .collect();
+
+    Ok(serde_json::json!({
+        "success": true,
+        "agent_id": agent_id,
+        "query": query_text,
+        "count": results.len(),
+        "results": results
+    })
+    .to_string())
+}

src/tools/store.rs

@@ -0,0 +1,66 @@
+//! Store Tool - Store memories with automatic embeddings
+
+use anyhow::{Context, Result, anyhow};
+use serde_json::Value;
+use std::sync::Arc;
+use tracing::info;
+
+use crate::embedding::extract_keywords;
+use crate::AppState;
+
+/// Execute the store tool
+pub async fn execute(state: &Arc<AppState>, arguments: Value) -> Result<String> {
+    // Get embedding engine, return error if not ready
+    let embedding_engine = state
+        .get_embedding()
+        .await
+        .ok_or_else(|| anyhow!("Embedding engine not ready - service is still initializing"))?;
+
+    // Extract parameters
+    let content = arguments
+        .get("content")
+        .and_then(|v| v.as_str())
+        .context("Missing required parameter: content")?;
+
+    let agent_id = arguments
+        .get("agent_id")
+        .and_then(|v| v.as_str())
+        .unwrap_or("default");
+
+    let metadata = arguments
+        .get("metadata")
+        .cloned()
+        .unwrap_or(serde_json::json!({}));
+
+    info!(
+        "Storing memory for agent '{}': {} chars",
+        agent_id,
+        content.len()
+    );
+
+    // Generate embedding using Arc<EmbeddingEngine>
+    let embedding = embedding_engine
+        .embed(content)
+        .context("Failed to generate embedding")?;
+
+    // Extract keywords
+    let keywords = extract_keywords(content, 10);
+
+    // Store in database
+    let id = state
+        .db
+        .store_memory(agent_id, content, &embedding, &keywords, metadata)
+        .await
+        .context("Failed to store memory")?;
+
+    info!("Memory stored with ID: {}", id);
+
+    Ok(serde_json::json!({
+        "success": true,
+        "id": id.to_string(),
+        "agent_id": agent_id,
+        "keywords": keywords,
+        "embedding_dimension": embedding.len()
+    })
+    .to_string())
+}