diff --git a/graph/graph.py b/graph/graph.py index f47d342..285c89a 100644 --- a/graph/graph.py +++ b/graph/graph.py @@ -361,3 +361,56 @@ class LocalGraph: cross_file_in.append(f"{rel_type}:{pred_type}:{pred_name}@{pred_file}") return cross_file_out, cross_file_in + + def get_node_by_id(self, node_id: str) -> Optional[dict]: + """Get node data by ID for graph queries.""" + if node_id in self.graph.nodes: + return dict(self.graph.nodes[node_id]) + return None + + def get_relationships_for_query(self, node_id: str, max_relationships: int = 10) -> dict: + """Get relationships in a format suitable for graph queries.""" + if node_id not in self.graph.nodes: + return {} + + result = { + 'node_id': node_id, + 'outgoing': [], + 'incoming': [] + } + + node_file = self.graph.nodes[node_id].get('file', '') + + # Outgoing relationships + for neighbor_id in list(self.graph.successors(node_id))[:max_relationships]: + 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', ''), + 'cross_file': neighbor_data.get('file', '') != node_file + } + result['outgoing'].append(rel_info) + + # Incoming relationships + for predecessor_id in list(self.graph.predecessors(node_id))[:max_relationships]: + 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', ''), + 'cross_file': predecessor_data.get('file', '') != node_file + } + result['incoming'].append(rel_info) + + return result diff --git a/mcp_codebase.py b/mcp_codebase.py index 2f02f20..b4eda94 100644 --- a/mcp_codebase.py +++ b/mcp_codebase.py @@ -4995,7 +4995,7 @@ def _build_enhanced_result(chunk_text: str, meta: dict, score: float, graph: Opt def _get_result_graph_context(meta: dict, graph: Optional[LocalGraph] = None) -> dict: - """Fixed graph context using enhanced graph methods.""" + """Enhanced graph context that provides queryable node IDs and relationships.""" if graph is None or graph.graph.number_of_nodes() == 0: return {} @@ -5005,153 +5005,134 @@ def _get_result_graph_context(meta: dict, graph: Optional[LocalGraph] = None) -> entity_type = meta.get('type', '') language = meta.get('language', '') - # Try multiple node ID formats to find the right one - node_candidates = [] + # 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) - # Format 1: Standard node ID - if entity_name and file_path and language and entity_type: - node_candidates.append(f"{language}::{entity_type}::{file_path}::{entity_name}") + if node_id and node_id in graph.graph: + context['node_id'] = node_id - # Format 2: Method with class name - if entity_type in ['method', 'function'] and meta.get('class'): - class_name = meta.get('class') - node_candidates.append(f"{language}::{entity_type}::{file_path}::{class_name}.{entity_name}") + # Get comprehensive relationships for graph queries + cross_file_out, cross_file_in = graph.get_cross_file_relationships(node_id, file_path) - # Format 3: Svelte component - if entity_type == 'component' and language == 'svelte': - node_candidates.append(f"svelte::component::{file_path}::{entity_name}") + # 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}") - # Try to find the node using graph's enhanced search - target_node_id = None - for candidate in node_candidates: - if candidate in graph.graph.nodes: - target_node_id = candidate - break + 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}") - # If still not found, try attribute-based search - if not target_node_id and entity_name and file_path: - matches = list(graph.find_nodes_by_attributes( - name=entity_name, - file=file_path, - type=entity_type if entity_type != 'code' else None - )) - if matches: - target_node_id = matches[0][0] # Take first match + # Same-file relationships for local graph exploration + same_file_rels = [] + 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}") - if not target_node_id: - return context - - # Use the enhanced graph methods to get relationships - cross_file_out, cross_file_in = graph.get_cross_file_relationships(target_node_id, file_path) - - if cross_file_out: - context['cross_file'] = cross_file_out - if cross_file_in: - context['used_by'] = cross_file_in - - # Get same-file relationships - same_file_rels = [] - for neighbor_id in graph.graph.successors(target_node_id): - if neighbor_id in graph.graph.nodes: - neighbor_file = graph.graph.nodes[neighbor_id].get('file', '') - if neighbor_file == file_path: # Same file - edge_data = graph.graph.edges[target_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}") - - if same_file_rels: - context['same_file'] = same_file_rels - - # Get meaningful file entities (filter out noise) - meaningful_types = {'function', 'method', 'class', 'struct', 'interface', 'enum', 'module', 'component', 'trait', 'impl'} - file_entities = [] - - file_node_id = f"file::{file_path}" - if file_node_id in graph.graph.nodes: - for neighbor_id, neighbor_data in graph.neighbors(file_node_id): - neighbor_type = neighbor_data.get('type', '') - neighbor_name = neighbor_data.get('name', '') - - # Filter criteria - if (neighbor_type in meaningful_types and - neighbor_name and - len(neighbor_name) > 1 and # No single letters - neighbor_name != entity_name and - not neighbor_name.startswith(('c:', 'b:', 'p:', 'r:', 't:', 'v:'))): # No import aliases - file_entities.append(f"{neighbor_type}:{neighbor_name}") - - if file_entities: - context['file_entities'] = file_entities + if same_file_rels: + context['same_file'] = same_file_rels[:10] return context +def _build_node_id(attrs: 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", "") + + # 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}" + else: + return f"{lang}::{node_type}::{fpath}::{name}" + def _format_enhanced_results(results: List[Dict], query: str, result_type: str) -> str: - """Comprehensive format with ALL information preserved - optimized for LLM analysis.""" + """Format results in TOON (Token-Oriented Object Notation) for LLM consumption.""" if not results: return f"# {result_type.title()}: {query}\nNo results found.\n" lines = [ f"# {result_type.title()}: {query}", - f"results_count={len(results)}" + f"results[{len(results)}]{{file,line,type,name,language,score,content,graph_context}}:" ] - for i, result in enumerate(results, 1): - # Extract ALL metadata + for result in results: + # Extract all fields file_path = result.get('file', '') line_num = result.get('line', '') entity_type = result.get('type', 'code') entity_name = result.get('name', '') language = result.get('language', '') score = result.get('score', 0.0) - content = result.get('content', '') # NO TRUNCATION - let LLM handle it - docstring = result.get('docstring', '') - + content = result.get('content', '') graph_context = result.get('graph_context', {}) - # Build comprehensive result block - lines.append(f"\n--- RESULT {i} ---") - lines.append(f"file:{file_path}:{line_num}") - lines.append(f"entity:{entity_type}:{entity_name}") - lines.append(f"language:{language}") - lines.append(f"score:{score:.3f}") + # Build compact graph context string + graph_parts = [] + if graph_context.get('node_id'): + graph_parts.append(f"node:{graph_context['node_id']}") - # Content (FULL, no truncation) - if content: - lines.append(f"content: {content}") + if graph_context.get('cross_file'): + graph_parts.append(f"xfile:{'|'.join(graph_context['cross_file'][:4])}") - # Docstring if available - if docstring: - lines.append(f"doc: {docstring}") + if graph_context.get('used_by'): + graph_parts.append(f"used:{'|'.join(graph_context['used_by'][:3])}") - # Graph context - PRESERVE EVERYTHING - if graph_context: - lines.append("graph_context:") + if graph_context.get('same_file'): + graph_parts.append(f"same:{'|'.join(graph_context['same_file'][:5])}") - # Cross-file dependencies (highest value) - if graph_context.get('cross_file'): - lines.append(f" cross_file: {' | '.join(graph_context['cross_file'])}") + graph_str = ";".join(graph_parts) - # Reverse dependencies - if graph_context.get('used_by'): - lines.append(f" used_by: {' | '.join(graph_context['used_by'])}") + # Build TOON row + row = [ + file_path, + str(line_num), + entity_type, + entity_name, + language, + f"{score:.2f}", + content, # No truncation - let LLM handle it + graph_str + ] - # Same-file relationships - if graph_context.get('same_file'): - lines.append(f" same_file_rels: {' | '.join(graph_context['same_file'])}") + # Escape fields according to TOON spec + escaped_row = [] + for field in row: + field_str = str(field) + if any(char in field_str for char in [',', '"', '\n', '\r']): + escaped = field_str.replace('"', '\\"') + escaped_row.append(f'"{escaped}"') + else: + escaped_row.append(field_str) - # All relationships (comprehensive) - if graph_context.get('all_rels'): - lines.append(f" all_relationships: {' | '.join(graph_context['all_rels'][:15])}") # Reasonable limit - - # File contents - if graph_context.get('file_contents'): - lines.append(f" file_entities: {' | '.join(graph_context['file_contents'])}") + lines.append(" " + ",".join(escaped_row)) return "\n".join(lines) - def _extract_clean_content(chunk_text: str) -> str: """Extract clean content WITHOUT truncation - preserve full context.""" lines = chunk_text.split('\n')