From 5a8b12ca8384eab478dbf24127c5d6d88ef0a635 Mon Sep 17 00:00:00 2001 From: Stephen Adamson Date: Sat, 15 Nov 2025 22:02:31 +0000 Subject: [PATCH] upgrade graph --- mcp_codebase.py | 1195 +++++++++++++++++++++++++++++++++++------------ 1 file changed, 893 insertions(+), 302 deletions(-) diff --git a/mcp_codebase.py b/mcp_codebase.py index c808889..34c9d4e 100644 --- a/mcp_codebase.py +++ b/mcp_codebase.py @@ -2,7 +2,7 @@ import os import signal import json from collections import defaultdict, deque - +import ast import re import subprocess @@ -339,207 +339,505 @@ def load_ragignore(base_path: Path) -> pathspec.PathSpec: # ----------------------------- @lru_cache(maxsize=1000) def parse_python_file_cached(filepath_str: str, file_hash: str) -> tuple: - """Cached Python AST parsing with enhanced metadata for graph relationships.""" + """Cached Python AST parsing with enhanced metadata for graph relationships. + ALWAYS returns a tuple of (text, metadata_tuple) entries. If parsing + fails or produces no chunks, return a single fallback chunk containing + the whole file text and minimal metadata. + """ filepath = Path(filepath_str) chunks = [] + + logger.info(f"🐍 Parsing Python file: {filepath}") + try: code = filepath.read_text(encoding="utf-8") - tree = ast.parse(code) + tree = ast.parse(code, filename=str(filepath)) lines = code.splitlines(keepends=True) # Track imports for module relationships imports = [] + module_imports = {} # Maps names to modules for better tracking + # First pass: Extract all imports for node in ast.walk(tree): - # Extract imports for module dependencies - if isinstance(node, (ast.Import, ast.ImportFrom)): - import_info = _extract_import_info(node) - if import_info: - imports.append(import_info) + if isinstance(node, ast.Import): + for alias in node.names: + import_name = alias.asname if alias.asname else alias.name + imports.append(alias.name) + module_imports[import_name] = alias.name + elif isinstance(node, ast.ImportFrom): + module = node.module or "" + for alias in node.names: + import_name = alias.asname if alias.asname else alias.name + full_import = f"{module}.{alias.name}" if module else alias.name + imports.append(full_import) + module_imports[import_name] = full_import - if isinstance(node, (ast.FunctionDef, ast.ClassDef, ast.AsyncFunctionDef)): + logger.info(f" Found {len(imports)} imports") + + # Second pass: Process declarations + decl_count = 0 + for node in ast.walk(tree): + chunk_text = None + metadata = None + + # FUNCTION DEFINITIONS + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): docstring = ast.get_docstring(node) or "" start_line = node.lineno - 1 end_line = getattr(node, "end_lineno", start_line + 1) + + # Guard the end_line + if end_line <= start_line: + end_line = start_line + 1 + chunk_code = "".join(lines[start_line:end_line]) - # Extract enhanced metadata metadata = { "file": str(filepath), - "name": getattr(node, "name", ""), - "type": type(node).__name__, + "name": node.name, + "type": "function" if isinstance(node, ast.FunctionDef) else "async_function", "line": node.lineno, "language": "python", } - # Add specific metadata based on node type - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): - metadata.update(_extract_function_metadata(node, code)) - elif isinstance(node, ast.ClassDef): - metadata.update(_extract_class_metadata(node, code)) + # Extract function metadata + func_meta = _extract_function_metadata(node, code) + metadata.update(func_meta) - # Include imports in the context + # Include imports if imports: metadata["imports"] = imports - chunk_text = f"File: {filepath}\nType: {type(node).__name__}\nName: {getattr(node,'name', '')}\nDocstring: {docstring}\n" + # Build chunk text + chunk_text = f"File: {filepath}\nType: {metadata['type']}\nName: {node.name}\n" - # Add enhanced information to text - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): - args = _format_arguments(node.args) - returns = _extract_return_annotation(node) - chunk_text += f"Arguments: {args}\n" - if returns: - chunk_text += f"Returns: {returns}\n" - elif isinstance(node, ast.ClassDef): - bases = [ast.unparse(base) for base in node.bases] if hasattr(ast, 'unparse') else [ast.dump(base) for base in node.bases] - if bases: - chunk_text += f"Base Classes: {', '.join(bases)}\n" + # Add decorators + if func_meta.get("decorators"): + chunk_text += f"Decorators: {', '.join(func_meta['decorators'])}\n" - chunk_text += f"\nCode:\n{chunk_code}" + # Add arguments + args = _format_arguments(node.args) + chunk_text += f"Arguments: {args}\n" + # Add return type + returns = _extract_return_annotation(node) + if returns: + chunk_text += f"Returns: {returns}\n" + + # Add async/generator info + if isinstance(node, ast.AsyncFunctionDef): + chunk_text += "Async: yes\n" + if func_meta.get("is_generator"): + chunk_text += "Generator: yes\n" + + # Add docstring + if docstring: + chunk_text += f"Docstring:\n{docstring}\n\n" + + chunk_text += f"Code:\n{chunk_code}" + + # CLASS DEFINITIONS + elif isinstance(node, ast.ClassDef): + docstring = ast.get_docstring(node) or "" + start_line = node.lineno - 1 + end_line = getattr(node, "end_lineno", start_line + 1) + + if end_line <= start_line: + end_line = start_line + 1 + + chunk_code = "".join(lines[start_line:end_line]) + + metadata = { + "file": str(filepath), + "name": node.name, + "type": "class", + "line": node.lineno, + "language": "python", + } + + # Extract class metadata + class_meta = _extract_class_metadata(node, code) + metadata.update(class_meta) + + # Include imports + if imports: + metadata["imports"] = imports + + # Build chunk text + chunk_text = f"File: {filepath}\nType: class\nName: {node.name}\n" + + # Add decorators + if class_meta.get("decorators"): + chunk_text += f"Decorators: {', '.join(class_meta['decorators'])}\n" + + # Add base classes + if class_meta.get("bases"): + chunk_text += f"Base Classes: {', '.join(class_meta['bases'])}\n" + + # Add methods + if class_meta.get("methods"): + chunk_text += f"Methods: {', '.join(class_meta['methods'])}\n" + + # Add properties + if class_meta.get("properties"): + chunk_text += f"Properties: {', '.join(class_meta['properties'])}\n" + + # Add class variables + if class_meta.get("class_variables"): + chunk_text += f"Class Variables: {', '.join(class_meta['class_variables'])}\n" + + # Add docstring + if docstring: + chunk_text += f"Docstring:\n{docstring}\n\n" + + chunk_text += f"Code:\n{chunk_code}" + + # MODULE-LEVEL ASSIGNMENTS (constants, globals) + elif isinstance(node, ast.Assign) and node.lineno: + # Only process top-level assignments (not inside functions/classes) + if _is_module_level(node, tree): + for target in node.targets: + if isinstance(target, ast.Name): + # Check if it looks like a constant (UPPER_CASE) + if target.id.isupper(): + start_line = node.lineno - 1 + end_line = getattr(node, "end_lineno", start_line + 1) + + if end_line <= start_line: + end_line = start_line + 1 + + chunk_code = "".join(lines[start_line:end_line]) + + metadata = { + "file": str(filepath), + "name": target.id, + "type": "constant", + "line": node.lineno, + "language": "python", + } + + # Try to extract value + try: + if hasattr(ast, 'unparse'): + value_str = ast.unparse(node.value) + metadata["value"] = value_str + except Exception: + pass + + chunk_text = f"File: {filepath}\nType: constant\nName: {target.id}\n" + chunk_text += f"Code:\n{chunk_code}" + + # ANNOTATED ASSIGNMENTS (type-annotated variables) + elif isinstance(node, ast.AnnAssign) and node.lineno: + if _is_module_level(node, tree) and isinstance(node.target, ast.Name): + start_line = node.lineno - 1 + end_line = getattr(node, "end_lineno", start_line + 1) + + if end_line <= start_line: + end_line = start_line + 1 + + chunk_code = "".join(lines[start_line:end_line]) + + metadata = { + "file": str(filepath), + "name": node.target.id, + "type": "variable", + "line": node.lineno, + "language": "python", + } + + # Extract type annotation + if node.annotation: + try: + if hasattr(ast, 'unparse'): + type_str = ast.unparse(node.annotation) + metadata["var_type"] = type_str + except Exception: + pass + + chunk_text = f"File: {filepath}\nType: variable\nName: {node.target.id}\n" + if metadata.get("var_type"): + chunk_text += f"Type: {metadata['var_type']}\n" + chunk_text += f"Code:\n{chunk_code}" + + # Add the chunk if we created one + if chunk_text and metadata: chunks.append({ "text": chunk_text, "metadata": metadata }) + decl_count += 1 + logger.info(f" Parsed {decl_count} declarations from {filepath}") + + except SyntaxError as e: + logger.warning(f"Python syntax error in {filepath}: {e}") except Exception as e: - logger.warning(f"parse_python_file: failed {filepath}: {e}") - return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks) + logger.warning(f"Failed to parse Python file {filepath}: {e}") + + # Return parsed chunks or fallback + if chunks: + logger.info(f"✅ Successfully parsed {len(chunks)} chunks from {filepath}") + return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks) + + # Fallback: emit a whole-file chunk + logger.warning(f"❌ No chunks created from Python parsing for {filepath}. Emitting fallback whole-file chunk.") + try: + full_code = filepath.read_text(encoding="utf-8") + except Exception: + full_code = "" + + fallback_meta = { + "file": str(filepath), + "name": filepath.stem, + "type": "module", + "line": 1, + "language": "python", + } + return ((full_code, tuple(fallback_meta.items())),) + + +def _is_module_level(node, tree): + """Check if a node is at module level (not inside a function or class).""" + for parent in ast.walk(tree): + if isinstance(parent, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + for child in ast.walk(parent): + if child is node: + return False + return True + def _extract_import_info(node): - """Extract import information for module dependencies.""" + """Extract import information from Import or ImportFrom nodes.""" if isinstance(node, ast.Import): - return { - "type": "import", - "modules": [alias.name for alias in node.names], - "level": 0 - } + return [alias.name for alias in node.names] elif isinstance(node, ast.ImportFrom): - return { - "type": "import_from", - "module": node.module, - "names": [alias.name for alias in node.names], - "level": node.level - } - return None + module = node.module or "" + return [f"{module}.{alias.name}" if module else alias.name for alias in node.names] + return [] + def _extract_function_metadata(node, code): - """Extract enhanced metadata for functions/methods.""" + """Extract enhanced metadata from function nodes.""" metadata = {} - # Extract arguments - args = node.args - arg_names = [arg.arg for arg in args.args] - if args.vararg: - arg_names.append(f"*{args.vararg.arg}") - if args.kwarg: - arg_names.append(f"**{args.kwarg.arg}") - - metadata["parameters"] = arg_names - # Extract decorators - decorators = [] - for decorator in node.decorator_list: - if isinstance(decorator, ast.Name): - decorators.append(decorator.id) - elif isinstance(decorator, ast.Attribute): - decorators.append(ast.unparse(decorator) if hasattr(ast, 'unparse') else ast.dump(decorator)) - elif isinstance(decorator, ast.Call): - decorators.append(ast.unparse(decorator.func) if hasattr(ast, 'unparse') else ast.dump(decorator.func)) + if node.decorator_list: + decorators = [] + for dec in node.decorator_list: + try: + if hasattr(ast, 'unparse'): + decorators.append(ast.unparse(dec)) + elif isinstance(dec, ast.Name): + decorators.append(dec.id) + elif isinstance(dec, ast.Call) and isinstance(dec.func, ast.Name): + decorators.append(dec.func.id) + except Exception: + pass + if decorators: + metadata["decorators"] = decorators - if decorators: - metadata["decorators"] = decorators + # Extract parameters with type hints + parameters = [] + if node.args.args: + for arg in node.args.args: + param_info = arg.arg + if arg.annotation: + try: + if hasattr(ast, 'unparse'): + param_info += f": {ast.unparse(arg.annotation)}" + except Exception: + pass + parameters.append(param_info) - # Try to extract return annotation - return_annotation = _extract_return_annotation(node) - if return_annotation: - metadata["return_type"] = return_annotation + if parameters: + metadata["parameters"] = parameters - # Detect if it's a method (has 'self' or 'cls' as first argument) - if arg_names and arg_names[0] in ('self', 'cls'): - metadata["is_method"] = True - # Try to find the containing class - metadata["method_type"] = "classmethod" if arg_names[0] == 'cls' else "instancemethod" - else: - metadata["is_method"] = False + # Check for special function types + if node.args.vararg: + metadata["has_varargs"] = True + if node.args.kwarg: + metadata["has_kwargs"] = True + + # Check if it's a generator (contains yield) + for child in ast.walk(node): + if isinstance(child, (ast.Yield, ast.YieldFrom)): + metadata["is_generator"] = True + break + + # Check for property decorator + if any(isinstance(dec, ast.Name) and dec.id == "property" for dec in node.decorator_list): + metadata["is_property"] = True + + # Check for staticmethod/classmethod + for dec in node.decorator_list: + if isinstance(dec, ast.Name): + if dec.id == "staticmethod": + metadata["is_static"] = True + elif dec.id == "classmethod": + metadata["is_classmethod"] = True + + # Extract return type annotation + if node.returns: + try: + if hasattr(ast, 'unparse'): + metadata["return_type"] = ast.unparse(node.returns) + except Exception: + pass return metadata + def _extract_class_metadata(node, code): - """Extract enhanced metadata for classes.""" + """Extract enhanced metadata from class nodes.""" metadata = {} + # Extract decorators + if node.decorator_list: + decorators = [] + for dec in node.decorator_list: + try: + if hasattr(ast, 'unparse'): + decorators.append(ast.unparse(dec)) + elif isinstance(dec, ast.Name): + decorators.append(dec.id) + except Exception: + pass + if decorators: + metadata["decorators"] = decorators + # Extract base classes bases = [] for base in node.bases: - if isinstance(base, ast.Name): - bases.append(base.id) - elif isinstance(base, ast.Attribute): - bases.append(ast.unparse(base) if hasattr(ast, 'unparse') else ast.dump(base)) - + try: + if hasattr(ast, 'unparse'): + bases.append(ast.unparse(base)) + elif isinstance(base, ast.Name): + bases.append(base.id) + except Exception: + pass if bases: - metadata["base_classes"] = bases + metadata["bases"] = bases - # Extract decorators - decorators = [] - for decorator in node.decorator_list: - if isinstance(decorator, ast.Name): - decorators.append(decorator.id) + # Extract methods and properties + methods = [] + properties = [] + static_methods = [] + class_methods = [] + class_variables = [] - if decorators: - metadata["decorators"] = decorators + for item in node.body: + if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)): + # Check decorators + is_property = any( + isinstance(dec, ast.Name) and dec.id == "property" + for dec in item.decorator_list + ) + is_static = any( + isinstance(dec, ast.Name) and dec.id == "staticmethod" + for dec in item.decorator_list + ) + is_classmethod = any( + isinstance(dec, ast.Name) and dec.id == "classmethod" + for dec in item.decorator_list + ) - # Try to detect class type from common patterns - class_code = ast.unparse(node) if hasattr(ast, 'unparse') else ast.dump(node) - if "metaclass" in class_code: - metadata["has_metaclass"] = True + if is_property: + properties.append(item.name) + elif is_static: + static_methods.append(item.name) + elif is_classmethod: + class_methods.append(item.name) + else: + methods.append(item.name) - # Look for common base class patterns - for base in bases: - if base in ("Exception", "BaseException"): - metadata["class_type"] = "exception" - break - elif base in ("Enum", "IntEnum", "StrEnum"): - metadata["class_type"] = "enum" - break - elif base in ("Model", "BaseModel"): - metadata["class_type"] = "model" - break + elif isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name): + # Class variable with type annotation + class_variables.append(item.target.id) + + elif isinstance(item, ast.Assign): + # Class variable without type annotation + for target in item.targets: + if isinstance(target, ast.Name): + class_variables.append(target.id) + + if methods: + metadata["methods"] = methods + if properties: + metadata["properties"] = properties + if static_methods: + metadata["static_methods"] = static_methods + if class_methods: + metadata["class_methods"] = class_methods + if class_variables: + metadata["class_variables"] = class_variables + + # Check for metaclass + for keyword in node.keywords: + if keyword.arg == "metaclass": + try: + if hasattr(ast, 'unparse'): + metadata["metaclass"] = ast.unparse(keyword.value) + elif isinstance(keyword.value, ast.Name): + metadata["metaclass"] = keyword.value.id + except Exception: + pass return metadata -def _format_arguments(args): - """Format function arguments in a readable way.""" - parts = [] - # Positional arguments +def _format_arguments(args): + """Format function arguments with type hints.""" + arg_strs = [] + + # Regular arguments for arg in args.args: - parts.append(arg.arg) + arg_str = arg.arg + if arg.annotation: + try: + if hasattr(ast, 'unparse'): + arg_str += f": {ast.unparse(arg.annotation)}" + except Exception: + pass + arg_strs.append(arg_str) # *args if args.vararg: - parts.append(f"*{args.vararg.arg}") - - # Keyword-only arguments - for arg in args.kwonlyargs: - parts.append(arg.arg) + vararg_str = f"*{args.vararg.arg}" + if args.vararg.annotation: + try: + if hasattr(ast, 'unparse'): + vararg_str += f": {ast.unparse(args.vararg.annotation)}" + except Exception: + pass + arg_strs.append(vararg_str) # **kwargs if args.kwarg: - parts.append(f"**{args.kwarg.arg}") + kwarg_str = f"**{args.kwarg.arg}" + if args.kwarg.annotation: + try: + if hasattr(ast, 'unparse'): + kwarg_str += f": {ast.unparse(args.kwarg.annotation)}" + except Exception: + pass + arg_strs.append(kwarg_str) + + return ", ".join(arg_strs) if arg_strs else "None" - return ", ".join(parts) def _extract_return_annotation(node): - """Extract return type annotation if available.""" - if hasattr(node, 'returns') and node.returns: - if isinstance(node.returns, ast.Name): - return node.returns.id - elif isinstance(node.returns, ast.Attribute): - return ast.unparse(node.returns) if hasattr(ast, 'unparse') else ast.dump(node.returns) - elif isinstance(node.returns, ast.Subscript): - return ast.unparse(node.returns) if hasattr(ast, 'unparse') else ast.dump(node.returns) + """Extract return type annotation from function node.""" + if not node.returns: + return None + + try: + if hasattr(ast, 'unparse'): + return ast.unparse(node.returns) + except Exception: + pass + return None def parse_python_file(filepath: Path) -> List[Dict]: @@ -2540,50 +2838,53 @@ def build_indexes(): # ------------------------------------------------- texts, metadatas = index_codebase(CODEBASE_PATH) if len(texts) != len(metadatas): - # Helpful debug output to find offending files quickly logger.error( "Indexing mismatch: %d texts vs %d metadata entries. Aborting index build.", len(texts), len(metadatas), ) - # raise so CI/dev runs fail fast and you can inspect logs raise ValueError(f"Indexing produced {len(texts)} texts but {len(metadatas)} metadata entries") logger.info("=== Sample metadata inspection ===") - for i, m in enumerate(metadatas[:5]): # Check first 5 + for i, m in enumerate(metadatas[:5]): logger.info(f"Chunk {i}: type={m.get('type')}, name={m.get('name')}") logger.info(f" imports={m.get('imports')}") logger.info(f" calls={m.get('calls')}") logger.info(f" parameters={m.get('parameters')}") - # FIX: Convert Path to string for LocalGraph + # Initialize graph graph = LocalGraph(root_path=str(CODEBASE_PATH)) graph.clear() # Track all created nodes for second-pass relationships created_nodes = {} - # Create graph nodes/edges from metadata with enhanced relationships + # Track language-specific features for cross-language relationships + component_registry = {} # Svelte components + interface_registry = {} # TypeScript/Java interfaces + trait_registry = {} # Rust traits + class_registry = {} # Classes across all languages + + logger.info("=== Phase 1: Creating nodes ===") + + # ------------------------------------------------- + # PHASE 1: Create all nodes + # ------------------------------------------------- for m in metadatas: fpath = m.get("file") lang = m.get("language", "unknown") node_type = m.get("type", "unknown") name = m.get("name", "") - block_type = m.get("block_type", "") - component_name = m.get("component_name") # Skip empty/invalid nodes if not name or node_type == "unknown" or not fpath: continue - # convert rust defs + # Normalize metadata if "implements_traits" in m and m["implements_traits"]: m["implements"] = m["implements_traits"] - # ---- File node --------------------------------------------- - # Keep the same ID (`file::`) but store the full set of - # attributes so that downstream code can rely on `name`, `lang`, - # and `file` (even though `name` will just be the file path). + # Create file node file_node_id = f"file::{fpath}" if file_node_id not in graph.graph.nodes: graph.add_node( @@ -2600,7 +2901,7 @@ def build_indexes(): "lang": lang, } - # Create node ID with language + # Create node ID with language prefix node_id = f"{lang}::{node_type}::{fpath}::{name}" # Base attributes for all nodes @@ -2608,203 +2909,434 @@ def build_indexes(): "name": name, "file": fpath, "line": m.get("line"), - "lang": lang # Always include language + "lang": lang } + # Add docstring/documentation if available + if m.get("docstring"): + base_attrs["docstring"] = m.get("docstring") + # Store for second-pass processing created_nodes[node_id] = base_attrs.copy() created_nodes[node_id]["type"] = node_type + created_nodes[node_id]["metadata"] = m # Keep full metadata for second pass - # Create appropriate node based on type - if node_type in ["function", "method", "constructor"]: - # Function-like entities + # ===== FUNCTIONS & METHODS ===== + if node_type in ["function", "method", "async_function", "constructor"]: if node_type == "method" and "class" in m: - # Method belongs to a class - use compound ID class_name = m['class'] node_id = f"{lang}::{node_type}::{fpath}::{class_name}.{name}" - # Update created_nodes with new ID created_nodes[node_id] = base_attrs.copy() created_nodes[node_id]["type"] = node_type created_nodes[node_id]["class_name"] = class_name + created_nodes[node_id]["metadata"] = m - graph.add_node(node_id, **base_attrs, type=node_type, - class_name=class_name, - return_type=m.get("return_type"), - parameters=m.get("parameters", [])) + # Enhanced method attributes + graph.add_node( + node_id, + **base_attrs, + type=node_type, + class_name=class_name, + return_type=m.get("return_type"), + parameters=m.get("parameters", []), + visibility=m.get("visibility"), + is_static=m.get("is_static", False), + is_abstract=m.get("is_abstract", False), + is_async=m.get("is_async", False), + decorators=m.get("decorators", []), + annotations=m.get("annotations", []) + ) - # Connect method to its class + # Connect method to class class_node_id = f"{lang}::class::{fpath}::{class_name}" graph.add_edge(class_node_id, node_id, "contains") - - # Connect return type if available - return_type = m.get("return_type") - if return_type and return_type not in ['void', '()', 'Result', '']: - for type_id, type_data in graph.find_nodes(name=return_type): - graph.add_edge(node_id, type_id, "returns") - - # Connect parameter types - parameters = m.get("parameters", []) - for param in parameters: - if ':' in param: - param_parts = param.split(':') - if len(param_parts) >= 2: - param_type = param_parts[-1].strip() - param_type = ( - param_type - .replace('&', '') - .replace('mut ', '') - .replace('<', '') - .replace('>', '') - .replace(',', '') - .strip() - ) - param_type = param_type.replace('&', '').replace('mut ', '').strip() - if (param_type and - param_type not in ['self', '&self', '&mut self', 'mut self'] and - not param_type.startswith('&') and - param_type not in ['i32', 'i64', 'f32', 'f64', 'bool', 'String', 'str']): - for type_id, type_data in graph.find_nodes(name=param_type): - graph.add_edge(node_id, type_id, "uses_parameter") - elif node_type == "impl": - # ---------- IMPL NODES ---------- - node_id = f"{lang}::{node_type}::{fpath}::{name}" - created_nodes[node_id] = base_attrs.copy() - created_nodes[node_id]["type"] = node_type - - graph.add_node(node_id, **base_attrs, type=node_type) - - # Guard against empty lists that can be returned by the Rust parser - target_name = None - if "implements_traits" in m and isinstance(m["implements_traits"], list) and m["implements_traits"]: - # assume first trait - target_name = m["implements_traits"][0] - elif "implements" in m and isinstance(m["implements"], list) and m["implements"]: - target_name = m["implements"][0] - - - if target_name: - target_node_id = f"{lang}::struct::{fpath}::{target_name}" - if target_node_id not in graph.graph: - graph.add_node(target_node_id, name=target_name, file=fpath, lang=lang, type="struct") - graph.add_edge(node_id, target_node_id, "implements") else: - # Regular function - graph.add_node(node_id, **base_attrs, type=node_type, - return_type=m.get("return_type"), - parameters=m.get("parameters", [])) - # Connect function to file + # Regular function or async function + graph.add_node( + node_id, + **base_attrs, + type=node_type, + return_type=m.get("return_type"), + parameters=m.get("parameters", []), + visibility=m.get("visibility"), + is_async=m.get("is_async", False), + is_generator=m.get("is_generator", False), + is_exported=m.get("is_exported", False), + decorators=m.get("decorators", []), + annotations=m.get("annotations", []) + ) + graph.add_edge(file_node_id, node_id, "contains") - # Connect return type if available - return_type = m.get("return_type") - if return_type and return_type not in ['void', '()', 'Result', '']: - for type_id, type_data in graph.find_nodes(name=return_type): - graph.add_edge(node_id, type_id, "returns") + # ===== CLASSES ===== + elif node_type == "class": + graph.add_node( + node_id, + **base_attrs, + type=node_type, + extends=m.get("extends"), + implements=m.get("implements", []), + methods=m.get("methods", []), + properties=m.get("properties", []), + fields=m.get("fields", []), + static_methods=m.get("static_methods", []), + class_methods=m.get("class_methods", []), + visibility=m.get("visibility"), + is_abstract=m.get("is_abstract", False), + decorators=m.get("decorators", []), + annotations=m.get("annotations", []) + ) - # Connect parameter types - parameters = m.get("parameters", []) - for param in parameters: - if ':' in param: - param_parts = param.split(':') - if len(param_parts) >= 2: - param_type = param_parts[-1].strip() - param_type = ( - param_type - .replace('&', '') - .replace('mut ', '') - .replace('<', '') - .replace('>', '') - .replace(',', '') - .strip() - ) - param_type = param_type.replace('&', '').replace('mut ', '').strip() - if (param_type and - param_type not in ['self', '&self', '&mut self', 'mut self'] and - not param_type.startswith('&') and - param_type not in ['i32', 'i64', 'f32', 'f64', 'bool', 'String', 'str']): - for type_id, type_data in graph.find_nodes(name=param_type): - graph.add_edge(node_id, type_id, "uses_parameter") - - elif node_type in ["class", "struct", "interface", "enum"]: - # Type definitions - graph.add_node(node_id, **base_attrs, type=node_type, - extends=m.get("extends"), - implements=m.get("implements", []), - fields=m.get("fields", []), - variants=m.get("variants", [])) - - # Connect type to file graph.add_edge(file_node_id, node_id, "contains") - # Store additional attributes for second-pass - created_nodes[node_id]["extends"] = m.get("extends") - created_nodes[node_id]["implements"] = m.get("implements", []) - created_nodes[node_id]["fields"] = m.get("fields", []) - created_nodes[node_id]["variants"] = m.get("variants", []) + # Register class for cross-references + class_registry[name] = { + "node_id": node_id, + "lang": lang, + "file": fpath, + "metadata": m + } + + # ===== STRUCTS ===== + elif node_type == "struct": + graph.add_node( + node_id, + **base_attrs, + type=node_type, + fields=m.get("fields", []), + visibility=m.get("visibility"), + generics=m.get("generics", []), + derives=m.get("derives", []) + ) + + graph.add_edge(file_node_id, node_id, "contains") + + # Register for type relationships + class_registry[name] = { + "node_id": node_id, + "lang": lang, + "file": fpath, + "metadata": m + } + + # ===== INTERFACES ===== + elif node_type == "interface": + graph.add_node( + node_id, + **base_attrs, + type=node_type, + extends=m.get("extends", []), + methods=m.get("methods", []), + properties=m.get("properties", []), + visibility=m.get("visibility"), + generics=m.get("generics", []) + ) + + graph.add_edge(file_node_id, node_id, "contains") + + # Register interface + interface_registry[name] = { + "node_id": node_id, + "lang": lang, + "file": fpath, + "metadata": m + } + + # ===== TRAITS ===== elif node_type == "trait": - # Traits - graph.add_node(node_id, **base_attrs, type=node_type, - methods=m.get("methods", [])) + graph.add_node( + node_id, + **base_attrs, + type=node_type, + methods=m.get("methods", []), + visibility=m.get("visibility"), + generics=m.get("generics", []) + ) + graph.add_edge(file_node_id, node_id, "contains") - created_nodes[node_id]["methods"] = m.get("methods", []) + # Register trait + trait_registry[name] = { + "node_id": node_id, + "lang": lang, + "file": fpath, + "metadata": m + } - elif node_type == "module": - # Modules + # ===== RUST IMPL BLOCKS ===== + elif node_type == "impl": + graph.add_node( + node_id, + **base_attrs, + type=node_type, + traits=m.get("traits", []), + methods=m.get("methods", []), + generics=m.get("generics", []) + ) + + graph.add_edge(file_node_id, node_id, "contains") + + # ===== ENUMS ===== + elif node_type == "enum": + graph.add_node( + node_id, + **base_attrs, + type=node_type, + variants=m.get("variants", []), + members=m.get("members", []), + visibility=m.get("visibility") + ) + + graph.add_edge(file_node_id, node_id, "contains") + + # ===== TYPE ALIASES ===== + elif node_type == "type": + graph.add_node( + node_id, + **base_attrs, + type=node_type, + type_definition=m.get("type_definition"), + generics=m.get("generics", []) + ) + + graph.add_edge(file_node_id, node_id, "contains") + + # ===== SVELTE COMPONENTS ===== + elif node_type == "markup" and lang == "svelte": + component_name = m.get("component_name", name) + component_node_id = f"svelte::component::{fpath}::{component_name}" + + graph.add_node( + component_node_id, + **base_attrs, + type="component", + component_name=component_name, + props=m.get("component_props", []), + slots=m.get("slots", []), + has_default_slot=m.get("has_default_slot", False), + events=m.get("event_handlers", []), + used_components=m.get("used_components", []), + stores=m.get("stores", []) + ) + + graph.add_edge(file_node_id, component_node_id, "contains") + + # Register component + component_registry[component_name] = { + "node_id": component_node_id, + "file": fpath, + "metadata": m + } + + node_id = component_node_id + + # ===== VARIABLES & CONSTANTS ===== + elif node_type in ["variable", "constant", "prop", "reactive"]: + graph.add_node( + node_id, + **base_attrs, + type=node_type, + var_type=m.get("var_type"), + var_kind=m.get("var_kind"), # const, let, var + is_exported=m.get("is_exported", False), + is_reactive=m.get("is_reactive", False), + is_store=m.get("is_store", False) + ) + + graph.add_edge(file_node_id, node_id, "contains") + + # ===== STYLES ===== + elif node_type == "style": + graph.add_node( + node_id, + **base_attrs, + type=node_type, + css_classes=m.get("css_classes", []), + css_ids=m.get("css_ids", []), + css_variables=m.get("css_variables", []), + is_scoped=m.get("is_scoped", True) + ) + + graph.add_edge(file_node_id, node_id, "contains") + + # ===== IMPORTS ===== + elif node_type == "import": + # Create import node + import_path = m.get("import_path") or m.get("full_name") or name + import_node_id = f"{lang}::import::{fpath}::{import_path}" + + graph.add_node( + import_node_id, + **base_attrs, + type=node_type, + import_path=import_path, + imported_names=m.get("imported_names", []) + ) + + graph.add_edge(file_node_id, import_node_id, "imports") + + # ===== GENERIC FALLBACK ===== + else: graph.add_node(node_id, **base_attrs, type=node_type) graph.add_edge(file_node_id, node_id, "contains") - else: - # Generic code element fallback - graph.add_node(node_id, **base_attrs, type=node_type) - graph.add_edge(file_node_id, node_id, "contains") + logger.info(f"Phase 1 complete: Created {len(created_nodes)} nodes") + logger.info(f" Classes: {len(class_registry)}, Interfaces: {len(interface_registry)}") + logger.info(f" Traits: {len(trait_registry)}, Components: {len(component_registry)}") - # ---------- IMPORTS ---------- - if node_type == "import" and "parameters" in m: - imports_raw = m["parameters"] # override metadata pathway - else: - imports_raw = m.get("imports", []) + # ------------------------------------------------- + # PHASE 2: Build relationships + # ------------------------------------------------- + logger.info("=== Phase 2: Building relationships ===") + + for node_id, node_data in created_nodes.items(): + m = node_data.get("metadata", {}) + if not m: + continue + + node_type = node_data.get("type") + fpath = node_data.get("file") + lang = node_data.get("lang") + + # ===== IMPORTS & DEPENDENCIES ===== + imports_raw = m.get("imports", []) if isinstance(imports_raw, str): - # Split the string back into a list imports = [imp.strip() for imp in imports_raw.split(",") if imp.strip()] else: - imports = imports_raw + imports = imports_raw if isinstance(imports_raw, list) else [] - graph._add_import_edges(file_node_id, imports) + if imports: + file_node_id = f"file::{fpath}" + graph._add_import_edges(file_node_id, imports) - # ---------- FUNCTION CALLS ---------- + # ===== FUNCTION CALLS ===== calls_raw = m.get("calls", []) if isinstance(calls_raw, str): - # Split the string back into a list calls = [c.strip() for c in calls_raw.split(",") if c.strip()] else: - calls = calls_raw + calls = calls_raw if isinstance(calls_raw, list) else [] - graph._add_call_edges(node_id, calls, fpath) + if calls: + graph._add_call_edges(node_id, calls, fpath) + # ===== CLASS INHERITANCE ===== + if node_type == "class": + # Extends relationships + extends = m.get("extends") + if extends: + for parent_name, parent_info in class_registry.items(): + if parent_name == extends: + graph.add_edge(node_id, parent_info["node_id"], "extends") + + # Implements relationships + implements = m.get("implements", []) + for interface_name in implements: + for iface_name, iface_info in interface_registry.items(): + if iface_name == interface_name: + graph.add_edge(node_id, iface_info["node_id"], "implements") + + # ===== INTERFACE INHERITANCE ===== + if node_type == "interface": + extends_list = m.get("extends", []) + if isinstance(extends_list, str): + extends_list = [extends_list] + + for parent_interface in extends_list: + for iface_name, iface_info in interface_registry.items(): + if iface_name == parent_interface: + graph.add_edge(node_id, iface_info["node_id"], "extends") + + # ===== RUST IMPL BLOCKS ===== + if node_type == "impl": + traits = m.get("traits", []) + if isinstance(traits, str): + traits = [traits] + + for trait_name in traits: + # Link to trait + for t_name, t_info in trait_registry.items(): + if t_name == trait_name: + graph.add_edge(node_id, t_info["node_id"], "implements") + + # Link to struct/type being implemented + for class_name, class_info in class_registry.items(): + if class_name in m.get("name", ""): + graph.add_edge(node_id, class_info["node_id"], "implements_for") + + # ===== METHOD RETURN TYPES & PARAMETERS ===== + if node_type in ["function", "method", "async_function"]: + return_type = m.get("return_type") + if return_type and return_type not in ['void', '()', 'Result', 'None', '', 'any']: + # Clean up return type + clean_return = _clean_type_string(return_type) + + # Find matching type definition + for type_name, type_info in {**class_registry, **interface_registry, **trait_registry}.items(): + if type_name == clean_return: + graph.add_edge(node_id, type_info["node_id"], "returns") + + # Parameter types + parameters = m.get("parameters", []) + if isinstance(parameters, list): + for param in parameters: + param_type = _extract_param_type(param) + if param_type: + for type_name, type_info in {**class_registry, **interface_registry}.items(): + if type_name == param_type: + graph.add_edge(node_id, type_info["node_id"], "uses_parameter") + + # ===== SVELTE COMPONENT RELATIONSHIPS ===== + if node_type == "markup" and lang == "svelte": + # Link to used components + used_components = m.get("used_components", []) + for comp_name in used_components: + if comp_name in component_registry: + comp_info = component_registry[comp_name] + graph.add_edge(node_id, comp_info["node_id"], "uses_component") + + # ===== FIELD TYPE RELATIONSHIPS ===== + if node_type in ["struct", "class"]: + fields = m.get("fields", []) + for field in fields: + field_type = _extract_field_type(field) + if field_type: + for type_name, type_info in {**class_registry, **interface_registry}.items(): + if type_name == field_type: + graph.add_edge(node_id, type_info["node_id"], "has_field_type") + + logger.info("Phase 2 complete: Built relationships") # ------------------------------------------------- - # Step 2: Second-pass relationship building + # PHASE 3: Advanced relationship building # ------------------------------------------------- - logger.info("Building second-pass relationships...") + logger.info("=== Phase 3: Advanced relationships ===") # Build type relationships type_connections = build_type_relationships(graph, created_nodes) logger.info(f"Added {type_connections} type relationships") - # Build Bevy-specific relationships + # Build framework-specific relationships bevy_connections = build_bevy_relationships(graph, created_nodes) logger.info(f"Added {bevy_connections} Bevy-specific relationships") + # Build cross-language relationships + cross_lang_connections = build_cross_language_relationships( + graph, created_nodes, component_registry, interface_registry, class_registry + ) + logger.info(f"Added {cross_lang_connections} cross-language relationships") + + # Save graph graph.save() graph.to_json() graph.to_toon() logger.info(f"Graph built with {graph.graph.number_of_nodes()} nodes and {graph.graph.number_of_edges()} edges.") # ------------------------------------------------- - # Step 2: embedding/index pipeline + # PHASE 4: Embedding/index pipeline # ------------------------------------------------- + logger.info("=== Phase 4: Building embeddings ===") + os.environ["OLLAMA_HOST"] = OLLAMA_BASE_URL embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL, base_url=OLLAMA_BASE_URL) @@ -2819,45 +3351,24 @@ def build_indexes(): # Create vector store with weighted embeddings logger.info("Creating vector store (Chroma) with weighted embeddings...") - # Create a wrapper that will return our pre-computed weighted embeddings class WeightedEmbeddingFunction: def __init__(self, texts, weighted_embs): - """ - texts: list[str] - weighted_embs: list[list[float]] - """ - - # Map each text to a queue of its embeddings (duplicate-safe) self.lookup = defaultdict(deque) - for t, emb in zip(texts, weighted_embs): self.lookup[t].append(emb) def embed_documents(self, docs): - """ - Chroma may pass docs in any size or order. - Always return correct per-document embeddings. - """ result = [] for d in docs: if not self.lookup[d]: - raise ValueError( - f"No embedding left for text {repr(d[:80])} — " - f"likely mismatch between texts and embeddings." - ) + raise ValueError(f"No embedding left for text {repr(d[:80])}") result.append(self.lookup[d].popleft()) return result def embed_query(self, q): - """ - Queries should use base model, not weighted context embeddings. - """ return embeddings.embed_query(q) - weighted_emb_func = WeightedEmbeddingFunction( - texts=texts, - weighted_embs=weighted_embeddings - ) + weighted_emb_func = WeightedEmbeddingFunction(texts=texts, weighted_embs=weighted_embeddings) vectorstore = Chroma.from_texts( texts=texts, @@ -2878,9 +3389,90 @@ def build_indexes(): json.dumps({"corpus": texts, "metadata": metadatas}, indent=2), encoding="utf-8" ) + index_build_time = time.time() elapsed = index_build_time - start_time - logger.info(f"Index build complete with contextual weighting applied. Total time: {elapsed:.1f}s") + logger.info(f"✅ Index build complete. Total time: {elapsed:.1f}s") + + +# ------------------------------------------------- +# Helper Functions +# ------------------------------------------------- + +def _clean_type_string(type_str: str) -> str: + """Clean up type strings for matching.""" + if not type_str: + return "" + + # Remove common type wrappers + type_str = type_str.strip() + type_str = re.sub(r'Option<(.+)>', r'\1', type_str) + type_str = re.sub(r'Result<(.+?)(?:,|>).*', r'\1', type_str) + type_str = re.sub(r'Vec<(.+)>', r'\1', type_str) + type_str = re.sub(r'Box<(.+)>', r'\1', type_str) + type_str = re.sub(r'Arc<(.+)>', r'\1', type_str) + type_str = re.sub(r'Rc<(.+)>', r'\1', type_str) + type_str = re.sub(r'Promise<(.+)>', r'\1', type_str) + type_str = re.sub(r'Array<(.+)>', r'\1', type_str) + + # Remove references and mutability + type_str = type_str.replace('&', '').replace('mut ', '').strip() + + # Remove generics for simpler matching + type_str = re.sub(r'<.*>', '', type_str).strip() + + return type_str + + +def _extract_param_type(param: str) -> str: + """Extract type from parameter string.""" + if not param or ':' not in param: + return "" + + parts = param.split(':') + if len(parts) < 2: + return "" + + param_type = parts[-1].strip() + return _clean_type_string(param_type) + + +def _extract_field_type(field: str) -> str: + """Extract type from field string.""" + if not field or ':' not in field: + return "" + + parts = field.split(':') + if len(parts) < 2: + return "" + + field_type = parts[-1].strip() + # Remove struct tags for languages like Go + field_type = re.sub(r'`.*`', '', field_type).strip() + return _clean_type_string(field_type) + + +def build_cross_language_relationships(graph, created_nodes, component_registry, + interface_registry, class_registry): + """Build relationships across different programming languages.""" + connections = 0 + + # Example: Connect TypeScript interfaces to Rust structs with same name + for ts_name, ts_info in interface_registry.items(): + if ts_info["lang"] == "typescript": + for rust_name, rust_info in class_registry.items(): + if rust_info["lang"] == "rust" and ts_name == rust_name: + graph.add_edge(ts_info["node_id"], rust_info["node_id"], "mirrors") + connections += 1 + + # Connect Svelte components to their TypeScript definitions + for comp_name, comp_info in component_registry.items(): + for class_name, class_info in class_registry.items(): + if class_name == comp_name or class_name == f"{comp_name}Props": + graph.add_edge(comp_info["node_id"], class_info["node_id"], "uses_type") + connections += 1 + + return connections def build_type_relationships(graph, created_nodes): """Build type relationships between nodes (parameter types, return types, etc.)""" @@ -3874,7 +4466,6 @@ Example: semantic_rag_search("user authentication", 5, True) → 5 results with if vectorstore is None or bm25_corpus is None: return "❌ Index not built. Please call init_repo() or rebuild_index() first." - # FIX: Convert Path to string graph = LocalGraph(root_path=str(CODEBASE_PATH)) graph_loaded = graph.load()