Skip to main content

do_memory_mcp/server/tools/embeddings/
search.rs

1//! Search by embedding tool handler for MCP server
2//!
3//! This module provides the tool for searching episodes by embedding
4//! similarity using a pre-computed embedding vector.
5
6use crate::server::MemoryMCPServer;
7use anyhow::Result;
8use serde_json::json;
9use std::sync::Arc;
10use tracing::debug;
11
12impl MemoryMCPServer {
13    /// Execute the search_by_embedding tool
14    ///
15    /// # Arguments
16    ///
17    /// * `input` - Parameters for embedding search
18    ///
19    /// # Returns
20    ///
21    /// Returns episodes similar to the provided embedding vector
22    pub async fn execute_search_by_embedding(
23        &self,
24        input: crate::mcp::tools::embeddings::SearchByEmbeddingInput,
25    ) -> Result<serde_json::Value> {
26        self.track_tool_usage("search_by_embedding").await;
27
28        // Start monitoring request
29        let request_id = format!(
30            "search_by_embedding_{}",
31            std::time::SystemTime::now()
32                .duration_since(std::time::UNIX_EPOCH)
33                .unwrap_or_default()
34                .as_nanos()
35        );
36        self.monitoring
37            .start_request(request_id.clone(), "search_by_embedding".to_string())
38            .await;
39
40        debug!(
41            "Searching by embedding (dimension: {}, limit: {}, threshold: {})",
42            input.embedding.len(),
43            input.limit,
44            input.similarity_threshold
45        );
46
47        let tool = crate::mcp::tools::embeddings::EmbeddingTools::new(Arc::clone(&self.memory));
48
49        let result = tool.execute_search_by_embedding(input).await;
50
51        // End monitoring request
52        self.monitoring
53            .end_request(&request_id, result.is_ok(), None)
54            .await;
55
56        let output = result?;
57
58        // Convert result to JSON
59        Ok(json!(output))
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    #[test]
68    #[allow(clippy::manual_async_fn)]
69    fn test_search_by_embedding_signature_compile() {
70        // This test ensures the method signature compiles correctly
71        use crate::mcp::tools::embeddings::SearchByEmbeddingInput;
72        fn method_signature(
73            _server: &MemoryMCPServer,
74            _input: SearchByEmbeddingInput,
75        ) -> impl std::future::Future<Output = Result<serde_json::Value>> {
76            async { Ok(json!({})) }
77        }
78        let _ = method_signature; // Use the function to avoid unused warnings
79    }
80}