From dffe1750591f8784b47aae658853493a51e54965 Mon Sep 17 00:00:00 2001 From: Stephen Adamson Date: Sat, 15 Nov 2025 23:27:43 +0000 Subject: [PATCH] better format --- graph/graph.py | 85 ++++++++++++-- mcp_codebase.py | 304 ++++++++++++++++++++++++++++++------------------ 2 files changed, 271 insertions(+), 118 deletions(-) diff --git a/graph/graph.py b/graph/graph.py index 96fe75f..0e9dd3d 100644 --- a/graph/graph.py +++ b/graph/graph.py @@ -16,16 +16,23 @@ class LocalGraph: # --- Creation --- def add_node(self, node_id: str | dict, **attrs): """ - If a dict is passed instead of a plain string, build the node_id - as: :::::: + Enhanced add_node that can return the generated node_id """ if isinstance(node_id, dict): - lang = node_id.get("lang", "unknown") - kind = node_id.get("type", "Symbol") - name = node_id.get("name", "unknown") - fpath = node_id.get("file", "") + lang = node_id.get("lang", "unknown") + kind = node_id.get("type", "Symbol") + name = node_id.get("name", "unknown") + fpath = node_id.get("file", "") node_id = f"{lang}::{kind}::{fpath}::{name}" - self.graph.add_node(node_id, **attrs) + + # Merge any additional attributes + final_attrs = {} + if isinstance(node_id, dict): + final_attrs.update(node_id) + final_attrs.update(attrs) + + self.graph.add_node(node_id, **final_attrs) + return node_id # Return the ID for reference def _add_import_edges(self, src: str, imports: list[str]): for imp in imports: @@ -213,3 +220,67 @@ class LocalGraph: except Exception: # swallow; graph should remain usable pass + + def get_node_relationships(self, node_id: str) -> dict: + """ + Get comprehensive relationships for a node in a structured format. + Returns: { + 'outgoing': [(neighbor_id, edge_type, neighbor_data)], + 'incoming': [(predecessor_id, edge_type, predecessor_data)], + 'cross_file': list of cross-file relationships, + 'same_file': list of same-file relationships + } + """ + if node_id not in self.graph: + return {} + + result = { + 'outgoing': [], + 'incoming': [], + 'cross_file': [], + 'same_file': [] + } + + node_file = self.graph.nodes[node_id].get('file', '') + + # Outgoing relationships + for neighbor_id in self.graph.successors(node_id): + edge_data = self.graph.edges[node_id, neighbor_id] + neighbor_data = self.graph.nodes[neighbor_id] + edge_type = edge_data.get('type', 'related') + + result['outgoing'].append((neighbor_id, edge_type, neighbor_data)) + + # Cross-file classification + neighbor_file = neighbor_data.get('file', '') + rel_info = f"{edge_type}:{neighbor_data.get('type', '')}:{neighbor_data.get('name', '')}" + + if neighbor_file and neighbor_file != node_file: + result['cross_file'].append(f"{rel_info}@{neighbor_file}") + else: + result['same_file'].append(rel_info) + + # Incoming relationships + for predecessor_id in self.graph.predecessors(node_id): + edge_data = self.graph.edges[predecessor_id, node_id] + predecessor_data = self.graph.nodes[predecessor_id] + edge_type = edge_data.get('type', 'related') + + result['incoming'].append((predecessor_id, edge_type, predecessor_data)) + + return result + + def get_file_entities(self, file_path: str) -> list: + """Get all entities in a file in compact format.""" + file_node_id = f"file::{file_path}" + if file_node_id not in self.graph: + return [] + + entities = [] + for neighbor_id, neighbor_data in self.neighbors(file_node_id): + entity_type = neighbor_data.get('type', '') + entity_name = neighbor_data.get('name', '') + if entity_type and entity_name: + entities.append(f"{entity_type}:{entity_name}") + + return entities diff --git a/mcp_codebase.py b/mcp_codebase.py index e603ae2..e1ff29a 100644 --- a/mcp_codebase.py +++ b/mcp_codebase.py @@ -5007,93 +5007,174 @@ 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: - """Get MOST relevant graph context - optimized for LLM tokens.""" - if graph is None or graph.graph.number_of_nodes() == 0: + """Ultra-compact graph context using LocalGraph helper methods.""" + if graph is None: return {} - context = {} file_path = meta.get('file', '') entity_name = meta.get('name', '') entity_type = meta.get('type', '') language = meta.get('language', '') - # Build node ID - node_id = None - if entity_name and file_path and language: - if entity_type in ['class', 'struct', 'interface', 'enum', 'impl']: - node_id = f"{language}::{entity_type}::{file_path}::{entity_name}" - elif entity_type in ['method', 'function', 'constructor']: - class_name = meta.get('class') - if class_name: - node_id = f"{language}::{entity_type}::{file_path}::{class_name}.{entity_name}" - else: - node_id = f"{language}::{entity_type}::{file_path}::{entity_name}" - elif entity_type == 'component' and language == 'svelte': - node_id = f"svelte::component::{file_path}::{entity_name}" + # Build node ID using graph's method + node_attrs = {"lang": language, "type": entity_type, "file": file_path, "name": entity_name} + node_id = graph.add_node(node_attrs, return_id=True) - if node_id and node_id in graph.graph: - # PRIORITY 1: Cross-file dependencies (highest value for LLM) - cross_file_rels = [] - incoming_cross_file = [] + if node_id not in graph.graph: + return {} - # Outgoing cross-file relationships - for neighbor_id, neighbor_data in graph.neighbors(node_id): - neighbor_file = neighbor_data.get('file', '') - if neighbor_file and neighbor_file != file_path: - edge_data = graph.graph.edges[node_id, neighbor_id] - rel_type = edge_data.get('type', 'related') - neighbor_name = neighbor_data.get('name', neighbor_id.split('::')[-1]) - neighbor_type = neighbor_data.get('type', '') + # Get all relationships in one call + relationships = graph.get_node_relationships(node_id) + context = {} - # Compress representation based on relationship importance - if rel_type in ['extends', 'implements', 'calls']: - # High-value relationships - keep full info - cross_file_rels.append(f"{rel_type}:{neighbor_type}:{neighbor_name}@{neighbor_file}") - else: - # Medium-value - compress type - cross_file_rels.append(f"{rel_type}:{neighbor_name}@{neighbor_file}") + if relationships.get('cross_file'): + context['cross_file'] = relationships['cross_file'] + if relationships.get('same_file'): + context['same_file'] = relationships['same_file'] - # PRIORITY 2: Reverse dependencies (what uses this entity) - for predecessor in graph.graph.predecessors(node_id): - pred_data = graph.graph.nodes[predecessor] + # Build reverse dependencies from incoming relationships + if relationships.get('incoming'): + reverse_deps = [] + node_file = graph.graph.nodes[node_id].get('file', '') + for pred_id, edge_type, pred_data in relationships['incoming']: pred_file = pred_data.get('file', '') pred_name = pred_data.get('name', '') + pred_type = pred_data.get('type', '') - # Only show cross-file incoming (same-file is less valuable) - if pred_file and pred_file != file_path: - edge_data = graph.graph.edges[predecessor, node_id] - rel_type = edge_data.get('type', 'related') + if pred_file and pred_file != node_file: + reverse_deps.append(f"{edge_type}:{pred_type}:{pred_name}@{pred_file}") + else: + reverse_deps.append(f"{edge_type}:{pred_type}:{pred_name}") - if rel_type in ['calls', 'extends', 'implements']: - incoming_cross_file.append(f"{rel_type}:{pred_name}@{pred_file}") + if reverse_deps: + context['used_by'] = reverse_deps - # Apply aggressive limits for token conservation - if cross_file_rels: - # Sort by relationship importance - priority_order = {'extends': 0, 'implements': 1, 'calls': 2} - cross_file_rels.sort(key=lambda x: priority_order.get(x.split(':')[0], 99)) - context['xfile'] = cross_file_rels[:6] # Only top 6 cross-file deps - - if incoming_cross_file: - context['used_by'] = incoming_cross_file[:4] # Only top 4 reverse deps - - # PRIORITY 3: Critical same-file context (only if cross-file is sparse) - if not context.get('xfile') and not context.get('used_by'): - file_node_id = f"file::{file_path}" - if file_node_id in graph.graph: - file_entities = [] - for neighbor_id, neighbor_data in graph.neighbors(file_node_id): - neighbor_type = neighbor_data.get('type', '') - neighbor_name = neighbor_data.get('name', '') - if (neighbor_type in ['class', 'struct', 'interface', 'component'] and - neighbor_name != entity_name): - file_entities.append(f"{neighbor_type}:{neighbor_name}") - - if file_entities: - context['file'] = file_entities[:3] # Only top 3 same-file entities + # Get file entities + file_entities = graph.get_file_entities(file_path) + if file_entities: + context['file_contents'] = file_entities return context +def _format_enhanced_results(results: List[Dict], query: str, result_type: str) -> str: + """Comprehensive format with ALL information preserved - optimized for LLM analysis.""" + if not results: + return f"# {result_type.title()}: {query}\nNo results found.\n" + + lines = [ + f"# {result_type.title()}: {query}", + f"results_count={len(results)}" + ] + + for i, result in enumerate(results, 1): + # Extract ALL metadata + 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', '') + + 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}") + + # Content (FULL, no truncation) + if content: + lines.append(f"content: {content}") + + # Docstring if available + if docstring: + lines.append(f"doc: {docstring}") + + # Graph context - PRESERVE EVERYTHING + if graph_context: + lines.append("graph_context:") + + # Cross-file dependencies (highest value) + if graph_context.get('cross_file'): + lines.append(f" cross_file: {' | '.join(graph_context['cross_file'])}") + + # Reverse dependencies + if graph_context.get('used_by'): + lines.append(f" used_by: {' | '.join(graph_context['used_by'])}") + + # Same-file relationships + if graph_context.get('same_file'): + lines.append(f" same_file_rels: {' | '.join(graph_context['same_file'])}") + + # 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'])}") + + 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') + content_lines = [] + + # Skip only the pure metadata header lines + skip_metadata = True + for line in lines: + stripped = line.strip() + if skip_metadata: + # Only skip actual metadata labels, not content + if any(stripped.startswith(prefix) for prefix in ['File:', 'Type:', 'Name:', 'Language:', 'Doc:']): + continue + # Stop skipping when we hit actual code/content + if stripped and not stripped.startswith(('File:', 'Type:', 'Name:', 'Language:', 'Doc:')): + skip_metadata = False + content_lines.append(stripped) + else: + content_lines.append(stripped) + + # Join with spaces but preserve ALL content + content = ' '.join(content_lines) + content = re.sub(r'\s+', ' ', content) # Only normalize whitespace + return content.strip() + + +def _build_enhanced_result(chunk_text: str, meta: dict, score: float, graph: Optional[LocalGraph] = None) -> dict: + """Build comprehensive result with ALL context preserved.""" + # Extract key metadata + file_path = meta.get('file', '') + line_num = meta.get('line', '') + entity_type = meta.get('type', 'code') + language = meta.get('language', '') + entity_name = meta.get('name', '') + + # Extract FULL content and docstring + clean_content = _extract_clean_content(chunk_text) + docstring = _extract_docstring(chunk_text) + + # Get COMPREHENSIVE graph context + graph_context = _get_result_graph_context(meta, graph) + + return { + 'file': file_path, + 'line': line_num, + 'type': entity_type, + 'language': language, + 'name': entity_name, + 'score': score, + 'content': clean_content, # NO TRUNCATION + 'docstring': docstring, + 'graph_context': graph_context + } + def _extract_docstring(chunk_text: str) -> str: """Extract docstring or comments from chunk text.""" lines = chunk_text.split('\n') @@ -5132,94 +5213,95 @@ def _extract_docstring(chunk_text: str) -> str: def _format_enhanced_results(results: List[Dict], query: str, result_type: str) -> str: - """Ultra-compact formatting optimized for LLM token usage.""" + """Comprehensive format with ALL information preserved - optimized for LLM analysis.""" if not results: return f"# {result_type.title()}: {query}\nNo results found.\n" lines = [ f"# {result_type.title()}: {query}", - f"Found {len(results)} results:" + f"results_count={len(results)}" ] for i, result in enumerate(results, 1): - # Extract core info + # Extract ALL metadata 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', '')[:350] # Strict content limit + content = result.get('content', '') # NO TRUNCATION - let LLM handle it + docstring = result.get('docstring', '') graph_context = result.get('graph_context', {}) - # Build ultra-compact context string - context_parts = [] + # 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}") - # Cross-file deps (highest priority) - if graph_context.get('xfile'): - # Compress further: remove file paths, keep only key info - compressed_deps = [] - for dep in graph_context['xfile'][:4]: # Only top 4 - parts = dep.split('@')[0] # Remove file path - compressed_deps.append(parts) - context_parts.append(f"deps:{'|'.join(compressed_deps)}") + # Content (FULL, no truncation) + if content: + lines.append(f"content: {content}") - # Reverse deps - if graph_context.get('used_by'): - compressed_used = [] - for used in graph_context['used_by'][:3]: # Only top 3 - parts = used.split('@')[0] # Remove file path - compressed_used.append(parts) - context_parts.append(f"used:{'|'.join(compressed_used)}") + # Docstring if available + if docstring: + lines.append(f"doc: {docstring}") - # Same-file context (lowest priority) - if graph_context.get('file') and not context_parts: - context_parts.append(f"file:{'|'.join(graph_context['file'][:2])}") + # Graph context - PRESERVE EVERYTHING + if graph_context: + lines.append("graph_context:") - context_str = " ".join(context_parts) + # Cross-file dependencies (highest value) + if graph_context.get('cross_file'): + lines.append(f" cross_file: {' | '.join(graph_context['cross_file'])}") - # Ultra-compact row format - row = f"{i}. {file_path}:{line_num} {entity_type}:{entity_name} (score:{score:.2f})" - if context_str: - row += f" [{context_str}]" + # Reverse dependencies + if graph_context.get('used_by'): + lines.append(f" used_by: {' | '.join(graph_context['used_by'])}") - lines.append(row) - lines.append(f" Content: {content}") + # Same-file relationships + if graph_context.get('same_file'): + lines.append(f" same_file_rels: {' | '.join(graph_context['same_file'])}") - # Add docstring only if highly relevant and short - docstring = result.get('docstring', '') - if docstring and len(docstring) < 100 and any(keyword in query.lower() for keyword in ['how', 'what', 'why', 'documentation']): - lines.append(f" Docs: {docstring}") + # All relationships (comprehensive) + if graph_context.get('all_rels'): + lines.append(f" all_relationships: {' | '.join(graph_context['all_rels'][:15])}") # Reasonable limit - lines.append("") # Empty line between results + # File contents + if graph_context.get('file_contents'): + lines.append(f" file_entities: {' | '.join(graph_context['file_contents'])}") return "\n".join(lines) + def _extract_clean_content(chunk_text: str) -> str: - """Extract clean content from chunk text by removing duplicate metadata.""" + """Extract clean content WITHOUT truncation - preserve full context.""" lines = chunk_text.split('\n') content_lines = [] - # Skip the first few metadata lines (File:, Type:, Name:, etc.) + # Skip only the pure metadata header lines skip_metadata = True for line in lines: stripped = line.strip() if skip_metadata: + # Only skip actual metadata labels, not content if any(stripped.startswith(prefix) for prefix in ['File:', 'Type:', 'Name:', 'Language:', 'Doc:']): continue - # Once we hit actual content, stop skipping - if stripped and not any(stripped.startswith(prefix) for prefix in ['File:', 'Type:', 'Name:', 'Language:', 'Doc:', 'SQL:', 'Code:']): + # Stop skipping when we hit actual code/content + if stripped and not stripped.startswith(('File:', 'Type:', 'Name:', 'Language:', 'Doc:')): skip_metadata = False content_lines.append(stripped) else: content_lines.append(stripped) - # Join and clean up + # Join with spaces but preserve ALL content content = ' '.join(content_lines) - content = re.sub(r'\s+', ' ', content) # Normalize whitespace + content = re.sub(r'\s+', ' ', content) # Only normalize whitespace return content.strip() - @mcp.tool() def find_code_references(symbol: str, top_k: int = 20) -> str: """What it does - Fast, unranked lookup of all file/line occurrences of a symbol.