reddb_server/storage/query/executors/
natural.rs

1//! Natural Language Query Executor
2//!
3//! Translates natural language queries to RQL and executes them,
4//! providing explanations of the translation process.
5//!
6//! # Features
7//!
8//! - Intent classification: list, find, path, count, show
9//! - Entity extraction: hosts, users, credentials, services, vulnerabilities
10//! - Query generation with confidence scoring
11//! - Execution explanation for user understanding
12
13use std::sync::Arc;
14
15use crate::storage::engine::graph_store::GraphStore;
16use crate::storage::query::unified::{
17    ExecutionError, MatchedNode, QueryStats, UnifiedRecord, UnifiedResult,
18};
19use reddb_rql::modes::natural::{
20    EntityType, ExtractedEntity, NaturalParser, NaturalQuery, QueryIntent,
21};
22
23/// Natural language executor with translation explanation
24pub struct NaturalExecutor {
25    graph: Arc<GraphStore>,
26}
27
28impl NaturalExecutor {
29    /// Create a new natural language executor
30    pub fn new(graph: Arc<GraphStore>) -> Self {
31        Self { graph }
32    }
33
34    /// Execute a natural language query and return explanation
35    pub fn execute_with_explanation(
36        &self,
37        query: &str,
38    ) -> Result<(UnifiedResult, String), ExecutionError> {
39        // Parse natural language
40        let parsed = NaturalParser::parse(query).map_err(|e| ExecutionError::new(e.to_string()))?;
41
42        // Generate explanation
43        let explanation = self.generate_explanation(&parsed, query);
44
45        // Execute
46        let result = self.execute_natural(&parsed)?;
47
48        Ok((result, explanation))
49    }
50
51    /// Execute a natural language query
52    pub fn execute(&self, query: &str) -> Result<UnifiedResult, ExecutionError> {
53        let parsed = NaturalParser::parse(query).map_err(|e| ExecutionError::new(e.to_string()))?;
54        self.execute_natural(&parsed)
55    }
56
57    /// Execute a parsed natural language query
58    fn execute_natural(&self, query: &NaturalQuery) -> Result<UnifiedResult, ExecutionError> {
59        let mut stats = QueryStats::default();
60        let mut result = UnifiedResult::empty();
61
62        match query.intent {
63            QueryIntent::Find => {
64                // Find handles both "find" and "list" semantics
65                self.execute_find(query, &mut result, &mut stats)?;
66            }
67            QueryIntent::Path => {
68                self.execute_path(query, &mut result, &mut stats)?;
69            }
70            QueryIntent::Count => {
71                self.execute_count(query, &mut result, &mut stats)?;
72            }
73            QueryIntent::Show => {
74                self.execute_show(query, &mut result, &mut stats)?;
75            }
76            QueryIntent::Check => {
77                self.execute_check(query, &mut result, &mut stats)?;
78            }
79        }
80
81        result.stats = stats;
82        Ok(result)
83    }
84
85    /// Execute FIND intent (also handles LIST semantics)
86    fn execute_find(
87        &self,
88        query: &NaturalQuery,
89        result: &mut UnifiedResult,
90        stats: &mut QueryStats,
91    ) -> Result<(), ExecutionError> {
92        let entity_label = self.primary_entity_label(query);
93
94        for node in self.graph.iter_nodes() {
95            stats.nodes_scanned += 1;
96
97            // Check label match — compare against the legacy enum's string form.
98            let type_matches = match entity_label {
99                Some(label) => node.node_type.as_str() == label,
100                None => true,
101            };
102
103            if !type_matches {
104                continue;
105            }
106
107            // Check entity filters
108            if !self.node_matches_filters(&node, &query.entities) {
109                continue;
110            }
111
112            // Check relationship constraints
113            let mut rel_match = true;
114            for entity in &query.entities {
115                if let Some(ref value) = entity.value {
116                    // Check if node has relationship to this entity
117                    if !self.has_relationship_to(&node.id, value, stats) {
118                        rel_match = false;
119                        break;
120                    }
121                }
122            }
123
124            if rel_match {
125                let mut record = UnifiedRecord::new();
126                record.set_node("_", MatchedNode::from_stored(&node));
127                result.push(record);
128            }
129        }
130
131        // Apply limit if specified
132        if let Some(limit) = query.limit {
133            if result.len() > limit as usize {
134                result.records.truncate(limit as usize);
135            }
136        }
137
138        Ok(())
139    }
140
141    /// Execute PATH intent
142    fn execute_path(
143        &self,
144        query: &NaturalQuery,
145        result: &mut UnifiedResult,
146        stats: &mut QueryStats,
147    ) -> Result<(), ExecutionError> {
148        // Extract source and target from entities
149        let (source, target) = self.extract_path_endpoints(query)?;
150
151        // BFS to find path
152        use crate::storage::query::unified::{GraphPath, MatchedEdge};
153        use std::collections::{HashSet, VecDeque};
154
155        let mut queue: VecDeque<(String, GraphPath)> = VecDeque::new();
156        let mut visited: HashSet<String> = HashSet::new();
157
158        queue.push_back((source.clone(), GraphPath::start(&source)));
159        visited.insert(source.clone());
160
161        let max_hops = query.limit.unwrap_or(10) as usize;
162
163        while let Some((current, path)) = queue.pop_front() {
164            if path.len() > max_hops {
165                continue;
166            }
167
168            if current == target {
169                let mut record = UnifiedRecord::new();
170                record.paths.push(path);
171                result.push(record);
172                break; // Found shortest path
173            }
174
175            for (edge_type, neighbor, weight) in self.graph.outgoing_edges(&current) {
176                stats.edges_scanned += 1;
177
178                if !visited.contains(&neighbor) {
179                    visited.insert(neighbor.clone());
180                    let edge = MatchedEdge::from_tuple(&current, edge_type, &neighbor, weight);
181                    let new_path = path.extend(edge, &neighbor);
182                    queue.push_back((neighbor, new_path));
183                }
184            }
185        }
186
187        if result.is_empty() {
188            return Err(ExecutionError::new(format!(
189                "No path found from {} to {}",
190                source, target
191            )));
192        }
193
194        Ok(())
195    }
196
197    /// Execute COUNT intent
198    fn execute_count(
199        &self,
200        query: &NaturalQuery,
201        result: &mut UnifiedResult,
202        stats: &mut QueryStats,
203    ) -> Result<(), ExecutionError> {
204        let entity_label = self.primary_entity_label(query);
205        let mut count = 0u64;
206
207        for node in self.graph.iter_nodes() {
208            stats.nodes_scanned += 1;
209
210            let type_matches = match entity_label {
211                Some(label) => node.node_type.as_str() == label,
212                None => true,
213            };
214
215            if type_matches && self.node_matches_filters(&node, &query.entities) {
216                count += 1;
217            }
218        }
219
220        let mut record = UnifiedRecord::new();
221        record.set("count", reddb_types::Value::Integer(count as i64));
222        result.push(record);
223        result.columns.push("count".to_string());
224
225        Ok(())
226    }
227
228    /// Execute SHOW intent
229    fn execute_show(
230        &self,
231        query: &NaturalQuery,
232        result: &mut UnifiedResult,
233        stats: &mut QueryStats,
234    ) -> Result<(), ExecutionError> {
235        // SHOW is like FIND but includes more details
236        self.execute_find(query, result, stats)?;
237
238        // Add neighbors for context
239        if result.len() == 1 {
240            if let Some(node) = result.records.first().and_then(|r| r.nodes.get("_")) {
241                // Add outgoing connections
242                for (edge_type, target, _) in self.graph.outgoing_edges(&node.id) {
243                    stats.edges_scanned += 1;
244                    if let Some(target_node) = self.graph.get_node(&target) {
245                        let mut record = UnifiedRecord::new();
246                        record.set_node("related", MatchedNode::from_stored(&target_node));
247                        record.set(
248                            "relationship",
249                            reddb_types::Value::text(format!("{:?}", edge_type)),
250                        );
251                        result.push(record);
252                    }
253                }
254            }
255        }
256
257        Ok(())
258    }
259
260    /// Execute CHECK intent - verify if a relationship exists
261    fn execute_check(
262        &self,
263        query: &NaturalQuery,
264        result: &mut UnifiedResult,
265        stats: &mut QueryStats,
266    ) -> Result<(), ExecutionError> {
267        // Check requires two entities with a relationship
268        let (source, target) = self.extract_path_endpoints(query)?;
269
270        // Check if direct connection exists
271        let mut found = false;
272        for (edge_type, neighbor, weight) in self.graph.outgoing_edges(&source) {
273            stats.edges_scanned += 1;
274            if neighbor == target || neighbor.contains(&target) {
275                found = true;
276                // Add the relationship to result
277                let mut record = UnifiedRecord::new();
278                if let Some(src_node) = self.graph.get_node(&source) {
279                    record.set_node("source", MatchedNode::from_stored(&src_node));
280                }
281                if let Some(tgt_node) = self.graph.get_node(&neighbor) {
282                    record.set_node("target", MatchedNode::from_stored(&tgt_node));
283                }
284                record.set(
285                    "relationship",
286                    reddb_types::Value::text(format!("{:?}", edge_type)),
287                );
288                record.set("exists", reddb_types::Value::Boolean(true));
289                record.set("weight", reddb_types::Value::Float(weight as f64));
290                result.push(record);
291                break;
292            }
293        }
294
295        if !found {
296            // Report that no relationship was found
297            let mut record = UnifiedRecord::new();
298            record.set("exists", reddb_types::Value::Boolean(false));
299            record.set("source", reddb_types::Value::text(source));
300            record.set("target", reddb_types::Value::text(target));
301            result.push(record);
302        }
303
304        result.columns = vec![
305            "source".into(),
306            "target".into(),
307            "relationship".into(),
308            "exists".into(),
309        ];
310        Ok(())
311    }
312
313    /// Get the primary entity label from query — canonical lower-snake-case
314    /// label string corresponding to the first matched entity.
315    fn primary_entity_label(&self, query: &NaturalQuery) -> Option<&'static str> {
316        for entity in &query.entities {
317            match entity.entity_type {
318                EntityType::Host => return Some("host"),
319                EntityType::User => return Some("user"),
320                EntityType::Credential => return Some("credential"),
321                EntityType::Service | EntityType::Port => return Some("service"),
322                EntityType::Vulnerability => return Some("vulnerability"),
323                EntityType::Technology => return Some("technology"),
324                EntityType::Domain => return Some("domain"),
325                EntityType::Certificate => return Some("certificate"),
326                // Network has no canonical legacy label.
327                EntityType::Network => continue,
328            }
329        }
330        None
331    }
332
333    /// Check if node matches entity filters
334    fn node_matches_filters(
335        &self,
336        node: &crate::storage::engine::graph_store::StoredNode,
337        entities: &[ExtractedEntity],
338    ) -> bool {
339        for entity in entities {
340            if let Some(ref value) = entity.value {
341                // Check if value matches node ID or label
342                let matches = node.id.contains(value)
343                    || node.label.to_lowercase().contains(&value.to_lowercase())
344                    || value.to_lowercase().contains(&node.label.to_lowercase());
345                if matches {
346                    return true;
347                }
348            }
349        }
350        // If no values to match, accept all
351        entities.iter().all(|e| e.value.is_none())
352    }
353
354    /// Check if node has relationship to target
355    fn has_relationship_to(&self, node_id: &str, target: &str, stats: &mut QueryStats) -> bool {
356        for (_, neighbor, _) in self.graph.outgoing_edges(node_id) {
357            stats.edges_scanned += 1;
358            if neighbor.contains(target) {
359                return true;
360            }
361            // Check neighbor's label
362            if let Some(neighbor_node) = self.graph.get_node(&neighbor) {
363                if neighbor_node
364                    .label
365                    .to_lowercase()
366                    .contains(&target.to_lowercase())
367                {
368                    return true;
369                }
370            }
371        }
372        false
373    }
374
375    /// Extract path endpoints from query
376    fn extract_path_endpoints(
377        &self,
378        query: &NaturalQuery,
379    ) -> Result<(String, String), ExecutionError> {
380        // Look for "from X to Y" pattern in entities
381        let mut source = None;
382        let mut target = None;
383
384        for entity in &query.entities {
385            if let Some(ref value) = entity.value {
386                // Find nodes matching this value
387                for node in self.graph.iter_nodes() {
388                    if node.id.contains(value)
389                        || node.label.to_lowercase().contains(&value.to_lowercase())
390                    {
391                        if source.is_none() {
392                            source = Some(node.id.clone());
393                        } else if target.is_none() && Some(&node.id) != source.as_ref() {
394                            target = Some(node.id.clone());
395                        }
396                    }
397                }
398            }
399        }
400
401        match (source, target) {
402            (Some(s), Some(t)) => Ok((s, t)),
403            (Some(s), None) => Err(ExecutionError::new(format!(
404                "Path query needs a target. Found source: {}",
405                s
406            ))),
407            _ => Err(ExecutionError::new(
408                "Path query needs source and target. Try: 'path from host X to host Y'",
409            )),
410        }
411    }
412
413    /// Generate explanation of query translation
414    fn generate_explanation(&self, query: &NaturalQuery, original: &str) -> String {
415        let mut explanation = Vec::new();
416
417        explanation.push(format!("Query: \"{}\"", original));
418        explanation.push(format!("Intent: {:?}", query.intent));
419
420        if !query.entities.is_empty() {
421            let entities: Vec<String> = query
422                .entities
423                .iter()
424                .map(|e| {
425                    if let Some(ref val) = e.value {
426                        format!("{:?}({})", e.entity_type, val)
427                    } else {
428                        format!("{:?}", e.entity_type)
429                    }
430                })
431                .collect();
432            explanation.push(format!("Entities: {}", entities.join(", ")));
433        }
434
435        // Generate equivalent RQL
436        let rql = self.to_rql(query);
437        explanation.push(format!("Equivalent RQL: {}", rql));
438
439        explanation.join("\n")
440    }
441
442    /// Convert natural query to RQL string
443    fn to_rql(&self, query: &NaturalQuery) -> String {
444        match query.intent {
445            QueryIntent::Find => {
446                let node_type = self.primary_entity_label(query).unwrap_or("*");
447
448                let filters: Vec<String> = query
449                    .entities
450                    .iter()
451                    .filter_map(|e| {
452                        e.value
453                            .as_ref()
454                            .map(|v| format!("n.label CONTAINS '{}'", v))
455                    })
456                    .collect();
457
458                if filters.is_empty() {
459                    format!("MATCH (n:{}) RETURN n", node_type)
460                } else {
461                    format!(
462                        "MATCH (n:{}) WHERE {} RETURN n",
463                        node_type,
464                        filters.join(" AND ")
465                    )
466                }
467            }
468            QueryIntent::Path => {
469                let endpoints: Vec<&str> = query
470                    .entities
471                    .iter()
472                    .filter_map(|e| e.value.as_deref())
473                    .collect();
474                if endpoints.len() >= 2 {
475                    format!("PATH FROM '{}' TO '{}'", endpoints[0], endpoints[1])
476                } else {
477                    "PATH FROM ? TO ?".to_string()
478                }
479            }
480            QueryIntent::Count => {
481                let node_type = self.primary_entity_label(query).unwrap_or("*");
482                format!("MATCH (n:{}) RETURN COUNT(n)", node_type)
483            }
484            QueryIntent::Show => {
485                let filters: Vec<String> = query
486                    .entities
487                    .iter()
488                    .filter_map(|e| e.value.as_ref().map(|v| format!("n.id = '{}'", v)))
489                    .collect();
490                if filters.is_empty() {
491                    "MATCH (n) RETURN n".to_string()
492                } else {
493                    format!(
494                        "MATCH (n) WHERE {} RETURN n, n.neighbors",
495                        filters.first().unwrap()
496                    )
497                }
498            }
499            QueryIntent::Check => {
500                let endpoints: Vec<&str> = query
501                    .entities
502                    .iter()
503                    .filter_map(|e| e.value.as_deref())
504                    .collect();
505                if endpoints.len() >= 2 {
506                    format!(
507                        "MATCH (a)-[r]->(b) WHERE a.id = '{}' AND b.id = '{}' RETURN EXISTS(r)",
508                        endpoints[0], endpoints[1]
509                    )
510                } else {
511                    "MATCH (a)-[r]->(b) RETURN EXISTS(r)".to_string()
512                }
513            }
514        }
515    }
516}
517
518#[cfg(test)]
519mod tests {
520    use super::*;
521    use crate::storage::query::test_support::service_graph_with_user;
522    use reddb_rql::ast::EdgeDirection;
523
524    fn create_test_graph() -> Arc<GraphStore> {
525        service_graph_with_user()
526    }
527
528    #[test]
529    fn test_list_hosts() {
530        let graph = create_test_graph();
531        let executor = NaturalExecutor::new(graph);
532
533        let (result, explanation) = executor.execute_with_explanation("list all hosts").unwrap();
534        assert_eq!(result.records.len(), 2);
535        // "list" maps to Find intent in NaturalParser
536        assert!(explanation.contains("Intent: Find"));
537    }
538
539    #[test]
540    fn test_find_services() {
541        let graph = create_test_graph();
542        let executor = NaturalExecutor::new(graph);
543
544        let (result, explanation) = executor.execute_with_explanation("find services").unwrap();
545        assert_eq!(result.records.len(), 2);
546        assert!(explanation.contains("Service"));
547    }
548
549    #[test]
550    fn test_count_hosts() {
551        let graph = create_test_graph();
552        let executor = NaturalExecutor::new(graph);
553
554        let (result, _) = executor.execute_with_explanation("how many hosts").unwrap();
555        assert_eq!(result.records.len(), 1);
556        let count = result.records[0].get("count");
557        assert!(count.is_some());
558    }
559
560    #[test]
561    fn test_explanation_includes_rql() {
562        let graph = create_test_graph();
563        let executor = NaturalExecutor::new(graph);
564
565        let (_, explanation) = executor
566            .execute_with_explanation("find hosts with SSH")
567            .unwrap();
568        assert!(explanation.contains("Equivalent RQL:"));
569        assert!(explanation.contains("MATCH"));
570    }
571
572    #[test]
573    fn test_path_query() {
574        let graph = create_test_graph();
575        let executor = NaturalExecutor::new(graph);
576
577        let (result, explanation) = executor
578            .execute_with_explanation("path from host 10.0.0.1 to host 10.0.0.2")
579            .unwrap();
580        assert!(!result.is_empty());
581        assert!(explanation.contains("Path"));
582    }
583}
reddb_server/storage/query/executors/natural.rs

reddb_server/storage/query/executors/
natural.rs
