Skip to main content

phago_rag/
code_query.rs

1//! Code-aware query interface.
2//!
3//! Queries the knowledge graph for code-related concepts like
4//! function names, type references, and file associations.
5
6use phago_core::topology::TopologyGraph;
7use phago_runtime::colony::Colony;
8
9/// A code query result.
10#[derive(Debug, Clone)]
11pub struct CodeQueryResult {
12    pub label: String,
13    pub score: f64,
14    pub related: Vec<String>,
15}
16
17/// Query the code knowledge graph for a term.
18///
19/// Performs a 2-3 hop BFS traversal from seed nodes with cumulative
20/// edge weight scoring. This surfaces related types, functions,
21/// and concepts beyond immediate neighbors.
22pub fn code_query(colony: &Colony, query: &str, max_results: usize) -> Vec<CodeQueryResult> {
23    let graph = colony.substrate().graph();
24
25    // Find seed nodes
26    let seed_ids = graph.find_nodes_by_label(query);
27    if seed_ids.is_empty() {
28        return Vec::new();
29    }
30
31    let max_depth: usize = 3;
32    let max_expansions: usize = 150;
33
34    // BFS with cumulative edge weight scoring
35    // (cumulative_weight, node_id, depth)
36    let mut frontier: Vec<(f64, phago_core::types::NodeId, usize)> = Vec::new();
37    let mut visited: std::collections::HashSet<phago_core::types::NodeId> = std::collections::HashSet::new();
38    let mut scored: Vec<(String, f64, phago_core::types::NodeId)> = Vec::new();
39
40    for seed_id in &seed_ids {
41        visited.insert(*seed_id);
42        if let Some(node) = graph.get_node(seed_id) {
43            scored.push((node.label.clone(), (node.access_count as f64 + 1.0) * 10.0, *seed_id));
44        }
45        frontier.push((1.0, *seed_id, 0));
46    }
47
48    let mut expansions = 0;
49    while !frontier.is_empty() && expansions < max_expansions {
50        // Sort ascending, pop best from end
51        frontier.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
52        let (weight, current, depth) = frontier.pop().unwrap();
53        expansions += 1;
54
55        if depth >= max_depth {
56            continue;
57        }
58
59        let neighbors = graph.neighbors(&current);
60        for (nid, edge) in &neighbors {
61            if visited.contains(nid) {
62                continue;
63            }
64            visited.insert(*nid);
65
66            if let Some(node) = graph.get_node(nid) {
67                let cumulative_weight = weight * edge.weight;
68                let score = cumulative_weight * (1.0 + (node.access_count as f64).ln().max(0.0));
69                scored.push((node.label.clone(), score, *nid));
70                frontier.push((cumulative_weight, *nid, depth + 1));
71            }
72        }
73    }
74
75    // Sort by score descending
76    scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
77
78    // Build results with related nodes
79    let mut results: Vec<CodeQueryResult> = Vec::new();
80    let mut seen = std::collections::HashSet::new();
81
82    for (label, score, nid) in &scored {
83        if seen.contains(label) {
84            continue;
85        }
86        seen.insert(label.clone());
87
88        // Get neighbors for "related" field
89        let mut neighbors: Vec<_> = graph.neighbors(nid)
90            .into_iter()
91            .filter_map(|(neighbor_id, edge)| {
92                let n = graph.get_node(&neighbor_id)?;
93                Some((n.label.clone(), edge.weight))
94            })
95            .collect();
96        neighbors.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
97
98        let related: Vec<String> = neighbors.iter()
99            .take(10)
100            .map(|(l, _)| l.clone())
101            .collect();
102
103        results.push(CodeQueryResult {
104            label: label.clone(),
105            score: *score,
106            related,
107        });
108
109        if results.len() >= max_results {
110            break;
111        }
112    }
113
114    results
115}