From cf44ea8ef98fe7360af24ac26869c47568cf5b1a Mon Sep 17 00:00:00 2001 From: Stephen Adamson Date: Sun, 16 Nov 2025 00:50:08 +0000 Subject: [PATCH] toon-er --- graph/graph.py | 109 ++++++++++++++++++++++++++++++++++++++++++++++ mcp_codebase.py | 112 ++++++++++++++++++++++++++---------------------- 2 files changed, 169 insertions(+), 52 deletions(-) diff --git a/graph/graph.py b/graph/graph.py index 285c89a..df8f913 100644 --- a/graph/graph.py +++ b/graph/graph.py @@ -414,3 +414,112 @@ class LocalGraph: result['incoming'].append(rel_info) return result + + def get_queryable_relationships(self, node_id: str) -> dict: + """Get relationships in a format optimized for graph queries.""" + if node_id not in self.graph.nodes: + return {} + + node_data = self.graph.nodes[node_id] + result = { + 'node_id': node_id, + 'node_type': node_data.get('type', ''), + 'node_name': node_data.get('name', ''), + 'file': node_data.get('file', ''), + 'outgoing': [], + 'incoming': [], + 'cross_file': [], + 'same_file': [] + } + + node_file = node_data.get('file', '') + + # Process all relationships + for neighbor_id in self.graph.successors(node_id): + if neighbor_id in self.graph.nodes: + edge_data = self.graph.edges[node_id, neighbor_id] + neighbor_data = self.graph.nodes[neighbor_id] + + rel_info = { + 'type': edge_data.get('type', 'related'), + 'target_id': neighbor_id, + 'target_name': neighbor_data.get('name', ''), + 'target_type': neighbor_data.get('type', ''), + 'target_file': neighbor_data.get('file', '') + } + + result['outgoing'].append(rel_info) + + # Classify by file relationship + if neighbor_data.get('file') and neighbor_data.get('file') != node_file: + result['cross_file'].append(rel_info) + else: + result['same_file'].append(rel_info) + + # Reverse relationships + for predecessor_id in self.graph.predecessors(node_id): + if predecessor_id in self.graph.nodes: + edge_data = self.graph.edges[predecessor_id, node_id] + predecessor_data = self.graph.nodes[predecessor_id] + + rel_info = { + 'type': edge_data.get('type', 'related'), + 'source_id': predecessor_id, + 'source_name': predecessor_data.get('name', ''), + 'source_type': predecessor_data.get('type', ''), + 'source_file': predecessor_data.get('file', '') + } + + result['incoming'].append(rel_info) + + return result + + def classify_relationships(self, node_id): + node_data = self.graph.nodes[node_id] + node_file = node_data.get('file', '') + + result = { + 'outgoing': [], + 'incoming': [], + 'same_file': [], + 'cross_file': [] + } + + for neighbor_id in self.graph.successors(node_id): + if neighbor_id in self.graph.nodes: + edge_data = self.graph.edges[node_id, neighbor_id] + neighbor_data = self.graph.nodes[neighbor_id] + + rel_info = { + 'type': edge_data.get('type', 'related'), + 'target_id': neighbor_id, + 'target_name': neighbor_data.get('name', ''), + 'target_type': neighbor_data.get('type', ''), + 'target_file': neighbor_data.get('file', '') + } + + result['outgoing'].append(rel_info) + + # Classify by file relationship + if neighbor_data.get('file') and neighbor_data.get('file') != node_file: + result['cross_file'].append(rel_info) + else: + result['same_file'].append(rel_info) + + # Reverse relationships + for predecessor_id in self.graph.predecessors(node_id): + if predecessor_id in self.graph.nodes: + edge_data = self.graph.edges[predecessor_id, node_id] + predecessor_data = self.graph.nodes[predecessor_id] + + rel_info = { + 'type': edge_data.get('type', 'related'), + 'source_id': predecessor_id, + 'source_name': predecessor_data.get('name', ''), + 'source_type': predecessor_data.get('type', ''), + 'source_file': predecessor_data.get('file', '') + } + + result['incoming'].append(rel_info) + + return result diff --git a/mcp_codebase.py b/mcp_codebase.py index b4eda94..ca7a7f8 100644 --- a/mcp_codebase.py +++ b/mcp_codebase.py @@ -4999,76 +4999,79 @@ def _get_result_graph_context(meta: dict, graph: Optional[LocalGraph] = None) -> if graph is None or graph.graph.number_of_nodes() == 0: return {} - context = {} + _format_enhanced_results = {} file_path = meta.get('file', '') entity_name = meta.get('name', '') entity_type = meta.get('type', '') language = meta.get('language', '') # Build node ID using the same logic as graph building - node_attrs = { - "lang": language, - "type": entity_type, - "file": file_path, - "name": entity_name - } - node_id = _build_node_id(node_attrs) + node_id = _build_node_id(language, entity_type, file_path, entity_name, meta) if node_id and node_id in graph.graph: context['node_id'] = node_id - # Get comprehensive relationships for graph queries - cross_file_out, cross_file_in = graph.get_cross_file_relationships(node_id, file_path) - - # Format for graph queries: "node_id:relationship_type" - if cross_file_out: - context['cross_file'] = [] - for rel in cross_file_out[:8]: # Limit to most important - # Extract target node info for graph queries - parts = rel.split('@') - relationship = parts[0] - target_file = parts[1] if len(parts) > 1 else "" - context['cross_file'].append(f"{relationship}->{target_file}") - - if cross_file_in: - context['used_by'] = [] - for rel in cross_file_in[:6]: # Limit to most important - parts = rel.split('@') - relationship = parts[0] - source_file = parts[1] if len(parts) > 1 else "" - context['used_by'].append(f"{relationship}<-{source_file}") - - # Same-file relationships for local graph exploration + # Get ALL relationships for comprehensive graph queries + cross_file_out = [] + cross_file_in = [] same_file_rels = [] + + # Outgoing relationships for neighbor_id in graph.graph.successors(node_id): if neighbor_id in graph.graph.nodes: - neighbor_file = graph.graph.nodes[neighbor_id].get('file', '') - if neighbor_file == file_path: - edge_data = graph.graph.edges[node_id, neighbor_id] - rel_type = edge_data.get('type', 'related') - neighbor_name = graph.graph.nodes[neighbor_id].get('name', '') - neighbor_type = graph.graph.nodes[neighbor_id].get('type', '') - same_file_rels.append(f"{rel_type}:{neighbor_type}:{neighbor_name}") + edge_data = graph.graph.edges[node_id, neighbor_id] + neighbor_data = graph.graph.nodes[neighbor_id] + rel_type = edge_data.get('type', 'related') + neighbor_name = neighbor_data.get('name', '') + neighbor_type = neighbor_data.get('type', '') + neighbor_file = neighbor_data.get('file', '') + + rel_info = f"{rel_type}:{neighbor_type}:{neighbor_name}" + + # Cross-file relationships (most valuable for queries) + if neighbor_file and neighbor_file != file_path: + cross_file_out.append(f"{rel_info}@{neighbor_file}") + else: + same_file_rels.append(rel_info) + + # Incoming relationships (reverse dependencies) + for predecessor_id in graph.graph.predecessors(node_id): + if predecessor_id in graph.graph.nodes: + edge_data = graph.graph.edges[predecessor_id, node_id] + predecessor_data = graph.graph.nodes[predecessor_id] + + rel_type = edge_data.get('type', 'related') + pred_name = predecessor_data.get('name', '') + pred_type = predecessor_data.get('type', '') + pred_file = predecessor_data.get('file', '') + + if pred_file and pred_file != file_path: + cross_file_in.append(f"{rel_type}:{pred_type}:{pred_name}@{pred_file}") + + # Store in context for TOON format + if cross_file_out: + context['cross_file'] = cross_file_out[:6] # Top 6 cross-file deps + if cross_file_in: + context['used_by'] = cross_file_in[:4] # Top 4 reverse deps if same_file_rels: - context['same_file'] = same_file_rels[:10] + context['same_file'] = same_file_rels[:8] # Top 8 same-file rels return context -def _build_node_id(attrs: dict) -> str: +def _build_node_id(lang: str, entity_type: str, file_path: str, entity_name: str, meta: dict) -> str: """Build node ID using the same logic as graph building.""" - lang = attrs.get("lang", "unknown") - node_type = attrs.get("type", "Symbol") - name = attrs.get("name", "unknown") - fpath = attrs.get("file", "") + if not entity_name or not file_path or not lang: + return None - # Handle special cases - if node_type == "component" and lang == "svelte": - return f"svelte::component::{fpath}::{name}" - elif node_type in ["method", "function"] and attrs.get("class_name"): - return f"{lang}::{node_type}::{fpath}::{attrs['class_name']}.{name}" + # Handle special cases to match graph building logic + if entity_type in ["method", "function"] and meta.get("class"): + class_name = meta.get("class") + return f"{lang}::{entity_type}::{file_path}::{class_name}.{entity_name}" + elif entity_type == "component" and lang == "svelte": + return f"svelte::component::{file_path}::{entity_name}" else: - return f"{lang}::{node_type}::{fpath}::{name}" + return f"{lang}::{entity_type}::{file_path}::{entity_name}" def _format_enhanced_results(results: List[Dict], query: str, result_type: str) -> str: """Format results in TOON (Token-Oriented Object Notation) for LLM consumption.""" @@ -5091,23 +5094,28 @@ def _format_enhanced_results(results: List[Dict], query: str, result_type: str) content = result.get('content', '') graph_context = result.get('graph_context', {}) - # Build compact graph context string + # Build compact graph context string for queries graph_parts = [] + + # Node ID for direct graph queries if graph_context.get('node_id'): graph_parts.append(f"node:{graph_context['node_id']}") + # Cross-file dependencies for import/call queries if graph_context.get('cross_file'): graph_parts.append(f"xfile:{'|'.join(graph_context['cross_file'][:4])}") + # Reverse dependencies for "used by" queries if graph_context.get('used_by'): graph_parts.append(f"used:{'|'.join(graph_context['used_by'][:3])}") + # Same-file relationships for file context queries if graph_context.get('same_file'): graph_parts.append(f"same:{'|'.join(graph_context['same_file'][:5])}") graph_str = ";".join(graph_parts) - # Build TOON row + # Build TOON row - compact format for LLM parsing row = [ file_path, str(line_num), @@ -5115,7 +5123,7 @@ def _format_enhanced_results(results: List[Dict], query: str, result_type: str) entity_name, language, f"{score:.2f}", - content, # No truncation - let LLM handle it + content, # Full content - LLM can handle truncation if needed graph_str ]