Skip to main content

do_memory_mcp/mcp/tools/embeddings/tool/execute/
status.rs

1//! Test embeddings and provider status tool implementations.
2
3use super::super::definitions::EmbeddingTools;
4use crate::mcp::tools::embeddings::types::{
5    EmbeddingProviderStatusInput, EmbeddingProviderStatusOutput, ProviderTestResult,
6    TestEmbeddingsOutput,
7};
8use anyhow::Result;
9use tracing::{debug, info, instrument};
10
11impl EmbeddingTools {
12    /// Execute the test_embeddings tool
13    #[instrument(skip(self))]
14    pub async fn execute_test_embeddings(&self) -> Result<TestEmbeddingsOutput> {
15        let start_time = std::time::Instant::now();
16
17        info!("Testing embedding provider connectivity");
18
19        // Check if semantic_service is configured
20        if let Some(semantic_service) = self.memory.semantic_service() {
21            // Call provider.embed_text("test") to validate
22            match semantic_service.provider.embed_text("test").await {
23                Ok(test_embedding) => {
24                    let test_time_ms = start_time.elapsed().as_millis() as u64;
25
26                    debug!("Embedding provider test completed in {}ms", test_time_ms);
27
28                    let config = semantic_service.config();
29                    let model_name = config.provider.model_name();
30                    let dimension = config.provider.effective_dimension();
31                    let provider = format!("{:?}", config.provider);
32                    let embedding_len = test_embedding.len();
33                    let model_name_for_msg = model_name.clone();
34
35                    return Ok(TestEmbeddingsOutput {
36                        available: true,
37                        provider: provider.clone(),
38                        model: model_name,
39                        dimension,
40                        test_time_ms,
41                        sample_embedding: test_embedding,
42                        message: format!(
43                            "Successfully tested {} provider with model {} (dimension: {}, embedding size: {})",
44                            provider, model_name_for_msg, dimension, embedding_len
45                        ),
46                        errors: vec![],
47                    });
48                }
49                Err(e) => {
50                    let test_time_ms = start_time.elapsed().as_millis() as u64;
51                    let config = semantic_service.config();
52                    let provider = format!("{:?}", config.provider);
53                    return Ok(TestEmbeddingsOutput {
54                        available: false,
55                        provider: provider.clone(),
56                        model: config.provider.model_name(),
57                        dimension: config.provider.effective_dimension(),
58                        test_time_ms,
59                        sample_embedding: vec![],
60                        message: format!("Embedding provider test failed: {}", e),
61                        errors: vec![format!("Failed to generate test embedding: {}", e)],
62                    });
63                }
64            }
65        }
66
67        // No semantic service configured
68        let test_time_ms = start_time.elapsed().as_millis() as u64;
69        Ok(TestEmbeddingsOutput {
70            available: false,
71            provider: "not-configured".to_string(),
72            model: "none".to_string(),
73            dimension: 384,
74            test_time_ms,
75            sample_embedding: vec![],
76            message: "Semantic service not yet configured. Use configure_embeddings first."
77                .to_string(),
78            errors: vec!["Semantic embeddings feature requires configuration".to_string()],
79        })
80    }
81
82    /// Execute the embedding_provider_status tool
83    #[instrument(skip(self))]
84    pub async fn execute_embedding_provider_status(
85        &self,
86        input: EmbeddingProviderStatusInput,
87    ) -> Result<EmbeddingProviderStatusOutput> {
88        info!("Getting embedding provider status");
89
90        let mut warnings = Vec::new();
91
92        // Check if semantic_service is configured
93        if let Some(semantic_service) = self.memory.semantic_service() {
94            let config = semantic_service.config();
95            let model_name = config.provider.model_name();
96            let dimension = config.provider.effective_dimension();
97            let provider = format!("{:?}", config.provider);
98            let similarity_threshold = config.similarity_threshold;
99            let batch_size = config.batch_size;
100            let cache_enabled = config.cache_embeddings;
101
102            // Get provider metadata
103            let metadata = semantic_service.provider.metadata();
104
105            // Perform connectivity test if requested
106            let test_result = if input.test_connectivity {
107                let start_time = std::time::Instant::now();
108                match semantic_service.provider.embed_text("test").await {
109                    Ok(embedding) => {
110                        let duration_ms = start_time.elapsed().as_millis() as u64;
111                        // Get first 5 values as sample
112                        let sample_embedding: Vec<f32> = embedding.into_iter().take(5).collect();
113
114                        Some(ProviderTestResult {
115                            success: true,
116                            duration_ms,
117                            sample_embedding,
118                            error: None,
119                        })
120                    }
121                    Err(e) => {
122                        let duration_ms = start_time.elapsed().as_millis() as u64;
123                        warnings.push(format!("Connectivity test failed: {}", e));
124                        Some(ProviderTestResult {
125                            success: false,
126                            duration_ms,
127                            sample_embedding: vec![],
128                            error: Some(e.to_string()),
129                        })
130                    }
131                }
132            } else {
133                None
134            };
135
136            // Check for potential configuration issues
137            if similarity_threshold < 0.5 {
138                warnings.push(format!(
139                    "Low similarity threshold ({}) may return many irrelevant results",
140                    similarity_threshold
141                ));
142            }
143            if similarity_threshold > 0.95 {
144                warnings.push(format!(
145                    "High similarity threshold ({}) may return very few results",
146                    similarity_threshold
147                ));
148            }
149            if batch_size > 100 {
150                warnings.push(format!(
151                    "Large batch size ({}) may cause timeout issues",
152                    batch_size
153                ));
154            }
155
156            // Determine availability based on test result or previous success
157            let available = test_result.as_ref().map(|t| t.success).unwrap_or_else(|| {
158                // If no test was requested, assume available since it was configured
159                true
160            });
161
162            return Ok(EmbeddingProviderStatusOutput {
163                configured: true,
164                available,
165                provider,
166                model: model_name,
167                dimension,
168                similarity_threshold,
169                batch_size,
170                cache_enabled,
171                metadata,
172                test_result,
173                warnings,
174            });
175        }
176
177        // No semantic service configured
178        Ok(EmbeddingProviderStatusOutput {
179            configured: false,
180            available: false,
181            provider: "not-configured".to_string(),
182            model: "none".to_string(),
183            dimension: 384,
184            similarity_threshold: 0.7,
185            batch_size: 32,
186            cache_enabled: false,
187            metadata: serde_json::json!({"status": "not_configured"}),
188            test_result: None,
189            warnings: vec![
190                "Semantic embeddings not configured. Use configure_embeddings to enable embedding features.".to_string()
191            ],
192        })
193    }
194}