diff --git a/mcp_codebase.py b/mcp_codebase.py index b0da96a..e960e71 100644 --- a/mcp_codebase.py +++ b/mcp_codebase.py @@ -41,7 +41,9 @@ OLLAMA_BASE_URL = "http://127.0.0.1:11434" os.environ["OLLAMA_HOST"] = OLLAMA_BASE_URL os.environ["OLLAMA_API_BASE"] = OLLAMA_BASE_URL -CODEBASE_PATH = Path(os.environ.get("CODEBASE_PATH", ".")) # change as needed +CODEBASE_PATH = Path("./working_repo") +CODEBASE_PATH.mkdir(exist_ok=True) + VECTOR_DB_PATH = Path(os.environ.get("VECTOR_DB_PATH", "./chroma_db")) BM25_INDEX_PATH = Path(os.environ.get("BM25_INDEX_PATH", "./bm25_index.json")) RAGIGNORE_PATH = CODEBASE_PATH / ".ragignore" @@ -329,23 +331,23 @@ def parse_python_file_cached(filepath_str: str, file_hash: str) -> tuple: code = filepath.read_text(encoding="utf-8") tree = ast.parse(code) lines = code.splitlines(keepends=True) - + # Track imports for module relationships 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.FunctionDef, ast.ClassDef, ast.AsyncFunctionDef)): docstring = ast.get_docstring(node) or "" start_line = node.lineno - 1 end_line = getattr(node, "end_lineno", start_line + 1) chunk_code = "".join(lines[start_line:end_line]) - + # Extract enhanced metadata metadata = { "file": str(filepath), @@ -354,19 +356,19 @@ def parse_python_file_cached(filepath_str: str, file_hash: str) -> tuple: "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)) - + # Include imports in the context if imports: metadata["imports"] = imports - + chunk_text = f"File: {filepath}\nType: {type(node).__name__}\nName: {getattr(node,'name', '')}\nDocstring: {docstring}\n" - + # Add enhanced information to text if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): args = _format_arguments(node.args) @@ -378,14 +380,14 @@ def parse_python_file_cached(filepath_str: str, file_hash: str) -> tuple: 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" - + chunk_text += f"\nCode:\n{chunk_code}" - + chunks.append({ "text": chunk_text, "metadata": metadata }) - + 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) @@ -410,7 +412,7 @@ def _extract_import_info(node): def _extract_function_metadata(node, code): """Extract enhanced metadata for functions/methods.""" metadata = {} - + # Extract arguments args = node.args arg_names = [arg.arg for arg in args.args] @@ -418,9 +420,9 @@ def _extract_function_metadata(node, code): 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: @@ -430,15 +432,15 @@ def _extract_function_metadata(node, code): 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 decorators: metadata["decorators"] = decorators - + # Try to extract return annotation return_annotation = _extract_return_annotation(node) if return_annotation: metadata["return_type"] = return_annotation - + # 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 @@ -446,13 +448,13 @@ def _extract_function_metadata(node, code): metadata["method_type"] = "classmethod" if arg_names[0] == 'cls' else "instancemethod" else: metadata["is_method"] = False - + return metadata def _extract_class_metadata(node, code): """Extract enhanced metadata for classes.""" metadata = {} - + # Extract base classes bases = [] for base in node.bases: @@ -460,24 +462,24 @@ def _extract_class_metadata(node, code): bases.append(base.id) elif isinstance(base, ast.Attribute): bases.append(ast.unparse(base) if hasattr(ast, 'unparse') else ast.dump(base)) - + if bases: metadata["base_classes"] = bases - + # Extract decorators decorators = [] for decorator in node.decorator_list: if isinstance(decorator, ast.Name): decorators.append(decorator.id) - + if decorators: metadata["decorators"] = decorators - + # 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 - + # Look for common base class patterns for base in bases: if base in ("Exception", "BaseException"): @@ -489,29 +491,29 @@ def _extract_class_metadata(node, code): elif base in ("Model", "BaseModel"): metadata["class_type"] = "model" break - + return metadata def _format_arguments(args): """Format function arguments in a readable way.""" parts = [] - + # Positional arguments for arg in args.args: parts.append(arg.arg) - + # *args if args.vararg: parts.append(f"*{args.vararg.arg}") - + # Keyword-only arguments for arg in args.kwonlyargs: parts.append(arg.arg) - + # **kwargs if args.kwarg: parts.append(f"**{args.kwarg.arg}") - + return ", ".join(parts) def _extract_return_annotation(node): @@ -690,7 +692,7 @@ def parse_sql_file_cached(filepath_str: str, file_hash: str) -> tuple: chunks.append({"text": chunk_text, "metadata": meta}) return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks) - + def parse_sql_file(filepath: Path) -> List[Dict[str, Any]]: """Parse SQL file with caching.""" file_hash = get_file_hash(filepath) @@ -827,7 +829,7 @@ def parse_go_file_cached(filepath_str: str, file_hash: str) -> tuple: "line": start + 1, "language": "go", } - + # Enhanced metadata for graph relationships if fields: metadata["fields"] = fields @@ -896,7 +898,7 @@ def parse_go_file_cached(filepath_str: str, file_hash: str) -> tuple: "line": start_line + 1, "language": "go", } - + # Receiver type for methods if receiver_type: metadata["receiver_type"] = receiver_type @@ -917,7 +919,7 @@ def parse_go_file_cached(filepath_str: str, file_hash: str) -> tuple: end_line = find_brace_block_end_go(lines, start_line) struct_code = "".join(lines[start_line:end_line + 1]) comments = extract_go_comments(lines, start_line) - + # Extract field names from struct field_names = [] struct_body_match = re.search(r'struct\s*\{([^}]+)\}', struct_code, re.DOTALL) @@ -938,11 +940,11 @@ def parse_go_file_cached(filepath_str: str, file_hash: str) -> tuple: metadata = { "file": str(filepath), "name": struct_name, - "type": "struct", + "type": "struct", "line": start_line + 1, "language": "go", } - + # Struct fields for graph relationships if field_names: metadata["fields"] = field_names @@ -961,7 +963,7 @@ def parse_go_file_cached(filepath_str: str, file_hash: str) -> tuple: end_line = find_brace_block_end_go(lines, start_line) interface_code = "".join(lines[start_line:end_line + 1]) comments = extract_go_comments(lines, start_line) - + # Extract method signatures from interface method_names = [] interface_body_match = re.search(r'interface\s*\{([^}]+)\}', interface_code, re.DOTALL) @@ -986,7 +988,7 @@ def parse_go_file_cached(filepath_str: str, file_hash: str) -> tuple: "line": start_line + 1, "language": "go", } - + # Interface methods for graph relationships if method_names: metadata["interface_methods"] = method_names @@ -1114,12 +1116,12 @@ def parse_java_file_cached(filepath_str: str) -> tuple: doc = get_context_with_comments(filepath, start_line + 1) extends_class = None implements_interfaces = [] - + if node.extends: extends_class = node.extends.name if hasattr(node, 'implements') and node.implements: implements_interfaces = [impl.name for impl in node.implements] - + chunks.append({ "text": f"File: {filepath}\nType: Class\nName: {node.name}\nExtends: {extends_class or 'None'}\nImplements: {', '.join(implements_interfaces) or 'None'}\nDoc:\n{doc}\nCode:\n{chunk_text}", "metadata": { @@ -1152,7 +1154,7 @@ def parse_java_file_cached(filepath_str: str) -> tuple: for param in node.parameters: param_type = param.type.name if param.type else "unknown" parameters.append(f"{param_type} {param.name}") - + chunks.append({ "text": f"File: {filepath}\nType: Method\nClass: {class_name or 'unknown'}\nName: {node.name}\nReturn: {return_type}\nParameters: {', '.join(parameters) or 'None'}\nDoc:\n{doc}\nCode:\n{chunk_text}", "metadata": { @@ -1176,7 +1178,7 @@ def parse_java_file_cached(filepath_str: str) -> tuple: end_line = start_line chunk_text = lines[start_line] doc = get_context_with_comments(filepath, start_line + 1) - field_type = node.type.name if node.type else "unknown" + field_type = node.type.name if node.type else "unknown" for declarator in node.declarators: chunks.append({ "text": f"File: {filepath}\nType: Field\nName: {declarator.name}\nField Type: {field_type}\nDoc:\n{doc}\nCode:\n{chunk_text}", @@ -1229,7 +1231,7 @@ def parse_java_file_cached(filepath_str: str) -> tuple: "language": "java" } }) - + elif isinstance(node, javalang.tree.InterfaceDeclaration): start_line = node.position.line - 1 if node.position else 0 end_line = find_brace_block_end(start_line) @@ -1284,28 +1286,28 @@ def parse_svelte_file_cached(filepath_str: str, file_hash: str) -> tuple: try: text = filepath.read_text(encoding="utf-8") lines = text.splitlines(keepends=True) - + # Track component props and context for graph relationships component_props = [] component_context = [] - + # Extract component name from filename for graph relationships component_name = filepath.stem if component_name[0].islower(): component_name = component_name[0].upper() + component_name[1:] # PascalCase - + script_matches = list(re.finditer(r"]*)?>(.*?)", text, flags=re.DOTALL)) idx = 0 - + for m in script_matches: script_content = m.group(1) script_start_line = text[:m.start(1)].count("\n") + 1 script_attrs = m.group(0).split('>')[0] # Get ", "", text, flags=re.DOTALL) markup = re.sub(r"]*>.*?", "", markup, flags=re.DOTALL).strip() - + if markup: # Extract component usage and bindings from markup used_components = re.findall(r'<([A-Z][a-zA-Z]*)', markup) prop_bindings = re.findall(r'(\w+)={([^}]+)}', markup) event_handlers = re.findall(r'on:(\w+)=', markup) - + chunk_text = f"File: {filepath}\nType: markup\nComponent: {component_name}\n" if used_components: chunk_text += f"Uses Components: {', '.join(set(used_components))}\n" @@ -1426,7 +1428,7 @@ def parse_svelte_file_cached(filepath_str: str, file_hash: str) -> tuple: if event_handlers: chunk_text += f"Event Handlers: {', '.join(set(event_handlers))}\n" chunk_text += f"Markup:\n{markup}" - + chunks.append({ "text": chunk_text, "metadata": { @@ -1434,7 +1436,7 @@ def parse_svelte_file_cached(filepath_str: str, file_hash: str) -> tuple: "name": "markup", "type": "markup", "line": 1, - "language": "svelte", + "language": "svelte", "block_type": "markup", "component_name": component_name, "used_components": list(set(used_components)), @@ -1456,7 +1458,7 @@ def _make_script_chunk(filepath, script_content, start_line, block_type, compone exports = re.findall(r'export\s+(?:let|const|function|class)\s+(\w+)', script_content) functions = re.findall(r'(?:export\s+)?function\s+(\w+)', script_content) imports = re.findall(r'import.*from\s+[\'"]([^\'"]+)[\'"]', script_content) - + chunk_text = f"File: {filepath}\nType: {block_type}\n" if component_name: chunk_text += f"Component: {component_name}\n" @@ -1465,7 +1467,7 @@ def _make_script_chunk(filepath, script_content, start_line, block_type, compone if functions: chunk_text += f"Functions: {', '.join(functions)}\n" chunk_text += f"Code:\n{script_content}" - + metadata = { "file": str(filepath), "name": f"{block_type}_chunk", @@ -1474,14 +1476,14 @@ def _make_script_chunk(filepath, script_content, start_line, block_type, compone "language": "typescript" if block_type == "script" else "css", "block_type": block_type } - + if component_name: metadata["component_name"] = component_name if exports: metadata["exports"] = exports if imports: metadata["imports"] = imports - + return {"text": chunk_text, "metadata": metadata} def parse_svelte_file(filepath: Path) -> List[Dict]: @@ -1538,7 +1540,7 @@ def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple: try: code = filepath.read_text(encoding="utf-8") lines = code.splitlines(keepends=True) - + # Try using rust-analyzer AST if available rust_helper = Path("./tools/parse_rust_ast") if rust_helper.exists(): @@ -1551,12 +1553,12 @@ def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple: timeout=20 ) decls = json.loads(proc.stdout) - + for d in decls: start_line = max(0, d.get("start_line", 1) - 1) end_line = d.get("end_line", start_line + 1) chunk_code = "".join(lines[start_line:end_line]) - + # Build comprehensive metadata metadata = { "file": str(filepath), @@ -1565,7 +1567,7 @@ def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple: "line": start_line + 1, "language": "rust", } - + # Add Rust-specific metadata if d.get("visibility"): metadata["visibility"] = d.get("visibility") @@ -1585,9 +1587,9 @@ def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple: metadata["return_type"] = d.get("return_type") if d.get("parameters"): metadata["parameters"] = d.get("parameters") - + chunk_text = f"File: {filepath}\nType: {d.get('type')}\nName: {d.get('name')}\n" - + # Add Rust-specific details to text if d.get("visibility"): chunk_text += f"Visibility: {d.get('visibility')}\n" @@ -1601,68 +1603,68 @@ def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple: chunk_text += f"Implements: {', '.join(d.get('traits', []))}\n" if d.get("fields"): chunk_text += f"Fields: {', '.join(d.get('fields', []))}\n" - + chunk_text += f"Code:\n{chunk_code}" - + chunks.append({ "text": chunk_text, "metadata": metadata }) - + if chunks: return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks) - + except Exception as e: logger.warning(f"Rust AST helper failed for {filepath}: {e}; falling back to regex") - + # Fallback: Regex-based Rust parsing chunks.extend(_parse_rust_with_regex(filepath, code, lines)) - + except Exception as e: logger.warning(f"parse_rust_file failed {filepath}: {e}") - + return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks) def _parse_rust_with_regex(filepath: Path, code: str, lines: list) -> list: """Fallback regex-based Rust parser.""" chunks = [] - + # Pattern for Rust functions (including async, unsafe, methods) func_pattern = re.compile( r'^(?:pub\s+)?(?:async\s+)?(?:unsafe\s+)?fn\s+(\w+)\s*[<(]', re.MULTILINE ) - + # Pattern for structs struct_pattern = re.compile( r'^(?:pub\s+)?struct\s+(\w+)\s*(?:<[^>]*>)?\s*\{', re.MULTILINE ) - + # Pattern for enums enum_pattern = re.compile( r'^(?:pub\s+)?enum\s+(\w+)\s*\{', re.MULTILINE ) - + # Pattern for traits trait_pattern = re.compile( r'^(?:pub\s+)?trait\s+(\w+)\s*\{', re.MULTILINE ) - + # Pattern for impl blocks impl_pattern = re.compile( r'^impl\s+(?:<[^>]*>)?\s*(\w+)\s*(?:<[^>]*>)?\s*\{', re.MULTILINE ) - + # Pattern for modules mod_pattern = re.compile( r'^(?:pub\s+)?mod\s+(\w+)\s*\{', re.MULTILINE ) - + # Parse functions for match in func_pattern.finditer(code): func_name = match.group(1) @@ -1670,12 +1672,12 @@ def _parse_rust_with_regex(filepath: Path, code: str, lines: list) -> list: start_line = code[:start_pos].count('\n') end_line = _find_rust_brace_block_end(lines, start_line) func_code = "".join(lines[start_line:end_line + 1]) - + # Enhanced function analysis is_async = 'async' in match.group(0) is_unsafe = 'unsafe' in match.group(0) is_pub = 'pub' in match.group(0) - + chunk_text = f"File: {filepath}\nType: function\nName: {func_name}\n" if is_async: chunk_text += "Async: yes\n" @@ -1684,7 +1686,7 @@ def _parse_rust_with_regex(filepath: Path, code: str, lines: list) -> list: if is_pub: chunk_text += "Visibility: pub\n" chunk_text += f"Code:\n{func_code}" - + metadata = { "file": str(filepath), "name": func_name, @@ -1695,9 +1697,9 @@ def _parse_rust_with_regex(filepath: Path, code: str, lines: list) -> list: "is_unsafe": is_unsafe, "visibility": "pub" if is_pub else "private" } - + chunks.append({"text": chunk_text, "metadata": metadata}) - + # Parse structs for match in struct_pattern.finditer(code): struct_name = match.group(1) @@ -1705,15 +1707,15 @@ def _parse_rust_with_regex(filepath: Path, code: str, lines: list) -> list: start_line = code[:start_pos].count('\n') end_line = _find_rust_brace_block_end(lines, start_line) struct_code = "".join(lines[start_line:end_line + 1]) - + # Extract fields from struct fields = _extract_rust_struct_fields(struct_code) - + chunk_text = f"File: {filepath}\nType: struct\nName: {struct_name}\n" if fields: chunk_text += f"Fields: {', '.join(fields)}\n" chunk_text += f"Code:\n{struct_code}" - + metadata = { "file": str(filepath), "name": struct_name, @@ -1722,9 +1724,9 @@ def _parse_rust_with_regex(filepath: Path, code: str, lines: list) -> list: "language": "rust", "fields": fields } - + chunks.append({"text": chunk_text, "metadata": metadata}) - + # Parse enums for match in enum_pattern.finditer(code): enum_name = match.group(1) @@ -1732,15 +1734,15 @@ def _parse_rust_with_regex(filepath: Path, code: str, lines: list) -> list: start_line = code[:start_pos].count('\n') end_line = _find_rust_brace_block_end(lines, start_line) enum_code = "".join(lines[start_line:end_line + 1]) - + # Extract variants from enum variants = _extract_rust_enum_variants(enum_code) - + chunk_text = f"File: {filepath}\nType: enum\nName: {enum_name}\n" if variants: chunk_text += f"Variants: {', '.join(variants)}\n" chunk_text += f"Code:\n{enum_code}" - + metadata = { "file": str(filepath), "name": enum_name, @@ -1749,9 +1751,9 @@ def _parse_rust_with_regex(filepath: Path, code: str, lines: list) -> list: "language": "rust", "variants": variants } - + chunks.append({"text": chunk_text, "metadata": metadata}) - + # Parse traits for match in trait_pattern.finditer(code): trait_name = match.group(1) @@ -1759,15 +1761,15 @@ def _parse_rust_with_regex(filepath: Path, code: str, lines: list) -> list: start_line = code[:start_pos].count('\n') end_line = _find_rust_brace_block_end(lines, start_line) trait_code = "".join(lines[start_line:end_line + 1]) - + # Extract method signatures from trait methods = _extract_rust_trait_methods(trait_code) - + chunk_text = f"File: {filepath}\nType: trait\nName: {trait_name}\n" if methods: chunk_text += f"Methods: {', '.join(methods)}\n" chunk_text += f"Code:\n{trait_code}" - + metadata = { "file": str(filepath), "name": trait_name, @@ -1776,9 +1778,9 @@ def _parse_rust_with_regex(filepath: Path, code: str, lines: list) -> list: "language": "rust", "methods": methods } - + chunks.append({"text": chunk_text, "metadata": metadata}) - + # Parse impl blocks for match in impl_pattern.finditer(code): impl_target = match.group(1) @@ -1786,15 +1788,15 @@ def _parse_rust_with_regex(filepath: Path, code: str, lines: list) -> list: start_line = code[:start_pos].count('\n') end_line = _find_rust_brace_block_end(lines, start_line) impl_code = "".join(lines[start_line:end_line + 1]) - + # Extract methods from impl block impl_methods = _extract_rust_impl_methods(impl_code) - + chunk_text = f"File: {filepath}\nType: impl\nTarget: {impl_target}\n" if impl_methods: chunk_text += f"Implementation Methods: {', '.join(impl_methods)}\n" chunk_text += f"Code:\n{impl_code}" - + metadata = { "file": str(filepath), "name": f"impl_{impl_target}", @@ -1804,9 +1806,9 @@ def _parse_rust_with_regex(filepath: Path, code: str, lines: list) -> list: "target": impl_target, "methods": impl_methods } - + chunks.append({"text": chunk_text, "metadata": metadata}) - + return chunks def _find_rust_brace_block_end(lines, start_line): @@ -2077,24 +2079,24 @@ def build_indexes(): name = m.get("name", "") block_type = m.get("block_type", "") component_name = m.get("component_name") - + # Ensure file node exists file_node_id = f"file::{fpath}" graph.add_node(file_node_id, type="File", path=fpath, lang=lang) - + # Create appropriate node based on type if node_type in ["function", "method", "constructor"]: # Function-like entities if node_type == "method" and "class" in m: # Method belongs to a class node_id = f"{lang}::{node_type}::{fpath}::{m['class']}.{name}" - graph.add_node(node_id, type=node_type, name=name, file=fpath, + graph.add_node(node_id, type=node_type, name=name, file=fpath, class_name=m.get("class"), line=m.get("line"), return_type=m.get("return_type"), parameters=m.get("parameters", [])) # Connect method to its class class_node_id = f"{lang}::class::{fpath}::{m['class']}" graph.add_edge(class_node_id, node_id, "contains") - + return_type = m.get("return_type") if return_type and return_type != "void": # Try to find return type class/interface @@ -2106,25 +2108,25 @@ def build_indexes(): graph.add_node(node_id, type=node_type, name=name, file=fpath, line=m.get("line"), return_type=m.get("return_type"), parameters=m.get("parameters", [])) graph.add_edge(file_node_id, node_id, "contains") - + elif node_type in ["class", "struct", "interface", "enum"]: # Type definitions node_id = f"{lang}::{node_type}::{fpath}::{name}" graph.add_node(node_id, type=node_type, name=name, file=fpath, line=m.get("line"), extends=m.get("extends"), implements=m.get("implements", [])) graph.add_edge(file_node_id, node_id, "contains") - + extends_class = m.get("extends") if extends_class: # Try to find the parent class node for parent_id, parent_data in graph.find_nodes(type="class", name=extends_class): graph.add_edge(node_id, parent_id, "extends") - + implements_interfaces = m.get("implements", []) for interface_name in implements_interfaces: for interface_id, interface_data in graph.find_nodes(type="interface", name=interface_name): graph.add_edge(node_id, interface_id, "implements") - + elif node_type == "field": # Class/struct fields node_id = f"{lang}::field::{fpath}::{name}" @@ -2136,32 +2138,32 @@ def build_indexes(): graph.add_edge(class_node_id, node_id, "contains") else: graph.add_edge(file_node_id, node_id, "contains") - + field_type = m.get("field_type") if field_type and field_type != "unknown": # Try to find the field type class for type_id, type_data in graph.find_nodes(name=field_type): graph.add_edge(node_id, type_id, "type_reference") - + elif node_type in ["style", "markup", "code_block"]: # Code blocks and stylistic elements block_type = m.get("block_type", node_type) node_id = f"{lang}::{block_type}::{fpath}::{name}" - graph.add_node(node_id, type=node_type, name=name, file=fpath, + graph.add_node(node_id, type=node_type, name=name, file=fpath, line=m.get("line"), block_type=block_type) graph.add_edge(file_node_id, node_id, "contains") - + else: # Generic code element fallback node_id = f"{lang}::{node_type}::{fpath}::{name}" graph.add_node(node_id, type=node_type, name=name, file=fpath, line=m.get("line")) graph.add_edge(file_node_id, node_id, "contains") - + # Handle imports if present for imp in m.get("imports", []): graph.add_edge(file_node_id, imp, "imports") - - # Handle function calls if present + + # Handle function calls if present for call in m.get("calls", []): graph.add_edge(node_id, call, "calls") @@ -2176,12 +2178,12 @@ def build_indexes(): graph.add_edge(type_id, node_id, "has_method") for type_id, type_data in graph.find_nodes(type="interface", name=receiver_type): graph.add_edge(type_id, node_id, "has_method") - + elif node_type == "struct": # Connect struct to its fields (if we had field type info) fields = m.get("fields", []) # Could potentially connect fields to their types if we extract field types - + elif node_type == "interface": # Connect interface to its required methods interface_methods = m.get("interface_methods", []) @@ -2191,18 +2193,18 @@ def build_indexes(): # Create component node if it doesn't exist component_node_id = f"svelte::component::{fpath}::{component_name}" if component_node_id not in graph.graph: - graph.add_node(component_node_id, type="component", name=component_name, + graph.add_node(component_node_id, type="component", name=component_name, file=fpath, language="svelte") graph.add_edge(file_node_id, component_node_id, "contains") - + # Connect different parts to the component if node_type in ["script", "script_module", "style", "markup"]: graph.add_edge(component_node_id, node_id, "contains") - + # Component props and interface if node_type == "prop": graph.add_edge(component_node_id, node_id, "has_prop") - + # Component usage relationships if node_type == "markup": used_components = m.get("used_components", []) @@ -2210,19 +2212,19 @@ def build_indexes(): # Try to find the used component in the graph for comp_id, comp_data in graph.find_nodes(type="component", name=used_comp): graph.add_edge(component_node_id, comp_id, "uses_component") - + # Event handler relationships event_handlers = m.get("event_handlers", []) for event in event_handlers: graph.add_edge(component_node_id, f"event::{event}", "handles_event") - + # Svelte script relationships (TypeScript) elif lang == "typescript" and block_type in ["script", "script_module"]: # Connect TypeScript declarations to their Svelte component if component_name: component_node_id = f"svelte::component::{fpath}::{component_name}" graph.add_edge(component_node_id, node_id, "contains") - + # Import relationships for imp in m.get("imports", []): # Try to resolve imports to other components/files @@ -2239,19 +2241,19 @@ def build_indexes(): fields = m.get("fields", []) for field in fields: graph.add_edge(node_id, f"rust::field::{fpath}::{name}.{field}", "has_field") - + elif node_type == "enum": # Connect enum to its variants variants = m.get("variants", []) for variant in variants: graph.add_edge(node_id, f"rust::variant::{fpath}::{name}::{variant}", "has_variant") - + elif node_type == "trait": # Connect trait to its required methods methods = m.get("methods", []) for method in methods: graph.add_edge(node_id, f"rust::trait_method::{name}::{method}", "requires_method") - + elif node_type == "impl": # Connect impl block to its target type target = m.get("target") @@ -2259,13 +2261,13 @@ def build_indexes(): # Try to find the target struct/enum for target_id, target_data in graph.find_nodes(name=target): graph.add_edge(node_id, target_id, "implements_for") - + # Connect impl block to implemented traits traits = m.get("traits", []) for trait_name in traits: for trait_id, trait_data in graph.find_nodes(type="trait", name=trait_name): graph.add_edge(node_id, trait_id, "implements") - + elif node_type == "function": # Connect function to its return type return_type = m.get("return_type") @@ -3006,28 +3008,45 @@ def find_brace_block_end(lines: List[str], start_idx: int) -> Optional[int]: # ----------------------------- @mcp.tool() def health_check() -> str: - """What it does - Quick JSON status of the RAG server (ready, indexed path, chunk count, age). -When to use - Before any other query, to confirm the index is current. + """What it does - Quick JSON status of the RAG server (ready, indexed path, chunk count, age, graph stats). +When to use - Before any other query, to confirm the index is current and see codebase structure. When not to use - After you already know the server is healthy; it adds no value. -Example - health_check() → { "status":"ready","total_chunks":11234,"index_age_hours":4.7 }""" - try: +Example - health_check() → { "status":"ready","total_chunks":11234,"graph_nodes":4567,"graph_edges":12345,"index_age_hours":4.7 }""" +try: with _startup_lock: + # Check if we have a git repo in working_repo + repo_info = "No repository loaded" + if (CODEBASE_PATH / ".git").exists(): + try: + branch = subprocess.check_output( + ["git", "-C", str(CODEBASE_PATH), "rev-parse", "--abbrev-ref", "HEAD"], + text=True + ).strip() + repo_info = f"Loaded: {branch}" + except: + repo_info = "Git repository (no branch info)" + status = { - "status": "ready" if vectorstore is not None else "initializing", - "codebase": str(CODEBASE_PATH), + "status": "ready" if vectorstore is not None else "no_index", + "working_repository": str(CODEBASE_PATH), + "repo_status": repo_info, "ollama_url": OLLAMA_BASE_URL, "config": { "embedding_model": EMBEDDING_MODEL, "vector_weight": VECTOR_WEIGHT, "bm25_weight": BM25_WEIGHT, "rerank_enabled": ENABLE_RERANK, - "batch_size": EMBEDDING_BATCH_SIZE + "batch_size": EMBEDDING_BATCH_SIZE, + "graph_enabled": True }, "Tools": { "NOTE": "THESE TOOLS ARE RESTRICTED BY .gitignore AS WELL AS .ragignore", - "search_codebase": search_codebase.__doc__, - "find_code_references": find_code_references.__doc__, - "read_file_lines_tool": read_file_lines_tool.__doc__ + "search_codebase": "Hybrid RAG search with graph context (cross-file deps, inheritance, calls)", + "find_code_references": "Find exact symbol references across codebase", + "read_file_lines_tool": "Read specific file lines with syntax highlighting", + "find_path": "Find files by glob patterns", + "grep": "Search file contents with regex", + "list_directory": "List directory contents" } } @@ -3040,14 +3059,85 @@ Example - health_check() → { "status":"ready","total_chunks":11234,"index_age_ "index_age_hours": round((time.time() - index_build_time) / 3600, 1) if index_build_time > 0 else None } + # Add graph statistics + if graph_loaded: + graph_stats = _calculate_graph_stats(graph) + status["graph"] = graph_stats + else: + status["graph"] = { + "status": "not_loaded", + "nodes": 0, + "edges": 0, + "message": "Graph will be built on next index rebuild" + } + return json.dumps(status, indent=2) except Exception as e: return json.dumps({"status": "error", "message": str(e)}, indent=2) +def _calculate_graph_stats(graph: LocalGraph) -> dict: + """Calculate comprehensive graph statistics.""" + if graph.graph.number_of_nodes() == 0: + return {"status": "empty", "nodes": 0, "edges": 0} + + # Basic counts + nodes = graph.graph.number_of_nodes() + edges = graph.graph.number_of_edges() + + # Count by node type + node_types = {} + for _, data in graph.graph.nodes(data=True): + node_type = data.get('type', 'unknown') + node_types[node_type] = node_types.get(node_type, 0) + 1 + + # Count by edge type + edge_types = {} + for _, _, data in graph.graph.edges(data=True): + edge_type = data.get('type', 'unknown') + edge_types[edge_type] = edge_types.get(edge_type, 0) + 1 + + # Language distribution + languages = {} + for _, data in graph.graph.nodes(data=True): + lang = data.get('language', 'unknown') + languages[lang] = languages.get(lang, 0) + 1 + + # File statistics + files = [n for n, d in graph.graph.nodes(data=True) if d.get('type') == 'File'] + + # Relationship density + avg_edges_per_node = edges / nodes if nodes > 0 else 0 + + # Most connected nodes (hubs) + degree_centrality = dict(graph.graph.degree()) + top_hubs = sorted(degree_centrality.items(), key=lambda x: x[1], reverse=True)[:5] + top_hubs_info = [] + for node_id, degree in top_hubs: + node_data = graph.graph.nodes[node_id] + top_hubs_info.append({ + "name": node_data.get('name', node_id), + "type": node_data.get('type', 'unknown'), + "file": node_data.get('file', ''), + "connections": degree + }) + + return { + "status": "loaded", + "nodes": nodes, + "edges": edges, + "node_types": node_types, + "edge_types": edge_types, + "languages": languages, + "files": len(files), + "avg_edges_per_node": round(avg_edges_per_node, 2), + "top_connected_nodes": top_hubs_info, + "relationship_density": "high" if avg_edges_per_node > 2.0 else "medium" if avg_edges_per_node > 1.0 else "low" + } + @mcp.tool() def search_codebase(query: str, top_k: int = 5, rerank: bool = True) -> str: """Hybrid RAG search: returns ranked code snippets + file/line + docstrings + graph context (cross-file deps, inheritance, callers/callees, type refs). - + Use for: Open questions ("how does auth work?"), unknown patterns, architecture exploration. Natural language queries. Don't use for: Single symbol/exact phrase (use find_code_references or grep then read_file_lines_tool or read), exploring known entity relationships. @@ -3056,8 +3146,14 @@ rerank=False: fast, 65% accuracy | rerank=True: slower, 95% accuracy Output includes: code content, docs, cross_file:[calls/extends/uses@filepath], used_by:[entity@filepath], file:[same-file entities] Example: search_codebase("user authentication", 5, True) → 5 results with code + "cross_file:calls:ValidateToken@auth/jwt.go|used_by:setupAuth@server/routes.go" """ + if vectorstore is None or bm25_corpus is None: + return "❌ Index not built. Please call init_repo() or rebuild_index() first." + graph = LocalGraph(CODEBASE_PATH) - graph.load() + graph_loaded = graph.load() + + if not graph_loaded: + logger.warning("Graph not loaded - search will proceed without graph context") try: hy = hybrid_search(query, k=max(top_k, RERANK_TOP_N)) @@ -3101,14 +3197,14 @@ def _build_enhanced_result(chunk_text: str, meta: dict, score: float, graph: Opt entity_type = meta.get('type', 'code') language = meta.get('language', '') entity_name = meta.get('name', '') - + # Extract clean content and docstring clean_content = _extract_clean_content(chunk_text) docstring = _extract_docstring(chunk_text) - + # Get graph context for this result graph_context = _get_result_graph_context(meta, graph) - + return { 'file': file_path, 'line': line_num, @@ -3126,13 +3222,13 @@ def _get_result_graph_context(meta: dict, graph: Optional[LocalGraph] = None) -> """Get relevant graph context for a search result with cross-file awareness.""" if graph is None or graph.graph.number_of_nodes() == 0: 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: @@ -3146,21 +3242,21 @@ def _get_result_graph_context(meta: dict, graph: Optional[LocalGraph] = None) -> 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}" - + if node_id and node_id in graph.graph: # Get direct relationships WITH file info neighbors = list(graph.neighbors(node_id)) if neighbors: context['relationships'] = [] context['cross_file_deps'] = [] # NEW: Track cross-file dependencies - + for neighbor_id, neighbor_data in neighbors[:10]: # Show more neighbors 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_file = neighbor_data.get('file', '') neighbor_type = neighbor_data.get('type', '') - + # Build compact relationship string with file info if neighbor_file and neighbor_file != file_path: # Cross-file relationship - show file @@ -3169,9 +3265,9 @@ def _get_result_graph_context(meta: dict, graph: Optional[LocalGraph] = None) -> else: # Same-file relationship - omit file for brevity rel_str = f"{rel_type}:{neighbor_name}" - + context['relationships'].append(rel_str) - + # NEW: Find reverse dependencies (what depends on THIS entity) incoming = [] for predecessor in graph.graph.predecessors(node_id): @@ -3180,14 +3276,14 @@ def _get_result_graph_context(meta: dict, graph: Optional[LocalGraph] = None) -> pred_name = pred_data.get('name', '') edge_data = graph.graph.edges[predecessor, node_id] rel_type = edge_data.get('type', 'related') - + # Only show cross-file incoming dependencies if pred_file and pred_file != file_path: incoming.append(f"{rel_type}:{pred_name}@{pred_file}") - + if incoming: context['used_by'] = incoming[:10] - + # Get file-level context (same as before) file_node_id = f"file::{file_path}" if file_node_id in graph.graph: @@ -3199,7 +3295,7 @@ def _get_result_graph_context(meta: dict, graph: Optional[LocalGraph] = None) -> neighbor_name = neighbor_data.get('name', '') if neighbor_type and neighbor_name and neighbor_name != entity_name: context['file_entities'].append(f"{neighbor_type}:{neighbor_name}") - + return context @@ -3207,12 +3303,12 @@ def _extract_docstring(chunk_text: str) -> str: """Extract docstring or comments from chunk text.""" lines = chunk_text.split('\n') doc_lines = [] - + # Look for docstring patterns in_docstring = False for line in lines: stripped = line.strip() - + # Java/Go style comments if stripped.startswith('//') or stripped.startswith('/*') or stripped.startswith('*'): doc_lines.append(stripped) @@ -3230,13 +3326,13 @@ def _extract_docstring(chunk_text: str) -> str: doc_content = stripped[4:].strip() if doc_content and doc_content not in ['None', '""']: doc_lines.append(doc_content) - + # Clean up docstring if doc_lines: docstring = ' '.join(doc_lines) docstring = re.sub(r'\s+', ' ', docstring) return docstring.strip()[:300] # Reasonable limit - + return "" @@ -3246,43 +3342,43 @@ def _format_enhanced_results(results: List[Dict], query: str, result_type: str) return f"# {result_type.title()}: {query}\nNo results found.\n" lines = [ - f"# {result_type.title()}: {query}", + f"# {result_type.title()}: {query}", f"results[{len(results)}]{{file,line,type,name,score,content,context}}:" ] for result in results: context_parts = [] - + # Graph context as dense key:value pairs graph_context = result.get('graph_context', {}) - + # Cross-file dependencies (highest value) if graph_context.get('cross_file_deps'): deps = '|'.join(graph_context['cross_file_deps'][:8]) context_parts.append(f"xfile:{deps}") - + # All relationships if graph_context.get('relationships'): rels = '|'.join(graph_context['relationships'][:12]) context_parts.append(f"rels:{rels}") - + # Reverse dependencies if graph_context.get('used_by'): used = '|'.join(graph_context['used_by'][:6]) context_parts.append(f"used:{used}") - + # File context if graph_context.get('file_entities'): entities = '|'.join(graph_context['file_entities'][:10]) context_parts.append(f"file:{entities}") - + # Docstring (if valuable) docstring = result.get('docstring', '') if docstring and len(docstring) > 10: # Only include substantial docs context_parts.append(f"doc:{docstring[:200]}") - + context_str = ";".join(context_parts) - + # Minimal row with essential info row = [ result.get('file', ''), @@ -3396,61 +3492,42 @@ Note: This tool treats / as '/home/popertots/Crussell/' so adjust paths accordin @mcp.tool() def init_repo(git_url: str) -> str: - """What it does - Initialise the RAG stack with a brand-new Git repository or update an existing one. - - The function will: - 1. Clone the repo if it does not already exist locally, - 2. Pull the latest changes if it does, - 3. Reset the global `CODEBASE_PATH` to the local clone, - 4. Remove any stale index artefacts, - 5. Re-build the vector/BM25 indexes with the same lock-protected flow used by `rebuild_index`, - 6. Return a human-readable status message. - - Next steps: - • Call `health_check` to confirm that the server is ready. - • Use `search_codebase` or `find_code_references` to explore the new code. - • Use `read_file_lines_tool` to inspect any hit in detail. - """ + """What it does - Initialise the RAG stack with a brand-new Git repository.""" import shutil from pathlib import Path import subprocess - # ----- 1. Determine where to clone --------------------------------- - repo_name = git_url.split('/')[-1] - if repo_name.endswith('.git'): - repo_name = repo_name[:-4] - clone_dir = Path.cwd() / repo_name + # ----- 1. Always wipe and recreate working_repo --------------------- + if CODEBASE_PATH.exists(): + shutil.rmtree(CODEBASE_PATH) + logger.info(f"Wiped existing working repo: {CODEBASE_PATH}") - # ----- 2. Clone / pull ----------------------------------------------- + CODEBASE_PATH.mkdir(exist_ok=True) + + # ----- 2. Clone fresh ----------------------------------------------- try: - if clone_dir.exists(): - # Existing repo → pull - subprocess.run( - ["git", "-C", str(clone_dir), "pull"], - check=True, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - ) - else: - # New repo → clone - subprocess.run( - ["git", "clone", git_url], - cwd=str(Path.cwd()), - check=True, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - ) + subprocess.run( + ["git", "clone", git_url, str(CODEBASE_PATH)], + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + logger.info(f"Cloned {git_url} to {CODEBASE_PATH}") except subprocess.CalledProcessError as e: - return f"❌ Failed to clone/pull {git_url}: {e.stderr.strip()}" + return f"❌ Failed to clone {git_url}: {e.stderr.strip()}" - # ----- 3. Update the global CODEBASE_PATH -------------------------------- - global CODEBASE_PATH - CODEBASE_PATH = clone_dir + # ----- 3. Reset ALL global state ------------------------------------ + global vectorstore, bm25, bm25_corpus, chunks_metadata, index_build_time - # ----- 4. Remove stale indices ---------------------------------------- - for artefact in ("chroma_db", "bm25_index.json", "embeddings"): + vectorstore = None + bm25 = None + bm25_corpus = None + chunks_metadata = None + index_build_time = 0 + + # ----- 4. Remove any stale indices ---------------------------------- + for artefact in ("chroma_db", "bm25_index.json", "embeddings", ".mcp_cache"): artefact_path = CODEBASE_PATH / artefact if artefact_path.exists(): try: @@ -3458,27 +3535,37 @@ def init_repo(git_url: str) -> str: shutil.rmtree(artefact_path) else: artefact_path.unlink() + logger.info(f"Removed stale artefact: {artefact_path}") except Exception as e: logger.warning(f"Could not delete {artefact_path}: {e}") - # ----- 5. Re-build the index (protected by the startup lock) ------------ + # ----- 5. Re-build the index ---------------------------------------- try: - rebuild_index() + build_indexes() + logger.info("Index rebuilt successfully") except Exception as e: + logger.exception("Index rebuild failed") return f"❌ Index rebuild failed: {e}" - # ----- 6. Report success ----------------------------------------------- - branch_name = subprocess.check_output( - ["git", "-C", str(CODEBASE_PATH), "rev-parse", "--abbrev-ref", "HEAD"], - text=True, - ).strip() - git_hash = subprocess.check_output( - ["git", "-C", str(CODEBASE_PATH), "rev-parse", "HEAD"], - text=True, - ).strip() + # ----- 6. Report success -------------------------------------------- + try: + branch_name = subprocess.check_output( + ["git", "-C", str(CODEBASE_PATH), "rev-parse", "--abbrev-ref", "HEAD"], + text=True, + ).strip() + git_hash = subprocess.check_output( + ["git", "-C", str(CODEBASE_PATH), "rev-parse", "HEAD"], + text=True, + ).strip() - return f"✅ Initialized repository at {CODEBASE_PATH} on branch {branch_name} with hash {git_hash}" + # Verify the index was actually built + if vectorstore is None or bm25_corpus is None: + return f"⚠️ Repository cloned but index failed to build properly" + return f"✅ Initialized repository at {CODEBASE_PATH}\n Branch: {branch_name}\n Commit: {git_hash[:8]}\n Indexed chunks: {len(bm25_corpus)}" + + except Exception as e: + return f"⚠️ Repository cloned but status check failed: {e}" @mcp.tool() def rebuild_index() -> str: